Model-Optimizer/tests/examples/diffusers/test_diffusers.py

# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from pathlib import Path
from typing import NamedTuple

import pytest
from _test_utils.examples.models import FLUX_SCHNELL_PATH, SD3_PATH, SDXL_PATH
from _test_utils.examples.run_command import run_example_command
from _test_utils.torch.misc import minimum_sm

# Calibrate from local prompts: the default Hub datasets (OpenVid-1M for Wan) are slow to load
# and add a network dependency without adding coverage.
CALIB_PROMPTS_FILE = Path(__file__).parent / "calib_prompts.txt"

# Run the example scripts in this process: ~20s of startup saved per step (see conftest).
pytestmark = pytest.mark.usefixtures("in_process_examples")


def assert_hf_ckpt_exported(hf_ckpt_dir: Path) -> None:
    assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}"

    config_files = list(hf_ckpt_dir.rglob("config.json"))
    assert len(config_files) > 0, f"No config.json found in {hf_ckpt_dir}"

    weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin"))
    assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}"


class DiffuserModel(NamedTuple):
    dtype: str
    name: str
    path: str
    format_type: str
    quant_algo: str
    collect_method: str
    # Also export an HF checkpoint from the quantize run, rather than calibrating again in a
    # separate export test.
    hf_export: bool = False

    def _run_cmd(self, script: str, *args: str) -> str:
        cmd_args = [
            "python",
            script,
            "--model",
            self.name,
            "--override-model-path",
            self.path,
        ]
        cmd_args.extend(args)
        return run_example_command(cmd_args, "diffusers/quantization")

    def _format_args(self) -> list[str]:
        return [
            "--calib-size",
            "4",
            "--percentile",
            "1.0",
            "--alpha",
            "0.8",
            "--n-steps",
            "2",
            "--batch-size",
            "2",
            "--format",
            self.format_type,
            "--collect-method",
            self.collect_method,
            "--quant-algo",
            self.quant_algo,
            "--prompts-file",
            str(CALIB_PROMPTS_FILE),
        ]

    def quantize(self, tmp_path: Path, *export_args: str) -> None:
        self._run_cmd(
            "quantize.py",
            *self._format_args(),
            "--trt-high-precision-dtype",
            self.dtype,
            "--quantized-torch-ckpt-save-path",
            str(tmp_path / f"{self.name}_{self.format_type}.pt"),
            *export_args,
        )

        assert any((tmp_path / f"{self.name}_{self.format_type}.pt").glob("*.pt"))

    def restore(self, tmp_path: Path, *export_args: str) -> None:
        output = self._run_cmd(
            "quantize.py",
            *self._format_args(),
            "--trt-high-precision-dtype",
            self.dtype,
            "--restore-from",
            str(tmp_path / f"{self.name}_{self.format_type}.pt"),
            *export_args,
        )
        assert f"Detected restored quantization format: {self.format_type}" in output


class Wan22Model(NamedTuple):
    model: str
    backbone: str | None
    format_type: str
    quant_algo: str
    collect_method: str
    # See DiffuserModel.hf_export.
    hf_export: bool = False

    def _ckpt_path(self, tmp_path: Path) -> str:
        stem = self.model.replace("wan2.2-t2v-", "")
        parts = [stem, *([self.backbone] if self.backbone else []), self.format_type]
        return str(tmp_path / f"wan22_{'_'.join(parts)}.pt")

    def _common_args(self, tiny_wan22_path: str) -> list[str]:
        cmd_args = [
            "python",
            "quantize.py",
            "--model",
            self.model,
            "--override-model-path",
            tiny_wan22_path,
            "--format",
            self.format_type,
            "--quant-algo",
            self.quant_algo,
            "--collect-method",
            self.collect_method,
            "--model-dtype",
            "BFloat16",
            "--trt-high-precision-dtype",
            "BFloat16",
            "--calib-size",
            "2",
            "--batch-size",
            "1",
            "--n-steps",
            "2",
            "--prompts-file",
            str(CALIB_PROMPTS_FILE),
            # Tiny video dims — override MODEL_DEFAULTS for fast CI.
            "--extra-param",
            "height=16",
            "--extra-param",
            "width=16",
            "--extra-param",
            "num_frames=5",
        ]
        if self.backbone is not None:
            cmd_args.extend(["--backbone", self.backbone])
        return cmd_args

    def quantize(self, tiny_wan22_path: str, tmp_path: Path, *export_args: str) -> None:
        run_example_command(
            [
                *self._common_args(tiny_wan22_path),
                "--quantized-torch-ckpt-save-path",
                self._ckpt_path(tmp_path),
                *export_args,
            ],
            "diffusers/quantization",
        )

    def restore(self, tiny_wan22_path: str, tmp_path: Path) -> None:
        run_example_command(
            [*self._common_args(tiny_wan22_path), "--restore-from", self._ckpt_path(tmp_path)],
            "diffusers/quantization",
        )


DIFFUSER_MODELS = [
    pytest.param(
        DiffuserModel(
            name="flux-schnell",
            path=FLUX_SCHNELL_PATH,
            dtype="BFloat16",
            format_type="int8",
            quant_algo="smoothquant",
            collect_method="min-mean",
        ),
        id="flux_schnell_bf16_int8_smoothquant_3.0_min_mean",
    ),
    pytest.param(
        DiffuserModel(
            name="sd3-medium",
            path=SD3_PATH,
            dtype="Half",
            format_type="int8",
            quant_algo="smoothquant",
            collect_method="min-mean",
        ),
        id="sd3_medium_fp16_int8_smoothquant_3.0_min_mean",
    ),
    pytest.param(
        DiffuserModel(
            name="sd3-medium",
            path=SD3_PATH,
            dtype="Half",
            format_type="fp8",
            quant_algo="max",
            collect_method="default",
        ),
        marks=minimum_sm(89),
        id="sd3_medium_fp16_fp8_max_3.0_default",
    ),
    pytest.param(
        DiffuserModel(
            name="sdxl-1.0",
            path=SDXL_PATH,
            dtype="Half",
            format_type="fp8",
            quant_algo="max",
            collect_method="default",
            hf_export=True,
        ),
        marks=minimum_sm(89),
        id="sdxl_1.0_fp16_fp8_max_3.0_default",
    ),
    pytest.param(
        DiffuserModel(
            name="sdxl-1.0",
            path=SDXL_PATH,
            dtype="Half",
            format_type="fp4",
            quant_algo="max",
            collect_method="default",
        ),
        marks=minimum_sm(100),
        id="sdxl_1.0_fp16_fp4_max_3.0_default",
    ),
    pytest.param(
        DiffuserModel(
            name="sdxl-1.0",
            path=SDXL_PATH,
            dtype="Half",
            format_type="int8",
            quant_algo="smoothquant",
            collect_method="min-mean",
            hf_export=True,
        ),
        id="sdxl_1.0_fp16_int8_smoothquant_3.0_min_mean",
    ),
]


# The VAE (``AutoencoderKLWan``) is shared between Wan 2.2 14B and 5B, so the
# Conv3D NVFP4 implicit-GEMM dispatch exercises the same kernel either way; we
# parametrize both ``--model`` values to also cover the ``quantize.py`` dispatch
# for each.
WAN22_MODELS = [
    pytest.param(
        Wan22Model("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean", hf_export=True),
        id="wan22_14b_transformer_int8_smoothquant",
    ),
    pytest.param(
        Wan22Model("wan2.2-t2v-14b", None, "fp8", "max", "default", hf_export=True),
        marks=minimum_sm(89),
        id="wan22_14b_transformer_fp8_max",
    ),
    pytest.param(
        Wan22Model("wan2.2-t2v-14b", None, "fp4", "max", "default"),
        marks=minimum_sm(89),
        id="wan22_14b_transformer_fp4_max",
    ),
    pytest.param(
        Wan22Model("wan2.2-t2v-14b", "vae", "fp8", "max", "default"),
        marks=minimum_sm(89),
        id="wan22_14b_vae_fp8_max",
    ),
    pytest.param(
        Wan22Model("wan2.2-t2v-14b", "vae", "fp4", "max", "default"),
        marks=minimum_sm(89),
        id="wan22_14b_vae_fp4_max",
    ),
    pytest.param(
        Wan22Model("wan2.2-t2v-5b", "vae", "fp8", "max", "default"),
        marks=minimum_sm(89),
        id="wan22_5b_vae_fp8_max",
    ),
    pytest.param(
        Wan22Model("wan2.2-t2v-5b", "vae", "fp4", "max", "default"),
        marks=minimum_sm(89),
        id="wan22_5b_vae_fp4_max",
    ),
]


@pytest.mark.parametrize("model", DIFFUSER_MODELS)
def test_diffusers_quantization(model: DiffuserModel, tmp_path: Path) -> None:
    hf_ckpt_dir = tmp_path / "hf_ckpt"
    model.quantize(tmp_path, *(["--hf-ckpt-dir", str(hf_ckpt_dir)] if model.hf_export else []))
    if model.hf_export:
        assert_hf_ckpt_exported(hf_ckpt_dir)
    model.restore(tmp_path)


@pytest.mark.parametrize("wan_model", WAN22_MODELS)
def test_wan22_quantization(wan_model: Wan22Model, tiny_wan22_path: str, tmp_path: Path) -> None:
    hf_ckpt_dir = tmp_path / "hf_ckpt"
    wan_model.quantize(
        tiny_wan22_path,
        tmp_path,
        *(["--hf-ckpt-dir", str(hf_ckpt_dir)] if wan_model.hf_export else []),
    )
    if wan_model.hf_export:
        assert_hf_ckpt_exported(hf_ckpt_dir)
    wan_model.restore(tiny_wan22_path, tmp_path)


@pytest.mark.parametrize(
    ("model_name", "model_path", "torch_compile"),
    [
        ("flux-schnell", FLUX_SCHNELL_PATH, False),
        ("flux-schnell", FLUX_SCHNELL_PATH, True),
        ("sd3-medium", SD3_PATH, False),
        ("sd3-medium", SD3_PATH, True),
        ("sdxl-1.0", SDXL_PATH, False),
        ("sdxl-1.0", SDXL_PATH, True),
    ],
    ids=[
        "flux_schnell_torch",
        "flux_schnell_torch_compile",
        "sd3_medium_torch",
        "sd3_medium_torch_compile",
        "sdxl_1.0_torch",
        "sdxl_1.0_torch_compile",
    ],
)
@pytest.mark.filterwarnings("ignore:.*env= given; running this step as a subprocess")
def test_diffusion_trt_torch(
    model_name: str,
    model_path: str,
    torch_compile: bool,
) -> None:
    cmd_args = [
        "python",
        "diffusion_trt.py",
        "--model",
        model_name,
        "--override-model-path",
        model_path,
        "--torch",
        "--num-inference-steps",
        "2",
    ]
    env = None
    if torch_compile:
        cmd_args.append("--torch-compile")
        # The script compiles with mode="max-autotune"; benchmarking Triton GEMM/conv templates is
        # a large share of the run on these tiny models and covers nothing in the example. Passing
        # env= also keeps this step in a subprocess, where the Inductor settings and cudagraphs
        # belong rather than in the shared pytest process.
        env = {
            **os.environ,
            "TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS": "ATEN",
            "TORCHINDUCTOR_MAX_AUTOTUNE_CONV_BACKENDS": "ATEN",
        }
    run_example_command(cmd_args, "diffusers/quantization", env=env)