hao-ai-lab-FastVideo/fastvideo/tests/inference/test_basic_fasth3_profile.py

# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import importlib.util
from pathlib import Path
from types import SimpleNamespace

import pytest

from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler


REPO_ROOT = Path(__file__).resolve().parents[3]
EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "basic_fasth3.py"


def _load_example():
    spec = importlib.util.spec_from_file_location("basic_fasth3_contract", EXAMPLE_PATH)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


fasth3 = _load_example()


def _args(*overrides: str):
    return fasth3.parse_args(["--prompt", "a test prompt", *overrides])


def test_default_all_profile_matches_fastest_contract(tmp_path):
    args = _args()
    config = fasth3.build_generator_config(args)
    request = fasth3.build_request(args, tmp_path / "result.mp4", args.seed)
    environment = fasth3.profile_environment(args)

    assert args.profile == "all"
    assert args.num_gpus == 4
    assert args.vsa_sparsity == 0.9
    assert args.vsa_tile_size == 64
    assert args.vsa_kernel == "sm100a"
    assert args.seed == 1000
    assert args.warmup_seed == 999
    assert args.warmup is True
    assert args.repeats == 3
    assert args.inference_torch_compile is True
    assert args.ulysses_a2a == "off"

    assert config.engine.num_gpus == 4
    assert config.engine.parallelism.tp_size == 1
    assert config.engine.parallelism.sp_size == 4
    assert config.engine.use_fsdp_inference is False
    assert config.engine.offload.dit is False
    assert config.engine.offload.text_encoder is True
    assert config.engine.offload.vae is True
    assert config.engine.offload.pin_cpu_memory is True
    assert config.engine.compile.enabled is False
    assert config.engine.compile.vae_enabled is True
    assert config.pipeline.experimental == {}
    assert config.engine.attention.backend == "VIDEO_SPARSE_ATTN_H3"
    assert config.engine.attention.vsa_sparsity == 0.9
    assert config.engine.attention.vsa_tile_size == 64
    assert config.engine.compile.regional is True
    assert config.pipeline.model.vae_parallel_decode is True
    assert config.pipeline.model.vae_parallel_decode_strategy == "gather"

    assert environment["FASTVIDEO_VSA_SM100A"] == "1"
    assert environment["FASTVIDEO_VSA_CUTEDSL"] == "0"
    assert environment["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] == "0"
    assert environment["FASTVIDEO_FA4"] == "1"
    assert environment["FASTVIDEO_NVFP4_FA4"] == "0"
    assert environment["FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN"] == "0"
    assert environment["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "all"
    assert environment["FASTVIDEO_INFERENCE_TORCH_COMPILE"] == "1"
    assert environment["FASTVIDEO_VAE_PARALLEL_DECODE"] == "1"
    assert environment["FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY"] == "gather"
    assert environment["FASTVIDEO_ULYSSES_A2A"] == "off"
    assert environment["FASTVIDEO_H3_VSA_PROBE"] is None

    assert request.sampling.num_inference_steps == 5
    assert request.sampling.seed == 1000
    assert request.sampling.guidance_scale == 1.0
    assert request.sampling.batch_cfg is False
    assert request.output.output_path == str(tmp_path / "result.mp4")
    assert config.engine.offload.lazy_module_load is None


def test_lazy_module_load_is_tri_state():
    config = fasth3.build_generator_config(_args("--num-gpus", "1"))
    assert config.engine.offload.lazy_module_load is None

    enabled = fasth3.build_generator_config(_args("--lazy-module-load"))
    assert enabled.engine.offload.lazy_module_load is True

    disabled = fasth3.build_generator_config(_args("--no-lazy-module-load"))
    assert disabled.engine.offload.lazy_module_load is False


@pytest.mark.parametrize("num_frames", (124, 243, 345))
def test_fast_profile_supports_measured_durations_without_separate_scripts(tmp_path, num_frames):
    args = _args("--num-frames", str(num_frames))
    config = fasth3.build_generator_config(args)
    request = fasth3.build_request(args, tmp_path / f"result_{num_frames}.mp4", args.seed)

    assert args.inference_torch_compile is True
    assert args.ulysses_a2a == "off"
    assert config.engine.compile.regional is True
    assert request.sampling.num_frames == num_frames


def test_compile_mode_requires_regional_compile_opt_out(capsys):
    with pytest.raises(SystemExit):
        _args("--compile-mode", "reduce-overhead")
    assert "--compile-mode cannot be combined with regional compile" in capsys.readouterr().err

    args = _args("--compile-mode", "reduce-overhead", "--no-inference-torch-compile")
    assert args.compile_mode == "reduce-overhead"
    assert args.inference_torch_compile is False


def test_default_sigma_grid_runs_exactly_four_dit_forwards():
    scheduler = MiniMaxH3Scheduler(shift=12.0)
    scheduler.set_timesteps(5)

    assert scheduler.sigmas is not None and len(scheduler.sigmas) == 5
    assert scheduler.timesteps is not None and len(scheduler.timesteps) == 4


def test_strict_profile_changes_only_non_parity_fusions():
    all_args = _args("--profile", "all")
    strict_args = _args("--profile", "strict")
    all_environment = fasth3.profile_environment(all_args)
    strict_environment = fasth3.profile_environment(strict_args)

    changed = {key for key in all_environment if all_environment[key] != strict_environment[key]}
    assert changed == {"FASTVIDEO_MINIMAX_H3_FUSIONS"}
    assert all_environment["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "all"
    assert strict_environment["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "0"
    assert fasth3.build_generator_config(all_args) == fasth3.build_generator_config(strict_args)


def test_h3_sequential_load_is_opt_in():
    default = fasth3.build_generator_config(_args())
    assert default.pipeline.model.sequential_load is None

    enabled = fasth3.build_generator_config(_args("--h3-sequential-load"))
    assert enabled.pipeline.model.sequential_load is True

    disabled = fasth3.build_generator_config(_args("--no-h3-sequential-load"))
    assert disabled.pipeline.model.sequential_load is False


def test_opt_outs_override_inherited_environment(monkeypatch):
    args = _args(
        "--vsa-kernel",
        "triton",
        "--no-fa4",
        "--no-h3-fusions",
        "--no-compile-vae",
        "--no-parallel-vae",
        "--no-replicated-dit",
        "--no-pin-cpu-memory",
        "--no-inference-torch-compile",
    )
    expected = fasth3.profile_environment(args)
    for name in expected:
        monkeypatch.setenv(name, "inherited-experiment-value")

    fasth3.configure_environment(args)
    config = fasth3.build_generator_config(args)

    assert "FASTVIDEO_H3_VSA_PROBE" not in fasth3.os.environ
    assert fasth3.os.environ["FASTVIDEO_VSA_SM100A"] == "0"
    assert fasth3.os.environ["FASTVIDEO_VSA_CUTEDSL"] == "0"
    assert fasth3.os.environ["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] == "0"
    assert fasth3.os.environ["FASTVIDEO_FA4"] == "0"
    assert fasth3.os.environ["FASTVIDEO_NVFP4_FA4"] == "0"
    assert fasth3.os.environ["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "0"
    assert fasth3.os.environ["FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN"] == "0"
    assert fasth3.os.environ["FASTVIDEO_INFERENCE_TORCH_COMPILE"] == "0"
    assert fasth3.os.environ["FASTVIDEO_VAE_PARALLEL_DECODE"] == "0"
    assert fasth3.os.environ["FASTVIDEO_ULYSSES_A2A"] == "off"
    assert config.engine.compile.vae_enabled is False
    assert config.pipeline.model.vae_parallel_decode is False
    assert config.engine.use_fsdp_inference is True
    assert config.engine.offload.pin_cpu_memory is False


def test_selected_fast_profile_requires_its_optional_routes(monkeypatch):
    monkeypatch.setattr(fasth3, "_fa4_is_installed", lambda: False)
    monkeypatch.setattr(fasth3, "_sm100a_kernel_is_installed", lambda: False)

    with pytest.raises(RuntimeError, match="flash-attn-4"):
        fasth3.validate_profile_dependencies(_args())
    with pytest.raises(RuntimeError, match="fastvideo-kernel 0.3.4"):
        fasth3.validate_profile_dependencies(_args("--no-fa4"))

    fasth3.validate_profile_dependencies(_args("--no-fa4", "--vsa-kernel", "triton"))

    monkeypatch.setattr(fasth3, "_fa4_is_installed", lambda: True)
    monkeypatch.setattr(fasth3, "_sm100a_kernel_is_installed", lambda: True)
    fasth3.validate_profile_dependencies(_args())


def test_run_excludes_warmup_and_uses_distinct_measured_outputs(monkeypatch, tmp_path, capsys):
    calls = []

    class FakeGenerator:
        shutdown_called = False

        def generate(self, request):
            calls.append(request)
            Path(request.output.output_path).touch()
            return SimpleNamespace(
                video_path=request.output.output_path,
                generation_time=1.25,
                logging_info=SimpleNamespace(stages={"denoising": {"execution_time": 2.5}}),
            )

        def shutdown(self):
            self.shutdown_called = True

    fake_generator = FakeGenerator()

    class FakeVideoGenerator:

        @classmethod
        def from_config(cls, config):
            assert config == fasth3.build_generator_config(args)
            return fake_generator

    for name in fasth3.profile_environment(_args()):
        monkeypatch.setenv(name, "inherited-experiment-value")
    monkeypatch.setattr(fasth3, "validate_profile_dependencies", lambda args: None)
    monkeypatch.setattr(fasth3, "VideoGenerator", FakeVideoGenerator)
    clock = iter((10.0, 15.0, 20.0, 26.0, 30.0, 37.0, 40.0, 48.0))
    monkeypatch.setattr(fasth3.time, "perf_counter", lambda: next(clock))
    args = _args("--output", str(tmp_path))

    measured = fasth3.run(args)

    assert measured == [6.0, 7.0, 8.0]
    assert [request.sampling.seed for request in calls] == [999, 1000, 1000, 1000]
    assert [Path(request.output.output_path).name for request in calls] == [
        "_fasth3_warmup.mp4",
        "fasth3_all_run_01.mp4",
        "fasth3_all_run_02.mp4",
        "fasth3_all_run_03.mp4",
    ]
    assert (tmp_path / "_fasth3_warmup.mp4").is_file()
    assert fake_generator.shutdown_called is True
    output = capsys.readouterr().out
    assert "[warmup] wall=5.000s (excluded)" in output
    assert f"Warmup output written to: {tmp_path / '_fasth3_warmup.mp4'}" in output
    assert output.count("Output written to:") == 3
    assert output.count("Denoising time: 2.500s") == 3
    assert "Measured E2E wall times (n=3, warmup excluded): [6.0, 7.0, 8.0]" in output
    assert "Median E2E wall time: 7.000s" in output
    assert "Median denoising time: 2.500s" in output


def test_taeh3_backend_is_opt_in():
    default = fasth3.build_generator_config(_args())
    taeh3 = fasth3.build_generator_config(_args("--video-decode-backend", "taeh3"))

    assert default.pipeline.model.video_decode_backend is None
    assert taeh3.pipeline.model.video_decode_backend == "taeh3"


EIGHT_STEP_EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "basic_fasth3_8step.py"


def _load_eight_step_example():
    import sys
    sys.path.insert(0, str(EXAMPLE_PATH.parent))
    try:
        spec = importlib.util.spec_from_file_location("basic_fasth3_8step_contract", EIGHT_STEP_EXAMPLE_PATH)
        assert spec is not None and spec.loader is not None
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        return module
    finally:
        sys.path.remove(str(EXAMPLE_PATH.parent))


def test_eight_step_example_pins_public_checkpoint_contract(tmp_path):
    """The 8-step example selects the public V2 checkpoint and its trained recipe by default
    without touching the four-forward preview example's defaults."""
    eight = _load_eight_step_example()
    args = eight.parse_args(["--prompt", "a test prompt"])
    assert args.model_path == "FastVideo/FastVideo-FastH3-8-Step-V2"
    assert args.steps == 9
    assert args.vsa_sparsity == 0.8
    assert args.vsa_tile_size == 64
    assert args.output == "outputs/fasth3_8step"
    request = fasth3.build_request(args, tmp_path / "result.mp4", args.seed)
    assert request.sampling.num_inference_steps == 9
    config = fasth3.build_generator_config(args)
    assert config.model_path == "FastVideo/FastVideo-FastH3-8-Step-V2"
    # Other flags still flow through the shared parser.
    args = eight.parse_args(["--prompt", "p", "--height", "480", "--width", "832", "--model-path", "/local/snap"])
    assert (args.height, args.width, args.model_path) == (480, 832, "/local/snap")
    # The preview example is unchanged.
    preview = _args()
    assert preview.model_path == fasth3.DEFAULT_MODEL and preview.steps == 5 and preview.vsa_sparsity == 0.9


def test_eight_step_example_rejects_other_grids():
    eight = _load_eight_step_example()
    with pytest.raises(SystemExit):
        eight.parse_args(["--prompt", "p", "--steps", "5"])
    with pytest.raises(SystemExit):
        eight.parse_args(["--prompt", "p", "--steps", "8"])