diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index 23ad2c725c00..b75a84dcf203 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/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: 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) 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..c7871b0620e2 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. + # 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): @@ -110,7 +108,16 @@ 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 == 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 "