Commit d96121b
committed
[MAX] Add Wan T2V diffusion pipeline with MoE support
## Summary
Add the Wan text-to-video (T2V) diffusion pipeline with MoE (Mixture of Experts) dual-transformer support.
## Description
- Implements the full Wan T2V pipeline: text encoding → latent preparation → denoising loop → VAE decode
- Supports **Wan 2.2 MoE models** (A14B) with dual transformers: high-noise expert for early steps, low-noise expert for later steps, with configurable boundary timestep
- Supports **Wan 2.1 single-transformer models** (14B) with the same code path
- LoRA support with automatic download from HuggingFace (e.g. Lightning turbo LoRAs for 4-step generation)
- Classifier-free guidance with batched forward pass (positive + negative in one call)
- On-device UniPC scheduler steps via compiled graphs — no Python-side numpy during denoising
- Architecture registration for `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, etc.
- Adds `guidance_scale_2` field to `VideoProviderOptions` for MoE boundary guidance control
- Minimal upstream changes: only `_weight_paths` storage in `DiffusionPipeline.__init__` and Wan registration in `pixel_tokenizer.py` / `registry.py`
## Dependencies
Depends on all previous PRs: modular#6298 (scheduler), modular#6299 (UMT5), modular#6300 (VAE), modular#6301 (transformer).
## Checklist
- [x] PR is small and focused
- [x] I ran `./bazelw run format` to format my changes
Assisted-by: Claude Code
Assisted-by: Claude Code1 parent 15eb852 commit d96121b
11 files changed
Lines changed: 1869 additions & 26 deletions
File tree
- max
- kernels/src/nn/conv
- python/max
- interfaces/provider_options/modality
- pipelines
- architectures
- wan
- core
- lib
- interfaces
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
5280 | 5280 | | |
5281 | 5281 | | |
5282 | 5282 | | |
5283 | | - | |
5284 | | - | |
5285 | | - | |
| 5283 | + | |
| 5284 | + | |
| 5285 | + | |
| 5286 | + | |
5286 | 5287 | | |
5287 | | - | |
5288 | | - | |
| 5288 | + | |
| 5289 | + | |
5289 | 5290 | | |
5290 | 5291 | | |
5291 | 5292 | | |
| |||
Lines changed: 6 additions & 0 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
66 | 66 | | |
67 | 67 | | |
68 | 68 | | |
| 69 | + | |
| 70 | + | |
| 71 | + | |
| 72 | + | |
| 73 | + | |
| 74 | + | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
82 | 82 | | |
83 | 83 | | |
84 | 84 | | |
| 85 | + | |
85 | 86 | | |
86 | 87 | | |
87 | 88 | | |
| |||
137 | 138 | | |
138 | 139 | | |
139 | 140 | | |
| 141 | + | |
| 142 | + | |
140 | 143 | | |
141 | 144 | | |
142 | 145 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
| 3 | + | |
| 4 | + | |
| 5 | + | |
| 6 | + | |
| 7 | + | |
| 8 | + | |
| 9 | + | |
| 10 | + | |
| 11 | + | |
| 12 | + | |
| 13 | + | |
| 14 | + | |
| 15 | + | |
| 16 | + | |
| 17 | + | |
| 18 | + | |
| 19 | + | |
| 20 | + | |
| 21 | + | |
| 22 | + | |
| 23 | + | |
| 24 | + | |
| 25 | + | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
| 3 | + | |
| 4 | + | |
| 5 | + | |
| 6 | + | |
| 7 | + | |
| 8 | + | |
| 9 | + | |
| 10 | + | |
| 11 | + | |
| 12 | + | |
| 13 | + | |
| 14 | + | |
| 15 | + | |
| 16 | + | |
| 17 | + | |
| 18 | + | |
| 19 | + | |
| 20 | + | |
| 21 | + | |
| 22 | + | |
| 23 | + | |
| 24 | + | |
| 25 | + | |
| 26 | + | |
| 27 | + | |
| 28 | + | |
| 29 | + | |
| 30 | + | |
| 31 | + | |
| 32 | + | |
| 33 | + | |
| 34 | + | |
| 35 | + | |
| 36 | + | |
| 37 | + | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
| 54 | + | |
| 55 | + | |
| 56 | + | |
| 57 | + | |
| 58 | + | |
| 59 | + | |
| 60 | + | |
| 61 | + | |
| 62 | + | |
| 63 | + | |
| 64 | + | |
| 65 | + | |
| 66 | + | |
| 67 | + | |
| 68 | + | |
| 69 | + | |
| 70 | + | |
| 71 | + | |
| 72 | + | |
| 73 | + | |
| 74 | + | |
| 75 | + | |
| 76 | + | |
| 77 | + | |
| 78 | + | |
| 79 | + | |
| 80 | + | |
| 81 | + | |
| 82 | + | |
| 83 | + | |
| 84 | + | |
| 85 | + | |
| 86 | + | |
| 87 | + | |
0 commit comments