diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index d1e376b9a13e..40d26bd6c06f 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -39,6 +39,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode | `Qwen/Qwen-Image-2512` | Text-to-Image | | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video | | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video | +| `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | Text-to-Video | Models are auto-detected from the checkpoint directory. Diffusers-format models are detected via `model_index.json`; LTX-2 monolithic safetensors checkpoints are detected via embedded metadata. The `AutoPipeline` registry selects the appropriate pipeline class automatically. @@ -51,6 +52,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models | **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | | **Wan 2.1 VSA** [^2] | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | Yes | | **Wan 2.2** | Yes | Yes | Yes [^3] | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | +| **FastWan 2.2** | Yes | Yes | No | No | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | | **LTX-2** | Yes | Yes | Yes [^4] | Yes | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | No | | **Qwen-Image** [^5] | Yes | Yes | No | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | | **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | @@ -65,6 +67,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models [^5]: Qwen-Image ships a native BF16 implementation with per-module numerical parity against `diffusers.QwenImagePipeline` (cosine similarity >= 0.999 on the full 20B transformer) and supports `trtllm-serve` / `/v1/images/generations`. VisualGen supports FP8 blockwise and NVFP4 dynamic quantization from BF16 checkpoints, as well as direct loading of statically quantized FP8 and NVFP4 ModelOpt checkpoints. + ## Quick Start Here is a simple example to generate a video with Wan 2.1: diff --git a/tensorrt_llm/_torch/visual_gen/models/__init__.py b/tensorrt_llm/_torch/visual_gen/models/__init__.py index c5d63ed88b23..68f271f6bbf1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/models/__init__.py @@ -37,7 +37,7 @@ from .flux import Flux2Pipeline, FluxPipeline from .ltx2 import LTX2Pipeline # noqa: F401 from .qwen_image import QwenImagePipeline -from .wan import WanImageToVideoPipeline, WanPipeline +from .wan import WanDMDPipeline, WanImageToVideoPipeline, WanPipeline __all__ = [ "AutoPipeline", @@ -46,6 +46,7 @@ "Flux2Pipeline", "QwenImagePipeline", "WanPipeline", + "WanDMDPipeline", "WanImageToVideoPipeline", "Cosmos3OmniMoTPipeline", "register_pipeline", diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/__init__.py b/tensorrt_llm/_torch/visual_gen/models/wan/__init__.py index 1eff9ff3ef2c..fd876e38bcb2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/__init__.py @@ -1,4 +1,5 @@ from .parallel_vae import ParallelVAE_Wan +from .pipeline_fastwan import WanDMDPipeline from .pipeline_wan import WanPipeline from .pipeline_wan_i2v import WanImageToVideoPipeline from .transformer_wan import WanTransformer3DModel @@ -6,6 +7,7 @@ __all__ = [ "WanPipeline", + "WanDMDPipeline", "WanImageToVideoPipeline", "WanTransformer3DModel", "ParallelVAE_Wan", diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/defaults.py b/tensorrt_llm/_torch/visual_gen/models/wan/defaults.py index 5ceacfe8af4c..f3beba76430e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/defaults.py @@ -131,6 +131,18 @@ def get_wan_default_params( return params +def get_fastwan_default_params() -> dict: + """Default generation params for FastWan 2.2 TI2V-5B. + + Same as the base Wan 2.2 TI2V-5B defaults, but distilled: a fixed 3-step + DMD schedule and CFG-free (guidance_scale 1.0). + """ + params = dict(_WAN22_5B_PARAMS) + params["num_inference_steps"] = 3 + params["guidance_scale"] = 1.0 + return params + + def get_wan_extra_param_specs(is_wan22_14b: bool) -> Dict[str, ExtraParamSchema]: """Return extra_param_specs for a Wan model. diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py new file mode 100644 index 000000000000..7ffa42478662 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py @@ -0,0 +1,192 @@ +# 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. +"""FastWan: distilled (DMD) Wan 2.2 TI2V-5B for 3-step text-to-video. + +FastWan shares Wan 2.2 TI2V-5B's architecture exactly — only the weights +(distilled) and the sampling recipe differ. So this pipeline subclasses +``WanPipeline`` and overrides only the denoising loop. +""" + +import time + +import PIL.Image +import torch +from diffusers.utils.torch_utils import randn_tensor + +from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput +from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline +from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.logger import logger + +from .defaults import get_fastwan_default_params +from .pipeline_wan import WanPipeline + + +@register_pipeline( + "WanDMDPipeline", + hf_ids=[ + "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", + ], + doc="FastWan 2.2 distilled (DMD) — 3-step Wan 2.2 TI2V-5B text-to-video.", +) +class WanDMDPipeline(WanPipeline): + """Wan 2.2 TI2V-5B with the DMD 3-step sampling loop. + + DMD inference (Distribution Matching Distillation): at each step the model + predicts the clean video ``x0``, then we re-noise it with FRESH noise to the + next (lower) noise level. The last step keeps ``x0``. CFG-free — one + transformer forward per step. Faithful to FastVideo's ``DmdDenoisingStage``. + """ + + # Fixed DMD schedule for FastWan2.2-TI2V-5B (see model card / FastVideo config). + DMD_TIMESTEPS = (1000, 757, 522) + NUM_TRAIN_TIMESTEPS = 1000 + + @property + def default_generation_params(self): + """Return FastWan-specific default generation parameters (steps, guidance, frame rate).""" + return get_fastwan_default_params() + + @nvtx_range("WanDMDPipeline.forward") + @torch.no_grad() + def forward( + self, + prompt: str | list[str], + seed: int, + negative_prompt: str | None = None, + height: int = 704, + width: int = 1280, + num_frames: int = 121, + num_inference_steps: int | None = None, + guidance_scale: float | None = None, + guidance_scale_2: float | None = None, + boundary_ratio: float | None = None, + max_sequence_length: int = 512, + image: PIL.Image.Image | torch.Tensor | str | None = None, + ) -> PipelineOutput: + """Generate a video from a text prompt using the 3-step DMD sampling loop. + + CFG-free: one transformer forward pass per step, no negative prompt support. + Unsupported parameters (num_inference_steps, guidance_scale, guidance_scale_2, + boundary_ratio) are accepted for base-class compatibility but logged as warnings + and ignored. + """ + if image is not None: + raise NotImplementedError( + "FastWan currently supports text-to-video only (no image conditioning)." + ) + if num_inference_steps not in (None, len(self.DMD_TIMESTEPS)): + logger.warning( + f"FastWan always uses {len(self.DMD_TIMESTEPS)} DMD steps; " + f"ignoring num_inference_steps={num_inference_steps}." + ) + if guidance_scale not in (None, 1.0): + logger.warning(f"FastWan is CFG-free; ignoring guidance_scale={guidance_scale}.") + if guidance_scale_2 is not None: + logger.warning( + f"FastWan does not support guidance_scale_2; ignoring guidance_scale_2={guidance_scale_2}." + ) + if boundary_ratio is not None: + logger.warning( + f"FastWan does not support boundary_ratio; ignoring boundary_ratio={boundary_ratio}." + ) + + pipeline_start = time.time() + timer = CudaPhaseTimer() + timer.mark_pre_start() + + if isinstance(prompt, str): + prompt = [prompt] + batch_size = len(prompt) + generator = torch.Generator(device=self.device).manual_seed(seed) + self.validate_resolution(height, width, num_frames) + + if negative_prompt: + raise ValueError("FastWan is CFG-free and does not support negative prompts.") + prompt_embeds, _ = self._encode_prompt(prompt, "", max_sequence_length) + + if self._fixed_latent is not None: + latents = self._fixed_latent.to(device=self.device, dtype=self.dtype) + else: + latents = self._prepare_latents(batch_size, height, width, num_frames, generator) + + timer.mark_denoise_start() + latents = self._denoise(latents, prompt_embeds, generator) + timer.mark_post_start() + + logger.info("Decoding video...") + decode_start = time.time() + video = self.decode_latents(latents, self._decode_latents) + if self.rank == 0: + logger.info(f"Video decoded in {time.time() - decode_start:.2f}s") + logger.info(f"Total pipeline time: {time.time() - pipeline_start:.2f}s") + + timer.mark_end() + frame_rate = self.default_generation_params.get("frame_rate", 24.0) + return timer.fill(PipelineOutput(video=video, frame_rate=frame_rate)) + + @nvtx_range("WanDMDPipeline._denoise", color="blue") + def _denoise( + self, latents: torch.Tensor, prompt_embeds: torch.Tensor, generator: torch.Generator + ) -> torch.Tensor: + """3-step DMD loop: predict x0, then re-noise with fresh noise. + + The per-step sigma reduces to ``sigma = t / NUM_TRAIN_TIMESTEPS`` because + FastVideo's scheduler keeps ``timesteps = sigmas * num_train_timesteps`` + (co-scaled), so the lookup returns ``matched_timestep / 1000`` and the + flow-shift cancels (verified numerically for shift in {1, 5, 8, 12}). + """ + timesteps = self.DMD_TIMESTEPS + num_steps = len(timesteps) + _, ph, pw = self.transformer.config.patch_size + + # The DMD ops are linear (mul/sub/add), so they run in the transformer's + # native (bf16) dtype to match the FastVideo reference; no fp32 needed. + latents = latents.to(self.dtype) + # Patch grid dimensions are fixed for the entire denoising loop. + nf = latents.shape[2] + nh = latents.shape[3] // ph + nw = latents.shape[4] // pw + start = time.time() + for i, t in enumerate(timesteps): + t_tensor = torch.full((latents.shape[0], nf * nh * nw), float(t), device=latents.device) + + pred_noise = self.transformer( + hidden_states=latents, + timestep=t_tensor / self.NUM_TRAIN_TIMESTEPS, + encoder_hidden_states=prompt_embeds, + ) + + sigma = t / self.NUM_TRAIN_TIMESTEPS + pred_video = latents - sigma * pred_noise # clean-video estimate (x0) + + if i < num_steps - 1: + sigma_next = timesteps[i + 1] / self.NUM_TRAIN_TIMESTEPS + noise = randn_tensor( + latents.shape, + generator=generator, + device=latents.device, + dtype=self.dtype, + ) + latents = (1.0 - sigma_next) * pred_video + sigma_next * noise + else: + latents = pred_video + + if self.rank == 0: + logger.info(f"DMD step {i + 1}/{num_steps} (t={t}, sigma={sigma:.3f})") + + if self.rank == 0: + logger.info(f"DMD denoising done in {time.time() - start:.2f}s") + return latents diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.json new file mode 100644 index 000000000000..25bef78765f5 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.json @@ -0,0 +1,16 @@ +{ + "video": "fastwan_lpips_golden_video.mp4", + "model": "FastWan2.2-TI2V-5B-FullAttn-Diffusers", + "source": "FastVideo VideoGenerator (trusted reference)", + "prompt": "A red fox walking through snow", + "negative_prompt": "", + "height": 704, + "width": 1280, + "num_frames": 121, + "num_inference_steps": 3, + "guidance_scale": 1.0, + "seed": 42, + "frame_rate": 24.0, + "lpips_net": "alex", + "lpips_threshold": 0.05 +} 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..83cb6cdb2735 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:1e6e0a6be5b9e219df7b881b601ce88f2dd77c031c7fae75925119b907bc69c2 +size 16803320 diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py index ed8a5978941d..b2e20f5525aa 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -120,6 +120,17 @@ COSMOS3_LPIPS_FRAME_RATE = 24.0 COSMOS3_LPIPS_THRESHOLD = 0.05 +# FastWan 2.2 TI2V-5B (DMD distilled) LPIPS configuration +FASTWAN_LPIPS_PROMPT = "A red fox walking through snow" +FASTWAN_LPIPS_NEGATIVE_PROMPT = "" +FASTWAN_LPIPS_HEIGHT = 704 +FASTWAN_LPIPS_WIDTH = 1280 +FASTWAN_LPIPS_NUM_FRAMES = 121 +FASTWAN_LPIPS_NUM_INFERENCE_STEPS = 3 +FASTWAN_LPIPS_GUIDANCE_SCALE = 1.0 +FASTWAN_LPIPS_SEED = 42 +FASTWAN_LPIPS_FRAME_RATE = 24.0 +FASTWAN_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" @@ -1134,6 +1145,42 @@ def test_qwenimage_cuda_graph_lpips_against_golden(tmp_path): _assert_lpips_below_threshold(score, QWENIMAGE_LPIPS_THRESHOLD) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_fastwan_lpips_against_golden(tmp_path): + generated_path = tmp_path / "fastwan_generated.mp4" + golden_path = _golden_media_path( + tmp_path, "fastwan_lpips_golden_video.mp4", "FastWan LPIPS golden video" + ) + _generate_wan_lpips_video( + _lpips_model_path("FastWan2.2-TI2V-5B-FullAttn-Diffusers"), + generated_path, + FASTWAN_LPIPS_PROMPT, + FASTWAN_LPIPS_NEGATIVE_PROMPT, + FASTWAN_LPIPS_HEIGHT, + FASTWAN_LPIPS_WIDTH, + FASTWAN_LPIPS_NUM_FRAMES, + FASTWAN_LPIPS_NUM_INFERENCE_STEPS, + FASTWAN_LPIPS_GUIDANCE_SCALE, + FASTWAN_LPIPS_SEED, + FASTWAN_LPIPS_FRAME_RATE, + ) + score = _run_lpips_eval( + tmp_path, + "fastwan", + "video", + FASTWAN_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, FASTWAN_LPIPS_THRESHOLD) + + +@pytest.fixture(scope="session") +def wan_trtllm_video_path(_visual_gen_deps, llm_venv, llm_root): + """Generate input video via models/wan_t2v.py and return path to trtllm_output.mp4.""" + return _generate_wan_video(llm_venv, llm_root, WAN_T2V_MODEL_SUBPATH, "wan") + + @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" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 464c3e553f91..8ffada48403f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -251,6 +251,8 @@ l0_b200: - unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py - unittest/_torch/visual_gen/test_wan_vsa_pipeline.py - unittest/_torch/visual_gen/test_wan_vae.py + - unittest/_torch/visual_gen/test_fastwan_dmd_math.py + - unittest/_torch/visual_gen/test_fastwan_pipeline.py - unittest/_torch/visual_gen/test_wan21_i2v_teacache.py - unittest/_torch/visual_gen/test_wan21_t2v_teacache.py - unittest/_torch/visual_gen/test_wan21_t2v_teacache_user_coefficients.py @@ -370,6 +372,7 @@ l0_b200: - 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.py::test_fastwan_lpips_against_golden TIMEOUT (10) - 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/unittest/_torch/visual_gen/test_fastwan_dmd_math.py b/tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py new file mode 100644 index 000000000000..ba49c140becc --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py @@ -0,0 +1,198 @@ +# 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. +"""DMD denoising loop math tests for WanDMDPipeline. + +Tests _denoise in isolation: no checkpoint, no text encoder, no VAE. +A stub pipeline (object.__new__) with a fake transformer supplies controlled +inputs so the DMD formulas can be verified numerically. + +GPU required (tensor ops + nvtx_range decorator). + +Run: + pytest tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py -v -s +""" + +import os +import unittest.mock as mock + +os.environ["TLLM_DISABLE_MPI"] = "1" + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.models.wan.pipeline_fastwan import WanDMDPipeline + + +@pytest.fixture(autouse=True, scope="module") +def _cleanup_mpi_env(): + yield + os.environ.pop("TLLM_DISABLE_MPI", None) + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +_PATCH_SIZE = (1, 2, 2) + +# Small latent (B=1, C=16, T=3, H=4, W=4) +# nf=3, nh=4//2=2, nw=4//2=2 → 12 patches per sample +_LATENT_SHAPE = (1, 16, 3, 4, 4) +_NUM_PATCHES = 3 * 2 * 2 # 12 + +# With pred_noise=0 and noise=0, re-noising reduces to: +# latents_{i+1} = (1 - sigma_next) * latents_i +# So final output = (1-sigma_1)*(1-sigma_2)*initial = 0.243*0.478*initial ≈ 0.1162 +_SIGMA_1 = WanDMDPipeline.DMD_TIMESTEPS[1] / WanDMDPipeline.NUM_TRAIN_TIMESTEPS # 0.757 +_SIGMA_2 = WanDMDPipeline.DMD_TIMESTEPS[2] / WanDMDPipeline.NUM_TRAIN_TIMESTEPS # 0.522 +_ZERO_NOISE_EXPECTED_SCALE = (1.0 - _SIGMA_1) * (1.0 - _SIGMA_2) # ≈ 0.1162 + +_RANDN_PATH = "tensorrt_llm._torch.visual_gen.models.wan.pipeline_fastwan.randn_tensor" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _RecordingTransformer: + """Fake transformer: records every timestep call, returns zeros. + + parameters() yields a single bfloat16 tensor so BasePipeline.dtype + returns torch.bfloat16 without needing a real nn.Module. + """ + + def __init__(self): + self.captured_timesteps = [] + self._dummy_param = torch.zeros(1, dtype=torch.bfloat16, device="cuda") + + class _Config: + patch_size = _PATCH_SIZE + + self.config = _Config() + + def parameters(self): + yield self._dummy_param + + def __call__(self, hidden_states, timestep, encoder_hidden_states): + self.captured_timesteps.append(timestep.detach().clone()) + return torch.zeros_like(hidden_states) + + @property + def call_count(self): + return len(self.captured_timesteps) + + +def _make_stub(transformer): + """WanDMDPipeline shell — no weights, no __init__. + + We bypass nn.Module.__setattr__ by writing directly into __dict__. + - transformer: accessed via __dict__, so no nn.Module submodule registration + - pipeline_config: WanPipeline.dtype reads self.pipeline_config.torch_dtype + - rank: read-only property on BasePipeline — returns 0 when MPI is disabled + """ + pipe = object.__new__(WanDMDPipeline) + pipe.__dict__["transformer"] = transformer + pipe.__dict__["pipeline_config"] = type("_Config", (), {"torch_dtype": torch.bfloat16})() + return pipe + + +def _run(pipe, latents, zero_noise=False): + gen = torch.Generator(device=latents.device).manual_seed(42) + embeds = torch.zeros(1, 8, 4096, device=latents.device, dtype=torch.bfloat16) + if zero_noise: + zeros = torch.zeros_like(latents) + with mock.patch(_RANDN_PATH, return_value=zeros): + return pipe._denoise(latents, embeds, gen) + return pipe._denoise(latents, embeds, gen) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestDMDTimesteps: + """sigma = t / 1000 for every step; tensor shape and call count correct.""" + + def test_sigma_values(self): + recorder = _RecordingTransformer() + _run(_make_stub(recorder), torch.ones(_LATENT_SHAPE, device="cuda", dtype=torch.bfloat16)) + + expected = [t / 1000.0 for t in WanDMDPipeline.DMD_TIMESTEPS] + for i, (captured, exp) in enumerate(zip(recorder.captured_timesteps, expected)): + assert torch.allclose(captured, torch.full_like(captured, exp)), ( + f"Step {i}: expected sigma {exp}, got {captured[0, 0].item():.6f}" + ) + + def test_timestep_tensor_shape(self): + recorder = _RecordingTransformer() + _run(_make_stub(recorder), torch.ones(_LATENT_SHAPE, device="cuda", dtype=torch.bfloat16)) + + for i, captured in enumerate(recorder.captured_timesteps): + assert captured.shape == (1, _NUM_PATCHES), ( + f"Step {i}: expected (1, {_NUM_PATCHES}), got {captured.shape}" + ) + + def test_transformer_called_three_times(self): + """One forward per step — CFG-free means no double-pass.""" + recorder = _RecordingTransformer() + _run(_make_stub(recorder), torch.ones(_LATENT_SHAPE, device="cuda", dtype=torch.bfloat16)) + assert recorder.call_count == 3 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestDMDMath: + def test_output_dtype_is_bfloat16(self): + output = _run( + _make_stub(_RecordingTransformer()), + torch.ones(_LATENT_SHAPE, device="cuda", dtype=torch.bfloat16), + ) + assert output.dtype == torch.bfloat16 + + def test_exact_output_with_zero_noise(self): + """pred_noise=0 and randn=0 → output = (1-sigma_1)*(1-sigma_2)*initial ≈ 0.1162*initial.""" + initial = torch.ones(_LATENT_SHAPE, device="cuda", dtype=torch.bfloat16) + output = _run(_make_stub(_RecordingTransformer()), initial, zero_noise=True) + + expected = torch.full( + _LATENT_SHAPE, _ZERO_NOISE_EXPECTED_SCALE, dtype=torch.float32, device="cuda" + ).to(torch.bfloat16) + assert torch.allclose(output.float(), expected.float(), atol=5e-3), ( + f"Expected ≈{_ZERO_NOISE_EXPECTED_SCALE:.4f}, got {output.float().mean().item():.4f}" + ) + + def test_final_step_does_not_call_randn_tensor(self): + """randn_tensor must be called for steps 0 and 1 only — never on the final step.""" + from diffusers.utils.torch_utils import randn_tensor as _orig + + call_count = [0] + + def _counting(*args, **kwargs): + call_count[0] += 1 + if call_count[0] >= 3: + raise AssertionError("randn_tensor called on the final step — must not happen") + return _orig(*args, **kwargs) + + pipe = _make_stub(_RecordingTransformer()) + gen = torch.Generator(device="cuda").manual_seed(42) + embeds = torch.zeros(1, 8, 4096, device="cuda", dtype=torch.bfloat16) + latents = torch.ones(_LATENT_SHAPE, device="cuda", dtype=torch.bfloat16) + + with mock.patch(_RANDN_PATH, side_effect=_counting): + pipe._denoise(latents, embeds, gen) + + assert call_count[0] == 2, f"Expected 2 randn_tensor calls, got {call_count[0]}" diff --git a/tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py b/tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py new file mode 100644 index 000000000000..1572c637c3e5 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py @@ -0,0 +1,136 @@ +# 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. +"""Pipeline-level tests for WanDMDPipeline. + +Run: + DIFFUSION_MODEL_PATH_FASTWAN=/path/to/checkpoint \\ + pytest tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py -v -s + +Override checkpoint path via DIFFUSION_MODEL_PATH_FASTWAN env var or place +the checkpoint at $LLM_MODELS_ROOT/FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers +""" + +import gc +import os + +os.environ["TLLM_DISABLE_MPI"] = "1" + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader +from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs + + +@pytest.fixture(autouse=True, scope="module") +def _cleanup_mpi_env(): + yield + os.environ.pop("TLLM_DISABLE_MPI", None) + + +def _checkpoint_path() -> str: + if "DIFFUSION_MODEL_PATH_FASTWAN" in os.environ: + return os.environ["DIFFUSION_MODEL_PATH_FASTWAN"] + root = os.environ.get( + "LLM_MODELS_ROOT", + "/home/scratch.trt_llm_data_ci/llm-models", + ) + return os.path.join(root, "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers") + + +FASTWAN_PATH = _checkpoint_path() + + +@pytest.fixture(scope="module") +def fastwan_pipeline(): + if not os.path.exists(FASTWAN_PATH): + pytest.skip(f"Checkpoint not found: {FASTWAN_PATH}") + args = VisualGenArgs( + model=FASTWAN_PATH, + torch_compile_config=TorchCompileConfig(enable=False), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + yield pipeline + del pipeline + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestFastWanForward: + """Run the full forward() pass and validate the output video shape.""" + + def test_single_prompt_shape(self, fastwan_pipeline): + """Single prompt returns (B, T, H, W, C) with B=1.""" + result = fastwan_pipeline.forward( + prompt="A red fox walking through snow", + height=704, + width=1280, + num_frames=9, + seed=42, + ) + assert result.video.dim() == 5, f"Expected 5D (B,T,H,W,C), got {result.video.dim()}D" + B, _T, H, W, C = result.video.shape + assert B == 1 and H == 704 and W == 1280 and C == 3 + + def test_batch_prompt_shape(self, fastwan_pipeline): + """List of prompts returns (B, T, H, W, C) with B=2.""" + result = fastwan_pipeline.forward( + prompt=["A red fox walking through snow", "A cat on a roof"], + height=704, + width=1280, + num_frames=9, + seed=42, + ) + assert result.video.dim() == 5, f"Expected 5D (B,T,H,W,C), got {result.video.dim()}D" + B, _T, H, W, C = result.video.shape + assert B == 2 and H == 704 and W == 1280 and C == 3 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestFastWanFP8: + """Confirm FP8 + TRTLLM attention runs end-to-end on the FastWan checkpoint. + + FastWan inherits FP8 support from WanPipeline.post_load_weights(), but that + method is only tested against the TI2V-5B checkpoint in CI. This test confirms + it also works with the FastWan distilled weights and the DMD denoising loop. + """ + + def test_fp8_trtllm(self): + if not os.path.exists(FASTWAN_PATH): + pytest.skip(f"Checkpoint not found: {FASTWAN_PATH}") + args = VisualGenArgs( + model=FASTWAN_PATH, + torch_compile_config=TorchCompileConfig(enable=False), + quant_config={"quant_algo": "FP8", "dynamic": True}, + attention_config=AttentionConfig(backend="TRTLLM"), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + with torch.no_grad(): + result = pipeline.forward( + prompt="A red fox walking through snow", + height=704, + width=1280, + num_frames=9, + seed=42, + ) + assert result.video.dim() == 5 + B, _T, H, W, C = result.video.shape + assert B == 1 and H == 704 and W == 1280 and C == 3 + finally: + del pipeline + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index e69c81d21f1b..3650439ae3b9 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -224,6 +224,18 @@ def test_wan22_5b_defaults(self): assert d["num_frames"] == 121 assert d["frame_rate"] == 24.0 + def test_fastwan_defaults(self): + """FastWan 2.2 TI2V-5B (DMD distilled): wan22-5b defaults, 3-step, CFG-free.""" + from tensorrt_llm._torch.visual_gen.models.wan.pipeline_fastwan import WanDMDPipeline + + d = WanDMDPipeline.default_generation_params.fget(None) + assert d["height"] == 704 + assert d["width"] == 1280 + assert d["num_inference_steps"] == 3 + assert d["guidance_scale"] == 1.0 + assert d["num_frames"] == 121 + assert d["frame_rate"] == 24.0 + def test_flux_defaults(self): from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux import FluxPipeline @@ -518,6 +530,22 @@ def test_wan22_5b_default_params(self): assert params.frame_rate == 24.0 assert params.extra_params is None + def test_fastwan_default_params(self): + """FastWan 2.2 TI2V-5B (DMD distilled): 3-step, CFG-free, no extra params.""" + from tensorrt_llm._torch.visual_gen.models.wan.pipeline_fastwan import WanDMDPipeline + + vg = self._make_visual_gen( + WanDMDPipeline, _wan_mock(is_wan22_14b=False, is_wan22_5b=True, num_heads=24) + ) + params = vg.default_params + assert params.height == 704 + assert params.width == 1280 + assert params.num_inference_steps == 3 + assert params.guidance_scale == 1.0 + assert params.num_frames == 121 + assert params.frame_rate == 24.0 + assert params.extra_params is None + def test_wan21_default_params(self): """Wan 2.1 small-model returns 480p defaults with no extra params.""" from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan import WanPipeline