[MInimax H3] Fix VAE decode

src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py

@@ -387,10 +387,12 @@ def __init__(
     def forward(
         self, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None
     ) -> torch.Tensor:
-        # The reference normalizes in float32 regardless of the compute dtype.
-        norm_hidden_states = self.norm1(hidden_states.float()).to(hidden_states.dtype)
+        # The reference normalizes in float32 regardless of the compute dtype. The residual stream is float32 too
+        # because the anchors are, so the normed activation follows the projections instead.
+        compute_dtype = get_parameter_dtype(self.attn.to_q)
+        norm_hidden_states = self.norm1(hidden_states.float()).to(compute_dtype)
         hidden_states = hidden_states + self.attn(norm_hidden_states, rotary_emb) * self.scale1
-        norm_hidden_states = self.norm2(hidden_states.float()).to(hidden_states.dtype)
+        norm_hidden_states = self.norm2(hidden_states.float()).to(compute_dtype)
         hidden_states = hidden_states + self.ff(norm_hidden_states) * self.scale2
         return hidden_states
 
@@ -453,7 +455,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = self.proj_in(hidden_states)
         num_patches = hidden_states.shape[1]
 
-        register_tokens = self.register_tokens.expand(batch_size, -1, -1)
+        # `proj_in` is pinned to float32 and `register_tokens` is not, so align rather than rely on promotion.
+        register_tokens = self.register_tokens.expand(batch_size, -1, -1).to(hidden_states.dtype)
         cls_token = torch.zeros_like(hidden_states[:, :1, :])
         hidden_states = torch.cat([hidden_states, register_tokens, cls_token], dim=1)
 
@@ -473,7 +476,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             else:
                 hidden_states = block(hidden_states, rotary_emb)
 
-        hidden_states = self.norm_out(hidden_states)
+        # `norm_out` is pinned to float32 and `proj_out` is not.
+        hidden_states = self.norm_out(hidden_states).to(get_parameter_dtype(self.proj_out))
         hidden_states = self.proj_out(hidden_states)
         hidden_states = hidden_states[:, :num_patches, :]
 
@@ -526,10 +530,20 @@ class AutoencoderKLMiniMaxH3(ModelMixin, ConfigMixin, AttentionMixin, Autoencode
     _no_split_modules = ["MiniMaxH3VideoResnetBlock3d", "MiniMaxH3VideoTransformerBlock"]
     _repeated_blocks = ["MiniMaxH3VideoTransformerBlock"]
     _skip_layerwise_casting_patterns = ["norm"]
-    # The released checkpoint is float32 and the verified decode recipe is float16 *autocast over float32 weights*
-    # (see `decode`). A pipeline-level `torch_dtype=torch.bfloat16` must therefore not downcast the weights, so every
-    # top-level module is pinned, mirroring the transformer's mixed-precision contract.
-    _keep_in_fp32_modules = ["encoder", "decoder", "quant_conv", "post_quant_conv"]
+    # The released checkpoint is float32; only the cast-sensitive modules are pinned so `torch_dtype` still reaches
+    # the decoder block stack. `proj_out` is left out on purpose: it sets the dtype of the decoded pixels, and
+    # pinning it holds the whole video in float32. Entries match whole segments of the parameter name.
+    _keep_in_fp32_modules = [
+        "encoder",
+        "quant_conv",
+        "post_quant_conv",
+        "proj_in",
+        "norm1",
+        "norm2",
+        "norm_out",
+        "scale1",
+        "scale2",
+    ]
 
     @register_to_config
     def __init__(
@@ -857,7 +871,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutp
             The latent distribution of the encoded videos. Note that MiniMax-H3 normalizes the sampled latents with
             `latents_mean` / `latents_std` afterwards.
         """
-        # Every module is pinned to float32 by `_keep_in_fp32_modules`, so a pipeline running in a lower `torch_dtype`
+        # The encoder is pinned to float32 by `_keep_in_fp32_modules`, so a pipeline running in a lower `torch_dtype`
         # hands over lower-precision pixels; align them with the weights, like the audio autoencoder does.
         x = x.to(get_parameter_dtype(self.encoder))
         if self.use_slicing and x.shape[0] > 1:
@@ -884,7 +898,8 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t
             [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
                 The decoded videos, shape `(batch_size, out_channels, num_frames, height, width)`.
         """
-        z = z.to(get_parameter_dtype(self.decoder))
+        # The decoder is mixed precision, so align with the first module that consumes the latents.
+        z = z.to(get_parameter_dtype(self.post_quant_conv))
         if self.use_slicing and z.shape[0] > 1:
             decoded = torch.cat([self._decode(z_slice) for z_slice in z.split(1)])
         else:

src/diffusers/modular_pipelines/minimax_h3/decoders.py

@@ -134,9 +134,8 @@ class MiniMaxH3VideoDecodeStep(ModularPipelineBlocks):
     def description(self) -> str:
         return (
             "Denormalizes the generated video latents and decodes them into video. The spatial tiling of the video "
-            "VAE covers the canvas exactly, so the decoded frames need no crop back, but the decode itself runs under "
-            "float16 autocast even though the VAE weights are float32, and the VAE produces ImageNet-normalized RGB "
-            "that is reverted here."
+            "VAE covers the canvas exactly, so the decoded frames need no crop back, and the VAE produces "
+            "ImageNet-normalized RGB that is reverted here."
         )
 
     @property
@@ -184,8 +183,7 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -
         latents_std = torch.tensor(components.vae.config.latents_std, device=device).view(1, -1, 1, 1, 1)
         latents = block_state.latents * latents_std + latents_mean
 
-        with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
-            video = components.vae.decode(latents, return_dict=False)[0]
+        video = components.vae.decode(latents, return_dict=False)[0]
         pixel_mean = torch.tensor(components.pixel_mean, device=device).view(1, -1, 1, 1, 1)
         pixel_std = torch.tensor(components.pixel_std, device=device).view(1, -1, 1, 1, 1)
         video = (video.float() * pixel_std + pixel_mean).clamp(0, 1)

tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py

@@ -96,21 +96,28 @@ class TestAutoencoderKLMiniMaxH3(AutoencoderKLMiniMaxH3TesterConfig, ModelTester
 
     @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
     def test_from_save_pretrained_dtype(self, tmp_path, dtype):
-        # `_keep_in_fp32_modules` pins every module: the released checkpoint is float32 and decoding through
-        # downcast weights degrades (the bfloat16 audio VAE decodes roughly 20 dB too quiet), so a requested
-        # `torch_dtype` cast at load time must be refused and the weights must stay float32.
+        # The requested dtype reaches the decoder; only `_keep_in_fp32_modules` stays float32.
         model = self.model_class(**self.get_init_dict())
         model.save_pretrained(tmp_path)
         new_model = self.model_class.from_pretrained(tmp_path, torch_dtype=dtype)
-        assert new_model.dtype == torch.float32
+        assert new_model.dtype == dtype
 
     @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
     def test_from_pretrained_dtype_alias(self, tmp_path, dtype):
         # Same contract through the `dtype` alias.
         model = self.model_class(**self.get_init_dict())
         model.save_pretrained(tmp_path)
         new_model = self.model_class.from_pretrained(tmp_path, dtype=dtype)
-        assert new_model.dtype == torch.float32
+        assert new_model.dtype == dtype
+
+    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
+    def test_decode_in_low_precision(self, tmp_path, dtype):
+        # Decode is mixed precision with no autocast, so every dtype seam has to hold on its own.
+        self.model_class(**self.get_init_dict()).save_pretrained(tmp_path)
+        model = self.model_class.from_pretrained(tmp_path, dtype=dtype).eval()
+        with torch.no_grad():
+            decoded = model.decode(torch.randn(1, 4, 7, HEIGHT // 4, WIDTH // 4), return_dict=False)[0]
+        assert decoded.dtype == dtype
 
     @pytest.mark.skip(
         "`forward` runs through the `apply_forward_hook`-decorated `encode` and `decode`, and that decorator's "