[None][feat] Add FastWan2.2 TI2V-5B DMD pipeline (3-step text-to-video)#16162
[None][feat] Add FastWan2.2 TI2V-5B DMD pipeline (3-step text-to-video)#16162Kambili wants to merge 15 commits into
Conversation
f37531e to
6e9eb69
Compare
|
Hi team, I'm adding FastWan 2.2 TI2V-5B support to VisualGen and want to check on the testing approach. Since HF Diffusers doesn't support FastWan, we can't use the standard diffusers-reference approach. Right now we're relying on two kinds of committed golden references:
We are noting this to confirm if this approach is acceptable. Thanks. |
📝 WalkthroughWalkthroughFastWan 2.2 TI2V-5B support adds a three-step DMD video pipeline, public exports and defaults, component and denoising tests, golden artifacts, FP8 checks, and LPIPS integration coverage. ChangesFastWan DMD pipeline
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant WanDMDPipeline
participant Transformer
participant VAE
Client->>WanDMDPipeline: forward(prompt, generation parameters)
WanDMDPipeline->>Transformer: run three DMD denoising steps
Transformer-->>WanDMDPipeline: predicted noise
WanDMDPipeline->>VAE: decode final latents
VAE-->>WanDMDPipeline: video tensor
WanDMDPipeline-->>Client: generation result
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.json (1)
3-5: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPin the trusted reference’s exact provenance.
Record the FastVideo commit/version, checkpoint revision, generation command, and video SHA-256. The committed golden approach is reasonable, but
"FastVideo VideoGenerator"alone is insufficient to reproduce or audit it.As per path instructions, coverage needs follow-up in
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.json.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.json` around lines 3 - 5, Update the golden metadata entry for the FastVideo trusted reference to include its exact FastVideo commit or version, checkpoint revision, generation command, and generated video SHA-256, while preserving the existing model, source, and prompt fields.Source: Path instructions
tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py (1)
17-22: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument reproducible golden generation and classify these as regression references.
The committed
.ptfiles are acceptable for regression detection, but because they were generated by this pipeline they do not independently establish component accuracy. Add a generator script ortests/unittest/_torch/visual_gen/golden/component_accuracy/README.mdrecording the exact checkpoint revision, software revision, command, tensor shapes/dtypes, and hashes. The independent FastVideo LPIPS golden can remain the end-to-end accuracy oracle.Coverage needs this provenance follow-up within the PR. As per path instructions, test feedback must identify concrete files and whether coverage is sufficient.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py` around lines 17 - 22, Document the provenance of the committed component-accuracy golden tensors under golden/component_accuracy/, either with a generator script or README.md. Record the exact checkpoint and software revisions, generation command, tensor shapes and dtypes, and file hashes; explicitly classify these files as regression references and retain FastVideo LPIPS as the independent end-to-end accuracy oracle.Source: Path instructions
tests/integration/defs/examples/visual_gen/test_visual_gen.py (1)
125-135: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the golden JSON as the single source of generation parameters.
These constants duplicate every relevant field in
fastwan_lpips_golden_video.json. Load that metadata for generation and threshold selection so the test cannot diverge from the artifact it evaluates.The LPIPS coverage is otherwise sufficient. As per path instructions, keep test metadata actionable and state coverage status.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py` around lines 125 - 135, Replace the duplicated FASTWAN_LPIPS_* generation and threshold constants in the visual generation test with values loaded from fastwan_lpips_golden_video.json. Use the golden metadata as the single source for prompt, dimensions, frame count, inference settings, seed, frame rate, and LPIPS threshold, while preserving the existing generation and evaluation flow; update the test metadata to state that LPIPS coverage is sufficient.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py`:
- Around line 72-81: Update the FastWan pipeline’s initialization or inference
validation around the visible parameters num_inference_steps, guidance_scale,
guidance_scale_2, and boundary_ratio to reject explicit values that conflict
with its fixed three-step, CFG-free contract instead of silently ignoring them;
preserve acceptance of omitted values and valid contract-compatible values. Add
focused coverage in test_fastwan_pipeline.py for each unsupported override and
the accepted defaults.
- Around line 58-78: Update WanDMDPipeline.forward to use built-in generic/union
syntax for its parameters where applicable and annotate its return type as
PipelineOutput. Add explicit tensor parameter annotations and a tensor return
annotation to WanDMDPipeline._denoise, preserving the existing behavior and
avoiding typing.Union/List where modern syntax applies.
In `@tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py`:
- Around line 27-41: Preserve the pre-existing TLLM_DISABLE_MPI value instead of
unconditionally removing it during teardown. In
tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py lines 27-41, update
the module-level _cleanup_mpi_env fixture to capture the value before assignment
and restore it afterward; apply the same save/restore pattern in
tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py lines 32-48.
- Around line 112-118: Update _run so the zero-noise mock returns a bfloat16
tensor instead of float32, matching the production dtype used by embeds and the
denoising path. Preserve the existing device, shape, and mocked _RANDN_PATH
behavior so the test exercises native bfloat16 re-noising arithmetic.
---
Nitpick comments:
In
`@tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.json`:
- Around line 3-5: Update the golden metadata entry for the FastVideo trusted
reference to include its exact FastVideo commit or version, checkpoint revision,
generation command, and generated video SHA-256, while preserving the existing
model, source, and prompt fields.
In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py`:
- Around line 125-135: Replace the duplicated FASTWAN_LPIPS_* generation and
threshold constants in the visual generation test with values loaded from
fastwan_lpips_golden_video.json. Use the golden metadata as the single source
for prompt, dimensions, frame count, inference settings, seed, frame rate, and
LPIPS threshold, while preserving the existing generation and evaluation flow;
update the test metadata to state that LPIPS coverage is sufficient.
In `@tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py`:
- Around line 17-22: Document the provenance of the committed component-accuracy
golden tensors under golden/component_accuracy/, either with a generator script
or README.md. Record the exact checkpoint and software revisions, generation
command, tensor shapes and dtypes, and file hashes; explicitly classify these
files as regression references and retain FastVideo LPIPS as the independent
end-to-end accuracy oracle.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a37fed43-8014-43ed-9874-275833e6f472
⛔ Files ignored due to path filters (1)
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zipis excluded by!**/*.zip
📒 Files selected for processing (13)
tensorrt_llm/_torch/visual_gen/models/__init__.pytensorrt_llm/_torch/visual_gen/models/wan/__init__.pytensorrt_llm/_torch/visual_gen/models/wan/defaults.pytensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.pytests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.jsontests/integration/defs/examples/visual_gen/test_visual_gen.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/visual_gen/golden/component_accuracy/text_encoder_reference.pttests/unittest/_torch/visual_gen/golden/component_accuracy/vae_input_latent.pttests/unittest/_torch/visual_gen/golden/component_accuracy/vae_reference.pttests/unittest/_torch/visual_gen/test_fastwan_dmd_math.pytests/unittest/_torch/visual_gen/test_fastwan_pipeline.pytests/unittest/_torch/visual_gen/test_visual_gen_params.py
zhenhuaw-me
left a comment
zhenhuaw-me
left a comment
There was a problem hiding this comment.
Thanks for the PR! I have some concerns about the uploaded torch tensors. If we don't have a decent solution for module level test, we may rely on the E2E tests for accuracy.
| version https://git-lfs.github.com/spec/v1 | ||
| oid sha256:6d38286a6ee60db331a89bb3563fcc0896309b6e21228d4867ab33990c5b47ed | ||
| size 14386487 | ||
| oid sha256:1e6e0a6be5b9e219df7b881b601ce88f2dd77c031c7fae75925119b907bc69c2 |
There was a problem hiding this comment.
(Out of the scope of this PR) @yibinl-nvidia I feel like we may need to improve this "zip" workflow in the long term - keep updating this zip is kind of strange.
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
LGTM — purely additive FastWan2.2 TI2V-5B DMD pipeline: new WanDMDPipeline(WanPipeline) subclass + registry/defaults wiring; existing WanPipeline and all other pipelines untouched, no public-API change.
|
Let's also add this model to the supported models in the doc |
… golden Introduces FastWanPipeline for the FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers checkpoint — a CFG-free DMD-distilled text+image-to-video model built on Wan 2.2. Changes: - pipeline_fastwan.py: FastWanPipeline with 3-step DMD denoising loop - wan/defaults.py, wan/__init__.py, models/__init__.py: registry wiring - test_fastwan_dmd_math.py: unit tests for DMD timestep math (no checkpoint) - test_fastwan_registry.py: unit tests for pipeline registration and defaults - test_fastwan_pipeline.py: component accuracy tests (text encoder, VAE) - test_visual_gen.py: LPIPS regression test against 121-frame golden video - visual_gen_lpips_golden_media.zip: updated with 121-frame FastWan golden - l0_b200.yml: register all FastWan tests in CI Signed-off-by: Nnamdi Wankwo <knwankwo@nvidia.com> Signed-off-by: Kambili Nwankwo <67324523+Kambili@users.noreply.github.com>
Remove stray 'NVIDIA#112' prefix from the SPDX copyright line. Signed-off-by: Nnamdi Wankwo <knwankwo@nvidia.com> Signed-off-by: Kambili Nwankwo <67324523+Kambili@users.noreply.github.com>
…nImageNotSupported removal Signed-off-by: knwankwo <knwankwo@nvidia.com>
Signed-off-by: knwankwo <knwankwo@nvidia.com>
Signed-off-by: knwankwo <knwankwo@nvidia.com>
Signed-off-by: knwankwo <knwankwo@nvidia.com>
…d param warnings - pipeline_fastwan.py: replace typing.Union/List/Optional with modern | syntax, add return type annotations to forward() and _denoise(), add docstrings to default_generation_params and forward(), replace hard ValueError for unsupported params (num_inference_steps, guidance_scale, guidance_scale_2, boundary_ratio) with logger.warning to match WanPipeline behaviour for the 5B model - fastwan_lpips_golden_video.json: remove FastVideo/ org prefix from model path - test_visual_gen.py: remove FastVideo/ org prefix from _lpips_model_path call - test_fastwan_dmd_math.py: use torch.zeros_like(latents) instead of explicit shape Signed-off-by: knwankwo <knwankwo@nvidia.com>
d31a6c3 to
5faa880
Compare
… fixes - Remove TestFastWanComponentAccuracy and golden/component_accuracy/ .pt reference tensors; E2E LPIPS test covers accuracy instead - Apply ruff-format fixes to pipeline_fastwan.py (logger.warning line collapsed, _denoise signature split across lines) Signed-off-by: knwankwo <knwankwo@nvidia.com>
|
/bot run |
|
PR_Github #60411 [ run ] triggered by Bot. Commit: |
|
PR_Github #60411 [ run ] completed with state |
Added FastWan 2.2 model to the visual generation documentation and updated the feature matrix accordingly. Signed-off-by: KAMBILI NWANKWO <67324523+Kambili@users.noreply.github.com>
|
/bot run |
|
PR_Github #60512 [ run ] triggered by Bot. Commit: |
|
PR_Github #60512 [ run ] completed with state
|
Simplified the description of FastVideo/FastWan2.2-TI2V-5B. Signed-off-by: KAMBILI NWANKWO <67324523+Kambili@users.noreply.github.com>
|
/bot run |
|
PR_Github #60530 [ run ] triggered by Bot. Commit: |
|
PR_Github #60530 [ run ] completed with state
|
zhenhuaw-me
left a comment
zhenhuaw-me
left a comment
There was a problem hiding this comment.
I take a look at the pipeline part. Could help to carefully refine the code? Thanks!
| 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}." | ||
| ) |
There was a problem hiding this comment.
I think you misunderstand the semantic of num_inference_steps
There was a problem hiding this comment.
same thing here as well as i noted below
| guidance_scale_2: float | None = None, | ||
| boundary_ratio: float | None = None, |
There was a problem hiding this comment.
What about disallow float type?
There was a problem hiding this comment.
its inherited from the wan base class even though its not used
| 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, |
There was a problem hiding this comment.
Should we update the signature?
There was a problem hiding this comment.
All four (num_inference_steps, guidance_scale, guidance_scale_2, boundary_ratio) are sent in by the inherited infer() but ignored by FastWan. I kept them explicit so all four are treated the same and the shape matches the base class; the warnings tell the user they're ignored. Can move them to **kwargs if you'd prefer that?
| version https://git-lfs.github.com/spec/v1 | ||
| oid sha256:6d38286a6ee60db331a89bb3563fcc0896309b6e21228d4867ab33990c5b47ed | ||
| size 14386487 | ||
| oid sha256:1e6e0a6be5b9e219df7b881b601ce88f2dd77c031c7fae75925119b907bc69c2 |
Co-authored-by: Zhenhua Wang <4936589+zhenhuaw-me@users.noreply.github.com> Signed-off-by: KAMBILI NWANKWO <67324523+Kambili@users.noreply.github.com>
Co-authored-by: Zhenhua Wang <4936589+zhenhuaw-me@users.noreply.github.com> Signed-off-by: KAMBILI NWANKWO <67324523+Kambili@users.noreply.github.com>
Signed-off-by: KAMBILI NWANKWO <67324523+Kambili@users.noreply.github.com>
|
/bot run |
|
PR_Github #61662 [ run ] triggered by Bot. Commit: |
|
PR_Github #61662 [ run ] completed with state |
8 participants
Description
Introduces
FastWanPipelinefor theFastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffuserscheckpoint — a CFG-free, DMD-distilled text-to-video model built on Wan 2.2 TI2V-5B.FastWan shares Wan 2.2 TI2V-5B's architecture exactly; only the distilled weights and the sampling recipe differ. The pipeline therefore subclasses
WanPipelineand overrides only the denoising loop, replacing the standard ~50-step diffusion with a 3-step DMD (Distribution Matching Distillation) schedule. This produces a 121×704×1280 video in ~3 denoising steps instead of ~50.Text-to-video only: FastWan2.2-5B's released weights are DMD-distilled for T2V, so the pipeline raises
NotImplementedErroron image input.Changes:
pipeline_fastwan.py: FastWanPipeline with the 3-step DMD denoising loopwan/defaults.py,wan/__init__.py,models/__init__.py: registry wiring for the WanDMDPipeline class and the HF checkpoint IDtest_fastwan_dmd_math.py: unit tests for the DMD timestep math (stubbed transformer, no checkpoint)test_visual_gen.py: LPIPS regression test against a 121-frame golden referencevisual_gen_lpips_golden_media.zip: adds the 121-frame FastWan golden videol0_b200.yml: registers the FastWan tests in CITest Coverage
test_fastwan_registry.py— verifies WanDMDPipeline resolves to FastWanPipeline, the HF ID is registered, defaults are correct (3 steps, CFG-free), and image input raises NotImplementedError. (No checkpoint; image input test requires GPU for the nvtx_range decorator.)test_fastwan_dmd_math.py— verifies the DMD loop math with a stubbed transformer: sigma = t/1000 giving {1.0, 0.757, 0.522}, correct re-noising on intermediate steps, x0 kept on the final step, and output dtype.test_visual_gen.py— end-to-end LPIPS regression: generates a 121-frame video with fixed settings and compares against the committed golden reference generated by FastVideo.l0_b200.yml— all of the above registered to run in CI on B200.Summary by CodeRabbit
New Features
Tests