diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 4c72f3495a99..e475f0698630 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -383,7 +383,7 @@ /tensorrt_llm/_torch/visual_gen/modules/attention.py @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-torch-visual-gen-devs /tensorrt_llm/media @NVIDIA/trt-llm-torch-visual-gen-devs /tensorrt_llm/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs @NVIDIA/trt-llm-runtime-devs -/tests/integration/defs/examples/test_visual_gen.py @NVIDIA/trt-llm-torch-visual-gen-devs +/tests/integration/defs/examples/visual_gen/ @NVIDIA/trt-llm-torch-visual-gen-devs /tests/integration/defs/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs /tests/scripts/perf-sanity/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs /tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit* @NVIDIA/trt-llm-torch-visual-gen-devs diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 23baf4fc8495..656ac85ff02a 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -671,23 +671,43 @@ "examples/test_ray.py::test_llm_inference_async_ray": 131.85468882175226, "examples/test_ray.py::test_llm_inference_distributed_ray[tep2]": 162.2443882352941, "examples/test_ray.py::test_llm_inference_distributed_ray[tp2]": 68.63462352941177, - "examples/visual_gen/test_visual_gen.py::test_cosmos3_example": 268.35481163434906, - "examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden": 273.255, - "examples/visual_gen/test_visual_gen.py::test_flux1_example": 221.66551968503938, - "examples/visual_gen/test_visual_gen.py::test_flux1_lpips_against_golden": 45.776666666666664, - "examples/visual_gen/test_visual_gen.py::test_flux2_example": 259.9329760638298, - "examples/visual_gen/test_visual_gen.py::test_flux2_lpips_against_golden": 97.27477777777779, - "examples/visual_gen/test_visual_gen.py::test_ltx2_cuda_graph_lpips_matches_eager": 250.33366666666666, - "examples/visual_gen/test_visual_gen.py::test_ltx2_example": 498.4825415282392, - "examples/visual_gen/test_visual_gen.py::test_ltx2_lpips_against_golden": 57.486583333333336, - "examples/visual_gen/test_visual_gen.py::test_qwen_image_example": 160.9602216358839, - "examples/visual_gen/test_visual_gen.py::test_qwenimage_cuda_graph_lpips_against_golden": 76.66260000000001, - "examples/visual_gen/test_visual_gen.py::test_qwenimage_lpips_against_golden": 54.41563636363636, - "examples/visual_gen/test_visual_gen.py::test_visual_gen_api_walkthrough": 300.98745974025974, - "examples/visual_gen/test_visual_gen.py::test_visual_gen_quickstart": 201.11261139896374, - "examples/visual_gen/test_visual_gen.py::test_wan21_t2v_lpips_against_golden": 16.938916666666668, - "examples/visual_gen/test_visual_gen.py::test_wan22_t2v_lpips_against_golden": 54.451727272727275, - "examples/visual_gen/test_visual_gen.py::test_wan_t2v_example": 647.1280441558441, + "examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_example": 268.35481163434906, + "examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_feature_accuracy_against_golden[fp8-blockwise]": 55.09, + "examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_feature_accuracy_against_golden[nvfp4]": 59.88, + "examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2v_lpips_against_golden": 273.255, + "examples/visual_gen/test_visual_gen_flux.py::test_flux1_example": 221.66551968503938, + "examples/visual_gen/test_visual_gen_flux.py::test_flux1_lpips_against_golden": 97.557, + "examples/visual_gen/test_visual_gen_flux.py::test_flux2_example": 259.9329760638298, + "examples/visual_gen/test_visual_gen_flux.py::test_flux2_lpips_against_golden": 250.544, + "examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-cuda-graph]": 12.84, + "examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-fp8-blockwise]": 119.97, + "examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-nvfp4]": 57.85, + "examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux2-cuda-graph]": 24.02, + "examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux2-fp8-blockwise]": 258.78, + "examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux2-nvfp4]": 277.28, + "examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_cuda_graph_lpips_matches_eager": 250.33366666666666, + "examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_example": 498.4825415282392, + "examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[cuda-graph]": 461.68, + "examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[fp8-blockwise]": 481.07, + "examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[nvfp4]": 496.24, + "examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_lpips_against_golden": 57.486583333333336, + "examples/visual_gen/test_visual_gen_qwen_image.py::test_qwen_image_example": 160.9602216358839, + "examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_cuda_graph_lpips_against_golden": 76.66260000000001, + "examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[cuda-graph]": 141.8, + "examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[fp8-blockwise]": 82.1, + "examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[nvfp4]": 104.22, + "examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_lpips_against_golden": 54.41563636363636, + "examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_api_walkthrough": 300.98745974025974, + "examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_quickstart": 201.11261139896374, + "examples/visual_gen/test_visual_gen_wan.py::test_wan21_t2v_lpips_against_golden": 16.938916666666668, + "examples/visual_gen/test_visual_gen_wan.py::test_wan22_t2v_lpips_against_golden": 54.451727272727275, + "examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan21-cuda-graph]": 8.49, + "examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan21-fp8-blockwise]": 65.48, + "examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan21-nvfp4]": 156.43, + "examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-cuda-graph]": 31.14, + "examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-fp8-blockwise]": 182.16, + "examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-nvfp4]": 161.28, + "examples/visual_gen/test_visual_gen_wan.py::test_wan_t2v_example": 647.1280441558441, "examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[attn2d_2x2]": 834.184, "examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2]": 963.1505, "examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_tp[cfg2_tp2]": 447.3106666666667, diff --git a/tests/integration/defs/examples/visual_gen/conftest.py b/tests/integration/defs/examples/visual_gen/conftest.py index 9d65d09ef545..ba7ebecc3ceb 100644 --- a/tests/integration/defs/examples/visual_gen/conftest.py +++ b/tests/integration/defs/examples/visual_gen/conftest.py @@ -12,16 +12,37 @@ # 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. +"""Fixtures shared by VisualGen example tests.""" -"""Shared fixtures for VisualGen example integration tests.""" +import shutil import pytest +from defs.examples.visual_gen.visual_gen_test_utils import _prepare_vbench_repo from defs.trt_test_alternative import check_call @pytest.fixture(scope="session") def _visual_gen_deps(llm_venv): - """Install optional media dependencies once per VisualGen test session.""" + """Install the Python and system media dependencies once per session.""" llm_venv.run_cmd(["-m", "pip", "install", "av"]) - check_call(["apt-get", "update", "-y"], shell=False) - check_call(["apt-get", "install", "-y", "ffmpeg"], shell=False) + if shutil.which("ffmpeg") is None: + check_call(["apt-get", "update", "-y"], shell=False) + check_call(["apt-get", "install", "-y", "ffmpeg"], shell=False) + + +@pytest.fixture(scope="session") +def _visual_gen_lpips_scorer(): + """Reuse one lazily initialized AlexNet LPIPS model for media tests.""" + from defs.examples.visual_gen.visual_gen_test_utils import ReusableLPIPSScorer + + scorer = ReusableLPIPSScorer() + try: + yield scorer + finally: + scorer.close() + + +@pytest.fixture(scope="session") +def vbench_repo_root(llm_venv): + """Prepare the pinned VBench checkout and return its repository root.""" + return _prepare_vbench_repo(llm_venv) diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_fp8_blockwise_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_fp8_blockwise_lpips_golden.json new file mode 100644 index 000000000000..0abf4430cc0c --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_fp8_blockwise_lpips_golden.json @@ -0,0 +1,27 @@ +{ + "image": "cosmos3_nano_fp8_blockwise_lpips_golden.png", + "model": "Cosmos3-Nano", + "source": "TensorRT-LLM VisualGen", + "prompt": "A serene mountain landscape with snow-capped peaks and a flowing river", + "height": 720, + "width": 1280, + "num_frames": 1, + "num_inference_steps": 35, + "guidance_scale": 6.0, + "frame_rate": 24.0, + "seed": 42, + "feature_config": { + "quantization": "FP8_BLOCK_SCALES", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "dbdfc83e3a6a038138f4c8e039fa09be8ab11cb308f1534405624523afed73a5" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_nvfp4_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_nvfp4_lpips_golden.json new file mode 100644 index 000000000000..99ece747507b --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_nvfp4_lpips_golden.json @@ -0,0 +1,27 @@ +{ + "image": "cosmos3_nano_nvfp4_lpips_golden.png", + "model": "Cosmos3-Nano", + "source": "TensorRT-LLM VisualGen", + "prompt": "A serene mountain landscape with snow-capped peaks and a flowing river", + "height": 720, + "width": 1280, + "num_frames": 1, + "num_inference_steps": 35, + "guidance_scale": 6.0, + "frame_rate": 24.0, + "seed": 42, + "feature_config": { + "quantization": "NVFP4", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "36a2927030c28eac1bcfee0361d13723b655b4d0db4a818537e71366360df85f" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_cuda_graph_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_cuda_graph_lpips_golden.json new file mode 100644 index 000000000000..251f1d41e53f --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_cuda_graph_lpips_golden.json @@ -0,0 +1,25 @@ +{ + "image": "flux1_cuda_graph_lpips_golden.png", + "model": "FLUX.1-dev", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 256, + "width": 256, + "num_inference_steps": 4, + "guidance_scale": 3.5, + "seed": 42, + "feature_config": { + "quantization": "none", + "cuda_graph": true + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "43c916aef97133281424e7e7f04bc6e58493e6e9535f5e5234f3fef51f0ab857" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_fp8_blockwise_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_fp8_blockwise_lpips_golden.json new file mode 100644 index 000000000000..22c3e0432fe9 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_fp8_blockwise_lpips_golden.json @@ -0,0 +1,25 @@ +{ + "image": "flux1_fp8_blockwise_lpips_golden.png", + "model": "FLUX.1-dev", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 256, + "width": 256, + "num_inference_steps": 4, + "guidance_scale": 3.5, + "seed": 42, + "feature_config": { + "quantization": "FP8_BLOCK_SCALES", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "9e71a2e12d8bd28278f95f02c1836033dccac4717578c30b6bbc27a3b4011022" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_nvfp4_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_nvfp4_lpips_golden.json new file mode 100644 index 000000000000..61f9d6fe7daa --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux1_nvfp4_lpips_golden.json @@ -0,0 +1,25 @@ +{ + "image": "flux1_nvfp4_lpips_golden.png", + "model": "FLUX.1-dev", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 256, + "width": 256, + "num_inference_steps": 4, + "guidance_scale": 3.5, + "seed": 42, + "feature_config": { + "quantization": "NVFP4", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "f34c21a23a8140f842a53a11937c08d6f034057daa89fb5825638a9b3a322566" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_cuda_graph_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_cuda_graph_lpips_golden.json new file mode 100644 index 000000000000..c592efd97e7a --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_cuda_graph_lpips_golden.json @@ -0,0 +1,25 @@ +{ + "image": "flux2_cuda_graph_lpips_golden.png", + "model": "FLUX.2-dev", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 256, + "width": 256, + "num_inference_steps": 4, + "guidance_scale": 3.5, + "seed": 42, + "feature_config": { + "quantization": "none", + "cuda_graph": true + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "0beac87c89fced52e5845b8de7b8dfa6837cca308de032a87f8dcf4d86c3dfee" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_fp8_blockwise_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_fp8_blockwise_lpips_golden.json new file mode 100644 index 000000000000..d277c1007fe2 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_fp8_blockwise_lpips_golden.json @@ -0,0 +1,25 @@ +{ + "image": "flux2_fp8_blockwise_lpips_golden.png", + "model": "FLUX.2-dev", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 256, + "width": 256, + "num_inference_steps": 4, + "guidance_scale": 3.5, + "seed": 42, + "feature_config": { + "quantization": "FP8_BLOCK_SCALES", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "8255196070843fe4a3d812f0880afcbfda5903537403581207c38bbf0ef615f5" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_nvfp4_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_nvfp4_lpips_golden.json new file mode 100644 index 000000000000..fc66f1d5ea50 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/flux2_nvfp4_lpips_golden.json @@ -0,0 +1,25 @@ +{ + "image": "flux2_nvfp4_lpips_golden.png", + "model": "FLUX.2-dev", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 256, + "width": 256, + "num_inference_steps": 4, + "guidance_scale": 3.5, + "seed": 42, + "feature_config": { + "quantization": "NVFP4", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "b290d7891340e259fa05f0b027ea543850aef453573c0384a7c4744d6d91d53b" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_cuda_graph_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_cuda_graph_lpips_golden_video.json new file mode 100644 index 000000000000..3c1432d103ed --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_cuda_graph_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "ltx2_cuda_graph_lpips_golden_video.mp4", + "model": "LTX-2 BF16 two-stage", + "source": "TensorRT-LLM VisualGen", + "prompt": "A woman with long brown hair and light skin smiles at the camera while standing in a sunlit park, her hair gently blowing in the breeze as she tilts her head slightly to the side.", + "negative_prompt": "worst quality, inconsistent motion, blurry, jittery, distorted", + "height": 512, + "width": 768, + "num_frames": 49, + "num_inference_steps": 8, + "guidance_scale": 4.0, + "frame_rate": 24.0, + "seed": 42, + "feature_config": { + "quantization": "none", + "cuda_graph": true + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "55b92e45eed7a6e795b669c2d3340a4e6094410cfe9af53163456fa4f84c36c2" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_fp8_blockwise_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_fp8_blockwise_lpips_golden_video.json new file mode 100644 index 000000000000..94d8128c800c --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_fp8_blockwise_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "ltx2_fp8_blockwise_lpips_golden_video.mp4", + "model": "LTX-2 BF16 two-stage", + "source": "TensorRT-LLM VisualGen", + "prompt": "A woman with long brown hair and light skin smiles at the camera while standing in a sunlit park, her hair gently blowing in the breeze as she tilts her head slightly to the side.", + "negative_prompt": "worst quality, inconsistent motion, blurry, jittery, distorted", + "height": 512, + "width": 768, + "num_frames": 49, + "num_inference_steps": 8, + "guidance_scale": 4.0, + "frame_rate": 24.0, + "seed": 42, + "feature_config": { + "quantization": "FP8_BLOCK_SCALES", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "fced64c757f3a7f1d6b6f1f0c68c40181fc11e880c9f272e62522182393b2484" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_nvfp4_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_nvfp4_lpips_golden_video.json new file mode 100644 index 000000000000..f29c899d6ab5 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx2_nvfp4_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "ltx2_nvfp4_lpips_golden_video.mp4", + "model": "LTX-2 BF16 two-stage", + "source": "TensorRT-LLM VisualGen", + "prompt": "A woman with long brown hair and light skin smiles at the camera while standing in a sunlit park, her hair gently blowing in the breeze as she tilts her head slightly to the side.", + "negative_prompt": "worst quality, inconsistent motion, blurry, jittery, distorted", + "height": 512, + "width": 768, + "num_frames": 49, + "num_inference_steps": 8, + "guidance_scale": 4.0, + "frame_rate": 24.0, + "seed": 42, + "feature_config": { + "quantization": "NVFP4", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "246651c299973a05f7bc305820b5a09e660289a1bd9ade17da21765edd61fd48" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_cuda_graph_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_cuda_graph_lpips_golden.json new file mode 100644 index 000000000000..d27a03e1fa2f --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_cuda_graph_lpips_golden.json @@ -0,0 +1,26 @@ +{ + "image": "qwenimage_cuda_graph_lpips_golden.png", + "model": "Qwen-Image", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "negative_prompt": "", + "height": 1328, + "width": 1328, + "num_inference_steps": 50, + "true_cfg_scale": 4.0, + "seed": 42, + "feature_config": { + "quantization": "none", + "cuda_graph": true + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "a7e05c09dd6fedbadc636e0962af02eab1c22ba5818451041c84c03a50ce9c76" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_fp8_blockwise_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_fp8_blockwise_lpips_golden.json new file mode 100644 index 000000000000..c2c5e95d4366 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_fp8_blockwise_lpips_golden.json @@ -0,0 +1,26 @@ +{ + "image": "qwenimage_fp8_blockwise_lpips_golden.png", + "model": "Qwen-Image", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "negative_prompt": "", + "height": 1328, + "width": 1328, + "num_inference_steps": 50, + "true_cfg_scale": 4.0, + "seed": 42, + "feature_config": { + "quantization": "FP8_BLOCK_SCALES", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "a6b6ba5917e9e5b83043d99a99e2f9b6199e24f3ca5f7d4a2c3ec83d62cd3bf2" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_nvfp4_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_nvfp4_lpips_golden.json new file mode 100644 index 000000000000..aa3cd843f4f2 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwenimage_nvfp4_lpips_golden.json @@ -0,0 +1,26 @@ +{ + "image": "qwenimage_nvfp4_lpips_golden.png", + "model": "Qwen-Image", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "negative_prompt": "", + "height": 1328, + "width": 1328, + "num_inference_steps": 50, + "true_cfg_scale": 4.0, + "seed": 42, + "feature_config": { + "quantization": "NVFP4", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "abb4854db2a842ad10d68555cba27c14824628a4c7afffe99e5c407bb91e483c" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip index 988c509f07b7..10edd4825dd6 100644 --- a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:78c7fde678df7b434b8f6da08d800d25bcfccb2509f73ab0c3d342eafe447100 -size 14648434 +oid sha256:038de10d5de9dc208c9d584efc6d1b64150042120b5e3157aaea994f7fb16950 +size 26850469 diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_cuda_graph_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_cuda_graph_lpips_golden_video.json new file mode 100644 index 000000000000..deff4e36e627 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_cuda_graph_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "wan21_cuda_graph_lpips_golden_video.mp4", + "model": "Wan2.1-T2V-1.3B-Diffusers", + "source": "TensorRT-LLM VisualGen", + "prompt": "A cat sitting on a windowsill", + "negative_prompt": null, + "height": 256, + "width": 256, + "num_frames": 5, + "num_inference_steps": 1, + "guidance_scale": 5.0, + "frame_rate": 16.0, + "seed": 42, + "feature_config": { + "quantization": "none", + "cuda_graph": true + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "e29016d3dbb594c2995d54111acdaf1a0ca4a8c5607d3ba71eab37407fdc6dc6" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_fp8_blockwise_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_fp8_blockwise_lpips_golden_video.json new file mode 100644 index 000000000000..c26a6f92e034 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_fp8_blockwise_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "wan21_fp8_blockwise_lpips_golden_video.mp4", + "model": "Wan2.1-T2V-1.3B-Diffusers", + "source": "TensorRT-LLM VisualGen", + "prompt": "A cat sitting on a windowsill", + "negative_prompt": null, + "height": 256, + "width": 256, + "num_frames": 5, + "num_inference_steps": 1, + "guidance_scale": 5.0, + "frame_rate": 16.0, + "seed": 42, + "feature_config": { + "quantization": "FP8_BLOCK_SCALES", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "f9864223aec27fed307c49faa4ee6a4cd5870891e9454ff141292cb2724da22f" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_nvfp4_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_nvfp4_lpips_golden_video.json new file mode 100644 index 000000000000..13bfe0d44ed6 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan21_nvfp4_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "wan21_nvfp4_lpips_golden_video.mp4", + "model": "Wan2.1-T2V-1.3B-Diffusers", + "source": "TensorRT-LLM VisualGen", + "prompt": "A cat sitting on a windowsill", + "negative_prompt": null, + "height": 256, + "width": 256, + "num_frames": 5, + "num_inference_steps": 1, + "guidance_scale": 5.0, + "frame_rate": 16.0, + "seed": 42, + "feature_config": { + "quantization": "NVFP4", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "5cd1a71e23264bbce786ea1ba7ba1196985399c668030cdba1d858d6167c8b88" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_cuda_graph_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_cuda_graph_lpips_golden_video.json new file mode 100644 index 000000000000..0a7493e0f94a --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_cuda_graph_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "wan22_cuda_graph_lpips_golden_video.mp4", + "model": "Wan2.2-T2V-A14B-Diffusers", + "source": "TensorRT-LLM VisualGen", + "prompt": "A cat sitting on a sunny windowsill watching birds outside.", + "negative_prompt": "", + "height": 480, + "width": 832, + "num_frames": 9, + "num_inference_steps": 4, + "guidance_scale": 4.0, + "frame_rate": 16.0, + "seed": 42, + "feature_config": { + "quantization": "none", + "cuda_graph": true + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "09c7138b04bb2b380f743dea1a82dc91868453a331a3dbb51348692f1e0969a5" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_fp8_blockwise_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_fp8_blockwise_lpips_golden_video.json new file mode 100644 index 000000000000..79ef0b9346ab --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_fp8_blockwise_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "wan22_fp8_blockwise_lpips_golden_video.mp4", + "model": "Wan2.2-T2V-A14B-Diffusers", + "source": "TensorRT-LLM VisualGen", + "prompt": "A cat sitting on a sunny windowsill watching birds outside.", + "negative_prompt": "", + "height": 480, + "width": 832, + "num_frames": 9, + "num_inference_steps": 4, + "guidance_scale": 4.0, + "frame_rate": 16.0, + "seed": 42, + "feature_config": { + "quantization": "FP8_BLOCK_SCALES", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "d0df505a3f0497997a7420d5091e429792a82b9c06563abbd0d06d5d1f297ff8" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_nvfp4_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_nvfp4_lpips_golden_video.json new file mode 100644 index 000000000000..ee6d74ecd26b --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_nvfp4_lpips_golden_video.json @@ -0,0 +1,28 @@ +{ + "video": "wan22_nvfp4_lpips_golden_video.mp4", + "model": "Wan2.2-T2V-A14B-Diffusers", + "source": "TensorRT-LLM VisualGen", + "prompt": "A cat sitting on a sunny windowsill watching birds outside.", + "negative_prompt": "", + "height": 480, + "width": 832, + "num_frames": 9, + "num_inference_steps": 4, + "guidance_scale": 4.0, + "frame_rate": 16.0, + "seed": 42, + "feature_config": { + "quantization": "NVFP4", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "b2131b181f5be6717cd302a0b53c22c6a70c65b3", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm@sha256:475999862b896593159f10f486e16a748e7b7abb5cee558932a59ea6fc690d6b", + "sha256": "1467b575093c1674060c3590e942295eac480c8ba5d6a9437614bb6f7775d648" +} diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py deleted file mode 100644 index ed8a5978941d..000000000000 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ /dev/null @@ -1,1883 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 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. -"""Integration tests for VisualGen examples and visual quality checks.""" - -import base64 -import contextlib -import gc -import glob -import json -import os -import random -import shutil -import subprocess -import sys -import time -import urllib.request -import zipfile - -import pytest -import torch -import torch._inductor.config as inductor_config -from defs.trt_test_alternative import check_call -from torch._inductor.async_compile import shutdown_compile_workers - -WAN_T2V_MODEL_SUBPATH = "Wan2.1-T2V-1.3B-Diffusers" -WAN22_A14B_FP8_MODEL_SUBPATH = "Wan2.2-T2V-A14B-Diffusers-FP8" -WAN22_A14B_NVFP4_MODEL_SUBPATH = "Wan2.2-T2V-A14B-Diffusers-NVFP4" -WAN22_I2V_A14B_NVFP4_MODEL_SUBPATH = "Wan2.2-I2V-A14B-Diffusers-NVFP4" -QWEN_IMAGE_MODEL_SUBPATH = "qwen-image" -VISUAL_GEN_OUTPUT_VIDEO = "trtllm_output.mp4" -DIFFUSERS_REFERENCE_VIDEO = "diffusers_reference.mp4" -WAN_T2V_PROMPT = "A cute cat playing piano" -WAN_T2V_HEIGHT = 480 -WAN_T2V_WIDTH = 832 -WAN_T2V_NUM_FRAMES = 165 - -# NB: this test file lives at tests/integration/defs/examples/visual_gen/, so the repo -# root is five levels up (the LPIPS eval script is referenced from /scripts/). -REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..")) -VISUAL_GEN_LPIPS_EVAL_SCRIPT = os.path.join( - REPO_ROOT, "scripts", "visualgen_eval", "visual_gen_lpips_score_eval.py" -) -VISUAL_GEN_LPIPS_GOLDEN_DIR = os.path.join(os.path.dirname(__file__), "golden", "visual_gen_lpips") -VISUAL_GEN_LPIPS_GOLDEN_MEDIA_ZIP = os.path.join( - VISUAL_GEN_LPIPS_GOLDEN_DIR, "visual_gen_lpips_golden_media.zip" -) - -FLUX_LPIPS_PROMPT = "a tiny astronaut hatching from an egg on the moon" -FLUX_LPIPS_HEIGHT = 256 -FLUX_LPIPS_WIDTH = 256 -FLUX_LPIPS_NUM_INFERENCE_STEPS = 4 -FLUX_LPIPS_GUIDANCE_SCALE = 3.5 -FLUX_LPIPS_SEED = 42 -FLUX_LPIPS_THRESHOLD = 0.05 - -LTX2_LPIPS_NUM_FRAMES = 49 -LTX2_LPIPS_NUM_INFERENCE_STEPS = 8 -LTX2_LPIPS_THRESHOLD = 0.05 -LTX2_CUDA_GRAPH_LPIPS_THRESHOLD = 0.01 - -WAN21_LPIPS_PROMPT = "A cat sitting on a windowsill" -WAN21_LPIPS_NEGATIVE_PROMPT = None -WAN21_LPIPS_HEIGHT = 256 -WAN21_LPIPS_WIDTH = 256 -WAN21_LPIPS_NUM_FRAMES = 5 -WAN21_LPIPS_NUM_INFERENCE_STEPS = 1 -WAN21_LPIPS_GUIDANCE_SCALE = 5.0 -WAN21_LPIPS_SEED = 42 -WAN_LPIPS_FRAME_RATE = 16.0 -WAN_LPIPS_THRESHOLD = 0.05 - -WAN22_LPIPS_PROMPT = "A cat sitting on a sunny windowsill watching birds outside." -WAN22_LPIPS_NEGATIVE_PROMPT = "" -WAN22_LPIPS_HEIGHT = 480 -WAN22_LPIPS_WIDTH = 832 -WAN22_LPIPS_NUM_FRAMES = 9 -WAN22_LPIPS_NUM_INFERENCE_STEPS = 4 -WAN22_LPIPS_GUIDANCE_SCALE = 4.0 -WAN22_LPIPS_SEED = 42 -WAN22_LPIPS_FRAME_RATE = 16.0 - -# QwenImage (text-to-image) — default-setting LPIPS golden. -# Params mirror the QwenImage 20B reference defaults (pipeline_qwen_image.py). -# NOTE: QwenImage's forward CFG knob is ``true_cfg_scale`` (not ``guidance_scale``), -# and real-CFG only engages when a negative prompt is supplied. -QWENIMAGE_MODEL_SUBPATH = "qwen-image" -QWENIMAGE_LPIPS_PROMPT = "a tiny astronaut hatching from an egg on the moon" -QWENIMAGE_LPIPS_NEGATIVE_PROMPT = "" -QWENIMAGE_LPIPS_HEIGHT = 1328 -QWENIMAGE_LPIPS_WIDTH = 1328 -QWENIMAGE_LPIPS_NUM_INFERENCE_STEPS = 50 -QWENIMAGE_LPIPS_TRUE_CFG_SCALE = 4.0 -QWENIMAGE_LPIPS_SEED = 42 -QWENIMAGE_LPIPS_THRESHOLD = 0.05 - -# Cosmos3-Nano (text-to-video + text-to-image) — default-setting LPIPS golden. -# Params are the Cosmos3 720P defaults (cosmos3/defaults.py:COSMOS3_720P_PARAMS). -# Cosmos3 requires VANILLA attention and guardrails disabled in CI. -COSMOS3_NANO_MODEL_SUBPATH = "Cosmos3-Nano" -COSMOS3_LPIPS_PROMPT = "A serene mountain landscape with snow-capped peaks and a flowing river" -COSMOS3_LPIPS_HEIGHT = 720 -COSMOS3_LPIPS_WIDTH = 1280 -COSMOS3_LPIPS_T2V_NUM_FRAMES = 189 -COSMOS3_LPIPS_T2I_NUM_FRAMES = 1 -COSMOS3_LPIPS_NUM_INFERENCE_STEPS = 35 -COSMOS3_LPIPS_GUIDANCE_SCALE = 6.0 -COSMOS3_LPIPS_SEED = 42 -COSMOS3_LPIPS_FRAME_RATE = 24.0 -COSMOS3_LPIPS_THRESHOLD = 0.05 - -# LTX-2 configuration -LTX2_MODEL_CHECKPOINT_PATH = "LTX-2/ltx-2-19b-dev.safetensors" -LTX2_TEXT_ENCODER_SUBPATH = "gemma-3-12b-it" -LTX2_T2V_PROMPT = ( - "A woman with long brown hair and light skin smiles at the camera while " - "standing in a sunlit park, her hair gently blowing in the breeze as she " - "tilts her head slightly to the side." -) -LTX2_T2V_HEIGHT = 512 -LTX2_T2V_WIDTH = 768 -LTX2_T2V_NUM_FRAMES = 121 -LTX2_T2V_STEPS = 40 -LTX2_T2V_GUIDANCE_SCALE = 4.0 -LTX2_T2V_MAX_SEQ_LEN = 1024 -LTX2_T2V_FRAME_RATE = 24.0 -LTX2_T2V_SEED = 42 -LTX2_T2V_NEGATIVE_PROMPT = "worst quality, inconsistent motion, blurry, jittery, distorted" - -# Dimensions to evaluate -VBENCH_DIMENSIONS = [ - "subject_consistency", - "background_consistency", - "motion_smoothness", - "dynamic_degree", - "aesthetic_quality", - "imaging_quality", -] - -# Golden VBench scores from HF reference video (WAN 2.1 1.3B); TRT-LLM is compared against these. -VBENCH_WAN_GOLDEN_SCORES = { - "subject_consistency": 0.8907, - "background_consistency": 0.9274, - "motion_smoothness": 0.9818, - "dynamic_degree": 1.0000, - "aesthetic_quality": 0.2928, - "imaging_quality": 0.3812, -} - - -# TODO: Reference scores from bf16 baseline runs -VBENCH_WAN22_BF16_GOLDEN_SCORES = { - "subject_consistency": 0.9103, - "background_consistency": 0.9516, - "motion_smoothness": 0.9693, - "dynamic_degree": 0.0000, - "aesthetic_quality": 0.6821, - "imaging_quality": 0.3993, -} -VBENCH_WAN22_A14B_FP8_GOLDEN_SCORES = { - "subject_consistency": 0.9173, - "background_consistency": 0.9717, - "motion_smoothness": 0.9865, - "dynamic_degree": 1.0000, - "aesthetic_quality": 0.5465, - "imaging_quality": 0.7142, -} -VBENCH_WAN22_A14B_NVFP4_GOLDEN_SCORES = { - "subject_consistency": 0.9173, - "background_consistency": 0.9717, - "motion_smoothness": 0.9865, - "dynamic_degree": 1.0000, - "aesthetic_quality": 0.5465, - "imaging_quality": 0.7142, -} - -# LTX-2 Two-Stage configuration -LTX2_UPSAMPLER_SUBPATH = "LTX-2/ltx-2-spatial-upscaler-x2-1.0.safetensors" -LTX2_DISTILLED_LORA_SUBPATH = "LTX-2/ltx-2-19b-distilled-lora-384.safetensors" -LTX2_TWO_STAGE_HEIGHT = 1024 -LTX2_TWO_STAGE_WIDTH = 1536 -LTX2_TWO_STAGE_NUM_FRAMES = 121 -LTX2_TWO_STAGE_STEPS = 40 -LTX2_TWO_STAGE_GUIDANCE_SCALE = 4.0 - -# Golden VBench scores for two-stage pipeline variants. -# Initially None — first CI run is a baseline that prints scores for capture. -VBENCH_LTX2_TWO_STAGE_BF16_GOLDEN_SCORES = { - "subject_consistency": 0.9877, - "background_consistency": 0.9601, - "motion_smoothness": 0.9952, - "dynamic_degree": 0.0, - "aesthetic_quality": 0.5839, - "imaging_quality": 0.5404, -} - -VBENCH_LTX2_TWO_STAGE_FP8_GOLDEN_SCORES = { - "subject_consistency": 0.9820, - "background_consistency": 0.9617, - "motion_smoothness": 0.9885, - "dynamic_degree": 1.0, - "aesthetic_quality": 0.6017, - "imaging_quality": 0.7136, -} - -VBENCH_REPO = "https://github.com/Vchitect/VBench.git" -# Pin to a fixed commit for reproducible shallow-fetch -VBENCH_COMMIT = "98b19513678e99c80d8377fda25ba53b81a491a6" - -DINO_REPO = "https://github.com/facebookresearch/dino.git" -DINO_HUB_DIR_NAME = "facebookresearch_dino_main" - -AESTHETIC_PREDICTOR_URL = ( - "https://raw.githubusercontent.com/LAION-AI/aesthetic-predictor/main/sa_0_4_vit_l_14_linear.pth" -) -AESTHETIC_PREDICTOR_FILENAME = "sa_0_4_vit_l_14_linear.pth" -AESTHETIC_PREDICTOR_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "emb_reader") - - -@pytest.fixture(scope="session") -def vbench_repo_root(llm_venv): - """Clone VBench repo into workspace and install; return repo root path.""" - workspace = llm_venv.get_working_directory() - repo_path = os.path.join(workspace, "VBench_repo") - _precache_dino_for_torch_hub() - _precache_aesthetic_predictor() - if os.path.exists(repo_path): - return repo_path - # Shallow-fetch only the pinned commit to avoid downloading full history (~350 MB) - os.makedirs(repo_path, exist_ok=True) - check_call(["git", "init", repo_path], shell=False) - check_call( - ["git", "-C", repo_path, "remote", "add", "origin", VBENCH_REPO], - shell=False, - ) - check_call( - ["git", "-C", repo_path, "fetch", "--depth", "1", "origin", VBENCH_COMMIT], - shell=False, - ) - check_call(["git", "-C", repo_path, "checkout", "FETCH_HEAD"], shell=False) - llm_venv.run_cmd( - [ - "-m", - "pip", - "install", - "tqdm>=4.60.0", - "openai-clip>=1.0", - "easydict", - "decord>=0.6.0", - "imageio", - ] - ) - llm_venv.run_cmd( - [ - "-m", - "pip", - "install", - "--no-deps", - "pyiqa>=0.1.0", - ] - ) - return repo_path - - -def _precache_dino_for_torch_hub(): - """Pre-clone facebookresearch/dino into torch.hub cache to avoid GitHub API rate limits. - - VBench's subject_consistency dimension calls - torch.hub.load('facebookresearch/dino:main', ...) which validates the repo - via the GitHub API and fails with HTTP 403 when rate-limited. Pre-cloning - the repo makes torch.hub skip the API validation for cached repos. - """ - hub_dir = torch.hub.get_dir() - os.makedirs(hub_dir, exist_ok=True) - - dino_cache = os.path.join(hub_dir, DINO_HUB_DIR_NAME) - if not os.path.isdir(dino_cache): - check_call( - ["git", "clone", "--depth", "1", "-b", "main", DINO_REPO, dino_cache], - shell=False, - ) - - -def _precache_aesthetic_predictor(): - """Pre-download LAION aesthetic predictor weights to avoid GitHub rate limits. - - VBench's aesthetic_quality dimension downloads sa_0_4_vit_l_14_linear.pth - from GitHub via wget at evaluation time. GitHub often returns HTTP 429 - (Too Many Requests) in CI environments. Pre-downloading with retries - and proper headers ensures the file is cached before VBench needs it. - """ - os.makedirs(AESTHETIC_PREDICTOR_CACHE_DIR, exist_ok=True) - cached_path = os.path.join(AESTHETIC_PREDICTOR_CACHE_DIR, AESTHETIC_PREDICTOR_FILENAME) - if os.path.isfile(cached_path): - return - - max_retries = 8 - for attempt in range(max_retries): - try: - req = urllib.request.Request( - AESTHETIC_PREDICTOR_URL, - headers={"User-Agent": "TensorRT-LLM-CI/1.0"}, - ) - with urllib.request.urlopen(req, timeout=120) as resp: - data = resp.read() - tmp_path = cached_path + ".tmp" - with open(tmp_path, "wb") as f: - f.write(data) - os.replace(tmp_path, cached_path) - return - except Exception as exc: - if attempt < max_retries - 1: - wait = min(10 * 2**attempt, 120) + random.uniform(0, 5) - print( - f"[precache] Aesthetic predictor download attempt {attempt + 1}/{max_retries} " - f"failed ({exc}), retrying in {wait:.0f}s..." - ) - time.sleep(wait) - else: - raise RuntimeError( - f"Failed to download aesthetic predictor after {max_retries} attempts: {exc}" - ) from exc - - -def _llm_models_root(): - # Imported lazily so that re-importing this module in a torch.multiprocessing.spawn - # child (a fresh interpreter) does not run a module-level `from defs import conftest`, - # which pulls in `tensorrt_llm.bindings` -- a compiled extension absent from the source - # tree the spawned child resolves, crashing the worker before the test runs. The parent - # process already imports conftest during collection, so this deferral is free. - from defs import conftest - - return conftest.llm_models_root() - - -def _venv_check_call(*args, **kwargs): - # Deferred like _llm_models_root above: defs.common does `from tensorrt_llm import - # LLM`, which pulls in tensorrt_llm.bindings. Importing it at module load would - # crash the torch.multiprocessing.spawn child processes used by the multi-GPU LPIPS - # tests, which re-import this module before the worker fixes sys.path. Only the - # single-GPU example tests call this, and only in the parent process. - from defs.common import venv_check_call - - return venv_check_call(*args, **kwargs) - - -def _lpips_model_path(*parts): - return os.path.join(_llm_models_root(), *parts) - - -def _skip_if_missing(path, label, is_dir=False): - exists = os.path.isdir(path) if is_dir else os.path.exists(path) - if not exists: - pytest.skip(f"{label} not found: {path}") - - -def _visual_gen_output_path(llm_venv, output_subdir): - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", output_subdir) - os.makedirs(out_dir, exist_ok=True) - return os.path.join(out_dir, VISUAL_GEN_OUTPUT_VIDEO) - - -def _extract_visual_gen_lpips_golden_media(tmp_path): - _skip_if_missing(VISUAL_GEN_LPIPS_GOLDEN_MEDIA_ZIP, "VisualGen LPIPS golden media zip") - extract_dir = tmp_path / "visual_gen_lpips_golden_media" - if extract_dir.exists(): - return extract_dir - - with zipfile.ZipFile(VISUAL_GEN_LPIPS_GOLDEN_MEDIA_ZIP) as archive: - for member in archive.namelist(): - if os.path.isabs(member) or ".." in member.split("/"): - raise ValueError(f"Unsafe golden media zip member: {member}") - archive.extractall(extract_dir) - return extract_dir - - -def _golden_media_path(tmp_path, media_name, label): - path = _extract_visual_gen_lpips_golden_media(tmp_path) / media_name - _skip_if_missing(path, label) - return path - - -def _ltx2_lpips_text_encoder_path(): - scratch_space = _llm_models_root() - candidates = [ - os.path.join(scratch_space, LTX2_TEXT_ENCODER_SUBPATH), - os.path.join(scratch_space, "gemma", LTX2_TEXT_ENCODER_SUBPATH), - ] - for candidate in candidates: - if os.path.isdir(candidate): - return candidate - return candidates[0] - - -def _disable_inductor_compile_worker_quiesce(): - # The quiesce timer thread can outlive shutdown_compile_workers() long - # enough for pytest-threadleak to report it as a leaked thread. - if hasattr(inductor_config, "quiesce_async_compile_pool"): - inductor_config.quiesce_async_compile_pool = False - - -def _cleanup_cuda(): - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - # torch.compile / unconditional @torch.compile decorators (e.g. on the TRT-LLM - # attention backend's _concat_qkv) lazily spawn an InductorSubproc worker - # pool the first time a compiled function runs. The pool's daemon threads - # outlive the test and trip pytest-threadleak. Tear them down explicitly. - shutdown_compile_workers() - - -@contextlib.contextmanager -def _lpips_deterministic_algorithms(*, fully_eager=False): - previous_deterministic = torch.are_deterministic_algorithms_enabled() - previous_warn_only = torch.is_deterministic_algorithms_warn_only_enabled() - previous_cublas_workspace_config = os.environ.get("CUBLAS_WORKSPACE_CONFIG") - - try: - os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") - torch.use_deterministic_algorithms(True) - compiler_context = ( - torch.compiler.set_stance("force_eager") if fully_eager else contextlib.nullcontext() - ) - with compiler_context: - yield - finally: - torch.use_deterministic_algorithms( - previous_deterministic, - warn_only=previous_warn_only, - ) - if previous_cublas_workspace_config is None: - os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None) - else: - os.environ["CUBLAS_WORKSPACE_CONFIG"] = previous_cublas_workspace_config - - -def _save_lpips_video_mp4(video, output_path, frame_rate): - from tensorrt_llm.media.encoding import save_video - - try: - save_video(video, output_path, frame_rate=frame_rate) - except RuntimeError as err: - if "MP4 format requires ffmpeg" not in str(err): - raise - # Never fall back to another codec here: the goldens are H.264/x264 and - # LPIPS compares decoded pixels, so a silent cv2/mp4v fallback compares - # codec artifacts, not model output (this produced a spurious 0.059 on - # wan22 while the generated frames were identical to the golden). - pytest.fail( - "ffmpeg is unavailable for LPIPS video encoding; refusing to fall back " - "to another codec because the golden comparison would measure codec " - f"artifacts instead of model output: {err}" - ) - assert os.path.isfile(output_path), f"Visual gen did not produce {output_path}" - - -def _run_lpips_eval(tmp_path, sample_id, media_type, prompt, reference_path, generated_path): - reference_key = "reference_video_path" if media_type == "video" else "reference_image_path" - generated_key = "generated_video_path" if media_type == "video" else "generated_image_path" - dataset_path = tmp_path / f"{sample_id}_dataset.json" - output_json = tmp_path / f"{sample_id}_lpips_results.json" - dataset_path.write_text( - json.dumps( - { - "samples": [ - { - "id": sample_id, - "media_type": media_type, - "prompt": prompt, - reference_key: str(reference_path), - generated_key: str(generated_path), - } - ] - }, - indent=2, - ) - + "\n", - encoding="utf-8", - ) - - env = os.environ.copy() - env.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") - env["PYTHONPATH"] = ( - f"{REPO_ROOT}{os.pathsep}{env['PYTHONPATH']}" if env.get("PYTHONPATH") else REPO_ROOT - ) - with _lpips_deterministic_algorithms(): - result = subprocess.run( - [ - sys.executable, - VISUAL_GEN_LPIPS_EVAL_SCRIPT, - "--dataset", - str(dataset_path), - "--output-json", - str(output_json), - "--json", - ], - cwd=REPO_ROOT, - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - check=False, - ) - if result.returncode != 0: - pytest.fail(f"LPIPS eval script failed for {sample_id}:\n{result.stdout}") - - scores = json.loads(output_json.read_text(encoding="utf-8")) - score = float(scores["mean_lpips_score"]) - print(f"\n[E2E {sample_id} LPIPS] score: {score:.6f}") - return score - - -def _assert_lpips_below_threshold(score, threshold): - assert score < threshold, f"LPIPS too high: {score:.6f} (expected < {threshold:.6f})" - - -def _preserve_lpips_candidate_on_failure(request, score, threshold, candidate_path, artifact_name): - """Copy the generated candidate into pytest's --output-dir when the LPIPS gate fails. - - CI archives the output dir per stage, so a threshold failure leaves behind the - exact CI-generated media needed to refresh the golden without guessing at - machine-to-machine kernel-stack drift (measured ~0.04 LPIPS across B200 hosts - on the same container for 1-step Wan2.1). - """ - if score < threshold: - return - output_dir = request.config.getoption("--output-dir", default=None) - if not output_dir: - return - dest_dir = os.path.join(str(output_dir), "lpips_failure_artifacts") - os.makedirs(dest_dir, exist_ok=True) - dest = os.path.join(dest_dir, artifact_name) - shutil.copy2(str(candidate_path), dest) - print(f"[LPIPS] candidate preserved for golden refresh: {dest}") - # CI's per-stage results tarball only collects flat files (results.xml etc.), not - # this subdirectory — and junit stdout IS collected. Small candidates therefore - # also go into stdout as base64 so a threshold failure is diagnosable from the - # archived results.xml alone (this is how wan22's spurious 0.059 was traced to a - # silent mpeg4 codec fallback rather than a generation difference). - size = os.path.getsize(dest) - if size <= 8 * 1024 * 1024: - with open(dest, "rb") as fh: - encoded = base64.b64encode(fh.read()).decode("ascii") - print(f"[LPIPS-B64-BEGIN {artifact_name} {size}]") - for i in range(0, len(encoded), 3072): - print(encoded[i : i + 3072]) - print(f"[LPIPS-B64-END {artifact_name}]") - - -def _generate_flux_lpips_image(model_path, output_path): - from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader - from tensorrt_llm.media.encoding import save_image - from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs - - _skip_if_missing(model_path, "FLUX checkpoint", is_dir=True) - _disable_inductor_compile_worker_quiesce() - with _lpips_deterministic_algorithms(): - args = VisualGenArgs( - model=model_path, - torch_compile_config=TorchCompileConfig(enable=False), - ) - pipeline = PipelineLoader(args).load(skip_warmup=True) - try: - result = pipeline.forward( - prompt=FLUX_LPIPS_PROMPT, - height=FLUX_LPIPS_HEIGHT, - width=FLUX_LPIPS_WIDTH, - num_inference_steps=FLUX_LPIPS_NUM_INFERENCE_STEPS, - guidance_scale=FLUX_LPIPS_GUIDANCE_SCALE, - seed=FLUX_LPIPS_SEED, - ) - generated_image = result.image[0].detach().cpu() - finally: - del pipeline - _cleanup_cuda() - - save_image(generated_image, output_path) - - -def _generate_ltx2_lpips_video(output_path, *, enable_cuda_graph=False): - from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader - from tensorrt_llm.visual_gen.args import CudaGraphConfig, TorchCompileConfig, VisualGenArgs - - checkpoint_path = _lpips_model_path("LTX-2", "ltx-2-19b-dev.safetensors") - text_encoder_path = _ltx2_lpips_text_encoder_path() - spatial_upsampler_path = _lpips_model_path("LTX-2", "ltx-2-spatial-upscaler-x2-1.0.safetensors") - distilled_lora_path = _lpips_model_path("LTX-2", "ltx-2-19b-distilled-lora-384.safetensors") - _skip_if_missing(checkpoint_path, "LTX-2 checkpoint") - _skip_if_missing(text_encoder_path, "LTX-2 text encoder", is_dir=True) - _skip_if_missing(spatial_upsampler_path, "LTX-2 spatial upsampler") - _skip_if_missing(distilled_lora_path, "LTX-2 distilled LoRA") - _disable_inductor_compile_worker_quiesce() - - # TorchCompileConfig(enable=False) does not suppress nested @torch.compile decorators. - # Wrapped here (not in the fixture) so the golden fixture and both sides of - # test_ltx2_cuda_graph_lpips_matches_eager run the same eager numerics. - with _lpips_deterministic_algorithms(), torch.compiler.set_stance("force_eager"): - args = VisualGenArgs( - model=checkpoint_path, - pipeline_config={ - "text_encoder_path": text_encoder_path, - "spatial_upsampler_path": spatial_upsampler_path, - "distilled_lora_path": distilled_lora_path, - }, - torch_compile_config=TorchCompileConfig(enable=False), - cuda_graph_config=CudaGraphConfig(enable=enable_cuda_graph), - ) - pipeline = PipelineLoader(args).load(skip_warmup=True) - try: - with torch.no_grad(): - result = pipeline.forward( - prompt=LTX2_T2V_PROMPT, - negative_prompt=LTX2_T2V_NEGATIVE_PROMPT, - height=LTX2_T2V_HEIGHT, - width=LTX2_T2V_WIDTH, - num_frames=LTX2_LPIPS_NUM_FRAMES, - num_inference_steps=LTX2_LPIPS_NUM_INFERENCE_STEPS, - guidance_scale=LTX2_T2V_GUIDANCE_SCALE, - seed=LTX2_T2V_SEED, - ) - generated_video = result.video.detach().cpu() - finally: - del pipeline - _cleanup_cuda() - - _save_lpips_video_mp4(generated_video, output_path, frame_rate=LTX2_T2V_FRAME_RATE) - - -def _generate_ltx2_cuda_graph_trtllm_backend_video(output_path): - from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams - from tensorrt_llm.visual_gen.args import ( - AttentionConfig, - CompilationConfig, - CudaGraphConfig, - ParallelConfig, - TorchCompileConfig, - ) - - scratch_space = _llm_models_root() - checkpoint_path = os.path.join(scratch_space, LTX2_MODEL_CHECKPOINT_PATH) - text_encoder_path = _ltx2_lpips_text_encoder_path() - spatial_upsampler_path = os.path.join(scratch_space, LTX2_UPSAMPLER_SUBPATH) - distilled_lora_path = os.path.join(scratch_space, LTX2_DISTILLED_LORA_SUBPATH) - _skip_if_missing(checkpoint_path, "LTX-2 checkpoint") - _skip_if_missing(text_encoder_path, "LTX-2 text encoder", is_dir=True) - _skip_if_missing(spatial_upsampler_path, "LTX-2 spatial upsampler") - _skip_if_missing(distilled_lora_path, "LTX-2 distilled LoRA") - _disable_inductor_compile_worker_quiesce() - - visual_gen_args = VisualGenArgs( - model=checkpoint_path, - quant_config={"quant_algo": "NVFP4", "dynamic": True}, - attention_config=AttentionConfig(backend="TRTLLM"), - parallel_config=ParallelConfig( - cfg_size=1, - ulysses_size=1, - parallel_vae_size=1, - ), - compilation_config=CompilationConfig( - resolutions=[ - ( - LTX2_T2V_HEIGHT, - LTX2_T2V_WIDTH, - ) - ], - num_frames=[LTX2_LPIPS_NUM_FRAMES], - ), - cuda_graph_config=CudaGraphConfig(enable=True), - torch_compile_config=TorchCompileConfig( - enable=True, - enable_fullgraph=False, - enable_autotune=True, - ), - pipeline_config={ - "text_encoder_path": text_encoder_path, - "spatial_upsampler_path": spatial_upsampler_path, - "distilled_lora_path": distilled_lora_path, - }, - ) - - visual_gen = VisualGen(model=checkpoint_path, args=visual_gen_args) - try: - params = VisualGenParams( - height=LTX2_T2V_HEIGHT, - width=LTX2_T2V_WIDTH, - num_frames=LTX2_LPIPS_NUM_FRAMES, - num_inference_steps=LTX2_LPIPS_NUM_INFERENCE_STEPS, - guidance_scale=LTX2_T2V_GUIDANCE_SCALE, - max_sequence_length=LTX2_T2V_MAX_SEQ_LEN, - seed=LTX2_T2V_SEED, - frame_rate=LTX2_T2V_FRAME_RATE, - negative_prompt=LTX2_T2V_NEGATIVE_PROMPT, - ) - output = visual_gen.generate(inputs=LTX2_T2V_PROMPT, params=params) - _save_lpips_video_mp4(output.video, output_path, frame_rate=LTX2_T2V_FRAME_RATE) - finally: - visual_gen.shutdown() - - assert os.path.isfile(output_path), f"LTX-2 TRTLLM backend did not produce {output_path}" - - -def _run_wan_lpips_pipeline( - model_path, - prompt, - negative_prompt, - height, - width, - num_frames, - num_inference_steps, - guidance_scale, - seed, - attention_backend="VANILLA", - parallel=None, - fully_eager=False, -): - from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader - from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs - - _skip_if_missing(model_path, "Wan checkpoint", is_dir=True) - _disable_inductor_compile_worker_quiesce() - args_kwargs = dict( - model=model_path, - attention_config=AttentionConfig(backend=attention_backend), - torch_compile_config=TorchCompileConfig(enable=False), - ) - if parallel is not None: - args_kwargs["parallel_config"] = parallel - with _lpips_deterministic_algorithms(fully_eager=fully_eager): - args = VisualGenArgs(**args_kwargs) - pipeline = PipelineLoader(args).load(skip_warmup=True) - try: - with torch.no_grad(): - result = pipeline.forward( - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - num_frames=num_frames, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - seed=seed, - ) - if result is None or result.video is None: - return None - return result.video.detach().cpu() - finally: - del pipeline - _cleanup_cuda() - - -def _generate_wan_lpips_video( - model_path, - output_path, - prompt, - negative_prompt, - height, - width, - num_frames, - num_inference_steps, - guidance_scale, - seed, - frame_rate, - attention_backend="VANILLA", - parallel=None, - fully_eager=False, -): - generated_video = _run_wan_lpips_pipeline( - model_path, - prompt, - negative_prompt, - height, - width, - num_frames, - num_inference_steps, - guidance_scale, - seed, - attention_backend=attention_backend, - parallel=parallel, - fully_eager=fully_eager, - ) - assert generated_video is not None, "Single-GPU Wan LPIPS run produced no video" - _save_lpips_video_mp4(generated_video, output_path, frame_rate=frame_rate) - - -@pytest.fixture(scope="session") -def wan21_bf16_video_path(_visual_gen_deps, llm_venv): - output_path = _visual_gen_output_path(llm_venv, "wan21_bf16") - if os.path.isfile(output_path): - return output_path - # TorchCompileConfig(enable=False) does not suppress nested @torch.compile decorators. - with torch.compiler.set_stance("force_eager"): - _generate_wan_lpips_video( - _lpips_model_path("Wan2.1-T2V-1.3B-Diffusers"), - output_path, - WAN21_LPIPS_PROMPT, - WAN21_LPIPS_NEGATIVE_PROMPT, - WAN21_LPIPS_HEIGHT, - WAN21_LPIPS_WIDTH, - WAN21_LPIPS_NUM_FRAMES, - WAN21_LPIPS_NUM_INFERENCE_STEPS, - WAN21_LPIPS_GUIDANCE_SCALE, - WAN21_LPIPS_SEED, - WAN_LPIPS_FRAME_RATE, - ) - return output_path - - -@pytest.fixture(scope="session") -def wan22_bf16_video_path(_visual_gen_deps, llm_venv): - output_path = _visual_gen_output_path(llm_venv, "wan22_bf16") - if os.path.isfile(output_path): - return output_path - # TorchCompileConfig(enable=False) does not suppress nested @torch.compile decorators. - with torch.compiler.set_stance("force_eager"): - _generate_wan_lpips_video( - _lpips_model_path("Wan2.2-T2V-A14B-Diffusers"), - output_path, - WAN22_LPIPS_PROMPT, - WAN22_LPIPS_NEGATIVE_PROMPT, - WAN22_LPIPS_HEIGHT, - WAN22_LPIPS_WIDTH, - WAN22_LPIPS_NUM_FRAMES, - WAN22_LPIPS_NUM_INFERENCE_STEPS, - WAN22_LPIPS_GUIDANCE_SCALE, - WAN22_LPIPS_SEED, - WAN22_LPIPS_FRAME_RATE, - ) - return output_path - - -def _generate_qwenimage_lpips_image(model_path, output_path, *, enable_cuda_graph=False): - """Generate the QwenImage text-to-image LPIPS sample.""" - from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader - from tensorrt_llm.media.encoding import save_image - from tensorrt_llm.visual_gen.args import CudaGraphConfig, TorchCompileConfig, VisualGenArgs - - _skip_if_missing(model_path, "QwenImage checkpoint", is_dir=True) - _disable_inductor_compile_worker_quiesce() - args = VisualGenArgs( - model=model_path, - torch_compile_config=TorchCompileConfig(enable=False), - cuda_graph_config=CudaGraphConfig(enable=enable_cuda_graph), - ) - pipeline = PipelineLoader(args).load(skip_warmup=True) - try: - with torch.no_grad(): - result = pipeline.forward( - prompt=QWENIMAGE_LPIPS_PROMPT, - negative_prompt=QWENIMAGE_LPIPS_NEGATIVE_PROMPT, - height=QWENIMAGE_LPIPS_HEIGHT, - width=QWENIMAGE_LPIPS_WIDTH, - num_inference_steps=QWENIMAGE_LPIPS_NUM_INFERENCE_STEPS, - true_cfg_scale=QWENIMAGE_LPIPS_TRUE_CFG_SCALE, - seed=QWENIMAGE_LPIPS_SEED, - ) - generated_image = result.image[0].detach().cpu() - finally: - del pipeline - _cleanup_cuda() - - save_image(generated_image, output_path) - - -def _run_cosmos3_lpips_pipeline(num_frames): - """Run the Cosmos3-Nano pipeline (default setting, VANILLA attn, compile-off). - - Returns the generated video tensor ``(B, T, H, W, C)`` (T == ``num_frames``), - or ``None`` if generation produced no video. ``num_frames=1`` yields the - single-frame text-to-image path. - """ - # Cosmos3 re-reads the guardrail flag in __init__; set it before the pipeline loads. - guardrails_env_key = "TRTLLM_DISABLE_COSMOS3_GUARDRAILS" - previous_guardrails_env = os.environ.get(guardrails_env_key) - os.environ[guardrails_env_key] = "1" - try: - from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader - from tensorrt_llm.visual_gen.args import ( - AttentionConfig, - CompilationConfig, - TorchCompileConfig, - VisualGenArgs, - ) - - model_path = _lpips_model_path(COSMOS3_NANO_MODEL_SUBPATH) - _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) - _disable_inductor_compile_worker_quiesce() - args = VisualGenArgs( - model=model_path, - compilation_config=CompilationConfig(skip_warmup=True), - torch_compile_config=TorchCompileConfig(enable=False), - attention_config=AttentionConfig(backend="VANILLA"), - ) - pipeline = PipelineLoader(args).load(skip_warmup=True) - try: - with torch.no_grad(): - result = pipeline.forward( - prompt=COSMOS3_LPIPS_PROMPT, - seed=COSMOS3_LPIPS_SEED, - height=COSMOS3_LPIPS_HEIGHT, - width=COSMOS3_LPIPS_WIDTH, - num_frames=num_frames, - num_inference_steps=COSMOS3_LPIPS_NUM_INFERENCE_STEPS, - guidance_scale=COSMOS3_LPIPS_GUIDANCE_SCALE, - frame_rate=COSMOS3_LPIPS_FRAME_RATE, - use_guardrails=False, - ) - if result is None or result.video is None: - return None - return result.video.detach().cpu() - finally: - del pipeline - _cleanup_cuda() - finally: - if previous_guardrails_env is None: - os.environ.pop(guardrails_env_key, None) - else: - os.environ[guardrails_env_key] = previous_guardrails_env - - -def _generate_cosmos3_lpips_video(output_path): - """Generate the Cosmos3-Nano text-to-video LPIPS sample.""" - video = _run_cosmos3_lpips_pipeline(COSMOS3_LPIPS_T2V_NUM_FRAMES) - assert video is not None, "Cosmos3-Nano T2V LPIPS run produced no video" - _save_lpips_video_mp4(video, output_path, frame_rate=COSMOS3_LPIPS_FRAME_RATE) - - -def _generate_cosmos3_lpips_image(output_path): - """Generate the Cosmos3-Nano text-to-image LPIPS sample (single frame).""" - from tensorrt_llm.media.encoding import save_image - - video = _run_cosmos3_lpips_pipeline(COSMOS3_LPIPS_T2I_NUM_FRAMES) - assert video is not None, "Cosmos3-Nano T2I LPIPS run produced no frame" - # video is (B, T, H, W, C); take the single frame -> (H, W, C) for save_image. - save_image(video[0, 0], output_path) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_flux1_lpips_against_golden(tmp_path): - generated_path = tmp_path / "flux1_generated.png" - golden_path = _golden_media_path( - tmp_path, "flux1_lpips_golden.png", "FLUX.1 LPIPS golden image" - ) - _generate_flux_lpips_image(_lpips_model_path("FLUX.1-dev"), generated_path) - score = _run_lpips_eval( - tmp_path, - "flux1", - "image", - FLUX_LPIPS_PROMPT, - golden_path, - generated_path, - ) - _assert_lpips_below_threshold(score, FLUX_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_flux2_lpips_against_golden(tmp_path): - generated_path = tmp_path / "flux2_generated.png" - golden_path = _golden_media_path( - tmp_path, "flux2_lpips_golden.png", "FLUX.2 LPIPS golden image" - ) - _generate_flux_lpips_image(_lpips_model_path("FLUX.2-dev"), generated_path) - score = _run_lpips_eval( - tmp_path, - "flux2", - "image", - FLUX_LPIPS_PROMPT, - golden_path, - generated_path, - ) - _assert_lpips_below_threshold(score, FLUX_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_ltx2_lpips_against_golden(request, tmp_path, ltx2_two_stage_bf16_video_path): - golden_path = _golden_media_path( - tmp_path, "ltx2_lpips_golden_video.mp4", "LTX-2 LPIPS golden video" - ) - score = _run_lpips_eval( - tmp_path, - "ltx2", - "video", - LTX2_T2V_PROMPT, - golden_path, - ltx2_two_stage_bf16_video_path, - ) - _preserve_lpips_candidate_on_failure( - request, - score, - LTX2_LPIPS_THRESHOLD, - ltx2_two_stage_bf16_video_path, - "ltx2_lpips_golden_video.mp4", - ) - _assert_lpips_below_threshold(score, LTX2_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_ltx2_cuda_graph_lpips_matches_eager(_visual_gen_deps, tmp_path): - eager_path = tmp_path / "ltx2_eager_generated.mp4" - cuda_graph_path = tmp_path / "ltx2_cuda_graph_generated.mp4" - - _generate_ltx2_lpips_video(eager_path, enable_cuda_graph=False) - _generate_ltx2_lpips_video(cuda_graph_path, enable_cuda_graph=True) - score = _run_lpips_eval( - tmp_path, - "ltx2_cuda_graph", - "video", - LTX2_T2V_PROMPT, - eager_path, - cuda_graph_path, - ) - _assert_lpips_below_threshold(score, LTX2_CUDA_GRAPH_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_ltx2_cuda_graph_trtllm_backend(request, _visual_gen_deps, tmp_path): - generated_path = tmp_path / "ltx2_cuda_graph_trtllm_backend_generated.mp4" - golden_path = _golden_media_path( - tmp_path, "ltx2_lpips_golden_video.mp4", "LTX-2 LPIPS golden video" - ) - _generate_ltx2_cuda_graph_trtllm_backend_video(generated_path) - score = _run_lpips_eval( - tmp_path, - "ltx2_cuda_graph_trtllm_backend", - "video", - LTX2_T2V_PROMPT, - golden_path, - generated_path, - ) - _preserve_lpips_candidate_on_failure( - request, - score, - LTX2_LPIPS_THRESHOLD, - generated_path, - "ltx2_cuda_graph_trtllm_backend_generated.mp4", - ) - _assert_lpips_below_threshold(score, LTX2_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_wan21_t2v_lpips_against_golden(request, tmp_path, wan21_bf16_video_path): - golden_path = _golden_media_path( - tmp_path, "wan21_t2v_lpips_golden_video.mp4", "Wan 2.1 LPIPS golden video" - ) - score = _run_lpips_eval( - tmp_path, - "wan21_t2v", - "video", - WAN21_LPIPS_PROMPT, - golden_path, - wan21_bf16_video_path, - ) - _preserve_lpips_candidate_on_failure( - request, - score, - WAN_LPIPS_THRESHOLD, - wan21_bf16_video_path, - "wan21_t2v_lpips_golden_video.mp4", - ) - _assert_lpips_below_threshold(score, WAN_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_wan22_t2v_lpips_against_golden(request, tmp_path, wan22_bf16_video_path): - golden_path = _golden_media_path( - tmp_path, "wan22_t2v_lpips_golden_video.mp4", "Wan 2.2 LPIPS golden video" - ) - score = _run_lpips_eval( - tmp_path, - "wan22_t2v", - "video", - WAN22_LPIPS_PROMPT, - golden_path, - wan22_bf16_video_path, - ) - _preserve_lpips_candidate_on_failure( - request, - score, - WAN_LPIPS_THRESHOLD, - wan22_bf16_video_path, - "wan22_t2v_lpips_golden_video.mp4", - ) - _assert_lpips_below_threshold(score, WAN_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_qwenimage_lpips_against_golden(tmp_path): - generated_path = tmp_path / "qwenimage_generated.png" - golden_path = _golden_media_path( - tmp_path, "qwenimage_lpips_golden.png", "QwenImage LPIPS golden image" - ) - _generate_qwenimage_lpips_image(_lpips_model_path(QWENIMAGE_MODEL_SUBPATH), generated_path) - score = _run_lpips_eval( - tmp_path, - "qwenimage", - "image", - QWENIMAGE_LPIPS_PROMPT, - golden_path, - generated_path, - ) - _assert_lpips_below_threshold(score, QWENIMAGE_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_qwenimage_cuda_graph_lpips_against_golden(tmp_path): - generated_path = tmp_path / "qwenimage_cuda_graph_generated.png" - golden_path = _golden_media_path( - tmp_path, "qwenimage_lpips_golden.png", "QwenImage LPIPS golden image" - ) - _generate_qwenimage_lpips_image( - _lpips_model_path(QWENIMAGE_MODEL_SUBPATH), - generated_path, - enable_cuda_graph=True, - ) - score = _run_lpips_eval( - tmp_path, - "qwenimage_cuda_graph", - "image", - QWENIMAGE_LPIPS_PROMPT, - golden_path, - generated_path, - ) - _assert_lpips_below_threshold(score, QWENIMAGE_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_cosmos3_nano_t2v_lpips_against_golden(_visual_gen_deps, tmp_path): - generated_path = tmp_path / "cosmos3_nano_t2v_generated.mp4" - golden_path = _golden_media_path( - tmp_path, - "cosmos3_nano_t2v_lpips_golden_video.mp4", - "Cosmos3-Nano T2V LPIPS golden video", - ) - _generate_cosmos3_lpips_video(generated_path) - score = _run_lpips_eval( - tmp_path, - "cosmos3_nano_t2v", - "video", - COSMOS3_LPIPS_PROMPT, - golden_path, - generated_path, - ) - _assert_lpips_below_threshold(score, COSMOS3_LPIPS_THRESHOLD) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_cosmos3_nano_t2i_lpips_against_golden(tmp_path): - generated_path = tmp_path / "cosmos3_nano_t2i_generated.png" - golden_path = _golden_media_path( - tmp_path, "cosmos3_nano_t2i_lpips_golden.png", "Cosmos3-Nano T2I LPIPS golden image" - ) - _generate_cosmos3_lpips_image(generated_path) - score = _run_lpips_eval( - tmp_path, - "cosmos3_nano_t2i", - "image", - COSMOS3_LPIPS_PROMPT, - golden_path, - generated_path, - ) - _assert_lpips_below_threshold(score, COSMOS3_LPIPS_THRESHOLD) - - -def _generate_wan_video(llm_venv, model_subpath, output_subdir): - """Generate a WAN video for a given model checkpoint. - - Returns the path to the generated .mp4, or calls pytest.skip if the model - is not found under LLM_MODELS_ROOT. - """ - from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams - - scratch_space = _llm_models_root() - model_path = os.path.join(scratch_space, model_subpath) - if not os.path.isdir(model_path): - pytest.skip( - f"Model not found: {model_path} " - f"(set LLM_MODELS_ROOT or place {model_subpath} under scratch)" - ) - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", output_subdir) - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, VISUAL_GEN_OUTPUT_VIDEO) - if os.path.isfile(output_path): - return output_path - - visual_gen_args = VisualGenArgs( - attention_config={"backend": "VANILLA"}, - parallel_config={"cfg_size": 2 if torch.cuda.device_count() >= 2 else 1}, - torch_compile_config={"enable": False}, - compilation_config={"skip_warmup": True}, - ) - visual_gen = VisualGen(model=model_path, args=visual_gen_args) - try: - frame_rate = visual_gen.default_params.frame_rate - output = visual_gen.generate( - inputs=WAN_T2V_PROMPT, - params=VisualGenParams( - height=WAN_T2V_HEIGHT, - width=WAN_T2V_WIDTH, - seed=42, - num_frames=WAN_T2V_NUM_FRAMES, - frame_rate=frame_rate, - ), - ) - assert output.error is None, f"unexpected error on WAN run: {output.error}" - assert output.video is not None - _save_lpips_video_mp4(output.video, output_path, frame_rate=frame_rate) - finally: - visual_gen.shutdown() - - assert os.path.isfile(output_path), f"Visual gen did not produce {output_path}" - return output_path - - -@pytest.fixture(scope="session") -def wan22_a14b_fp8_video_path(_visual_gen_deps, llm_venv): - """Generate video with Wan 2.2 A14B FP8 checkpoint.""" - return _generate_wan_video(llm_venv, WAN22_A14B_FP8_MODEL_SUBPATH, "wan22_fp8") - - -@pytest.fixture(scope="session") -def wan22_a14b_nvfp4_video_path(_visual_gen_deps, llm_venv): - """Generate video with Wan 2.2 A14B NVFP4 checkpoint.""" - return _generate_wan_video(llm_venv, WAN22_A14B_NVFP4_MODEL_SUBPATH, "wan22_nvfp4") - - -def _linear_type_to_quant_config(linear_type): - """Map linear_type shortcut to quant_config dict for VisualGenArgs.""" - mapping = { - "trtllm-fp8-per-tensor": {"quant_algo": "FP8", "dynamic": True}, - "trtllm-fp8-blockwise": {"quant_algo": "FP8_BLOCK_SCALES", "dynamic": True}, - "trtllm-nvfp4": {"quant_algo": "NVFP4", "dynamic": True}, - } - return mapping.get(linear_type) - - -def _generate_ltx2_two_stage_video(llm_venv, output_subdir, linear_type="default"): - """Generate a two-stage LTX-2 video using the Python API. - - Requires the main checkpoint, text encoder, spatial upsampler, and - distilled LoRA. Returns the path to the generated .mp4, or calls - pytest.skip if any asset is missing. - """ - from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams - - scratch_space = _llm_models_root() - model_path = os.path.join(scratch_space, LTX2_MODEL_CHECKPOINT_PATH) - text_encoder_path = os.path.join(scratch_space, LTX2_TEXT_ENCODER_SUBPATH) - upsampler_path = os.path.join(scratch_space, LTX2_UPSAMPLER_SUBPATH) - lora_path = os.path.join(scratch_space, LTX2_DISTILLED_LORA_SUBPATH) - - for label, path, is_file in [ - ("LTX-2 checkpoint", model_path, True), - ("text encoder", text_encoder_path, False), - ("spatial upsampler", upsampler_path, True), - ("distilled LoRA", lora_path, True), - ]: - exists = os.path.isfile(path) if is_file else os.path.isdir(path) - if not exists: - pytest.skip(f"Two-stage {label} not found: {path}") - - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", output_subdir) - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, VISUAL_GEN_OUTPUT_VIDEO) - if os.path.isfile(output_path): - return output_path - - vg_kwargs = dict( - pipeline_config={ - "text_encoder_path": text_encoder_path, - "spatial_upsampler_path": upsampler_path, - "distilled_lora_path": lora_path, - }, - ) - quant_config = _linear_type_to_quant_config(linear_type) - if quant_config is not None: - vg_kwargs["quant_config"] = quant_config - if torch.cuda.device_count() >= 2: - vg_kwargs["parallel_config"] = {"cfg_size": 2} - - visual_gen_args = VisualGenArgs(**vg_kwargs) - visual_gen = VisualGen(model=model_path, args=visual_gen_args) - - try: - params = VisualGenParams( - height=LTX2_TWO_STAGE_HEIGHT, - width=LTX2_TWO_STAGE_WIDTH, - num_frames=LTX2_TWO_STAGE_NUM_FRAMES, - num_inference_steps=LTX2_TWO_STAGE_STEPS, - guidance_scale=LTX2_TWO_STAGE_GUIDANCE_SCALE, - max_sequence_length=LTX2_T2V_MAX_SEQ_LEN, - seed=LTX2_T2V_SEED, - frame_rate=LTX2_T2V_FRAME_RATE, - negative_prompt=LTX2_T2V_NEGATIVE_PROMPT, - ) - output = visual_gen.generate(inputs=LTX2_T2V_PROMPT, params=params) - assert output.error is None - assert output.video is not None - assert output.frame_rate == LTX2_T2V_FRAME_RATE - assert output.metrics is not None - assert output.metrics.generation > 0 - assert output.metrics.denoise > 0 - assert output.metrics.post_denoise >= 0 - _save_lpips_video_mp4(output.video, output_path, frame_rate=LTX2_T2V_FRAME_RATE) - finally: - visual_gen.shutdown() - - assert os.path.isfile(output_path), f"LTX-2 two-stage did not produce {output_path}" - return output_path - - -@pytest.fixture(scope="session") -def ltx2_two_stage_bf16_video_path(_visual_gen_deps, llm_venv): - """Generate LTX-2 two-stage BF16 video with the LPIPS config and return path.""" - output_path = _visual_gen_output_path(llm_venv, "ltx2_two_stage_bf16") - if os.path.isfile(output_path): - return output_path - _generate_ltx2_lpips_video(output_path) - return output_path - - -@pytest.fixture(scope="session") -def ltx2_two_stage_fp8_video_path(_visual_gen_deps, llm_venv): - """Generate LTX-2 two-stage FP8 T2V video and return path.""" - return _generate_ltx2_two_stage_video( - llm_venv, "ltx2_two_stage_fp8", linear_type="trtllm-fp8-per-tensor" - ) - - -def _normalize_score(val): - """Normalize to 0-1 scale (e.g. imaging_quality can be 0-100).""" - if isinstance(val, bool): - return float(val) - if isinstance(val, (int, float)) and val > 1.5: - return val / 100.0 - return float(val) - - -def _get_per_video_scores(results, video_path_substr): - """From VBench results, get per-dimension score for the video whose path contains video_path_substr.""" - scores = {} - for dim in VBENCH_DIMENSIONS: - dim_result = results[dim] - assert isinstance(dim_result, list) and len(dim_result) >= 2, ( - f"Dimension '{dim}' result must be [overall_score, video_results]; got {type(dim_result)}" - ) - video_results = dim_result[1] - for entry in video_results: - if video_path_substr in entry.get("video_path", ""): - raw = entry.get("video_results") - scores[dim] = _normalize_score(raw) - break - else: - raise AssertionError( - f"No video matching '{video_path_substr}' in dimension '{dim}'; " - f"paths: {[e.get('video_path') for e in video_results]}" - ) - return scores - - -def test_vbench_dimension_score_wan(vbench_repo_root, wan21_bf16_video_path, llm_venv): - """Run VBench on WAN 2.1 BF16 video generated with the LPIPS config.""" - videos_dir = os.path.dirname(wan21_bf16_video_path) - assert os.path.isfile(wan21_bf16_video_path), "WAN 2.1 BF16 video must exist" - _run_vbench_and_report( - vbench_repo_root, - videos_dir, - VISUAL_GEN_OUTPUT_VIDEO, - llm_venv, - title="WAN 2.1 BF16", - golden_scores=VBENCH_WAN_GOLDEN_SCORES, - max_score_diff=0.05, - ) - - -def _run_vbench_and_report( - vbench_repo_root, - videos_dir, - trtllm_filename, - llm_venv, - title, - golden_scores=None, - max_score_diff=0.10, -): - """Run VBench, print scores, and optionally assert against golden values. - - When *golden_scores* is None the test is a baseline run: scores are printed - for the operator to capture but no assertion is made. Once golden values - are populated, the comparison kicks in automatically. - """ - output_path = os.path.join( - llm_venv.get_working_directory(), "vbench_eval_output", title.replace(" ", "_").lower() - ) - os.makedirs(output_path, exist_ok=True) - evaluate_script = os.path.join(vbench_repo_root, "evaluate.py") - cmd = [ - evaluate_script, - "--videos_path", - videos_dir, - "--output_path", - output_path, - "--mode", - "custom_input", - ] - cmd.extend(["--dimension"] + VBENCH_DIMENSIONS) - _venv_check_call(llm_venv, cmd) - - pattern = os.path.join(output_path, "*_eval_results.json") - result_files = glob.glob(pattern) - assert result_files, ( - f"No eval results found matching {pattern}; output dir: {os.listdir(output_path)}" - ) - with open(result_files[0], "r") as f: - results = json.load(f) - for dim in VBENCH_DIMENSIONS: - assert dim in results, ( - f"Expected dimension '{dim}' in results; keys: {list(results.keys())}" - ) - - scores_trtllm = _get_per_video_scores(results, trtllm_filename) - max_len = max(len(d) for d in VBENCH_DIMENSIONS) - - if golden_scores is not None: - header = f"{'Dimension':<{max_len}} | {'TRT-LLM':>10} | {'Golden':>10} | {'Diff':>8}" - else: - header = f"{'Dimension':<{max_len}} | {'TRT-LLM':>10}" - sep = "-" * len(header) - print("\n" + "=" * len(header)) - print(f"VBench dimension scores ({title})") - print("=" * len(header)) - print(header) - print(sep) - - for dim in VBENCH_DIMENSIONS: - t = scores_trtllm[dim] - if golden_scores is not None: - r = golden_scores[dim] - print(f"{dim:<{max_len}} | {t:>10.4f} | {r:>10.4f} | {abs(t - r):>8.4f}") - else: - print(f"{dim:<{max_len}} | {t:>10.4f}") - print(sep) - - if golden_scores is None: - print("\n** Baseline run — no golden scores to compare against. **") - print("** Copy the values above into the GOLDEN_SCORES dict. **\n") - return scores_trtllm - - max_diff_val = max(abs(scores_trtllm[d] - golden_scores[d]) for d in VBENCH_DIMENSIONS) - print(f"max_diff={max_diff_val:.4f} (threshold={max_score_diff})") - print("=" * len(header) + "\n") - for dim in VBENCH_DIMENSIONS: - diff = abs(scores_trtllm[dim] - golden_scores[dim]) - assert diff < max_score_diff or scores_trtllm[dim] >= golden_scores[dim], ( - f"Dimension '{dim}' score difference {diff:.4f} >= {max_score_diff} " - f"(TRT-LLM={scores_trtllm[dim]:.4f}, golden={golden_scores[dim]:.4f})" - ) - return scores_trtllm - - -def test_vbench_dimension_score_wan22_bf16(vbench_repo_root, wan22_bf16_video_path, llm_venv): - """VBench accuracy for Wan 2.2 A14B BF16 generated with the LPIPS config.""" - videos_dir = os.path.dirname(wan22_bf16_video_path) - assert os.path.isfile(wan22_bf16_video_path), "WAN 2.2 BF16 video must exist" - _run_vbench_and_report( - vbench_repo_root, - videos_dir, - VISUAL_GEN_OUTPUT_VIDEO, - llm_venv, - title="WAN 2.2 A14B BF16", - golden_scores=VBENCH_WAN22_BF16_GOLDEN_SCORES, - max_score_diff=0.05, - ) - - -def test_vbench_dimension_score_wan22_a14b_fp8( - vbench_repo_root, wan22_a14b_fp8_video_path, llm_venv -): - """VBench accuracy for Wan 2.2 A14B FP8 — full generate + evaluate.""" - videos_dir = os.path.dirname(wan22_a14b_fp8_video_path) - assert os.path.isfile(wan22_a14b_fp8_video_path), "FP8 video must exist" - _run_vbench_and_report( - vbench_repo_root, - videos_dir, - VISUAL_GEN_OUTPUT_VIDEO, - llm_venv, - title="WAN 2.2 A14B FP8", - golden_scores=VBENCH_WAN22_A14B_FP8_GOLDEN_SCORES, - max_score_diff=0.06, - ) - - -def test_vbench_dimension_score_wan22_a14b_nvfp4( - vbench_repo_root, wan22_a14b_nvfp4_video_path, llm_venv -): - """VBench accuracy for Wan 2.2 A14B NVFP4 — full generate + evaluate.""" - videos_dir = os.path.dirname(wan22_a14b_nvfp4_video_path) - assert os.path.isfile(wan22_a14b_nvfp4_video_path), "NVFP4 video must exist" - _run_vbench_and_report( - vbench_repo_root, - videos_dir, - VISUAL_GEN_OUTPUT_VIDEO, - llm_venv, - title="WAN 2.2 A14B NVFP4", - golden_scores=VBENCH_WAN22_A14B_NVFP4_GOLDEN_SCORES, - max_score_diff=0.05, - ) - - -def test_vbench_dimension_score_ltx2_two_stage_bf16( - vbench_repo_root, ltx2_two_stage_bf16_video_path, llm_venv -): - """VBench accuracy for LTX-2 two-stage BF16 T2V.""" - videos_dir = os.path.dirname(ltx2_two_stage_bf16_video_path) - assert os.path.isfile(ltx2_two_stage_bf16_video_path), "LTX-2 two-stage BF16 video must exist" - _run_vbench_and_report( - vbench_repo_root, - videos_dir, - VISUAL_GEN_OUTPUT_VIDEO, - llm_venv, - title="LTX-2 Two-Stage BF16", - golden_scores=VBENCH_LTX2_TWO_STAGE_BF16_GOLDEN_SCORES, - max_score_diff=0.05, - ) - - -def test_vbench_dimension_score_ltx2_two_stage_fp8( - vbench_repo_root, ltx2_two_stage_fp8_video_path, llm_venv -): - """VBench accuracy for LTX-2 two-stage FP8 T2V.""" - videos_dir = os.path.dirname(ltx2_two_stage_fp8_video_path) - assert os.path.isfile(ltx2_two_stage_fp8_video_path), "LTX-2 two-stage FP8 video must exist" - _run_vbench_and_report( - vbench_repo_root, - videos_dir, - VISUAL_GEN_OUTPUT_VIDEO, - llm_venv, - title="LTX-2 Two-Stage FP8", - golden_scores=VBENCH_LTX2_TWO_STAGE_FP8_GOLDEN_SCORES, - max_score_diff=0.05, - ) - - -def test_visual_gen_quickstart(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/quickstart_example.py end-to-end.""" - scratch_space = _llm_models_root() - model_src = os.path.join(scratch_space, WAN_T2V_MODEL_SUBPATH) - if not os.path.isdir(model_src): - pytest.skip( - f"Model not found: {model_src} " - f"(set LLM_MODELS_ROOT or place {WAN_T2V_MODEL_SUBPATH} under scratch)" - ) - - model_dst = os.path.join(llm_venv.get_working_directory(), "Wan-AI", WAN_T2V_MODEL_SUBPATH) - if not os.path.islink(model_dst): - os.makedirs(os.path.dirname(model_dst), exist_ok=True) - os.symlink(model_src, model_dst, target_is_directory=True) - - script_path = os.path.join(llm_root, "examples", "visual_gen", "quickstart_example.py") - _venv_check_call(llm_venv, [script_path]) - - output_path = os.path.join(llm_venv.get_working_directory(), "output.avi") - assert os.path.isfile(output_path), f"Quickstart did not produce output.avi at {output_path}" - - -def test_visual_gen_api_walkthrough(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/api_walkthrough.py end-to-end.""" - scratch_space = _llm_models_root() - model_src = os.path.join(scratch_space, WAN_T2V_MODEL_SUBPATH) - if not os.path.isdir(model_src): - pytest.skip( - f"Model not found: {model_src} " - f"(set LLM_MODELS_ROOT or place {WAN_T2V_MODEL_SUBPATH} under scratch)" - ) - - model_dst = os.path.join(llm_venv.get_working_directory(), "Wan-AI", WAN_T2V_MODEL_SUBPATH) - if not os.path.islink(model_dst): - os.makedirs(os.path.dirname(model_dst), exist_ok=True) - os.symlink(model_src, model_dst, target_is_directory=True) - - script_path = os.path.join(llm_root, "examples", "visual_gen", "api_walkthrough.py") - _venv_check_call(llm_venv, [script_path]) - - output_path = os.path.join(llm_venv.get_working_directory(), "api_walkthrough_output.avi") - assert os.path.isfile(output_path), f"API walkthrough did not produce {output_path}" - - -# ============================================================================= -# Core example tests — run per-model scripts from examples/visual_gen/models/ -# with shared YAML configs from examples/visual_gen/configs/. -# ============================================================================= - - -def test_wan_t2v_example(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/models/wan_t2v.py with NVFP4 config end-to-end. - - This is a core example test: it validates that the per-model example script - and the shared YAML config work together as documented in the README. - Uses the pre-quantized Wan 2.2 T2V A14B NVFP4 checkpoint and the shared - ``configs/wan2.2-t2v-fp4-1gpu.yaml`` (NVFP4 dynamic quant). The closest - overlapping test is ``test_vbench_dimension_score_wan22_a14b_nvfp4``, - which runs the same script but with a no-quant YAML synthesized at - runtime and additionally evaluates VBench scores. - """ - scratch_space = _llm_models_root() - model_path = os.path.join(scratch_space, WAN22_A14B_NVFP4_MODEL_SUBPATH) - assert os.path.isdir(model_path), ( - f"Model not found: {model_path} " - f"(set LLM_MODELS_ROOT or place {WAN22_A14B_NVFP4_MODEL_SUBPATH} under models root)" - ) - - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "wan_t2v_example") - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, "wan_t2v_output.mp4") - - script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "wan_t2v.py") - config_path = os.path.join( - llm_root, "examples", "visual_gen", "configs", "wan2.2-t2v-fp4-1gpu.yaml" - ) - assert os.path.isfile(script_path), f"Example script not found: {script_path}" - assert os.path.isfile(config_path), f"Config not found: {config_path}" - - _venv_check_call( - llm_venv, - [ - script_path, - "--model", - model_path, - "--visual_gen_args", - config_path, - "--output_path", - output_path, - ], - ) - assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" - - -def test_flux1_example(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/models/flux1.py with NVFP4 config end-to-end. - - Validates that the FLUX.1-dev example script and ``configs/flux1-dev-fp4-1gpu.yaml`` - work together as documented. Uses the local FLUX.1-dev checkpoint and the shared - NVFP4 dynamic-quant config. - """ - model_path = _lpips_model_path("FLUX.1-dev") - _skip_if_missing(model_path, "FLUX.1-dev checkpoint", is_dir=True) - - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "flux1_example") - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, "flux1_output.png") - - script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "flux1.py") - config_path = os.path.join( - llm_root, "examples", "visual_gen", "configs", "flux1-dev-fp4-1gpu.yaml" - ) - assert os.path.isfile(script_path), f"Example script not found: {script_path}" - assert os.path.isfile(config_path), f"Config not found: {config_path}" - - _venv_check_call( - llm_venv, - [ - script_path, - "--model", - model_path, - "--visual_gen_args", - config_path, - "--output_path", - output_path, - ], - ) - assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" - - -def test_flux2_example(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/models/flux2.py with NVFP4 config end-to-end. - - Validates that the FLUX.2-dev example script and ``configs/flux2-dev-fp4-1gpu.yaml`` - work together as documented. Uses the local FLUX.2-dev checkpoint and the shared - NVFP4 dynamic-quant config. - """ - model_path = _lpips_model_path("FLUX.2-dev") - _skip_if_missing(model_path, "FLUX.2-dev checkpoint", is_dir=True) - - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "flux2_example") - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, "flux2_output.png") - - script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "flux2.py") - config_path = os.path.join( - llm_root, "examples", "visual_gen", "configs", "flux2-dev-fp4-1gpu.yaml" - ) - assert os.path.isfile(script_path), f"Example script not found: {script_path}" - assert os.path.isfile(config_path), f"Config not found: {config_path}" - - _venv_check_call( - llm_venv, - [ - script_path, - "--model", - model_path, - "--visual_gen_args", - config_path, - "--output_path", - output_path, - ], - ) - assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" - - -def test_ltx2_example(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/models/ltx2.py with NVFP4 config end-to-end. - - Validates that the LTX-2 example script and ``configs/ltx2-t2v-fp4-1gpu.yaml`` - work together as documented. The Gemma3 text encoder is passed separately via - ``--text_encoder_path`` because the shared YAML intentionally omits it to keep - the config model-path-agnostic. - """ - model_path = _lpips_model_path("LTX-2", "ltx-2-19b-dev.safetensors") - _skip_if_missing(model_path, "LTX-2 checkpoint") - text_encoder_path = _ltx2_lpips_text_encoder_path() - _skip_if_missing(text_encoder_path, "LTX-2 text encoder (gemma-3-12b-it)", is_dir=True) - - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "ltx2_example") - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, "ltx2_output.mp4") - - script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "ltx2.py") - config_path = os.path.join( - llm_root, "examples", "visual_gen", "configs", "ltx2-t2v-fp4-1gpu.yaml" - ) - assert os.path.isfile(script_path), f"Example script not found: {script_path}" - assert os.path.isfile(config_path), f"Config not found: {config_path}" - - _venv_check_call( - llm_venv, - [ - script_path, - "--model", - model_path, - "--visual_gen_args", - config_path, - "--text_encoder_path", - text_encoder_path, - "--output_path", - output_path, - ], - ) - assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" - - -def test_wan_i2v_example(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/models/wan_i2v.py with NVFP4 config end-to-end. - - Validates that the Wan I2V example script and ``configs/wan2.2-i2v-fp4-1gpu.yaml`` - work together as documented. Uses the pre-quantized Wan 2.2 I2V A14B NVFP4 - checkpoint and the default input image (cat_piano.png) bundled with the examples. - """ - scratch_space = _llm_models_root() - model_path = os.path.join(scratch_space, WAN22_I2V_A14B_NVFP4_MODEL_SUBPATH) - if not os.path.isdir(model_path): - pytest.skip( - f"Model not found: {model_path} " - f"(set LLM_MODELS_ROOT or place {WAN22_I2V_A14B_NVFP4_MODEL_SUBPATH} under models root)" - ) - - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "wan_i2v_example") - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, "wan_i2v_output.mp4") - - script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "wan_i2v.py") - config_path = os.path.join( - llm_root, "examples", "visual_gen", "configs", "wan2.2-i2v-fp4-1gpu.yaml" - ) - assert os.path.isfile(script_path), f"Example script not found: {script_path}" - assert os.path.isfile(config_path), f"Config not found: {config_path}" - - _venv_check_call( - llm_venv, - [ - script_path, - "--model", - model_path, - "--visual_gen_args", - config_path, - "--output_path", - output_path, - ], - ) - assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" - - -def test_qwen_image_example(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/models/qwen_image.py with FP8 config end-to-end. - - Validates that the Qwen-Image example script and - ``configs/qwen-image-fp8-1gpu.yaml`` work together as documented. Uses the - local Qwen-Image checkpoint and the shared FP8 blockwise dynamic-quant config. - """ - scratch_space = _llm_models_root() - model_path = os.path.join(scratch_space, QWEN_IMAGE_MODEL_SUBPATH) - _skip_if_missing(model_path, "Qwen-Image checkpoint", is_dir=True) - model_index_path = os.path.join(model_path, "model_index.json") - if not os.path.isfile(model_index_path): - pytest.skip( - f"Qwen-Image checkpoint is incomplete: {model_path} (missing {model_index_path})" - ) - - out_dir = os.path.join( - llm_venv.get_working_directory(), "visual_gen_output", "qwen_image_example" - ) - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, "qwen_image_output.png") - - script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "qwen_image.py") - config_path = os.path.join( - llm_root, "examples", "visual_gen", "configs", "qwen-image-fp8-1gpu.yaml" - ) - assert os.path.isfile(script_path), f"Example script not found: {script_path}" - assert os.path.isfile(config_path), f"Config not found: {config_path}" - - _venv_check_call( - llm_venv, - [ - script_path, - "--model", - model_path, - "--visual_gen_args", - config_path, - "--output_path", - output_path, - ], - ) - assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" - - -def test_cosmos3_example(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/models/cosmos3/cosmos3.py with FP8 config end-to-end. - - Validates that the Cosmos3-Nano example script and ``configs/cosmos3-nano-1gpu.yaml`` - work together as documented. Uses the local Cosmos3-Nano checkpoint and - the shared FP8 dynamic-quant config. - """ - model_path = _lpips_model_path("Cosmos3-Nano") - _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) - - out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_example") - os.makedirs(out_dir, exist_ok=True) - output_path = os.path.join(out_dir, "cosmos3_output.mp4") - - script_path = os.path.join( - llm_root, "examples", "visual_gen", "models", "cosmos3", "cosmos3.py" - ) - config_path = os.path.join( - llm_root, "examples", "visual_gen", "configs", "cosmos3-nano-1gpu.yaml" - ) - assert os.path.isfile(script_path), f"Example script not found: {script_path}" - assert os.path.isfile(config_path), f"Config not found: {config_path}" - - _venv_check_call( - llm_venv, - [ - script_path, - "--model", - model_path, - "--visual_gen_args", - config_path, - "--prompt", - "A serene mountain landscape with snow-capped peaks and a flowing river", - "--output_path", - output_path, - ], - env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, - ) - assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py new file mode 100644 index 000000000000..51bca7f78dde --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py @@ -0,0 +1,352 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""Single-GPU integration and accuracy tests for Cosmos3.""" + +import os +from dataclasses import dataclass + +import pytest +import torch +from defs.common import venv_check_call +from defs.examples.visual_gen.visual_gen_test_utils import ( + FeatureConfigState, + _assert_feature_quantization_installed, + _assert_lpips_below_threshold, + _assert_resolved_single_device_feature_config, + _assert_single_device_feature_executed, + _build_single_device_feature_args, + _cleanup_cuda, + _cleanup_single_device_feature_pipeline, + _disable_inductor_compile_worker_quiesce, + _fixed_nvfp4_quantization_backend, + _golden_media_path, + _lpips_deterministic_algorithms, + _lpips_model_path, + _preserve_lpips_candidate_on_failure, + _run_lpips_eval, + _run_reusable_image_lpips_eval, + _run_single_device_feature_generator, + _save_lpips_video_mp4, + _skip_if_missing, + _validate_single_feature_config, +) + +# Cosmos3-Nano (text-to-video + text-to-image) — default-setting LPIPS golden. +# Params are the Cosmos3 720P defaults (cosmos3/defaults.py:COSMOS3_720P_PARAMS). +# Cosmos3 requires VANILLA attention and guardrails disabled in CI. +COSMOS3_NANO_MODEL_SUBPATH = "Cosmos3-Nano" +COSMOS3_LPIPS_PROMPT = "A serene mountain landscape with snow-capped peaks and a flowing river" +COSMOS3_LPIPS_HEIGHT = 720 +COSMOS3_LPIPS_WIDTH = 1280 +COSMOS3_LPIPS_T2V_NUM_FRAMES = 189 +COSMOS3_LPIPS_T2I_NUM_FRAMES = 1 +COSMOS3_LPIPS_NUM_INFERENCE_STEPS = 35 +COSMOS3_LPIPS_GUIDANCE_SCALE = 6.0 +COSMOS3_LPIPS_SEED = 42 +COSMOS3_LPIPS_FRAME_RATE = 24.0 +COSMOS3_LPIPS_THRESHOLD = 0.05 + +COSMOS3_FEATURE_LPIPS_THRESHOLD = 0.05 +COSMOS3_QUANTIZATION_IGNORE = [ + "language_model.*", + "vae2llm", + "llm2vae", + "time_embedder.*", +] +COSMOS3_SUPPORTED_FEATURES = frozenset({"fp8-blockwise", "nvfp4"}) + + +@dataclass(frozen=True) +class Cosmos3AccuracyCase: + id: str + golden_file: str + features: FeatureConfigState + lpips_threshold: float + + +# CUDA graph is not included yet: Cosmos3 reads a CUDA scalar with ``.item()`` +# inside the captured transformer forward, which CUDA rejects during capture. +COSMOS3_FEATURE_PROFILES = ( + ("fp8-blockwise", FeatureConfigState(quantization="FP8_BLOCK_SCALES")), + ("nvfp4", FeatureConfigState(quantization="NVFP4")), +) + + +def _build_cosmos3_accuracy_cases(): + cases = [] + for profile_id, features in COSMOS3_FEATURE_PROFILES: + _validate_single_feature_config( + features, + COSMOS3_SUPPORTED_FEATURES, + "Cosmos3", + ) + cases.append( + pytest.param( + Cosmos3AccuracyCase( + id=profile_id, + golden_file=(f"cosmos3_nano_{profile_id.replace('-', '_')}_lpips_golden.png"), + features=features, + lpips_threshold=COSMOS3_FEATURE_LPIPS_THRESHOLD, + ), + id=profile_id, + ) + ) + return cases + + +COSMOS3_ACCURACY_CASES = _build_cosmos3_accuracy_cases() + + +def _run_cosmos3_lpips_pipeline(num_frames): + """Run the Cosmos3-Nano pipeline (default setting, VANILLA attn, compile-off). + + Returns the generated video tensor ``(B, T, H, W, C)`` (T == ``num_frames``), + or ``None`` if generation produced no video. ``num_frames=1`` yields the + single-frame text-to-image path. + """ + # Cosmos3 re-reads the guardrail flag in __init__; set it before the pipeline loads. + guardrails_env_key = "TRTLLM_DISABLE_COSMOS3_GUARDRAILS" + previous_guardrails_env = os.environ.get(guardrails_env_key) + os.environ[guardrails_env_key] = "1" + try: + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + CompilationConfig, + TorchCompileConfig, + VisualGenArgs, + ) + + model_path = _lpips_model_path(COSMOS3_NANO_MODEL_SUBPATH) + _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + args = VisualGenArgs( + model=model_path, + compilation_config=CompilationConfig(skip_warmup=True), + torch_compile_config=TorchCompileConfig(enable=False), + attention_config=AttentionConfig(backend="VANILLA"), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + with torch.no_grad(): + result = pipeline.forward( + prompt=COSMOS3_LPIPS_PROMPT, + seed=COSMOS3_LPIPS_SEED, + height=COSMOS3_LPIPS_HEIGHT, + width=COSMOS3_LPIPS_WIDTH, + num_frames=num_frames, + num_inference_steps=COSMOS3_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=COSMOS3_LPIPS_GUIDANCE_SCALE, + frame_rate=COSMOS3_LPIPS_FRAME_RATE, + use_guardrails=False, + ) + if result is None or result.video is None: + return None + return result.video.detach().cpu() + finally: + del pipeline + _cleanup_cuda() + finally: + if previous_guardrails_env is None: + os.environ.pop(guardrails_env_key, None) + else: + os.environ[guardrails_env_key] = previous_guardrails_env + + +def _generate_cosmos3_lpips_video(output_path): + """Generate the Cosmos3-Nano text-to-video LPIPS sample.""" + video = _run_cosmos3_lpips_pipeline(COSMOS3_LPIPS_T2V_NUM_FRAMES) + assert video is not None, "Cosmos3-Nano T2V LPIPS run produced no video" + _save_lpips_video_mp4(video, output_path, frame_rate=COSMOS3_LPIPS_FRAME_RATE) + + +def _generate_cosmos3_lpips_image(output_path): + """Generate the Cosmos3-Nano text-to-image LPIPS sample (single frame).""" + from tensorrt_llm.media.encoding import save_image + + video = _run_cosmos3_lpips_pipeline(COSMOS3_LPIPS_T2I_NUM_FRAMES) + assert video is not None, "Cosmos3-Nano T2I LPIPS run produced no frame" + # video is (B, T, H, W, C); take the single frame -> (H, W, C) for save_image. + save_image(video[0, 0], output_path) + + +def _generate_cosmos3_feature_image(case, output_path): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.media.encoding import save_image + + guardrails_env_key = "TRTLLM_DISABLE_COSMOS3_GUARDRAILS" + previous_guardrails_env = os.environ.get(guardrails_env_key) + os.environ[guardrails_env_key] = "1" + pipeline = None + try: + model_path = _lpips_model_path(COSMOS3_NANO_MODEL_SUBPATH) + _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + with _lpips_deterministic_algorithms(), _fixed_nvfp4_quantization_backend(case.features): + args = _build_single_device_feature_args( + model_path, + case.features, + resolution=(COSMOS3_LPIPS_HEIGHT, COSMOS3_LPIPS_WIDTH), + num_frames=1, + quantization_kwargs={"ignore": COSMOS3_QUANTIZATION_IGNORE}, + ) + try: + pipeline = PipelineLoader(args).load(skip_warmup=False) + _assert_resolved_single_device_feature_config( + pipeline, + case.features, + resolution=(COSMOS3_LPIPS_HEIGHT, COSMOS3_LPIPS_WIDTH), + num_frames=1, + ) + _assert_feature_quantization_installed(pipeline, case.features) + result = pipeline.forward( + prompt=COSMOS3_LPIPS_PROMPT, + seed=COSMOS3_LPIPS_SEED, + height=COSMOS3_LPIPS_HEIGHT, + width=COSMOS3_LPIPS_WIDTH, + num_frames=1, + num_inference_steps=COSMOS3_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=COSMOS3_LPIPS_GUIDANCE_SCALE, + frame_rate=COSMOS3_LPIPS_FRAME_RATE, + use_guardrails=False, + ) + _assert_single_device_feature_executed(pipeline, case.features) + generated_image = result.video[0, 0].detach().cpu() + finally: + try: + if pipeline is not None: + _cleanup_single_device_feature_pipeline(pipeline) + del pipeline + finally: + _cleanup_cuda() + save_image(generated_image, output_path) + finally: + if previous_guardrails_env is None: + os.environ.pop(guardrails_env_key, None) + else: + os.environ[guardrails_env_key] = previous_guardrails_env + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("case", COSMOS3_ACCURACY_CASES) +def test_cosmos3_feature_accuracy_against_golden( + request, + tmp_path, + case, + _visual_gen_lpips_scorer, +): + generated_path = tmp_path / f"cosmos3_{case.id}_generated.png" + golden_path = _golden_media_path( + tmp_path, + case.golden_file, + f"Cosmos3-Nano {case.id} LPIPS golden image", + ) + _run_single_device_feature_generator( + case.features, _generate_cosmos3_feature_image, case, generated_path + ) + score = _run_reusable_image_lpips_eval( + f"cosmos3-{case.id}", + golden_path, + generated_path, + _visual_gen_lpips_scorer, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + case.lpips_threshold, + generated_path, + f"cosmos3_{case.id}_generated.png", + ) + _assert_lpips_below_threshold(score, case.lpips_threshold) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_nano_t2v_lpips_against_golden(_visual_gen_deps, tmp_path): + generated_path = tmp_path / "cosmos3_nano_t2v_generated.mp4" + golden_path = _golden_media_path( + tmp_path, + "cosmos3_nano_t2v_lpips_golden_video.mp4", + "Cosmos3-Nano T2V LPIPS golden video", + ) + _generate_cosmos3_lpips_video(generated_path) + score = _run_lpips_eval( + tmp_path, + "cosmos3_nano_t2v", + "video", + COSMOS3_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, COSMOS3_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_nano_t2i_lpips_against_golden(_visual_gen_deps, tmp_path): + generated_path = tmp_path / "cosmos3_nano_t2i_generated.png" + golden_path = _golden_media_path( + tmp_path, "cosmos3_nano_t2i_lpips_golden.png", "Cosmos3-Nano T2I LPIPS golden image" + ) + _generate_cosmos3_lpips_image(generated_path) + score = _run_lpips_eval( + tmp_path, + "cosmos3_nano_t2i", + "image", + COSMOS3_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, COSMOS3_LPIPS_THRESHOLD) + + +def test_cosmos3_example(_visual_gen_deps, llm_root, llm_venv): + """Run examples/visual_gen/models/cosmos3/cosmos3.py with FP8 config end-to-end. + + Validates that the Cosmos3-Nano example script and ``configs/cosmos3-nano-1gpu.yaml`` + work together as documented. Uses the local Cosmos3-Nano checkpoint and + the shared FP8 dynamic-quant config. + """ + model_path = _lpips_model_path("Cosmos3-Nano") + _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) + + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_example") + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "cosmos3_output.mp4") + + script_path = os.path.join( + llm_root, "examples", "visual_gen", "models", "cosmos3", "cosmos3.py" + ) + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "cosmos3-nano-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--prompt", + "A serene mountain landscape with snow-capped peaks and a flowing river", + "--output_path", + output_path, + ], + env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_flux.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_flux.py new file mode 100644 index 000000000000..65c57e9b3165 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_flux.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. +"""Single-device FLUX visual-quality and example regression tests.""" + +import os +from dataclasses import dataclass + +import pytest +import torch +from defs.common import venv_check_call +from defs.examples.visual_gen.visual_gen_test_utils import ( + FeatureConfigState, + _assert_feature_quantization_installed, + _assert_lpips_below_threshold, + _assert_resolved_single_device_feature_config, + _assert_single_device_feature_executed, + _build_single_device_feature_args, + _cleanup_cuda, + _cleanup_single_device_feature_pipeline, + _disable_inductor_compile_worker_quiesce, + _fixed_nvfp4_quantization_backend, + _golden_media_path, + _lpips_deterministic_algorithms, + _lpips_model_path, + _preserve_lpips_candidate_on_failure, + _run_lpips_eval, + _run_reusable_image_lpips_eval, + _run_single_device_feature_generator, + _skip_if_missing, + _validate_single_feature_config, +) + +FLUX_LPIPS_PROMPT = "a tiny astronaut hatching from an egg on the moon" +FLUX_LPIPS_HEIGHT = 256 +FLUX_LPIPS_WIDTH = 256 +FLUX_LPIPS_NUM_INFERENCE_STEPS = 4 +FLUX_LPIPS_GUIDANCE_SCALE = 3.5 +FLUX_LPIPS_SEED = 42 +FLUX_LPIPS_THRESHOLD = 0.05 +FLUX_FEATURE_LPIPS_THRESHOLD = 0.05 +FLUX_SUPPORTED_FEATURES = frozenset({"fp8-blockwise", "nvfp4", "cuda-graph"}) + + +@dataclass(frozen=True) +class FluxModelSpec: + id: str + checkpoint_subdir: str + + +@dataclass(frozen=True) +class FluxFeatureProfile: + id: str + features: FeatureConfigState + + +@dataclass(frozen=True) +class FluxAccuracyCase: + id: str + checkpoint_subdir: str + golden_file: str + features: FeatureConfigState + lpips_threshold: float + + +FLUX_MODEL_SPECS = ( + FluxModelSpec( + id="flux1", + checkpoint_subdir="FLUX.1-dev", + ), + FluxModelSpec( + id="flux2", + checkpoint_subdir="FLUX.2-dev", + ), +) + +FLUX_FEATURE_PROFILES = ( + FluxFeatureProfile( + id="fp8-blockwise", + features=FeatureConfigState(quantization="FP8_BLOCK_SCALES"), + ), + FluxFeatureProfile( + id="nvfp4", + features=FeatureConfigState(quantization="NVFP4"), + ), + FluxFeatureProfile( + id="cuda-graph", + features=FeatureConfigState(cuda_graph=True), + ), +) + + +def _build_flux_accuracy_cases(): + cases = [] + for model in FLUX_MODEL_SPECS: + for profile in FLUX_FEATURE_PROFILES: + _validate_single_feature_config( + profile.features, + FLUX_SUPPORTED_FEATURES, + "FLUX", + ) + case_id = f"{model.id}-{profile.id}" + cases.append( + pytest.param( + FluxAccuracyCase( + id=case_id, + checkpoint_subdir=model.checkpoint_subdir, + golden_file=(f"{model.id}_{profile.id.replace('-', '_')}_lpips_golden.png"), + features=profile.features, + lpips_threshold=FLUX_FEATURE_LPIPS_THRESHOLD, + ), + id=case_id, + ) + ) + return cases + + +FLUX_ACCURACY_CASES = _build_flux_accuracy_cases() + + +def _generate_flux_lpips_image(model_path, output_path): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.media.encoding import save_image + from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs + + _skip_if_missing(model_path, "FLUX checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + with _lpips_deterministic_algorithms(): + args = VisualGenArgs( + model=model_path, + torch_compile_config=TorchCompileConfig(enable=False), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + result = pipeline.forward( + prompt=FLUX_LPIPS_PROMPT, + height=FLUX_LPIPS_HEIGHT, + width=FLUX_LPIPS_WIDTH, + num_inference_steps=FLUX_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=FLUX_LPIPS_GUIDANCE_SCALE, + seed=FLUX_LPIPS_SEED, + ) + generated_image = result.image[0].detach().cpu() + finally: + del pipeline + _cleanup_cuda() + + save_image(generated_image, output_path) + + +def _generate_flux_image(case, output_path): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.media.encoding import save_image + + model_path = _lpips_model_path(case.checkpoint_subdir) + _skip_if_missing(model_path, f"{case.checkpoint_subdir} checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + pipeline = None + with _lpips_deterministic_algorithms(), _fixed_nvfp4_quantization_backend(case.features): + args = _build_single_device_feature_args( + model_path, + case.features, + resolution=(FLUX_LPIPS_HEIGHT, FLUX_LPIPS_WIDTH), + num_frames=1, + ) + try: + pipeline = PipelineLoader(args).load(skip_warmup=False) + _assert_resolved_single_device_feature_config( + pipeline, + case.features, + resolution=(FLUX_LPIPS_HEIGHT, FLUX_LPIPS_WIDTH), + num_frames=1, + ) + _assert_feature_quantization_installed(pipeline, case.features) + result = pipeline.forward( + prompt=FLUX_LPIPS_PROMPT, + height=FLUX_LPIPS_HEIGHT, + width=FLUX_LPIPS_WIDTH, + num_inference_steps=FLUX_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=FLUX_LPIPS_GUIDANCE_SCALE, + seed=FLUX_LPIPS_SEED, + ) + _assert_single_device_feature_executed(pipeline, case.features) + generated_image = result.image[0].detach().cpu() + finally: + try: + if pipeline is not None: + _cleanup_single_device_feature_pipeline(pipeline) + del pipeline + finally: + _cleanup_cuda() + + save_image(generated_image, output_path) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_flux1_lpips_against_golden(tmp_path): + generated_path = tmp_path / "flux1_generated.png" + golden_path = _golden_media_path( + tmp_path, "flux1_lpips_golden.png", "FLUX.1 LPIPS golden image" + ) + _generate_flux_lpips_image(_lpips_model_path("FLUX.1-dev"), generated_path) + score = _run_lpips_eval( + tmp_path, + "flux1", + "image", + FLUX_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, FLUX_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_flux2_lpips_against_golden(tmp_path): + generated_path = tmp_path / "flux2_generated.png" + golden_path = _golden_media_path( + tmp_path, "flux2_lpips_golden.png", "FLUX.2 LPIPS golden image" + ) + _generate_flux_lpips_image(_lpips_model_path("FLUX.2-dev"), generated_path) + score = _run_lpips_eval( + tmp_path, + "flux2", + "image", + FLUX_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, FLUX_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("case", FLUX_ACCURACY_CASES) +def test_flux_accuracy_against_golden(request, tmp_path, case, _visual_gen_lpips_scorer): + generated_path = tmp_path / f"{case.id}_generated.png" + reference_path = _golden_media_path( + tmp_path, + case.golden_file, + f"{case.checkpoint_subdir} LPIPS golden image", + ) + + _run_single_device_feature_generator(case.features, _generate_flux_image, case, generated_path) + score = _run_reusable_image_lpips_eval( + case.id, + reference_path, + generated_path, + _visual_gen_lpips_scorer, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + case.lpips_threshold, + generated_path, + f"{case.id}_generated.png", + ) + _assert_lpips_below_threshold(score, case.lpips_threshold) + + +def test_flux1_example(_visual_gen_deps, llm_root, llm_venv): + """Run the FLUX.1 example with the supported single-GPU NVFP4 config.""" + model_path = _lpips_model_path("FLUX.1-dev") + _skip_if_missing(model_path, "FLUX.1-dev checkpoint", is_dir=True) + + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "flux1_example") + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "flux1_output.png") + script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "flux1.py") + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "flux1-dev-fp4-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--output_path", + output_path, + ], + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + + +def test_flux2_example(_visual_gen_deps, llm_root, llm_venv): + """Run the FLUX.2 example with the supported single-GPU NVFP4 config.""" + model_path = _lpips_model_path("FLUX.2-dev") + _skip_if_missing(model_path, "FLUX.2-dev checkpoint", is_dir=True) + + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "flux2_example") + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "flux2_output.png") + script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "flux2.py") + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "flux2-dev-fp4-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--output_path", + output_path, + ], + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py new file mode 100644 index 000000000000..3ce77550021e --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py @@ -0,0 +1,620 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""Single-GPU integration and accuracy tests for LTX-2.""" + +import os +from dataclasses import dataclass + +import pytest +import torch +from defs import conftest +from defs.common import venv_check_call +from defs.examples.visual_gen.visual_gen_test_utils import ( + VISUAL_GEN_OUTPUT_VIDEO, + FeatureConfigState, + _assert_feature_quantization_installed, + _assert_lpips_below_threshold, + _assert_resolved_single_device_feature_config, + _assert_single_device_feature_executed, + _build_single_device_feature_args, + _cleanup_cuda, + _cleanup_single_device_feature_pipeline, + _disable_inductor_compile_worker_quiesce, + _fixed_nvfp4_quantization_backend, + _golden_media_path, + _lpips_deterministic_algorithms, + _lpips_model_path, + _preserve_lpips_candidate_on_failure, + _run_lpips_eval, + _run_reusable_video_lpips_eval, + _run_single_device_feature_generator, + _run_vbench_and_report, + _save_lpips_video_mp4, + _skip_if_missing, + _validate_single_feature_config, + _visual_gen_output_path, +) + +LTX2_LPIPS_NUM_FRAMES = 49 +LTX2_LPIPS_NUM_INFERENCE_STEPS = 8 +LTX2_LPIPS_THRESHOLD = 0.05 +LTX2_CUDA_GRAPH_LPIPS_THRESHOLD = 0.01 + +LTX2_FEATURE_LPIPS_THRESHOLD = 0.05 +LTX2_SUPPORTED_FEATURES = frozenset({"fp8-blockwise", "nvfp4", "cuda-graph"}) + + +@dataclass(frozen=True) +class LTX2AccuracyCase: + id: str + golden_file: str + features: FeatureConfigState + lpips_threshold: float + + +LTX2_FEATURE_PROFILES = ( + ("fp8-blockwise", FeatureConfigState(quantization="FP8_BLOCK_SCALES")), + ("nvfp4", FeatureConfigState(quantization="NVFP4")), + ("cuda-graph", FeatureConfigState(cuda_graph=True)), +) + + +def _build_ltx2_accuracy_cases(): + cases = [] + for profile_id, features in LTX2_FEATURE_PROFILES: + _validate_single_feature_config( + features, + LTX2_SUPPORTED_FEATURES, + "LTX-2", + ) + cases.append( + pytest.param( + LTX2AccuracyCase( + id=profile_id, + golden_file=(f"ltx2_{profile_id.replace('-', '_')}_lpips_golden_video.mp4"), + features=features, + lpips_threshold=LTX2_FEATURE_LPIPS_THRESHOLD, + ), + id=profile_id, + ) + ) + return cases + + +LTX2_ACCURACY_CASES = _build_ltx2_accuracy_cases() + + +# LTX-2 configuration +LTX2_MODEL_CHECKPOINT_PATH = "LTX-2/ltx-2-19b-dev.safetensors" +LTX2_TEXT_ENCODER_SUBPATH = "gemma-3-12b-it" +LTX2_T2V_PROMPT = ( + "A woman with long brown hair and light skin smiles at the camera while " + "standing in a sunlit park, her hair gently blowing in the breeze as she " + "tilts her head slightly to the side." +) +LTX2_T2V_HEIGHT = 512 +LTX2_T2V_WIDTH = 768 +LTX2_T2V_NUM_FRAMES = 121 +LTX2_T2V_STEPS = 40 +LTX2_T2V_GUIDANCE_SCALE = 4.0 +LTX2_T2V_MAX_SEQ_LEN = 1024 +LTX2_T2V_FRAME_RATE = 24.0 +LTX2_T2V_SEED = 42 +LTX2_T2V_NEGATIVE_PROMPT = "worst quality, inconsistent motion, blurry, jittery, distorted" + + +# LTX-2 Two-Stage configuration +LTX2_UPSAMPLER_SUBPATH = "LTX-2/ltx-2-spatial-upscaler-x2-1.0.safetensors" +LTX2_DISTILLED_LORA_SUBPATH = "LTX-2/ltx-2-19b-distilled-lora-384.safetensors" +LTX2_TWO_STAGE_HEIGHT = 1024 +LTX2_TWO_STAGE_WIDTH = 1536 +LTX2_TWO_STAGE_NUM_FRAMES = 121 +LTX2_TWO_STAGE_STEPS = 40 +LTX2_TWO_STAGE_GUIDANCE_SCALE = 4.0 + +# Golden VBench scores for two-stage pipeline variants, captured from the +# initial baseline runs. +VBENCH_LTX2_TWO_STAGE_BF16_GOLDEN_SCORES = { + "subject_consistency": 0.9877, + "background_consistency": 0.9601, + "motion_smoothness": 0.9952, + "dynamic_degree": 0.0, + "aesthetic_quality": 0.5839, + "imaging_quality": 0.5404, +} + +VBENCH_LTX2_TWO_STAGE_FP8_GOLDEN_SCORES = { + "subject_consistency": 0.9820, + "background_consistency": 0.9617, + "motion_smoothness": 0.9885, + "dynamic_degree": 1.0, + "aesthetic_quality": 0.6017, + "imaging_quality": 0.7136, +} + + +def _ltx2_lpips_text_encoder_path(): + scratch_space = conftest.llm_models_root() + candidates = [ + os.path.join(scratch_space, LTX2_TEXT_ENCODER_SUBPATH), + os.path.join(scratch_space, "gemma", LTX2_TEXT_ENCODER_SUBPATH), + ] + for candidate in candidates: + if os.path.isdir(candidate): + return candidate + return candidates[0] + + +def _generate_ltx2_feature_video(case, output_path): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + + checkpoint_path = _lpips_model_path("LTX-2", "ltx-2-19b-dev.safetensors") + text_encoder_path = _ltx2_lpips_text_encoder_path() + spatial_upsampler_path = _lpips_model_path("LTX-2", "ltx-2-spatial-upscaler-x2-1.0.safetensors") + distilled_lora_path = _lpips_model_path("LTX-2", "ltx-2-19b-distilled-lora-384.safetensors") + _skip_if_missing(checkpoint_path, "LTX-2 checkpoint") + _skip_if_missing(text_encoder_path, "LTX-2 text encoder", is_dir=True) + _skip_if_missing(spatial_upsampler_path, "LTX-2 spatial upsampler") + _skip_if_missing(distilled_lora_path, "LTX-2 distilled LoRA") + _disable_inductor_compile_worker_quiesce() + pipeline = None + with ( + _lpips_deterministic_algorithms(), + torch.compiler.set_stance("force_eager"), + _fixed_nvfp4_quantization_backend(case.features), + ): + args = _build_single_device_feature_args( + checkpoint_path, + case.features, + resolution=(LTX2_T2V_HEIGHT, LTX2_T2V_WIDTH), + num_frames=LTX2_LPIPS_NUM_FRAMES, + pipeline_config={ + "text_encoder_path": text_encoder_path, + "spatial_upsampler_path": spatial_upsampler_path, + "distilled_lora_path": distilled_lora_path, + }, + ) + try: + pipeline = PipelineLoader(args).load(skip_warmup=False) + _assert_resolved_single_device_feature_config( + pipeline, + case.features, + resolution=(LTX2_T2V_HEIGHT, LTX2_T2V_WIDTH), + num_frames=LTX2_LPIPS_NUM_FRAMES, + ) + _assert_feature_quantization_installed(pipeline, case.features) + result = pipeline.forward( + prompt=LTX2_T2V_PROMPT, + negative_prompt=LTX2_T2V_NEGATIVE_PROMPT, + height=LTX2_T2V_HEIGHT, + width=LTX2_T2V_WIDTH, + num_frames=LTX2_LPIPS_NUM_FRAMES, + num_inference_steps=LTX2_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=LTX2_T2V_GUIDANCE_SCALE, + seed=LTX2_T2V_SEED, + ) + assert result.video is not None, "LTX-2 feature run produced no video" + _assert_single_device_feature_executed(pipeline, case.features) + generated_video = result.video.detach().cpu() + finally: + try: + if pipeline is not None: + _cleanup_single_device_feature_pipeline(pipeline) + del pipeline + finally: + _cleanup_cuda() + + _save_lpips_video_mp4(generated_video, output_path, frame_rate=LTX2_T2V_FRAME_RATE) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("case", LTX2_ACCURACY_CASES) +def test_ltx2_feature_accuracy_against_golden( + request, + tmp_path, + case, + _visual_gen_deps, + _visual_gen_lpips_scorer, +): + generated_path = tmp_path / f"ltx2_{case.id}_generated.mp4" + golden_path = _golden_media_path( + tmp_path, + case.golden_file, + f"LTX-2 {case.id} LPIPS golden video", + ) + _run_single_device_feature_generator( + case.features, _generate_ltx2_feature_video, case, generated_path + ) + score = _run_reusable_video_lpips_eval( + f"ltx2-{case.id}", + golden_path, + generated_path, + _visual_gen_lpips_scorer, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + case.lpips_threshold, + generated_path, + f"ltx2_{case.id}_generated.mp4", + ) + _assert_lpips_below_threshold(score, case.lpips_threshold) + + +def _generate_ltx2_lpips_video(output_path, *, enable_cuda_graph=False): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.visual_gen.args import CudaGraphConfig, TorchCompileConfig, VisualGenArgs + + checkpoint_path = _lpips_model_path("LTX-2", "ltx-2-19b-dev.safetensors") + text_encoder_path = _ltx2_lpips_text_encoder_path() + spatial_upsampler_path = _lpips_model_path("LTX-2", "ltx-2-spatial-upscaler-x2-1.0.safetensors") + distilled_lora_path = _lpips_model_path("LTX-2", "ltx-2-19b-distilled-lora-384.safetensors") + _skip_if_missing(checkpoint_path, "LTX-2 checkpoint") + _skip_if_missing(text_encoder_path, "LTX-2 text encoder", is_dir=True) + _skip_if_missing(spatial_upsampler_path, "LTX-2 spatial upsampler") + _skip_if_missing(distilled_lora_path, "LTX-2 distilled LoRA") + _disable_inductor_compile_worker_quiesce() + + # TorchCompileConfig(enable=False) does not suppress nested @torch.compile decorators. + # Wrapped here (not in the fixture) so the golden fixture and both sides of + # test_ltx2_cuda_graph_lpips_matches_eager run the same eager numerics. + with _lpips_deterministic_algorithms(), torch.compiler.set_stance("force_eager"): + args = VisualGenArgs( + model=checkpoint_path, + pipeline_config={ + "text_encoder_path": text_encoder_path, + "spatial_upsampler_path": spatial_upsampler_path, + "distilled_lora_path": distilled_lora_path, + }, + torch_compile_config=TorchCompileConfig(enable=False), + cuda_graph_config=CudaGraphConfig(enable=enable_cuda_graph), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + with torch.no_grad(): + result = pipeline.forward( + prompt=LTX2_T2V_PROMPT, + negative_prompt=LTX2_T2V_NEGATIVE_PROMPT, + height=LTX2_T2V_HEIGHT, + width=LTX2_T2V_WIDTH, + num_frames=LTX2_LPIPS_NUM_FRAMES, + num_inference_steps=LTX2_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=LTX2_T2V_GUIDANCE_SCALE, + seed=LTX2_T2V_SEED, + ) + generated_video = result.video.detach().cpu() + finally: + del pipeline + _cleanup_cuda() + + _save_lpips_video_mp4(generated_video, output_path, frame_rate=LTX2_T2V_FRAME_RATE) + + +def _generate_ltx2_cuda_graph_trtllm_backend_video(output_path): + from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams + from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + CompilationConfig, + CudaGraphConfig, + ParallelConfig, + TorchCompileConfig, + ) + + scratch_space = conftest.llm_models_root() + checkpoint_path = os.path.join(scratch_space, LTX2_MODEL_CHECKPOINT_PATH) + text_encoder_path = _ltx2_lpips_text_encoder_path() + spatial_upsampler_path = os.path.join(scratch_space, LTX2_UPSAMPLER_SUBPATH) + distilled_lora_path = os.path.join(scratch_space, LTX2_DISTILLED_LORA_SUBPATH) + _skip_if_missing(checkpoint_path, "LTX-2 checkpoint") + _skip_if_missing(text_encoder_path, "LTX-2 text encoder", is_dir=True) + _skip_if_missing(spatial_upsampler_path, "LTX-2 spatial upsampler") + _skip_if_missing(distilled_lora_path, "LTX-2 distilled LoRA") + _disable_inductor_compile_worker_quiesce() + + visual_gen_args = VisualGenArgs( + model=checkpoint_path, + quant_config={"quant_algo": "NVFP4", "dynamic": True}, + attention_config=AttentionConfig(backend="TRTLLM"), + parallel_config=ParallelConfig( + cfg_size=1, + ulysses_size=1, + parallel_vae_size=1, + ), + compilation_config=CompilationConfig( + resolutions=[ + ( + LTX2_T2V_HEIGHT, + LTX2_T2V_WIDTH, + ) + ], + num_frames=[LTX2_LPIPS_NUM_FRAMES], + ), + cuda_graph_config=CudaGraphConfig(enable=True), + torch_compile_config=TorchCompileConfig( + enable=True, + enable_fullgraph=False, + enable_autotune=True, + ), + pipeline_config={ + "text_encoder_path": text_encoder_path, + "spatial_upsampler_path": spatial_upsampler_path, + "distilled_lora_path": distilled_lora_path, + }, + ) + + visual_gen = VisualGen(model=checkpoint_path, args=visual_gen_args) + try: + params = VisualGenParams( + height=LTX2_T2V_HEIGHT, + width=LTX2_T2V_WIDTH, + num_frames=LTX2_LPIPS_NUM_FRAMES, + num_inference_steps=LTX2_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=LTX2_T2V_GUIDANCE_SCALE, + max_sequence_length=LTX2_T2V_MAX_SEQ_LEN, + seed=LTX2_T2V_SEED, + frame_rate=LTX2_T2V_FRAME_RATE, + negative_prompt=LTX2_T2V_NEGATIVE_PROMPT, + ) + output = visual_gen.generate(inputs=LTX2_T2V_PROMPT, params=params) + _save_lpips_video_mp4(output.video, output_path, frame_rate=LTX2_T2V_FRAME_RATE) + finally: + visual_gen.shutdown() + + assert os.path.isfile(output_path), f"LTX-2 TRTLLM backend did not produce {output_path}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_ltx2_lpips_against_golden(request, tmp_path, ltx2_two_stage_bf16_video_path): + golden_path = _golden_media_path( + tmp_path, "ltx2_lpips_golden_video.mp4", "LTX-2 LPIPS golden video" + ) + score = _run_lpips_eval( + tmp_path, + "ltx2", + "video", + LTX2_T2V_PROMPT, + golden_path, + ltx2_two_stage_bf16_video_path, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + LTX2_LPIPS_THRESHOLD, + ltx2_two_stage_bf16_video_path, + "ltx2_lpips_golden_video.mp4", + ) + _assert_lpips_below_threshold(score, LTX2_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_ltx2_cuda_graph_lpips_matches_eager(_visual_gen_deps, tmp_path): + eager_path = tmp_path / "ltx2_eager_generated.mp4" + cuda_graph_path = tmp_path / "ltx2_cuda_graph_generated.mp4" + + _generate_ltx2_lpips_video(eager_path, enable_cuda_graph=False) + _generate_ltx2_lpips_video(cuda_graph_path, enable_cuda_graph=True) + score = _run_lpips_eval( + tmp_path, + "ltx2_cuda_graph", + "video", + LTX2_T2V_PROMPT, + eager_path, + cuda_graph_path, + ) + _assert_lpips_below_threshold(score, LTX2_CUDA_GRAPH_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_ltx2_cuda_graph_trtllm_backend(request, _visual_gen_deps, tmp_path): + generated_path = tmp_path / "ltx2_cuda_graph_trtllm_backend_generated.mp4" + golden_path = _golden_media_path( + tmp_path, "ltx2_lpips_golden_video.mp4", "LTX-2 LPIPS golden video" + ) + _generate_ltx2_cuda_graph_trtllm_backend_video(generated_path) + score = _run_lpips_eval( + tmp_path, + "ltx2_cuda_graph_trtllm_backend", + "video", + LTX2_T2V_PROMPT, + golden_path, + generated_path, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + LTX2_LPIPS_THRESHOLD, + generated_path, + "ltx2_cuda_graph_trtllm_backend_generated.mp4", + ) + _assert_lpips_below_threshold(score, LTX2_LPIPS_THRESHOLD) + + +def _linear_type_to_quant_config(linear_type): + """Map linear_type shortcut to quant_config dict for VisualGenArgs.""" + mapping = { + "trtllm-fp8-per-tensor": {"quant_algo": "FP8", "dynamic": True}, + "trtllm-fp8-blockwise": {"quant_algo": "FP8_BLOCK_SCALES", "dynamic": True}, + "trtllm-nvfp4": {"quant_algo": "NVFP4", "dynamic": True}, + } + return mapping.get(linear_type) + + +def _generate_ltx2_two_stage_video(llm_venv, output_subdir, linear_type="default"): + """Generate a two-stage LTX-2 video using the Python API. + + Requires the main checkpoint, text encoder, spatial upsampler, and + distilled LoRA. Returns the path to the generated .mp4, or calls + pytest.skip if any asset is missing. + """ + from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams + + scratch_space = conftest.llm_models_root() + model_path = os.path.join(scratch_space, LTX2_MODEL_CHECKPOINT_PATH) + text_encoder_path = _ltx2_lpips_text_encoder_path() + upsampler_path = os.path.join(scratch_space, LTX2_UPSAMPLER_SUBPATH) + lora_path = os.path.join(scratch_space, LTX2_DISTILLED_LORA_SUBPATH) + + for label, path, is_file in [ + ("LTX-2 checkpoint", model_path, True), + ("text encoder", text_encoder_path, False), + ("spatial upsampler", upsampler_path, True), + ("distilled LoRA", lora_path, True), + ]: + exists = os.path.isfile(path) if is_file else os.path.isdir(path) + if not exists: + pytest.skip(f"Two-stage {label} not found: {path}") + + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", output_subdir) + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, VISUAL_GEN_OUTPUT_VIDEO) + if os.path.isfile(output_path): + return output_path + + vg_kwargs = dict( + pipeline_config={ + "text_encoder_path": text_encoder_path, + "spatial_upsampler_path": upsampler_path, + "distilled_lora_path": lora_path, + }, + ) + quant_config = _linear_type_to_quant_config(linear_type) + if quant_config is not None: + vg_kwargs["quant_config"] = quant_config + if torch.cuda.device_count() >= 2: + vg_kwargs["parallel_config"] = {"cfg_size": 2} + + visual_gen_args = VisualGenArgs(**vg_kwargs) + visual_gen = VisualGen(model=model_path, args=visual_gen_args) + + try: + params = VisualGenParams( + height=LTX2_TWO_STAGE_HEIGHT, + width=LTX2_TWO_STAGE_WIDTH, + num_frames=LTX2_TWO_STAGE_NUM_FRAMES, + num_inference_steps=LTX2_TWO_STAGE_STEPS, + guidance_scale=LTX2_TWO_STAGE_GUIDANCE_SCALE, + max_sequence_length=LTX2_T2V_MAX_SEQ_LEN, + seed=LTX2_T2V_SEED, + frame_rate=LTX2_T2V_FRAME_RATE, + negative_prompt=LTX2_T2V_NEGATIVE_PROMPT, + ) + output = visual_gen.generate(inputs=LTX2_T2V_PROMPT, params=params) + assert output.error is None + assert output.video is not None + assert output.frame_rate == LTX2_T2V_FRAME_RATE + assert output.metrics is not None + assert output.metrics.generation > 0 + assert output.metrics.denoise > 0 + assert output.metrics.post_denoise >= 0 + _save_lpips_video_mp4(output.video, output_path, frame_rate=LTX2_T2V_FRAME_RATE) + finally: + visual_gen.shutdown() + + assert os.path.isfile(output_path), f"LTX-2 two-stage did not produce {output_path}" + return output_path + + +@pytest.fixture(scope="session") +def ltx2_two_stage_bf16_video_path(_visual_gen_deps, llm_venv): + """Generate LTX-2 two-stage BF16 video with the LPIPS config and return path.""" + output_path = _visual_gen_output_path(llm_venv, "ltx2_two_stage_bf16") + if os.path.isfile(output_path): + return output_path + _generate_ltx2_lpips_video(output_path) + return output_path + + +@pytest.fixture(scope="session") +def ltx2_two_stage_fp8_video_path(_visual_gen_deps, llm_venv): + """Generate LTX-2 two-stage FP8 T2V video and return path.""" + return _generate_ltx2_two_stage_video( + llm_venv, "ltx2_two_stage_fp8", linear_type="trtllm-fp8-per-tensor" + ) + + +def test_vbench_dimension_score_ltx2_two_stage_bf16( + vbench_repo_root, ltx2_two_stage_bf16_video_path, llm_venv +): + """VBench accuracy for LTX-2 two-stage BF16 T2V.""" + videos_dir = os.path.dirname(ltx2_two_stage_bf16_video_path) + assert os.path.isfile(ltx2_two_stage_bf16_video_path), "LTX-2 two-stage BF16 video must exist" + _run_vbench_and_report( + vbench_repo_root, + videos_dir, + VISUAL_GEN_OUTPUT_VIDEO, + llm_venv, + title="LTX-2 Two-Stage BF16", + golden_scores=VBENCH_LTX2_TWO_STAGE_BF16_GOLDEN_SCORES, + max_score_diff=0.05, + ) + + +def test_vbench_dimension_score_ltx2_two_stage_fp8( + vbench_repo_root, ltx2_two_stage_fp8_video_path, llm_venv +): + """VBench accuracy for LTX-2 two-stage FP8 T2V.""" + videos_dir = os.path.dirname(ltx2_two_stage_fp8_video_path) + assert os.path.isfile(ltx2_two_stage_fp8_video_path), "LTX-2 two-stage FP8 video must exist" + _run_vbench_and_report( + vbench_repo_root, + videos_dir, + VISUAL_GEN_OUTPUT_VIDEO, + llm_venv, + title="LTX-2 Two-Stage FP8", + golden_scores=VBENCH_LTX2_TWO_STAGE_FP8_GOLDEN_SCORES, + max_score_diff=0.05, + ) + + +def test_ltx2_example(_visual_gen_deps, llm_root, llm_venv): + """Run examples/visual_gen/models/ltx2.py with NVFP4 config end-to-end. + + Validates that the LTX-2 example script and ``configs/ltx2-t2v-fp4-1gpu.yaml`` + work together as documented. The Gemma3 text encoder is passed separately via + ``--text_encoder_path`` because the shared YAML intentionally omits it to keep + the config model-path-agnostic. + """ + model_path = _lpips_model_path("LTX-2", "ltx-2-19b-dev.safetensors") + _skip_if_missing(model_path, "LTX-2 checkpoint") + text_encoder_path = _ltx2_lpips_text_encoder_path() + _skip_if_missing(text_encoder_path, "LTX-2 text encoder (gemma-3-12b-it)", is_dir=True) + + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "ltx2_example") + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "ltx2_output.mp4") + + script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "ltx2.py") + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "ltx2-t2v-fp4-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--text_encoder_path", + text_encoder_path, + "--output_path", + output_path, + ], + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 1c2de13e3b06..75c4642dff5d 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -23,7 +23,7 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp -from defs.examples.visual_gen.test_visual_gen import ( +from defs.examples.visual_gen.visual_gen_test_utils import ( WAN22_LPIPS_FRAME_RATE, WAN22_LPIPS_GUIDANCE_SCALE, WAN22_LPIPS_HEIGHT, diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_qwen_image.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_qwen_image.py new file mode 100644 index 000000000000..dadd4a4a764b --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_qwen_image.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""Single-GPU integration and accuracy tests for Qwen-Image.""" + +import os +from dataclasses import dataclass + +import pytest +import torch +from defs import conftest +from defs.common import venv_check_call +from defs.examples.visual_gen.visual_gen_test_utils import ( + FeatureConfigState, + _assert_feature_quantization_installed, + _assert_lpips_below_threshold, + _assert_resolved_single_device_feature_config, + _assert_single_device_feature_executed, + _build_single_device_feature_args, + _cleanup_cuda, + _cleanup_single_device_feature_pipeline, + _disable_inductor_compile_worker_quiesce, + _fixed_nvfp4_quantization_backend, + _golden_media_path, + _lpips_deterministic_algorithms, + _lpips_model_path, + _preserve_lpips_candidate_on_failure, + _run_lpips_eval, + _run_reusable_image_lpips_eval, + _run_single_device_feature_generator, + _skip_if_missing, + _validate_single_feature_config, +) + +# QwenImage (text-to-image) — default-setting LPIPS golden. +# Params mirror the QwenImage 20B reference defaults (pipeline_qwen_image.py). +# NOTE: QwenImage's forward CFG knob is ``true_cfg_scale`` (not ``guidance_scale``), +# and real-CFG only engages when a negative prompt is supplied. +QWEN_IMAGE_MODEL_SUBPATH = "qwen-image" +QWENIMAGE_LPIPS_PROMPT = "a tiny astronaut hatching from an egg on the moon" +QWENIMAGE_LPIPS_NEGATIVE_PROMPT = "" +QWENIMAGE_LPIPS_HEIGHT = 1328 +QWENIMAGE_LPIPS_WIDTH = 1328 +QWENIMAGE_LPIPS_NUM_INFERENCE_STEPS = 50 +QWENIMAGE_LPIPS_TRUE_CFG_SCALE = 4.0 +QWENIMAGE_LPIPS_SEED = 42 +QWENIMAGE_LPIPS_THRESHOLD = 0.05 + +QWENIMAGE_FEATURE_LPIPS_THRESHOLD = 0.05 +QWENIMAGE_SUPPORTED_FEATURES = frozenset({"fp8-blockwise", "nvfp4", "cuda-graph"}) + + +@dataclass(frozen=True) +class QwenImageAccuracyCase: + id: str + golden_file: str + features: FeatureConfigState + lpips_threshold: float + + +QWENIMAGE_FEATURE_PROFILES = ( + ("fp8-blockwise", FeatureConfigState(quantization="FP8_BLOCK_SCALES")), + ("nvfp4", FeatureConfigState(quantization="NVFP4")), + ("cuda-graph", FeatureConfigState(cuda_graph=True)), +) + + +def _build_qwenimage_accuracy_cases(): + cases = [] + for profile_id, features in QWENIMAGE_FEATURE_PROFILES: + _validate_single_feature_config( + features, + QWENIMAGE_SUPPORTED_FEATURES, + "Qwen-Image", + ) + cases.append( + pytest.param( + QwenImageAccuracyCase( + id=profile_id, + golden_file=(f"qwenimage_{profile_id.replace('-', '_')}_lpips_golden.png"), + features=features, + lpips_threshold=QWENIMAGE_FEATURE_LPIPS_THRESHOLD, + ), + id=profile_id, + ) + ) + return cases + + +QWENIMAGE_ACCURACY_CASES = _build_qwenimage_accuracy_cases() + + +def _generate_qwenimage_lpips_image(model_path, output_path): + """Generate the QwenImage text-to-image LPIPS sample.""" + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.media.encoding import save_image + from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs + + _skip_if_missing(model_path, "QwenImage checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + args = VisualGenArgs( + model=model_path, + torch_compile_config=TorchCompileConfig(enable=False), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + with torch.no_grad(): + result = pipeline.forward( + prompt=QWENIMAGE_LPIPS_PROMPT, + negative_prompt=QWENIMAGE_LPIPS_NEGATIVE_PROMPT, + height=QWENIMAGE_LPIPS_HEIGHT, + width=QWENIMAGE_LPIPS_WIDTH, + num_inference_steps=QWENIMAGE_LPIPS_NUM_INFERENCE_STEPS, + true_cfg_scale=QWENIMAGE_LPIPS_TRUE_CFG_SCALE, + seed=QWENIMAGE_LPIPS_SEED, + ) + generated_image = result.image[0].detach().cpu() + finally: + del pipeline + _cleanup_cuda() + + save_image(generated_image, output_path) + + +def _generate_qwenimage_feature_image(case, output_path): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.media.encoding import save_image + + model_path = _lpips_model_path(QWEN_IMAGE_MODEL_SUBPATH) + _skip_if_missing(model_path, "QwenImage checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + pipeline = None + with _lpips_deterministic_algorithms(), _fixed_nvfp4_quantization_backend(case.features): + args = _build_single_device_feature_args( + model_path, + case.features, + resolution=(QWENIMAGE_LPIPS_HEIGHT, QWENIMAGE_LPIPS_WIDTH), + num_frames=1, + ) + try: + pipeline = PipelineLoader(args).load(skip_warmup=False) + _assert_resolved_single_device_feature_config( + pipeline, + case.features, + resolution=(QWENIMAGE_LPIPS_HEIGHT, QWENIMAGE_LPIPS_WIDTH), + num_frames=1, + ) + _assert_feature_quantization_installed(pipeline, case.features) + result = pipeline.forward( + prompt=QWENIMAGE_LPIPS_PROMPT, + negative_prompt=QWENIMAGE_LPIPS_NEGATIVE_PROMPT, + height=QWENIMAGE_LPIPS_HEIGHT, + width=QWENIMAGE_LPIPS_WIDTH, + num_inference_steps=QWENIMAGE_LPIPS_NUM_INFERENCE_STEPS, + true_cfg_scale=QWENIMAGE_LPIPS_TRUE_CFG_SCALE, + seed=QWENIMAGE_LPIPS_SEED, + ) + _assert_single_device_feature_executed(pipeline, case.features) + generated_image = result.image[0].detach().cpu() + finally: + try: + if pipeline is not None: + _cleanup_single_device_feature_pipeline(pipeline) + del pipeline + finally: + _cleanup_cuda() + + save_image(generated_image, output_path) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("case", QWENIMAGE_ACCURACY_CASES) +def test_qwenimage_feature_accuracy_against_golden( + request, + tmp_path, + case, + _visual_gen_lpips_scorer, +): + generated_path = tmp_path / f"qwenimage_{case.id}_generated.png" + golden_path = _golden_media_path( + tmp_path, + case.golden_file, + f"QwenImage {case.id} LPIPS golden image", + ) + _run_single_device_feature_generator( + case.features, _generate_qwenimage_feature_image, case, generated_path + ) + score = _run_reusable_image_lpips_eval( + f"qwenimage-{case.id}", + golden_path, + generated_path, + _visual_gen_lpips_scorer, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + case.lpips_threshold, + generated_path, + f"qwenimage_{case.id}_generated.png", + ) + _assert_lpips_below_threshold(score, case.lpips_threshold) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_qwenimage_lpips_against_golden(_visual_gen_deps, tmp_path): + generated_path = tmp_path / "qwenimage_generated.png" + golden_path = _golden_media_path( + tmp_path, "qwenimage_lpips_golden.png", "QwenImage LPIPS golden image" + ) + _generate_qwenimage_lpips_image(_lpips_model_path(QWEN_IMAGE_MODEL_SUBPATH), generated_path) + score = _run_lpips_eval( + tmp_path, + "qwenimage", + "image", + QWENIMAGE_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, QWENIMAGE_LPIPS_THRESHOLD) + + +def test_qwen_image_example(_visual_gen_deps, llm_root, llm_venv): + """Run examples/visual_gen/models/qwen_image.py with FP8 config end-to-end. + + Validates that the Qwen-Image example script and + ``configs/qwen-image-fp8-1gpu.yaml`` work together as documented. Uses the + local Qwen-Image checkpoint and the shared FP8 blockwise dynamic-quant config. + """ + scratch_space = conftest.llm_models_root() + model_path = os.path.join(scratch_space, QWEN_IMAGE_MODEL_SUBPATH) + _skip_if_missing(model_path, "Qwen-Image checkpoint", is_dir=True) + model_index_path = os.path.join(model_path, "model_index.json") + if not os.path.isfile(model_index_path): + pytest.skip( + f"Qwen-Image checkpoint is incomplete: {model_path} (missing {model_index_path})" + ) + + out_dir = os.path.join( + llm_venv.get_working_directory(), "visual_gen_output", "qwen_image_example" + ) + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "qwen_image_output.png") + + script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "qwen_image.py") + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "qwen-image-fp8-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--output_path", + output_path, + ], + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_wan.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_wan.py new file mode 100644 index 000000000000..82f812cab4d1 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_wan.py @@ -0,0 +1,665 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""Single-GPU integration and accuracy tests for Wan models and examples.""" + +import os +from dataclasses import dataclass + +import pytest +import torch +from defs import conftest +from defs.common import venv_check_call +from defs.examples.visual_gen.visual_gen_test_utils import ( + VISUAL_GEN_OUTPUT_VIDEO, + WAN21_LPIPS_GUIDANCE_SCALE, + WAN21_LPIPS_HEIGHT, + WAN21_LPIPS_NEGATIVE_PROMPT, + WAN21_LPIPS_NUM_FRAMES, + WAN21_LPIPS_NUM_INFERENCE_STEPS, + WAN21_LPIPS_PROMPT, + WAN21_LPIPS_SEED, + WAN21_LPIPS_WIDTH, + WAN22_LPIPS_FRAME_RATE, + WAN22_LPIPS_GUIDANCE_SCALE, + WAN22_LPIPS_HEIGHT, + WAN22_LPIPS_NEGATIVE_PROMPT, + WAN22_LPIPS_NUM_FRAMES, + WAN22_LPIPS_NUM_INFERENCE_STEPS, + WAN22_LPIPS_PROMPT, + WAN22_LPIPS_SEED, + WAN22_LPIPS_WIDTH, + WAN_LPIPS_FRAME_RATE, + WAN_LPIPS_THRESHOLD, + FeatureConfigState, + _assert_feature_quantization_installed, + _assert_lpips_below_threshold, + _assert_resolved_single_device_feature_config, + _assert_single_device_feature_executed, + _build_single_device_feature_args, + _cleanup_cuda, + _cleanup_single_device_feature_pipeline, + _disable_inductor_compile_worker_quiesce, + _fixed_nvfp4_quantization_backend, + _generate_wan_lpips_video, + _golden_media_path, + _lpips_deterministic_algorithms, + _lpips_model_path, + _preserve_lpips_candidate_on_failure, + _run_lpips_eval, + _run_reusable_video_lpips_eval, + _run_single_device_feature_generator, + _run_vbench_and_report, + _save_lpips_video_mp4, + _skip_if_missing, + _validate_single_feature_config, + _visual_gen_output_path, +) + +WAN_T2V_MODEL_SUBPATH = "Wan2.1-T2V-1.3B-Diffusers" +WAN22_T2V_MODEL_SUBPATH = "Wan2.2-T2V-A14B-Diffusers" +WAN22_A14B_FP8_MODEL_SUBPATH = "Wan2.2-T2V-A14B-Diffusers-FP8" +WAN22_A14B_NVFP4_MODEL_SUBPATH = "Wan2.2-T2V-A14B-Diffusers-NVFP4" +WAN22_I2V_A14B_NVFP4_MODEL_SUBPATH = "Wan2.2-I2V-A14B-Diffusers-NVFP4" +WAN_FEATURE_LPIPS_THRESHOLD = 0.05 +WAN_STANDARD_SUPPORTED_FEATURES = frozenset({"fp8-blockwise", "nvfp4", "cuda-graph"}) + + +@dataclass(frozen=True) +class WanModelSpec: + id: str + checkpoint_subdir: str + prompt: str + negative_prompt: str | None + height: int + width: int + num_frames: int + num_inference_steps: int + guidance_scale: float + seed: int + + +@dataclass(frozen=True) +class WanAccuracyCase: + id: str + model_id: str + checkpoint_subdir: str + golden_file: str + prompt: str + negative_prompt: str | None + height: int + width: int + num_frames: int + num_inference_steps: int + guidance_scale: float + seed: int + features: FeatureConfigState + lpips_threshold: float + + +WAN_MODEL_SPECS = ( + WanModelSpec( + id="wan21", + checkpoint_subdir=WAN_T2V_MODEL_SUBPATH, + prompt=WAN21_LPIPS_PROMPT, + negative_prompt=WAN21_LPIPS_NEGATIVE_PROMPT, + height=WAN21_LPIPS_HEIGHT, + width=WAN21_LPIPS_WIDTH, + num_frames=WAN21_LPIPS_NUM_FRAMES, + num_inference_steps=WAN21_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=WAN21_LPIPS_GUIDANCE_SCALE, + seed=WAN21_LPIPS_SEED, + ), + WanModelSpec( + id="wan22", + checkpoint_subdir=WAN22_T2V_MODEL_SUBPATH, + prompt=WAN22_LPIPS_PROMPT, + negative_prompt=WAN22_LPIPS_NEGATIVE_PROMPT, + height=WAN22_LPIPS_HEIGHT, + width=WAN22_LPIPS_WIDTH, + num_frames=WAN22_LPIPS_NUM_FRAMES, + num_inference_steps=WAN22_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=WAN22_LPIPS_GUIDANCE_SCALE, + seed=WAN22_LPIPS_SEED, + ), +) + +WAN_FEATURE_PROFILES = ( + ("fp8-blockwise", FeatureConfigState(quantization="FP8_BLOCK_SCALES")), + ("nvfp4", FeatureConfigState(quantization="NVFP4")), + ("cuda-graph", FeatureConfigState(cuda_graph=True)), +) + + +def _build_wan_accuracy_cases(): + cases = [] + for model in WAN_MODEL_SPECS: + for profile_id, features in WAN_FEATURE_PROFILES: + _validate_single_feature_config( + features, + WAN_STANDARD_SUPPORTED_FEATURES, + model.id, + ) + case_id = f"{model.id}-{profile_id}" + cases.append( + pytest.param( + WanAccuracyCase( + id=case_id, + model_id=model.id, + checkpoint_subdir=model.checkpoint_subdir, + golden_file=( + f"{model.id}_{profile_id.replace('-', '_')}_lpips_golden_video.mp4" + ), + prompt=model.prompt, + negative_prompt=model.negative_prompt, + height=model.height, + width=model.width, + num_frames=model.num_frames, + num_inference_steps=model.num_inference_steps, + guidance_scale=model.guidance_scale, + seed=model.seed, + features=features, + lpips_threshold=WAN_FEATURE_LPIPS_THRESHOLD, + ), + id=case_id, + ) + ) + + return cases + + +WAN_ACCURACY_CASES = _build_wan_accuracy_cases() + + +WAN_T2V_PROMPT = "A cute cat playing piano" +WAN_T2V_HEIGHT = 480 +WAN_T2V_WIDTH = 832 +WAN_T2V_NUM_FRAMES = 165 + + +# Golden VBench scores from HF reference video (WAN 2.1 1.3B); TRT-LLM is compared against these. +VBENCH_WAN_GOLDEN_SCORES = { + "subject_consistency": 0.8907, + "background_consistency": 0.9274, + "motion_smoothness": 0.9818, + "dynamic_degree": 1.0000, + "aesthetic_quality": 0.2928, + "imaging_quality": 0.3812, +} + + +# TODO: Reference scores from bf16 baseline runs +VBENCH_WAN22_BF16_GOLDEN_SCORES = { + "subject_consistency": 0.9103, + "background_consistency": 0.9516, + "motion_smoothness": 0.9693, + "dynamic_degree": 0.0000, + "aesthetic_quality": 0.6821, + "imaging_quality": 0.3993, +} +VBENCH_WAN22_A14B_FP8_GOLDEN_SCORES = { + "subject_consistency": 0.9173, + "background_consistency": 0.9717, + "motion_smoothness": 0.9865, + "dynamic_degree": 1.0000, + "aesthetic_quality": 0.5465, + "imaging_quality": 0.7142, +} +VBENCH_WAN22_A14B_NVFP4_GOLDEN_SCORES = { + "subject_consistency": 0.9173, + "background_consistency": 0.9717, + "motion_smoothness": 0.9865, + "dynamic_degree": 1.0000, + "aesthetic_quality": 0.5465, + "imaging_quality": 0.7142, +} + + +@pytest.fixture(scope="session") +def wan21_bf16_video_path(_visual_gen_deps, llm_venv): + output_path = _visual_gen_output_path(llm_venv, "wan21_bf16") + if os.path.isfile(output_path): + return output_path + # TorchCompileConfig(enable=False) does not suppress nested @torch.compile decorators. + with torch.compiler.set_stance("force_eager"): + _generate_wan_lpips_video( + _lpips_model_path(WAN_T2V_MODEL_SUBPATH), + output_path, + WAN21_LPIPS_PROMPT, + WAN21_LPIPS_NEGATIVE_PROMPT, + WAN21_LPIPS_HEIGHT, + WAN21_LPIPS_WIDTH, + WAN21_LPIPS_NUM_FRAMES, + WAN21_LPIPS_NUM_INFERENCE_STEPS, + WAN21_LPIPS_GUIDANCE_SCALE, + WAN21_LPIPS_SEED, + WAN_LPIPS_FRAME_RATE, + ) + return output_path + + +@pytest.fixture(scope="session") +def wan22_bf16_video_path(_visual_gen_deps, llm_venv): + output_path = _visual_gen_output_path(llm_venv, "wan22_bf16") + if os.path.isfile(output_path): + return output_path + # TorchCompileConfig(enable=False) does not suppress nested @torch.compile decorators. + with torch.compiler.set_stance("force_eager"): + _generate_wan_lpips_video( + _lpips_model_path(WAN22_T2V_MODEL_SUBPATH), + output_path, + WAN22_LPIPS_PROMPT, + WAN22_LPIPS_NEGATIVE_PROMPT, + WAN22_LPIPS_HEIGHT, + WAN22_LPIPS_WIDTH, + WAN22_LPIPS_NUM_FRAMES, + WAN22_LPIPS_NUM_INFERENCE_STEPS, + WAN22_LPIPS_GUIDANCE_SCALE, + WAN22_LPIPS_SEED, + WAN22_LPIPS_FRAME_RATE, + ) + return output_path + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_wan21_t2v_lpips_against_golden(request, tmp_path, wan21_bf16_video_path): + golden_path = _golden_media_path( + tmp_path, "wan21_t2v_lpips_golden_video.mp4", "Wan 2.1 LPIPS golden video" + ) + score = _run_lpips_eval( + tmp_path, + "wan21_t2v", + "video", + WAN21_LPIPS_PROMPT, + golden_path, + wan21_bf16_video_path, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + WAN_LPIPS_THRESHOLD, + wan21_bf16_video_path, + "wan21_t2v_lpips_golden_video.mp4", + ) + _assert_lpips_below_threshold(score, WAN_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_wan22_t2v_lpips_against_golden(request, tmp_path, wan22_bf16_video_path): + golden_path = _golden_media_path( + tmp_path, "wan22_t2v_lpips_golden_video.mp4", "Wan 2.2 LPIPS golden video" + ) + score = _run_lpips_eval( + tmp_path, + "wan22_t2v", + "video", + WAN22_LPIPS_PROMPT, + golden_path, + wan22_bf16_video_path, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + WAN_LPIPS_THRESHOLD, + wan22_bf16_video_path, + "wan22_t2v_lpips_golden_video.mp4", + ) + _assert_lpips_below_threshold(score, WAN_LPIPS_THRESHOLD) + + +def _generate_wan_feature_video(case, output_path): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + + model_path = _lpips_model_path(case.checkpoint_subdir) + _skip_if_missing(model_path, f"{case.checkpoint_subdir} checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + pipeline = None + with ( + _lpips_deterministic_algorithms(), + torch.compiler.set_stance("force_eager"), + _fixed_nvfp4_quantization_backend(case.features), + ): + args = _build_single_device_feature_args( + model_path, + case.features, + resolution=(case.height, case.width), + num_frames=case.num_frames, + ) + try: + pipeline = PipelineLoader(args).load(skip_warmup=False) + _assert_resolved_single_device_feature_config( + pipeline, + case.features, + resolution=(case.height, case.width), + num_frames=case.num_frames, + ) + _assert_feature_quantization_installed(pipeline, case.features) + result = pipeline.forward( + prompt=case.prompt, + negative_prompt=case.negative_prompt, + height=case.height, + width=case.width, + num_frames=case.num_frames, + num_inference_steps=case.num_inference_steps, + guidance_scale=case.guidance_scale, + seed=case.seed, + ) + assert result.video is not None, f"{case.id} produced no video" + _assert_single_device_feature_executed(pipeline, case.features) + generated_video = result.video.detach().cpu() + finally: + try: + if pipeline is not None: + _cleanup_single_device_feature_pipeline(pipeline) + del pipeline + finally: + _cleanup_cuda() + + _save_lpips_video_mp4(generated_video, output_path, frame_rate=WAN_LPIPS_FRAME_RATE) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("case", WAN_ACCURACY_CASES) +def test_wan_feature_accuracy_against_golden( + request, + tmp_path, + case, + _visual_gen_deps, + _visual_gen_lpips_scorer, +): + generated_path = tmp_path / f"{case.id}_generated.mp4" + golden_path = _golden_media_path( + tmp_path, + case.golden_file, + f"{case.id} LPIPS golden video", + ) + _run_single_device_feature_generator( + case.features, _generate_wan_feature_video, case, generated_path + ) + score = _run_reusable_video_lpips_eval( + case.id, + golden_path, + generated_path, + _visual_gen_lpips_scorer, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + case.lpips_threshold, + generated_path, + f"{case.id}_generated.mp4", + ) + _assert_lpips_below_threshold(score, case.lpips_threshold) + + +def _generate_wan_video(llm_venv, model_subpath, output_subdir): + """Generate a WAN video for a given model checkpoint. + + Returns the path to the generated .mp4, or calls pytest.skip if the model + is not found under LLM_MODELS_ROOT. + """ + from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams + + scratch_space = conftest.llm_models_root() + model_path = os.path.join(scratch_space, model_subpath) + if not os.path.isdir(model_path): + pytest.skip( + f"Model not found: {model_path} " + f"(set LLM_MODELS_ROOT or place {model_subpath} under scratch)" + ) + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", output_subdir) + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, VISUAL_GEN_OUTPUT_VIDEO) + if os.path.isfile(output_path): + return output_path + + visual_gen_args = VisualGenArgs( + attention_config={"backend": "VANILLA"}, + parallel_config={"cfg_size": 2 if torch.cuda.device_count() >= 2 else 1}, + torch_compile_config={"enable": False}, + compilation_config={"skip_warmup": True}, + ) + visual_gen = VisualGen(model=model_path, args=visual_gen_args) + try: + frame_rate = visual_gen.default_params.frame_rate + output = visual_gen.generate( + inputs=WAN_T2V_PROMPT, + params=VisualGenParams( + height=WAN_T2V_HEIGHT, + width=WAN_T2V_WIDTH, + seed=42, + num_frames=WAN_T2V_NUM_FRAMES, + frame_rate=frame_rate, + ), + ) + assert output.error is None, f"unexpected error on WAN run: {output.error}" + assert output.video is not None + _save_lpips_video_mp4(output.video, output_path, frame_rate=frame_rate) + finally: + visual_gen.shutdown() + + assert os.path.isfile(output_path), f"Visual gen did not produce {output_path}" + return output_path + + +@pytest.fixture(scope="session") +def wan22_a14b_fp8_video_path(_visual_gen_deps, llm_venv): + """Generate video with Wan 2.2 A14B FP8 checkpoint.""" + return _generate_wan_video(llm_venv, WAN22_A14B_FP8_MODEL_SUBPATH, "wan22_fp8") + + +@pytest.fixture(scope="session") +def wan22_a14b_nvfp4_video_path(_visual_gen_deps, llm_venv): + """Generate video with Wan 2.2 A14B NVFP4 checkpoint.""" + return _generate_wan_video(llm_venv, WAN22_A14B_NVFP4_MODEL_SUBPATH, "wan22_nvfp4") + + +def test_vbench_dimension_score_wan(vbench_repo_root, wan21_bf16_video_path, llm_venv): + """Run VBench on WAN 2.1 BF16 video generated with the LPIPS config.""" + videos_dir = os.path.dirname(wan21_bf16_video_path) + assert os.path.isfile(wan21_bf16_video_path), "WAN 2.1 BF16 video must exist" + _run_vbench_and_report( + vbench_repo_root, + videos_dir, + VISUAL_GEN_OUTPUT_VIDEO, + llm_venv, + title="WAN 2.1 BF16", + golden_scores=VBENCH_WAN_GOLDEN_SCORES, + max_score_diff=0.05, + ) + + +def test_vbench_dimension_score_wan22_bf16(vbench_repo_root, wan22_bf16_video_path, llm_venv): + """VBench accuracy for Wan 2.2 A14B BF16 generated with the LPIPS config.""" + videos_dir = os.path.dirname(wan22_bf16_video_path) + assert os.path.isfile(wan22_bf16_video_path), "WAN 2.2 BF16 video must exist" + _run_vbench_and_report( + vbench_repo_root, + videos_dir, + VISUAL_GEN_OUTPUT_VIDEO, + llm_venv, + title="WAN 2.2 A14B BF16", + golden_scores=VBENCH_WAN22_BF16_GOLDEN_SCORES, + max_score_diff=0.05, + ) + + +def test_vbench_dimension_score_wan22_a14b_fp8( + vbench_repo_root, wan22_a14b_fp8_video_path, llm_venv +): + """VBench accuracy for Wan 2.2 A14B FP8 — full generate + evaluate.""" + videos_dir = os.path.dirname(wan22_a14b_fp8_video_path) + assert os.path.isfile(wan22_a14b_fp8_video_path), "FP8 video must exist" + _run_vbench_and_report( + vbench_repo_root, + videos_dir, + VISUAL_GEN_OUTPUT_VIDEO, + llm_venv, + title="WAN 2.2 A14B FP8", + golden_scores=VBENCH_WAN22_A14B_FP8_GOLDEN_SCORES, + max_score_diff=0.06, + ) + + +def test_vbench_dimension_score_wan22_a14b_nvfp4( + vbench_repo_root, wan22_a14b_nvfp4_video_path, llm_venv +): + """VBench accuracy for Wan 2.2 A14B NVFP4 — full generate + evaluate.""" + videos_dir = os.path.dirname(wan22_a14b_nvfp4_video_path) + assert os.path.isfile(wan22_a14b_nvfp4_video_path), "NVFP4 video must exist" + _run_vbench_and_report( + vbench_repo_root, + videos_dir, + VISUAL_GEN_OUTPUT_VIDEO, + llm_venv, + title="WAN 2.2 A14B NVFP4", + golden_scores=VBENCH_WAN22_A14B_NVFP4_GOLDEN_SCORES, + max_score_diff=0.05, + ) + + +def test_visual_gen_quickstart(_visual_gen_deps, llm_root, llm_venv): + """Run examples/visual_gen/quickstart_example.py end-to-end.""" + scratch_space = conftest.llm_models_root() + model_src = os.path.join(scratch_space, WAN_T2V_MODEL_SUBPATH) + if not os.path.isdir(model_src): + pytest.skip( + f"Model not found: {model_src} " + f"(set LLM_MODELS_ROOT or place {WAN_T2V_MODEL_SUBPATH} under scratch)" + ) + + model_dst = os.path.join(llm_venv.get_working_directory(), "Wan-AI", WAN_T2V_MODEL_SUBPATH) + if not os.path.islink(model_dst): + os.makedirs(os.path.dirname(model_dst), exist_ok=True) + os.symlink(model_src, model_dst, target_is_directory=True) + + script_path = os.path.join(llm_root, "examples", "visual_gen", "quickstart_example.py") + venv_check_call(llm_venv, [script_path]) + + output_path = os.path.join(llm_venv.get_working_directory(), "output.avi") + assert os.path.isfile(output_path), f"Quickstart did not produce output.avi at {output_path}" + + +def test_visual_gen_api_walkthrough(_visual_gen_deps, llm_root, llm_venv): + """Run examples/visual_gen/api_walkthrough.py end-to-end.""" + scratch_space = conftest.llm_models_root() + model_src = os.path.join(scratch_space, WAN_T2V_MODEL_SUBPATH) + if not os.path.isdir(model_src): + pytest.skip( + f"Model not found: {model_src} " + f"(set LLM_MODELS_ROOT or place {WAN_T2V_MODEL_SUBPATH} under scratch)" + ) + + model_dst = os.path.join(llm_venv.get_working_directory(), "Wan-AI", WAN_T2V_MODEL_SUBPATH) + if not os.path.islink(model_dst): + os.makedirs(os.path.dirname(model_dst), exist_ok=True) + os.symlink(model_src, model_dst, target_is_directory=True) + + script_path = os.path.join(llm_root, "examples", "visual_gen", "api_walkthrough.py") + venv_check_call(llm_venv, [script_path]) + + output_path = os.path.join(llm_venv.get_working_directory(), "api_walkthrough_output.avi") + assert os.path.isfile(output_path), f"API walkthrough did not produce {output_path}" + + +# ============================================================================= +# Core example tests — run per-model scripts from examples/visual_gen/models/ +# with shared YAML configs from examples/visual_gen/configs/. +# ============================================================================= + + +def test_wan_t2v_example(_visual_gen_deps, llm_root, llm_venv): + """Run examples/visual_gen/models/wan_t2v.py with NVFP4 config end-to-end. + + This is a core example test: it validates that the per-model example script + and the shared YAML config work together as documented in the README. + Uses the pre-quantized Wan 2.2 T2V A14B NVFP4 checkpoint and the shared + ``configs/wan2.2-t2v-fp4-1gpu.yaml`` (NVFP4 dynamic quant). The closest + overlapping test is ``test_vbench_dimension_score_wan22_a14b_nvfp4``, + which runs the same script but with a no-quant YAML synthesized at + runtime and additionally evaluates VBench scores. + """ + scratch_space = conftest.llm_models_root() + model_path = os.path.join(scratch_space, WAN22_A14B_NVFP4_MODEL_SUBPATH) + assert os.path.isdir(model_path), ( + f"Model not found: {model_path} " + f"(set LLM_MODELS_ROOT or place {WAN22_A14B_NVFP4_MODEL_SUBPATH} under models root)" + ) + + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "wan_t2v_example") + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "wan_t2v_output.mp4") + + script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "wan_t2v.py") + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "wan2.2-t2v-fp4-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--output_path", + output_path, + ], + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + + +def test_wan_i2v_example(_visual_gen_deps, llm_root, llm_venv): + """Run examples/visual_gen/models/wan_i2v.py with NVFP4 config end-to-end. + + Validates that the Wan I2V example script and ``configs/wan2.2-i2v-fp4-1gpu.yaml`` + work together as documented. Uses the pre-quantized Wan 2.2 I2V A14B NVFP4 + checkpoint and the default input image (cat_piano.png) bundled with the examples. + """ + scratch_space = conftest.llm_models_root() + model_path = os.path.join(scratch_space, WAN22_I2V_A14B_NVFP4_MODEL_SUBPATH) + if not os.path.isdir(model_path): + pytest.skip( + f"Model not found: {model_path} " + f"(set LLM_MODELS_ROOT or place {WAN22_I2V_A14B_NVFP4_MODEL_SUBPATH} under models root)" + ) + + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "wan_i2v_example") + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "wan_i2v_output.mp4") + + script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "wan_i2v.py") + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "wan2.2-i2v-fp4-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--output_path", + output_path, + ], + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" diff --git a/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py b/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py new file mode 100644 index 000000000000..e6721a5316c8 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py @@ -0,0 +1,989 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. +"""Shared helpers for VisualGen accuracy tests.""" + +import base64 +import contextlib +import gc +import glob +import json +import multiprocessing +import os +import random +import shutil +import subprocess +import sys +import time +import traceback +import urllib.request +import zipfile +from dataclasses import dataclass +from typing import Collection, Literal + +import pytest +import torch +import torch._inductor.config as inductor_config +from defs.trt_test_alternative import check_call +from torch._inductor.async_compile import shutdown_compile_workers + +# ============================================================================= +# Shared paths and feature configuration +# ============================================================================= + +# This file lives at tests/integration/defs/examples/visual_gen/, so the repo +# root is five levels up. +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..")) +VISUAL_GEN_LPIPS_EVAL_SCRIPT = os.path.join( + REPO_ROOT, "scripts", "visualgen_eval", "visual_gen_lpips_score_eval.py" +) +VISUAL_GEN_LPIPS_GOLDEN_DIR = os.path.join(os.path.dirname(__file__), "golden", "visual_gen_lpips") +VISUAL_GEN_LPIPS_GOLDEN_MEDIA_ZIP = os.path.join( + VISUAL_GEN_LPIPS_GOLDEN_DIR, "visual_gen_lpips_golden_media.zip" +) +VISUAL_GEN_OUTPUT_VIDEO = "trtllm_output.mp4" + +QuantizationMode = Literal["none", "FP8_BLOCK_SCALES", "NVFP4"] + + +@dataclass(frozen=True) +class FeatureConfigState: + """Normalized single-device feature state requested by an accuracy case. + + ``quantization="none"`` means the unquantized BF16 VisualGen path. Model + test modules remain responsible for validating which combinations they + support and for mapping this state to ``VisualGenArgs``. ``torch.compile`` + is intentionally omitted because all feature tests disable it to match the + existing upstream LPIPS baseline configuration. + """ + + quantization: QuantizationMode = "none" + cuda_graph: bool = False + + +def _enabled_feature_ids(features): + """Return the normalized IDs for features enabled in ``features``.""" + enabled = [] + if features.quantization != "none": + enabled.append( + { + "FP8_BLOCK_SCALES": "fp8-blockwise", + "NVFP4": "nvfp4", + }[features.quantization] + ) + if features.cuda_graph: + enabled.append("cuda-graph") + return tuple(enabled) + + +def _validate_single_feature_config(features, supported_features: Collection[str], model_name): + """Validate an exactly-one-feature accuracy configuration.""" + enabled = _enabled_feature_ids(features) + if len(enabled) != 1: + raise ValueError(f"{model_name} accuracy cases must enable exactly one feature: {enabled}") + unsupported = set(enabled).difference(supported_features) + if unsupported: + raise ValueError(f"{model_name} does not support feature(s): {sorted(unsupported)}") + return len(enabled) + + +def _build_single_device_feature_args( + model_path, + features, + *, + resolution, + num_frames, + pipeline_config=None, + quantization_kwargs=None, +): + """Map normalized feature state to a one-GPU ``VisualGenArgs`` object.""" + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + CompilationConfig, + CudaGraphConfig, + ParallelConfig, + TorchCompileConfig, + VisualGenArgs, + ) + + if features.quantization == "none": + quant_config = QuantConfig() + else: + quant_config = { + "quant_algo": features.quantization, + "dynamic": True, + **(quantization_kwargs or {}), + } + kwargs = dict( + model=model_path, + quant_config=quant_config, + compilation_config=CompilationConfig( + resolutions=[resolution], + num_frames=[num_frames], + ), + attention_config=AttentionConfig(backend="VANILLA"), + parallel_config=ParallelConfig( + parallel_vae_size=1, + cfg_size=1, + ulysses_size=1, + async_ulysses=False, + ring_size=1, + attn2d_size=(1, 1), + tp_size=1, + ), + cuda_graph_config=CudaGraphConfig(enable=features.cuda_graph), + torch_compile_config=TorchCompileConfig(enable=False), + ) + if pipeline_config is not None: + kwargs["pipeline_config"] = pipeline_config + return VisualGenArgs(**kwargs) + + +def _assert_resolved_single_device_feature_config( + pipeline, + features, + *, + resolution, + num_frames, +): + """Assert that every normalized feature reached the resolved pipeline config.""" + from tensorrt_llm.quantization.mode import QuantAlgo + + config = pipeline.pipeline_config + expected_quant_algo = { + "none": None, + "FP8_BLOCK_SCALES": QuantAlgo.FP8_BLOCK_SCALES, + "NVFP4": QuantAlgo.NVFP4, + }[features.quantization] + assert config.quant_config.quant_algo == expected_quant_algo + assert config.dynamic_weight_quant is (expected_quant_algo is not None) + if features.quantization == "NVFP4": + assert config.force_dynamic_quantization is True + + assert config.cache_backend is None + assert pipeline.cache_accelerator is None + assert config.attention.backend == "VANILLA" + assert config.attention.sparse_attention_config is None + + assert config.cuda_graph.enable is features.cuda_graph + assert config.torch_compile.enable is False + assert config.parallel.total_parallel_size == 1 + assert resolution in config.compilation.resolutions + assert num_frames in config.compilation.num_frames + + +def _transformer_components(pipeline): + components = [] + for component_name in pipeline.transformer_components: + component = getattr(pipeline, component_name, None) + if component is not None: + components.append(component) + assert components, "No VisualGen transformer components were found" + return components + + +def _assert_feature_quantization_installed(pipeline, features): + """Verify that requested dynamic quantization changed real Linear modules.""" + from tensorrt_llm._torch.modules.linear import Linear + + linears = [ + module + for component in _transformer_components(pipeline) + for module in component.modules() + if isinstance(module, Linear) and getattr(module, "_weights_created", False) + ] + assert linears, "No initialized VisualGen Linear modules were found" + if features.quantization == "none": + assert not [module for module in linears if module.has_any_quant] + return + + predicate = { + "FP8_BLOCK_SCALES": lambda module: module.has_fp8_block_scales, + "NVFP4": lambda module: module.has_nvfp4, + }[features.quantization] + quantized_linears = [module for module in linears if predicate(module)] + assert quantized_linears, f"No {features.quantization} VisualGen Linear modules were created" + sample = quantized_linears[0] + assert getattr(sample, "weight", None) is not None + assert getattr(sample, "weight_scale", None) is not None + if features.quantization == "FP8_BLOCK_SCALES": + assert sample.weight.dtype == torch.float8_e4m3fn + assert sample.weight_scale.ndim == 2 + else: + assert sample.weight.shape[-1] * 2 == sample.in_features + assert sample.scaling_vector_size == 16 + assert getattr(sample, "weight_scale_2", None) is not None + + +def _assert_single_device_feature_executed(pipeline, features): + """Assert that runtime-only features did real work during generation.""" + runners = pipeline._cuda_graph_runners + if features.cuda_graph: + assert runners, "CUDA graph was requested but no runner was installed" + for name, runner in runners.items(): + assert runner.enabled, f"CUDA graph runner {name} is disabled" + assert runner.graphs, f"CUDA graph runner {name} captured no graphs" + else: + assert not runners, f"CUDA graph runners exist with cuda_graph=False: {list(runners)}" + + +# ============================================================================= +# LPIPS, golden-media, and runtime helpers +# ============================================================================= + + +class ReusableLPIPSScorer: + """Lazily load one LPIPS model and reuse it across media test cases.""" + + def __init__(self, net="alex", device=None): + self.net = net + self.device = str(device or ("cuda" if torch.cuda.is_available() else "cpu")) + self._model = None + + def _get_model(self): + # Keep model construction aligned with _make_lpips_model() in + # visual_gen_lpips_score_eval.py. + if self._model is None: + if self.device.startswith("cuda") and not torch.cuda.is_available(): + raise RuntimeError(f"LPIPS device is {self.device}, but CUDA is not available.") + import lpips + + self._model = lpips.LPIPS(net=self.net, verbose=False).to(self.device).eval() + return self._model + + def _to_lpips_tensor(self, image): + # Keep this preprocessing aligned with _to_lpips_tensor() and + # _score_images() in visual_gen_lpips_score_eval.py. + import numpy as np + + tensor = torch.from_numpy(np.array(image.convert("RGB"))) + tensor = tensor.to(device=self.device, dtype=torch.float32) + if tensor.max() > 2.0: + tensor = tensor / 255.0 + tensor = tensor.permute(2, 0, 1).unsqueeze(0) + return tensor * 2.0 - 1.0 + + def score(self, reference_path, generated_path): + from PIL import Image + + with Image.open(reference_path) as reference_image: + reference_tensor = self._to_lpips_tensor(reference_image.convert("RGB")) + with Image.open(generated_path) as generated_image: + generated_tensor = self._to_lpips_tensor(generated_image.convert("RGB")) + + if generated_tensor.shape != reference_tensor.shape: + raise ValueError( + "Generated image and reference image must have the same LPIPS tensor shape: " + f"{tuple(generated_tensor.shape)} vs {tuple(reference_tensor.shape)}." + ) + with torch.no_grad(): + return float( + self._get_model()(generated_tensor, reference_tensor).reshape(-1).mean().item() + ) + + def _decode_video_to_lpips_batch(self, video_path): + import cv2 + import numpy as np + + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + raise RuntimeError(f"Failed to open video for LPIPS comparison: {video_path}") + frames = [] + try: + while True: + ok, frame_bgr = cap.read() + if not ok: + break + frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)) + finally: + cap.release() + if not frames: + raise ValueError(f"Decoded video contains no frames: {video_path}") + batch = torch.from_numpy(np.stack(frames, axis=0).copy()) + batch = batch.permute(0, 3, 1, 2).to(device=self.device, dtype=torch.float32) + return batch / 127.5 - 1.0 + + def score_video(self, reference_path, generated_path): + reference = self._decode_video_to_lpips_batch(reference_path) + generated = self._decode_video_to_lpips_batch(generated_path) + if generated.shape[1:] != reference.shape[1:]: + raise ValueError( + "Generated and reference video frames must have the same LPIPS tensor shape: " + f"{tuple(generated.shape[1:])} vs {tuple(reference.shape[1:])}." + ) + paired_frame_count = min(generated.shape[0], reference.shape[0]) + with torch.no_grad(): + scores = self._get_model()( + generated[:paired_frame_count], + reference[:paired_frame_count], + ).flatten() + return float(scores.mean().item()) + + def close(self): + model = self._model + self._model = None + del model + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _llm_models_root(): + # Keep imports that load TRT-LLM bindings out of module initialization. + # torch.multiprocessing.spawn workers import this utility before the + # installed-wheel path is placed ahead of the source checkout. + from defs import conftest + + return conftest.llm_models_root() + + +def _venv_check_call(*args, **kwargs): + # Deferred for the same spawn-worker reason as _llm_models_root(). + from defs.common import venv_check_call + + return venv_check_call(*args, **kwargs) + + +def _lpips_model_path(*parts): + return os.path.join(_llm_models_root(), *parts) + + +def _skip_if_missing(path, label, is_dir=False): + exists = os.path.isdir(path) if is_dir else os.path.exists(path) + if not exists: + pytest.skip(f"{label} not found: {path}") + + +def _extract_visual_gen_lpips_golden_media(tmp_path): + _skip_if_missing(VISUAL_GEN_LPIPS_GOLDEN_MEDIA_ZIP, "VisualGen LPIPS golden media zip") + extract_dir = tmp_path / "visual_gen_lpips_golden_media" + if extract_dir.exists(): + return extract_dir + + with zipfile.ZipFile(VISUAL_GEN_LPIPS_GOLDEN_MEDIA_ZIP) as archive: + for member in archive.namelist(): + if os.path.isabs(member) or ".." in member.split("/"): + raise ValueError(f"Unsafe golden media zip member: {member}") + archive.extractall(extract_dir) + return extract_dir + + +def _golden_media_path(tmp_path, media_name, label): + path = _extract_visual_gen_lpips_golden_media(tmp_path) / media_name + _skip_if_missing(path, label) + return path + + +def _disable_inductor_compile_worker_quiesce(): + # The quiesce timer thread can outlive shutdown_compile_workers() long + # enough for pytest-threadleak to report it as a leaked thread. + if hasattr(inductor_config, "quiesce_async_compile_pool"): + inductor_config.quiesce_async_compile_pool = False + + +def _cleanup_single_device_feature_pipeline(pipeline): + """Release feature wrappers before dropping a single-device pipeline.""" + if pipeline is None: + return + try: + cache_accelerator = getattr(pipeline, "cache_accelerator", None) + if cache_accelerator is not None: + cache_accelerator.unwrap() + finally: + pipeline.cleanup() + + +def _cleanup_cuda(): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + # torch.compile and unconditional @torch.compile helpers can lazily spawn + # an Inductor worker pool whose daemon threads otherwise outlive the test. + shutdown_compile_workers() + + +@contextlib.contextmanager +def _lpips_deterministic_algorithms(*, fully_eager=False): + previous_deterministic = torch.are_deterministic_algorithms_enabled() + previous_warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + previous_cublas_workspace_config = os.environ.get("CUBLAS_WORKSPACE_CONFIG") + + try: + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.use_deterministic_algorithms(True) + compiler_context = ( + torch.compiler.set_stance("force_eager") if fully_eager else contextlib.nullcontext() + ) + with compiler_context: + yield + finally: + torch.use_deterministic_algorithms( + previous_deterministic, + warn_only=previous_warn_only, + ) + if previous_cublas_workspace_config is None: + os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None) + else: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = previous_cublas_workspace_config + + +@contextlib.contextmanager +def _fixed_nvfp4_quantization_backend(features): + """Pin NVFP4 golden runs to deterministic TRT-LLM and CUTLASS paths. + + VisualGen normally autotunes between the TRT-LLM and FlashInfer activation + quantizers and among several GEMM backends. Those performance choices can + vary between processes, while fixed accuracy goldens require a stable + numerical path. Backend selection and parity are covered by the NVFP4 + operator unit tests. + """ + if features.quantization != "NVFP4": + yield + return + + from tensorrt_llm._torch.custom_ops import torch_custom_ops + from tensorrt_llm._torch.utils import get_model_extra_attrs, model_extra_attrs + + previous_flashinfer_available = torch_custom_ops.IS_FLASHINFER_AVAILABLE + extra_attrs = dict(get_model_extra_attrs() or {}) + extra_attrs["nvfp4_gemm_allowed_backends"] = ["cutlass"] + try: + torch_custom_ops.IS_FLASHINFER_AVAILABLE = False + with model_extra_attrs(extra_attrs): + yield + finally: + torch_custom_ops.IS_FLASHINFER_AVAILABLE = previous_flashinfer_available + + +def _feature_generator_process_entry(connection, generator, args): + try: + generator(*args) + except pytest.skip.Exception as error: + connection.send(("skip", str(error))) + except BaseException: + connection.send(("error", traceback.format_exc())) + else: + connection.send(("ok", "")) + finally: + connection.close() + + +def _run_single_device_feature_generator(features, generator, *args): + """Run NVFP4 generation with process-local compiler and backend state. + + Dynamo caches compiled code by Python code object, but the cache guards do + not cover every NVFP4 backend choice. A pytest worker that runs multiple + feature profiles can therefore reuse code compiled for an earlier profile. + A spawned process gives NVFP4 the same clean state used to create its fixed + golden without leaking that state into the following profile. + """ + if features.quantization != "NVFP4": + generator(*args) + return + + context = multiprocessing.get_context("spawn") + parent_connection, child_connection = context.Pipe(duplex=False) + process = context.Process( + target=_feature_generator_process_entry, + args=(child_connection, generator, args), + ) + process.start() + child_connection.close() + + try: + status, detail = parent_connection.recv() + except EOFError: + status, detail = "error", "" + finally: + parent_connection.close() + + process.join() + exitcode = process.exitcode + process.close() + + if status == "skip": + pytest.skip(detail) + if status != "ok" or exitcode != 0: + pytest.fail( + detail or f"NVFP4 generator exited with code {exitcode} without a result", + pytrace=False, + ) + + +def _run_lpips_eval(tmp_path, sample_id, media_type, prompt, reference_path, generated_path): + reference_key = "reference_video_path" if media_type == "video" else "reference_image_path" + generated_key = "generated_video_path" if media_type == "video" else "generated_image_path" + dataset_path = tmp_path / f"{sample_id}_dataset.json" + output_json = tmp_path / f"{sample_id}_lpips_results.json" + dataset_path.write_text( + json.dumps( + { + "samples": [ + { + "id": sample_id, + "media_type": media_type, + "prompt": prompt, + reference_key: str(reference_path), + generated_key: str(generated_path), + } + ] + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + + env = os.environ.copy() + env.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + env["PYTHONPATH"] = ( + f"{REPO_ROOT}{os.pathsep}{env['PYTHONPATH']}" if env.get("PYTHONPATH") else REPO_ROOT + ) + with _lpips_deterministic_algorithms(): + result = subprocess.run( + [ + sys.executable, + VISUAL_GEN_LPIPS_EVAL_SCRIPT, + "--dataset", + str(dataset_path), + "--output-json", + str(output_json), + "--json", + ], + cwd=REPO_ROOT, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + ) + if result.returncode != 0: + pytest.fail(f"LPIPS eval script failed for {sample_id}:\n{result.stdout}") + + scores = json.loads(output_json.read_text(encoding="utf-8")) + score = float(scores["mean_lpips_score"]) + print(f"\n[E2E {sample_id} LPIPS] score: {score:.6f}") + return score + + +def _run_reusable_image_lpips_eval(sample_id, reference_path, generated_path, scorer): + try: + with _lpips_deterministic_algorithms(): + score = scorer.score(reference_path, generated_path) + except Exception as error: + pytest.fail(f"LPIPS image eval failed for {sample_id}: {error}") + print(f"\n[E2E {sample_id} LPIPS] score: {score:.6f}") + return score + + +def _run_reusable_video_lpips_eval(sample_id, reference_path, generated_path, scorer): + try: + with _lpips_deterministic_algorithms(): + score = scorer.score_video(reference_path, generated_path) + except Exception as error: + pytest.fail(f"LPIPS video eval failed for {sample_id}: {error}") + print(f"\n[E2E {sample_id} LPIPS] score: {score:.6f}") + return score + + +def _assert_lpips_below_threshold(score, threshold): + assert score < threshold, f"LPIPS too high: {score:.6f} (expected < {threshold:.6f})" + + +def _preserve_lpips_candidate_on_failure(request, score, threshold, candidate_path, artifact_name): + """Copy a failed LPIPS candidate into pytest's archived output directory.""" + if score < threshold: + return + output_dir = request.config.getoption("--output-dir", default=None) + if not output_dir: + return + dest_dir = os.path.join(str(output_dir), "lpips_failure_artifacts") + os.makedirs(dest_dir, exist_ok=True) + dest = os.path.join(dest_dir, artifact_name) + shutil.copy2(str(candidate_path), dest) + print(f"[LPIPS] candidate preserved for golden refresh: {dest}") + + size = os.path.getsize(dest) + if size <= 8 * 1024 * 1024: + with open(dest, "rb") as fh: + encoded = base64.b64encode(fh.read()).decode("ascii") + print(f"[LPIPS-B64-BEGIN {artifact_name} {size}]") + for i in range(0, len(encoded), 3072): + print(encoded[i : i + 3072]) + print(f"[LPIPS-B64-END {artifact_name}]") + + +def _visual_gen_output_path(llm_venv, output_subdir): + out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", output_subdir) + os.makedirs(out_dir, exist_ok=True) + return os.path.join(out_dir, VISUAL_GEN_OUTPUT_VIDEO) + + +def _save_lpips_video_mp4(video, output_path, frame_rate): + from tensorrt_llm.media.encoding import save_video + + try: + save_video(video, output_path, frame_rate=frame_rate) + except RuntimeError as err: + if "MP4 format requires ffmpeg" not in str(err): + raise + # Never fall back to another codec here: the goldens are H.264/x264 and + # LPIPS compares decoded pixels, so a silent cv2/mp4v fallback compares + # codec artifacts, not model output. + pytest.fail( + "ffmpeg is unavailable for LPIPS video encoding; refusing to fall back " + "to another codec because the golden comparison would measure codec " + f"artifacts instead of model output: {err}" + ) + assert os.path.isfile(output_path), f"Visual gen did not produce {output_path}" + + +# ============================================================================= +# VBench setup and score helpers +# ============================================================================= + +VBENCH_DIMENSIONS = [ + "subject_consistency", + "background_consistency", + "motion_smoothness", + "dynamic_degree", + "aesthetic_quality", + "imaging_quality", +] + +VBENCH_REPO = "https://github.com/Vchitect/VBench.git" +# Pin to a fixed commit for reproducible shallow-fetch. +VBENCH_COMMIT = "98b19513678e99c80d8377fda25ba53b81a491a6" + +DINO_REPO = "https://github.com/facebookresearch/dino.git" +DINO_HUB_DIR_NAME = "facebookresearch_dino_main" + +AESTHETIC_PREDICTOR_URL = ( + "https://raw.githubusercontent.com/LAION-AI/aesthetic-predictor/main/sa_0_4_vit_l_14_linear.pth" +) +AESTHETIC_PREDICTOR_FILENAME = "sa_0_4_vit_l_14_linear.pth" +AESTHETIC_PREDICTOR_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "emb_reader") + + +def _prepare_vbench_repo(llm_venv): + """Prepare the pinned VBench checkout and dependencies for this venv.""" + workspace = llm_venv.get_working_directory() + repo_path = os.path.join(workspace, "VBench_repo") + _precache_dino_for_torch_hub() + _precache_aesthetic_predictor() + if not os.path.exists(repo_path): + # Shallow-fetch only the pinned commit to avoid downloading full history (~350 MB). + os.makedirs(repo_path, exist_ok=True) + check_call(["git", "init", repo_path], shell=False) + check_call( + ["git", "-C", repo_path, "remote", "add", "origin", VBENCH_REPO], + shell=False, + ) + check_call( + ["git", "-C", repo_path, "fetch", "--depth", "1", "origin", VBENCH_COMMIT], + shell=False, + ) + check_call(["git", "-C", repo_path, "checkout", "FETCH_HEAD"], shell=False) + + # The checkout persists across sessions, while llm_venv may be new. + llm_venv.run_cmd( + [ + "-m", + "pip", + "install", + "tqdm>=4.60.0", + "openai-clip>=1.0", + "easydict", + "decord>=0.6.0", + "imageio", + ] + ) + llm_venv.run_cmd(["-m", "pip", "install", "--no-deps", "pyiqa>=0.1.0"]) + return repo_path + + +def _precache_dino_for_torch_hub(): + """Pre-clone DINO into the torch.hub cache to avoid GitHub API rate limits.""" + hub_dir = torch.hub.get_dir() + os.makedirs(hub_dir, exist_ok=True) + + dino_cache = os.path.join(hub_dir, DINO_HUB_DIR_NAME) + if not os.path.isdir(dino_cache): + check_call( + ["git", "clone", "--depth", "1", "-b", "main", DINO_REPO, dino_cache], + shell=False, + ) + + +def _precache_aesthetic_predictor(): + """Pre-download LAION aesthetic predictor weights with retry handling.""" + os.makedirs(AESTHETIC_PREDICTOR_CACHE_DIR, exist_ok=True) + cached_path = os.path.join(AESTHETIC_PREDICTOR_CACHE_DIR, AESTHETIC_PREDICTOR_FILENAME) + if os.path.isfile(cached_path): + return + + max_retries = 8 + for attempt in range(max_retries): + try: + req = urllib.request.Request( + AESTHETIC_PREDICTOR_URL, + headers={"User-Agent": "TensorRT-LLM-CI/1.0"}, + ) + with urllib.request.urlopen(req, timeout=120) as resp: + data = resp.read() + tmp_path = cached_path + ".tmp" + with open(tmp_path, "wb") as f: + f.write(data) + os.replace(tmp_path, cached_path) + return + except Exception as exc: + if attempt < max_retries - 1: + wait = min(10 * 2**attempt, 120) + random.uniform(0, 5) + print( + "[precache] Aesthetic predictor download attempt " + f"{attempt + 1}/{max_retries} failed ({exc}), retrying in {wait:.0f}s..." + ) + time.sleep(wait) + else: + raise RuntimeError( + f"Failed to download aesthetic predictor after {max_retries} attempts: {exc}" + ) from exc + + +def _normalize_score(val): + """Normalize a VBench score to a zero-to-one scale.""" + if isinstance(val, bool): + return float(val) + if isinstance(val, (int, float)) and val > 1.5: + return val / 100.0 + return float(val) + + +def _get_per_video_scores(results, video_path_substr): + """Get per-dimension scores for a matching video from VBench results.""" + scores = {} + for dim in VBENCH_DIMENSIONS: + dim_result = results[dim] + assert isinstance(dim_result, list) and len(dim_result) >= 2, ( + f"Dimension '{dim}' result must be [overall_score, video_results]; " + f"got {type(dim_result)}" + ) + video_results = dim_result[1] + for entry in video_results: + if video_path_substr in entry.get("video_path", ""): + scores[dim] = _normalize_score(entry.get("video_results")) + break + else: + raise AssertionError( + f"No video matching '{video_path_substr}' in dimension '{dim}'; " + f"paths: {[e.get('video_path') for e in video_results]}" + ) + return scores + + +def _run_vbench_and_report( + vbench_repo_root, + videos_dir, + trtllm_filename, + llm_venv, + title, + golden_scores=None, + max_score_diff=0.10, +): + """Run VBench, report scores, and optionally compare against golden values.""" + output_path = os.path.join( + llm_venv.get_working_directory(), "vbench_eval_output", title.replace(" ", "_").lower() + ) + os.makedirs(output_path, exist_ok=True) + evaluate_script = os.path.join(vbench_repo_root, "evaluate.py") + cmd = [ + evaluate_script, + "--videos_path", + videos_dir, + "--output_path", + output_path, + "--mode", + "custom_input", + ] + cmd.extend(["--dimension"] + VBENCH_DIMENSIONS) + _venv_check_call(llm_venv, cmd) + + pattern = os.path.join(output_path, "*_eval_results.json") + result_files = glob.glob(pattern) + assert result_files, ( + f"No eval results found matching {pattern}; output dir: {os.listdir(output_path)}" + ) + result_file = max(result_files, key=os.path.getmtime) + with open(result_file, "r") as f: + results = json.load(f) + for dim in VBENCH_DIMENSIONS: + assert dim in results, ( + f"Expected dimension '{dim}' in results; keys: {list(results.keys())}" + ) + + scores_trtllm = _get_per_video_scores(results, trtllm_filename) + max_len = max(len(d) for d in VBENCH_DIMENSIONS) + + if golden_scores is not None: + header = f"{'Dimension':<{max_len}} | {'TRT-LLM':>10} | {'Golden':>10} | {'Diff':>8}" + else: + header = f"{'Dimension':<{max_len}} | {'TRT-LLM':>10}" + separator = "-" * len(header) + print("\n" + "=" * len(header)) + print(f"VBench dimension scores ({title})") + print("=" * len(header)) + print(header) + print(separator) + + for dim in VBENCH_DIMENSIONS: + trtllm_score = scores_trtllm[dim] + if golden_scores is not None: + golden_score = golden_scores[dim] + print( + f"{dim:<{max_len}} | {trtllm_score:>10.4f} | {golden_score:>10.4f} | " + f"{abs(trtllm_score - golden_score):>8.4f}" + ) + else: + print(f"{dim:<{max_len}} | {trtllm_score:>10.4f}") + print(separator) + + if golden_scores is None: + print("\n** Baseline run — no golden scores to compare against. **") + print("** Copy the values above into the GOLDEN_SCORES dict. **\n") + return scores_trtllm + + max_diff_val = max(abs(scores_trtllm[dim] - golden_scores[dim]) for dim in VBENCH_DIMENSIONS) + print(f"max_diff={max_diff_val:.4f} (threshold={max_score_diff})") + print("=" * len(header) + "\n") + for dim in VBENCH_DIMENSIONS: + diff = abs(scores_trtllm[dim] - golden_scores[dim]) + assert diff < max_score_diff or scores_trtllm[dim] >= golden_scores[dim], ( + f"Dimension '{dim}' score difference {diff:.4f} >= {max_score_diff} " + f"(TRT-LLM={scores_trtllm[dim]:.4f}, golden={golden_scores[dim]:.4f})" + ) + return scores_trtllm + + +# ============================================================================= +# WAN LPIPS configuration and generation helpers +# ============================================================================= + +# Keep WAN LPIPS parameters here because the generation helpers and single-GPU +# tests share them; the WAN 2.2 values are also reused by the mechanically +# preserved multi-GPU tests. One source keeps those runs aligned with the +# committed baseline media. +WAN21_LPIPS_PROMPT = "A cat sitting on a windowsill" +WAN21_LPIPS_NEGATIVE_PROMPT = None +WAN21_LPIPS_HEIGHT = 256 +WAN21_LPIPS_WIDTH = 256 +WAN21_LPIPS_NUM_FRAMES = 5 +WAN21_LPIPS_NUM_INFERENCE_STEPS = 1 +WAN21_LPIPS_GUIDANCE_SCALE = 5.0 +WAN21_LPIPS_SEED = 42 +WAN_LPIPS_FRAME_RATE = 16.0 +WAN_LPIPS_THRESHOLD = 0.05 + +WAN22_LPIPS_PROMPT = "A cat sitting on a sunny windowsill watching birds outside." +WAN22_LPIPS_NEGATIVE_PROMPT = "" +WAN22_LPIPS_HEIGHT = 480 +WAN22_LPIPS_WIDTH = 832 +WAN22_LPIPS_NUM_FRAMES = 9 +WAN22_LPIPS_NUM_INFERENCE_STEPS = 4 +WAN22_LPIPS_GUIDANCE_SCALE = 4.0 +WAN22_LPIPS_SEED = 42 +WAN22_LPIPS_FRAME_RATE = 16.0 + + +def _run_wan_lpips_pipeline( + model_path, + prompt, + negative_prompt, + height, + width, + num_frames, + num_inference_steps, + guidance_scale, + seed, + attention_backend="VANILLA", + parallel=None, + fully_eager=False, +): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs + + _skip_if_missing(model_path, "Wan checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + args_kwargs = dict( + model=model_path, + attention_config=AttentionConfig(backend=attention_backend), + torch_compile_config=TorchCompileConfig(enable=False), + ) + if parallel is not None: + args_kwargs["parallel_config"] = parallel + with _lpips_deterministic_algorithms(fully_eager=fully_eager): + args = VisualGenArgs(**args_kwargs) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + with torch.no_grad(): + result = pipeline.forward( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + seed=seed, + ) + if result is None or result.video is None: + return None + return result.video.detach().cpu() + finally: + del pipeline + _cleanup_cuda() + + +def _generate_wan_lpips_video( + model_path, + output_path, + prompt, + negative_prompt, + height, + width, + num_frames, + num_inference_steps, + guidance_scale, + seed, + frame_rate, + attention_backend="VANILLA", + parallel=None, + fully_eager=False, +): + generated_video = _run_wan_lpips_pipeline( + model_path, + prompt, + negative_prompt, + height, + width, + num_frames, + num_inference_steps, + guidance_scale, + seed, + attention_backend=attention_backend, + parallel=parallel, + fully_eager=fully_eager, + ) + assert generated_video is not None, "Single-GPU Wan LPIPS run produced no video" + _save_lpips_video_mp4(generated_video, output_path, frame_rate=frame_rate) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index d812dcf5d84e..0b7711d94a81 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -228,14 +228,13 @@ l0_b200: - unittest/_torch/visual_gen/test_wan_transformer.py - unittest/_torch/visual_gen/test_cosmos3_transformer.py - unittest/_torch/visual_gen/test_cosmos3_pipeline.py - - examples/visual_gen/test_visual_gen.py::test_wan_t2v_example - - examples/visual_gen/test_visual_gen.py::test_flux1_example - - examples/visual_gen/test_visual_gen.py::test_flux2_example - - examples/visual_gen/test_visual_gen.py::test_ltx2_example - - examples/visual_gen/test_visual_gen.py::test_wan_i2v_example - - examples/visual_gen/test_visual_gen.py::test_cosmos3_example - - examples/visual_gen/test_visual_gen.py::test_qwen_image_example - # - examples/visual_gen/test_visual_gen.py + - examples/visual_gen/test_visual_gen_wan.py::test_wan_t2v_example + - examples/visual_gen/test_visual_gen_flux.py::test_flux1_example + - examples/visual_gen/test_visual_gen_flux.py::test_flux2_example + - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_example + - examples/visual_gen/test_visual_gen_wan.py::test_wan_i2v_example + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_example + - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwen_image_example # ------------- Host perf module regression tests (6 representative scenarios) --------------- - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_gen_only_bs8] - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_mixed_32gen_4ctx] @@ -322,21 +321,40 @@ l0_b200: - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[bf16] - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[nvfp4] # ------------- VisualGen single-GPU tests --------------- - - examples/visual_gen/test_visual_gen.py::test_visual_gen_quickstart - - examples/visual_gen/test_visual_gen.py::test_visual_gen_api_walkthrough - - examples/visual_gen/test_visual_gen.py::test_flux1_lpips_against_golden - - examples/visual_gen/test_visual_gen.py::test_flux2_lpips_against_golden - - examples/visual_gen/test_visual_gen.py::test_ltx2_lpips_against_golden - - examples/visual_gen/test_visual_gen.py::test_ltx2_cuda_graph_lpips_matches_eager - - examples/visual_gen/test_visual_gen.py::test_ltx2_cuda_graph_trtllm_backend - - examples/visual_gen/test_visual_gen.py::test_wan21_t2v_lpips_against_golden - - examples/visual_gen/test_visual_gen.py::test_wan22_t2v_lpips_against_golden + - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_quickstart + - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_api_walkthrough + - examples/visual_gen/test_visual_gen_flux.py::test_flux1_lpips_against_golden + - examples/visual_gen/test_visual_gen_flux.py::test_flux2_lpips_against_golden + - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-fp8-blockwise] + - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-nvfp4] + - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-cuda-graph] + - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux2-fp8-blockwise] + - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux2-nvfp4] + - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux2-cuda-graph] + - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[fp8-blockwise] + - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[nvfp4] + - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[cuda-graph] + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_feature_accuracy_against_golden[fp8-blockwise] + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_feature_accuracy_against_golden[nvfp4] + - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[fp8-blockwise] + - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[nvfp4] + - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[cuda-graph] + - examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan21-fp8-blockwise] + - examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan21-nvfp4] + - examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan21-cuda-graph] + - examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-fp8-blockwise] + - examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-nvfp4] + - examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-cuda-graph] + - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_lpips_against_golden + - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_cuda_graph_lpips_matches_eager + - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_cuda_graph_trtllm_backend + - examples/visual_gen/test_visual_gen_wan.py::test_wan21_t2v_lpips_against_golden + - examples/visual_gen/test_visual_gen_wan.py::test_wan22_t2v_lpips_against_golden # Measured on B200 with TRT-LLM commit 85665f5f from the staging main image: # QwenImage ~5 min, Cosmos3-Nano T2I ~3 min, and T2V ~5 min. - - examples/visual_gen/test_visual_gen.py::test_qwenimage_lpips_against_golden TIMEOUT (10) - - examples/visual_gen/test_visual_gen.py::test_qwenimage_cuda_graph_lpips_against_golden TIMEOUT (10) - - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden TIMEOUT (10) - - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden TIMEOUT (15) + - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_lpips_against_golden TIMEOUT (10) + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2i_lpips_against_golden TIMEOUT (10) + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2v_lpips_against_golden TIMEOUT (15) - visual_gen/test_visual_gen_benchmark.py::test_offline_benchmark - visual_gen/test_visual_gen_benchmark.py::test_online_benchmark[openai-videos] # ---- moved to post-merge (MoE CI optimization) ---- diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index e64a30ba11c0..0058c5bea298 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -327,7 +327,7 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=2-ep4-fp8kv=False-attention_dp=True-cuda_graph=False-overlap_scheduler=False-low_precision_combine=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=False-enable_gemm_allreduce_fusion=False] - - examples/visual_gen/test_visual_gen.py::test_wan_t2v_example + - examples/visual_gen/test_visual_gen_wan.py::test_wan_t2v_example - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[ulysses4] - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2] - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[attn2d_2x2] diff --git a/tests/integration/test_lists/test-db/l0_gh200.yml b/tests/integration/test_lists/test-db/l0_gh200.yml index 6cf8033c2d6b..7072b7605606 100644 --- a/tests/integration/test_lists/test-db/l0_gh200.yml +++ b/tests/integration/test_lists/test-db/l0_gh200.yml @@ -17,5 +17,5 @@ l0_gh200: - unittest/bindings - unittest/llmapi/test_llm_quant.py - llmapi/test_llm_examples.py::test_llmapi_quickstart_atexit - - examples/visual_gen/test_visual_gen.py::test_visual_gen_quickstart - - examples/visual_gen/test_visual_gen.py::test_visual_gen_api_walkthrough + - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_quickstart + - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_api_walkthrough diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 133a34e15300..4351d1bccf5d 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -294,8 +294,8 @@ l0_h100: backend: pytorch tests: - unittest/llmapi/test_llm_quant.py # 5.5 mins on H100 - - examples/visual_gen/test_visual_gen.py::test_visual_gen_quickstart - - examples/visual_gen/test_visual_gen.py::test_visual_gen_api_walkthrough + - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_quickstart + - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_api_walkthrough - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_l40s.yml b/tests/integration/test_lists/test-db/l0_l40s.yml index bf71ba62fc8c..6d14548d1a3c 100644 --- a/tests/integration/test_lists/test-db/l0_l40s.yml +++ b/tests/integration/test_lists/test-db/l0_l40s.yml @@ -68,8 +68,8 @@ l0_l40s: backend: pytorch tests: - llmapi/test_llm_examples.py::test_llmapi_quickstart - - examples/visual_gen/test_visual_gen.py::test_visual_gen_quickstart - - examples/visual_gen/test_visual_gen.py::test_visual_gen_api_walkthrough + - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_quickstart + - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_api_walkthrough - llmapi/test_llm_examples.py::test_llmapi_example_inference - llmapi/test_llm_examples.py::test_llmapi_example_inference_async - llmapi/test_llm_examples.py::test_llmapi_example_inference_async_streaming diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 5038092f4bf2..40f1339e1592 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -141,9 +141,9 @@ examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_accep examples/test_ray.py::test_llm_inference_distributed_ray[pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_llm_inference_distributed_ray[tp2pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_ray_disaggregated_serving[tp2] SKIP (https://nvbugs/5612502) -examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815) -examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341) -examples/visual_gen/test_visual_gen.py::test_ltx2_cuda_graph_trtllm_backend SKIP (https://nvbugs/6463822) +examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815) +examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341) +examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_cuda_graph_trtllm_backend SKIP (https://nvbugs/6463822) full:A100/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6479837) full:A100X/llmapi/test_llm_examples.py::test_llmapi_speculative_decoding_mtp SKIP (https://nvbugs/6287561) full:B200/accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] SKIP (https://nvbugs/6327718)