h3_quantization.py

"""Owner-opt-in INT8 loading. Conversion is at startup, never inside a GPU job.

Based on the official MiniMax-H3 Diffusers memory recipe. Kept separate because
GPU/PEFT/ZeroGPU compatibility must be measured on the owner's actual runtime.
"""


def load_int8_transformer(repo):
    import torch
    try:
        from diffusers import MiniMaxH3Transformer3DModel, TorchAoConfig
        from torchao.quantization import Int8WeightOnlyConfig
    except ImportError as error:
        raise RuntimeError(
            "INT8 needs torchao==0.17.0 in the main requirements.txt. "
            "Use H3_QUANTIZATION=bf16 for the existing supported path."
        ) from error
    model = MiniMaxH3Transformer3DModel.from_pretrained(
        repo, subfolder="transformer_ref", dtype=torch.bfloat16,
        quantization_config=TorchAoConfig(
            Int8WeightOnlyConfig(version=2),
            modules_to_not_convert=[
                "proj_in", "audio_proj_in", "context_embedder", "time_embedder", "time_proj",
                "token_refiner", "norm_out", "proj_out", "audio_proj_out",
            ],
        ),
        # The pinned Diffusers ModelMixin rejects False with a quantizer.
        # Accelerate is already pinned in requirements.txt for this load path.
        low_cpu_mem_usage=True,
    )
    model.requires_grad_(False)
    return model