diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index 0d8367c6..c08fff20 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -21,9 +21,10 @@ nvfp4 — with dequantized-matmul fallbacks for GPUs without the fast kernels), and the fp16/fp32 single-file VAEs. Files are resolved under ``MODELS_PATH`` (checked first, both at the repo-relative location and flat at the root) and downloaded from the hub into ``MODELS_PATH`` when missing. Individual files -can be overridden via ``model_kwargs``: ``dit_path``, ``text_encoder_path``, -``video_vae_path``, ``audio_vae_path``; ``model_kwargs.partition`` picks -``fl2va`` (default) or ``ref2va``. +can be overridden via ``model_kwargs``: ``dit__path``, +``text_encoder_path``, ``video_vae_path``, ``audio_vae_path``; +``model_kwargs.partition`` picks ``fl2va``, ``fl2va_pruned`` (default), +``ref2va``, or ``ref2va_pruned``. Conventions bridged to ai-toolkit: - the model consumes t = 1 - sigma in [0, 1] (t=1 clean) and predicts the @@ -93,8 +94,10 @@ scheduler_config = { # missing. COMFY_REPO = "Comfy-Org/MiniMax-H3" COMFY_FILES = { - "dit_fl2va": "diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors", - "dit_ref2va": "diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors", + "dit_fl2va": "diffusion_models/minimax_h3_fl2va_int8_convrot.safetensors", + "dit_fl2va_pruned": "diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors", + "dit_ref2va": "diffusion_models/minimax_h3_ref2va_int8_convrot.safetensors", + "dit_ref2va_pruned": "diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors", "text_encoder": "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "video_vae": "vae/minimax_h3_video_vae_fp16.safetensors", "audio_vae": "vae/minimax_h3_audio_vae_fp32.safetensors", @@ -268,11 +271,12 @@ class MinimaxH3Model(BaseModel): def _dit_component(self) -> str: partition = str( - self.model_config.model_kwargs.get("partition", "fl2va") + self.model_config.model_kwargs.get("partition", "fl2va_pruned") ).lower() - if partition not in ("fl2va", "ref2va"): + if partition not in ("fl2va", "fl2va_pruned", "ref2va", "ref2va_pruned"): raise ValueError( - f"model_kwargs.partition must be fl2va or ref2va, got {partition}" + "model_kwargs.partition must be fl2va, fl2va_pruned, ref2va, " + f"or ref2va_pruned, got {partition}" ) return f"dit_{partition}" diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py b/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py index a5a51387..f958268e 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py @@ -62,13 +62,19 @@ class MiniMaxH3TransformerParams: # MLP with a small lookup table: ``adaln_t_table`` of shape # (adaln_t_table_size, time_embed_dim) sampled by linear interpolation at # t * (size - 1), consumed by the AdaLN projections WITHOUT the SiLU. - # These checkpoints also shrink time_embed_dim (8 in the released files). + # These checkpoints also shrink time_embed_dim (8 in the released files) + # and add a bias to the block AdaLN linears, which the original weights + # lack (the final layer's AdaLN carries a bias in both variants). adaln_t_table_size: Optional[int] = None @property def adaln_apply_silu(self) -> bool: return self.adaln_t_table_size is None + @property + def adaln_bias(self) -> bool: + return self.adaln_t_table_size is not None + class MiniMaxH3Rope(nn.Module): """3-axis rotary embedding over the packed (t, h, w) coordinates. @@ -206,13 +212,14 @@ class MiniMaxH3AdalnProj(nn.Module): expand: int, modalities: int, apply_silu: bool = True, + bias: bool = True, ): super().__init__() self.expand = expand self.modalities = modalities self.hidden = hidden self.apply_silu = apply_silu - self.linear = nn.Linear(t_dim, expand * hidden * modalities, bias=True) + self.linear = nn.Linear(t_dim, expand * hidden * modalities, bias=bias) def forward(self, temb: torch.Tensor): if self.apply_silu: @@ -282,6 +289,7 @@ class MiniMaxH3Block(nn.Module): expand=6, modalities=MODALITY_NUM, apply_silu=p.adaln_apply_silu, + bias=p.adaln_bias, ) def forward( @@ -328,6 +336,7 @@ class MiniMaxH3FinalLayer(nn.Module): expand=2, modalities=1, apply_silu=p.adaln_apply_silu, + bias=True, ) self.video_out = nn.Linear(p.hidden_size, video_patch_dim, bias=True) self.audio_out = nn.Linear(p.hidden_size, p.audio_latents_dim, bias=True)