[MInimax H3] Fix VAE decode

[asomoza]

What does this PR do?

Drops the fp16 autocast in the video decode block and replaces decoder in _keep_in_fp32_modules with just the sensitive layers, so it runs at the dtype you pass instead of always fp32 weights with fp16 compute.

Error vs a true fp32 decode with no autocast:

	weights	peak	mean err	max err	pixels >1/255
main (fp32 weights + fp16 autocast)	9.70 GiB	16.69 GiB	9.64e-5	0.0182	0.0010%
fp32	9.70 GiB	14.58 GiB	- (reference)	-	-
fp16	5.19 GiB	7.66 GiB	7.34e-5	0.0060	0
bf16	5.19 GiB	7.66 GiB	5.86e-4	0.0514	0.057%

Reference links here and here.

bf16 is the worst of them but I don't see it visually when generating the video, still the reference doesn't use it and other libraries don't either, so this PR forces it to fp16 anyway it's respected now.

There's a user facing change with this, I now respect the dtype they set on the pipeline, so the generation changes for everyone except fp32 on CPU/MPS, including the bf16 recipe in the docs. Probably imperceptible to most eyes but still a difference..

videos:

Fp32

decoded_fp32.mp4

fp16

decoded_fp16.mp4

bf16

decoded_bf16.mp4

Fixes #14746

Who can review?

@yiyixuxu @sayakpaul

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[asomoza]

/diffusers-bot pytest tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py

[github-actions]

❌ pytest tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py failed on GPU — view logs.

[asomoza]

/diffusers-bot pytest tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py

[github-actions]

✅ pytest tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py passed on GPU — view logs.

[yiyixuxu]

Suggested change

	@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)

[asomoza]

/diffusers-bot pytest tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py

[github-actions]

✅ pytest tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py passed on GPU — view logs.

[sayakpaul]

Drops the fp16 autocast in the video decode block and narrows the VAE's _keep_in_fp32_modules

I am having difficulty understanding this. Where was this autocast block? I would have expected to see torch.autocast(...) but I didn't see that being removed in the diff.

How can we set up testing for this in a better manner?

Also, we're not trimming _keep_in_fp32_modules, we're adding more names to that list, no? How is that narrowing the list down?

[vladmandic]

Also, we're not trimming _keep_in_fp32_modules, we're adding more names to that list, no? How is that narrowing the list down?

it drops entire decoder in favor of individual layers/groups

[asomoza]

yeah, it narrows what's pinned inside the decoder, specific layers instead of the whole module, but I see how it can be misinterpreted so I changed the wording.

about the tests, what do you think about this (proposed by claude):

    @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, 2, 2), return_dict=False)[0]
        assert decoded.dtype == dtype

[sayakpaul]

@asomoza the test looks neat to me! Let's add it.