From 7c55d67d79675ea234d4c547956bcb1611f56b44 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Somoza?= Date: Thu, 10 Sep 2026 20:53:16 -0300 Subject: [PATCH 1/6] intiial fix --- .../autoencoders/autoencoder_kl_minimax_h3.py | 52 +++++++++++++++---- .../modular_pipelines/minimax_h3/decoders.py | 8 ++- 2 files changed, 44 insertions(+), 16 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index 23ad2c725c00..c5d01c9113c9 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py @@ -13,6 +13,7 @@ # limitations under the License. import math +import os import torch import torch.nn as nn @@ -387,10 +388,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 +456,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 +477,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 +531,34 @@ 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", + ] + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs): + r""" + Load a pretrained MiniMax-H3 video autoencoder. + + Takes the same arguments as [`ModelMixin.from_pretrained`], except that a requested `bfloat16` is loaded as + `float16`. Pipelines apply one dtype to every component, and this decoder degrades in `bfloat16` without + using any less memory than `float16`. + """ + for key in ("dtype", "torch_dtype"): + if kwargs.get(key) == torch.bfloat16: + kwargs[key] = torch.float16 + return super().from_pretrained(pretrained_model_name_or_path, **kwargs) @register_to_config def __init__( @@ -857,7 +886,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 +913,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: diff --git a/src/diffusers/modular_pipelines/minimax_h3/decoders.py b/src/diffusers/modular_pipelines/minimax_h3/decoders.py index 44e9b2034f11..e5b624cdb6c4 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/decoders.py +++ b/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) From c6a25e4a13dfe5ae5aa54506c4c537ff5dde0ddf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Somoza?= Date: Thu, 10 Sep 2026 21:36:23 -0300 Subject: [PATCH 2/6] fix code quality --- .../models/autoencoders/autoencoder_kl_minimax_h3.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index c5d01c9113c9..6eb6f893aecb 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py @@ -552,8 +552,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None Load a pretrained MiniMax-H3 video autoencoder. Takes the same arguments as [`ModelMixin.from_pretrained`], except that a requested `bfloat16` is loaded as - `float16`. Pipelines apply one dtype to every component, and this decoder degrades in `bfloat16` without - using any less memory than `float16`. + `float16`. Pipelines apply one dtype to every component, and this decoder degrades in `bfloat16` without using + any less memory than `float16`. """ for key in ("dtype", "torch_dtype"): if kwargs.get(key) == torch.bfloat16: From f416bab1482fb9e75380c6417b732e6e58c82030 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Somoza?= Date: Thu, 10 Sep 2026 21:55:33 -0300 Subject: [PATCH 3/6] fix failing pr test --- .../autoencoders/test_models_autoencoder_kl_minimax_h3.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index 2ffcd43af216..93ea0729b4aa 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -96,13 +96,11 @@ 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. + # A requested bfloat16 is loaded as float16; the modules in `_keep_in_fp32_modules` stay 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 == torch.float16 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) def test_from_pretrained_dtype_alias(self, tmp_path, dtype): @@ -110,7 +108,7 @@ def test_from_pretrained_dtype_alias(self, tmp_path, dtype): 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 == torch.float16 @pytest.mark.skip( "`forward` runs through the `apply_forward_hook`-decorated `encode` and `decode`, and that decorator's " From 06a0c5bcfb0630794e0390310d1e83c842c3724b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Somoza?= Date: Thu, 10 Sep 2026 22:36:33 -0300 Subject: [PATCH 4/6] fix gpu test --- .../test_models_autoencoder_kl_minimax_h3.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index 93ea0729b4aa..7772943bf477 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -19,7 +19,7 @@ from diffusers import AutoencoderKLMiniMaxH3 from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, require_accelerator, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -110,6 +110,18 @@ def test_from_pretrained_dtype_alias(self, tmp_path, dtype): new_model = self.model_class.from_pretrained(tmp_path, dtype=dtype) assert new_model.dtype == torch.float16 + @require_accelerator + @pytest.mark.skipif( + torch_device not in ["cuda", "xpu"], + reason="float16 and bfloat16 can only be use for inference with an accelerator", + ) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + @torch.no_grad() + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype, atol=1e-4, rtol=0): + # The shared test builds its own reference model, so it needs the dtype bfloat16 is loaded as. + loaded_dtype = torch.float16 if dtype == torch.bfloat16 else dtype + super().test_from_save_pretrained_dtype_inference(tmp_path, loaded_dtype, atol=atol, rtol=rtol) + @pytest.mark.skip( "`forward` runs through the `apply_forward_hook`-decorated `encode` and `decode`, and that decorator's " "`pre_forward` call clears the input device accelerate's `AlignDevicesHook` recorded for the caller, so the " From 28a3062c8f3ce997396973145e85812717d663c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Somoza?= Date: Fri, 11 Sep 2026 02:43:16 -0300 Subject: [PATCH 5/6] don't override the requested dtype --- .../autoencoders/autoencoder_kl_minimax_h3.py | 15 -------------- .../test_models_autoencoder_kl_minimax_h3.py | 20 ++++--------------- 2 files changed, 4 insertions(+), 31 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index 6eb6f893aecb..b75a84dcf203 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py @@ -13,7 +13,6 @@ # limitations under the License. import math -import os import torch import torch.nn as nn @@ -546,20 +545,6 @@ class AutoencoderKLMiniMaxH3(ModelMixin, ConfigMixin, AttentionMixin, Autoencode "scale2", ] - @classmethod - def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs): - r""" - Load a pretrained MiniMax-H3 video autoencoder. - - Takes the same arguments as [`ModelMixin.from_pretrained`], except that a requested `bfloat16` is loaded as - `float16`. Pipelines apply one dtype to every component, and this decoder degrades in `bfloat16` without using - any less memory than `float16`. - """ - for key in ("dtype", "torch_dtype"): - if kwargs.get(key) == torch.bfloat16: - kwargs[key] = torch.float16 - return super().from_pretrained(pretrained_model_name_or_path, **kwargs) - @register_to_config def __init__( self, diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index 7772943bf477..ff67c61a7bda 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -19,7 +19,7 @@ from diffusers import AutoencoderKLMiniMaxH3 from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, require_accelerator, torch_device +from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -96,11 +96,11 @@ 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): - # A requested bfloat16 is loaded as float16; the modules in `_keep_in_fp32_modules` 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.float16 + 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): @@ -108,19 +108,7 @@ def test_from_pretrained_dtype_alias(self, tmp_path, dtype): 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.float16 - - @require_accelerator - @pytest.mark.skipif( - torch_device not in ["cuda", "xpu"], - reason="float16 and bfloat16 can only be use for inference with an accelerator", - ) - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - @torch.no_grad() - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype, atol=1e-4, rtol=0): - # The shared test builds its own reference model, so it needs the dtype bfloat16 is loaded as. - loaded_dtype = torch.float16 if dtype == torch.bfloat16 else dtype - super().test_from_save_pretrained_dtype_inference(tmp_path, loaded_dtype, atol=atol, rtol=rtol) + assert new_model.dtype == dtype @pytest.mark.skip( "`forward` runs through the `apply_forward_hook`-decorated `encode` and `decode`, and that decorator's " From 6ac16a89ab78a5a74da9fc380e3a74ed8eb80f7c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Somoza?= Date: Fri, 11 Sep 2026 16:28:26 -0300 Subject: [PATCH 6/6] add test --- .../test_models_autoencoder_kl_minimax_h3.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index ff67c61a7bda..c7871b0620e2 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -110,6 +110,15 @@ def test_from_pretrained_dtype_alias(self, tmp_path, dtype): new_model = self.model_class.from_pretrained(tmp_path, dtype=dtype) 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 " "`pre_forward` call clears the input device accelerate's `AlignDevicesHook` recorded for the caller, so the "