MiniMax-H3 VAE decoder loads in fp32 but CUDA decode runs it in fp16

Describe the bug

The VAE keeps everything in fp32 no matter what dtype you pass, but the decode block wraps vae.decode in a fp16 autocast, so the decoder gets downcast to fp16 anyway on every call. On CUDA this just costs VRAM, RAM and disk space.

Decoding 192 frames at 1344x768:

As seen on the table, this prevents to use this model on 16GB GPUs for that resolution and duration.

Reproduction

Just normal inference on a cuda GPU will reproduce it

import torch

from diffusers import ComponentsManager, ModularPipeline


manager = ComponentsManager()
manager.enable_auto_cpu_offload(device="cuda")

pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3", components_manager=manager)
pipe.load_components(workflow="t2va", dtype=torch.bfloat16)

results = pipe(
    prompt="An astronaut hiking through the mountains, humming a tune",
    num_frames=192,
    output=["videos", "audio", "sampling_rate"],
)

code to test a PoC fix:

import torch

from diffusers import AutoencoderKLMiniMaxH3


class Fp16DecoderVae(AutoencoderKLMiniMaxH3):
    _keep_in_fp32_modules = [
        "encoder",
        "quant_conv",
        "post_quant_conv",
        "norm1",
        "norm2",
        "norm_out",
        "scale1",
        "scale2",
    ]

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        kwargs.pop("torch_dtype", None)
        kwargs["dtype"] = torch.float16
        return super().from_pretrained(*args, **kwargs)


vae = Fp16DecoderVae.from_pretrained("MiniMaxAI/MiniMax-H3", subfolder="vae", dtype=torch.bfloat16)

main

https://github.com/user-attachments/assets/5bfbfedb-cfed-4645-87b7-58d9c8e8fa9f

PoC fix

https://github.com/user-attachments/assets/72429905-03b6-4590-bee0-f4fd2e4a39d4

System Info

diffusers from main
any cuda platform

Who can help?

Commencez par l’implémentation de AutoencoderKLMiniMaxH3 et suivez le comportement d’autocast du bloc de decode pendant la reproduction d’inférence CUDA fournie. Comparez-le avec la PoC de Fp16DecoderVae, puis vérifiez que le decoder utilise le dtype prévu et que le cas de 192 frames en 1344x768 n’entraîne plus l’utilisation mémoire fp32 signalée.