FL2VA/video_vae/minimax_h3_video_vae.py
5.0 KB · 123 lines · python Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # Remote entry: self-contained MiniMax H3 visual VAE (3D CNN encoder + ViT3D decoder).
3 # Loaded via config.json:auto_map with trust_remote_code.
4 from __future__ import annotations
5
6 import json
7 from pathlib import Path
8
9 import safetensors.torch
10 import torch.nn as nn
11
12 # --- dependency manifest ---
13 # diffusers' dynamic-module loader only copies ONE level of relative
14 # imports into its cache; list every bundle module here so all files
15 # are copied, letting their own second-level imports resolve.
16 from .attention import Attention as _dep_attention # noqa: F401
17 from .base_module import FeedForward as _dep_base_module # noqa: F401
18 from .conv import SpatialParallelConv3d as _dep_conv # noqa: F401
19 from .flash import make_block_causal_mask_mod as _dep_flash # noqa: F401
20 from .func import create_token_ids as _dep_func # noqa: F401
21 from .klvae import AutoencoderKL as _dep_klvae # noqa: F401
22 from .norm import FusedGroupNorm3D as _dep_norm # noqa: F401
23 from .normalize import get_norm_constants as _dep_normalize # noqa: F401
24 from .parallel import get_parallel_state as _dep_parallel # noqa: F401
25 from .utils import apply_spatial_parallel as _dep_utils # noqa: F401
26 from .vae_cnn import EncoderFCN3D as _dep_vae_cnn # noqa: F401
27 from .vae_module import DiagonalGaussianDistribution as _dep_vae_module # noqa: F401
28 from .vae_processor import VAEProcessor as _dep_vae_processor # noqa: F401
29 from .vae_vit import ViTBase as _dep_vae_vit # noqa: F401
30 # --- end dependency manifest ---
31
32 from .klvae import AutoencoderKLLegacy
33 from .parallel import get_parallel_state
34
35 _SOURCE_CLASSES = {
36 "AutoencoderKLLegacy": AutoencoderKLLegacy,
37 }
38
39
40 def _ensure_vae_parallel_state() -> None:
41 """Seed the bundled VAE parallel state for single-process inference."""
42 state = get_parallel_state()
43 if not isinstance(state, dict):
44 raise TypeError("get_parallel_state() must return a dict")
45 if state:
46 return
47 state.update(
48 {
49 "group_size": 1,
50 "group_rank": 0,
51 "local_process_group": None,
52 "sp_size": 1,
53 "sp_rank": 0,
54 "sp_enabled": False,
55 "sp_process_group": None,
56 "tp_size": 1,
57 "tp_rank": 0,
58 }
59 )
60
61
62 class MiniMaxH3VideoVAE(nn.Module):
63 def __init__(self, model: nn.Module) -> None:
64 super().__init__()
65 self.model = model
66
67 @classmethod
68 def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
69 component_dir = Path(pretrained_model_name_or_path)
70 with (component_dir / "config.json").open("r", encoding="utf-8") as f:
71 config = json.load(f)
72 source_path = component_dir / config["source_path"]
73 source_class_name = config["source_class_name"]
74 if source_class_name not in _SOURCE_CLASSES:
75 raise ValueError(
76 f"unsupported source_class_name {source_class_name!r}; "
77 f"bundled: {sorted(_SOURCE_CLASSES)}"
78 )
79 source_cls = _SOURCE_CLASSES[source_class_name]
80 if "source_safetensors_path" not in config:
81 raise ValueError(
82 "source_safetensors_path is required; pickle checkpoints are "
83 "not supported"
84 )
85 weights_path = source_path / config["source_safetensors_path"]
86 if not weights_path.is_file():
87 raise FileNotFoundError(f"source weights not found: {weights_path}")
88 if bool(config["vae_parallel_tiling"]):
89 _ensure_vae_parallel_state()
90 load_kwargs = {
91 "clip_length": int(config["vae_clip_length"]),
92 "token_drop": int(config["vae_token_drop"]),
93 "encoder_tiling": int(config["vae_encoder_tiling"]),
94 "decoder_tiling": int(config["vae_decoder_tiling"]),
95 "parallel_tiling": int(config["vae_parallel_tiling"]),
96 "tile_size": int(config["vae_tile_size"]),
97 "tile_overlap_min": int(config["vae_tile_overlap_min"]),
98 "encoder_parallel": int(config["vae_encoder_parallel"]),
99 "decoder_parallel": int(config["vae_decoder_parallel"]),
100 "chunk_dim": int(config["vae_chunk_dim"]),
101 }
102 # Mirror diffusers ModelMixin.from_pretrained instantiation semantics
103 # (config-driven init via from_config) but load the state dict from an
104 # explicitly named safetensors file instead of the diffusers default
105 # weight filename.
106 source_config = source_cls.load_config(str(source_path))
107 model, _unused = source_cls.from_config(
108 source_config, return_unused_kwargs=True, **load_kwargs
109 )
110 state_dict = safetensors.torch.load_file(str(weights_path))
111 model.load_state_dict(state_dict, strict=True)
112 model.eval()
113 return cls(model)
114
115 def forward(self, *args, **kwargs):
116 return self.model(*args, **kwargs)
117
118 def __getattr__(self, name: str):
119 try:
120 return super().__getattr__(name)
121 except AttributeError:
122 return getattr(self.model, name)
123