Wan-DMD

4-step distillation of Wan2.1-T2V — decoupled-DMD + GAN on real data, for both 1.3B and 14B.

Wan-DMD turns the pretrained Wan2.1 text-to-video models into 4-step generators using decoupled distribution-matching distillation (decoupled-DMD) with an adversarial (GAN) term — trained entirely on real video data. One small codebase, one recipe, both the 1.3B and 14B checkpoints.

Highlights

• 🎞️ Distilled from real videos, not simulated trajectories. The student's training states are real-video latents run through the forward (noising) process — no teacher-simulated ODE trajectories and no self-generated rollout to synthesize first. DMD/DMD2's distribution-matching term is data-free (prompts + teacher score only), and CausVid/AccVideo precompute a dataset of teacher-simulated ODE pairs (DMD2, arXiv:2405.14867; CausVid, arXiv:2412.07772). Wan-DMD skips that synthesis cost, anchors training to real data, and can (hypothesis) adapt more directly to a specific target video set.
• 🥊 GAN on real-data latents. On top of the DMD terms, an adversarial loss discriminates student latents against the same real-video latents — DMD2-style, with no ODE-pair initialization.
• 🧪 Decoupled-DMD (CA + DM). The generator loss combines a CFG-Augmentation real-score term (CA) with a distribution-matching term (DM) between separate real-score and fake-score networks. Separate real/fake networks follow DMD2 (arXiv:2405.14867); the CA/DM decomposition follows Decoupled DMD (arXiv:2511.22677).
• 🧊 Frozen real score as the GAN backbone. The discriminator reads features from the frozen real score, so no spurious gradients leak into the trainable networks — a clean, stable video-GAN setup. The conv discriminator design is from Flash-Diffusion (arXiv:2406.02347).
• 🧠 1.3B and 14B in one codebase. FSDP ZeRO-3 + CPU parameter offload + T5 text-encoder CPU offload + gradient checkpointing bring 14B distillation down to modest hardware — see the measured memory table below.

Where this fits

Few-step Wan distillation is an active area, but most open training recipes target the 1.3B model (e.g. CausVid, Self-Forcing), and the public 14B artifacts tend to be distilled weights or LoRA adapters rather than a full training recipe. Wan-DMD is a compact, readable codebase that ships open full-parameter training for both sizes. Its main design choice is to distill from real video data — real latents run through the forward process supply the student's training states — instead of teacher-simulated ODE trajectories or self-generated rollouts, which removes the data-synthesis stage and keeps the pipeline single-stage. It aims to be easy to read, and adapt.

The method is documented in docs/method.md.

Visual examples

video1785747558305.mp4

Todo

• Phase-0 consistency distillation. Use consistency distillation (CD) instead of DMD at the highest-noise phase. Expected to improve motion dynamics and sample diversity.
• LoRA fine-tuning. Add a LoRA path for the student so the 4-step generator can be adapted cheaply without full-parameter training.

---

Install

git clone <this repo> && cd Wan-DMD
# install a CUDA build of torch matching your driver first, then:
pip install -r requirements.txt
pip install -e .
# strongly recommended for speed/memory:
pip install flash-attn --no-build-isolation

Weights

Wan-DMD builds on the official Wan2.1-T2V base weights (the T5 text encoder and VAE are bundled inside the 1.3B release). Download and point to them:

• Wan2.1-T2V-1.3B: https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B
• Wan2.1-T2V-14B: https://huggingface.co/Wan-AI/Wan2.1-T2V-14B

export WAN_1_3B_PATH=/abs/path/to/Wan2.1-T2V-1.3B   # also provides T5 + VAE
export WAN_14B_PATH=/abs/path/to/Wan2.1-T2V-14B

(Individual paths can be overridden — see wan_dmd/paths.py.)

Data

Training reads raw videos + captions from a CSV (videoID,caption); clips are VAE-encoded on the fly. See examples/data/README.md for the format. (Optional latent pre-encoding: scripts/encode_latents.py.)

Train

# 1.3B — single node, 8 GPUs
torchrun --nnodes 1 --nproc_per_node 8 scripts/train.py --config configs/wan1.3b.yaml

# 14B — multi-node, e.g. 2x8 = 16x H20
torchrun --nnodes 2 --nproc_per_node 8 --rdzv_backend=c10d \
         --rdzv_endpoint <MASTER_IP>:29500 \
         scripts/train.py --config configs/wan14b.yaml

Outputs (logs, TensorBoard, checkpoints, validation clips) go to OUTPUT_DIR/<timestamp>-<EXP_NAME>/. Resume by setting RESUME_CKPT. Set VERBOSE: true for per-step loss printing and a banner when the GAN turns on. Validate a config's wiring (and the few-step schedule) without GPUs/weights:

python scripts/train.py --config configs/wan14b.yaml --dry-run

Data: on-the-fly video vs pre-encoded latents

DATA_MODE: video (default) decodes raw videos and VAE-encodes them each step. DATA_MODE: latent reads pre-encoded latents (faster; skips decode + VAE):

python scripts/encode_latents.py --csv train.csv --video-root /videos \
    --out-dir /latents --sequence-length 81 --resolution 480 832
# then set in the config:  DATA_MODE: latent   LATENT_DIR: [/latents]

Memory switches

Three independent switches trade speed for memory: CPU_OFFLOAD (FSDP CPU param offload), SHARD_REAL (shard the frozen real score), TEXT_ENCODER_OFFLOAD (keep T5 on CPU). Measured peak GPU memory (4×H200, FlashAttention-2, batch_size=1):

model	CPU_OFFLOAD	SHARD_REAL	TE_OFFLOAD	peak/GPU
1.3B	✗	✗	✗	~42 GB
1.3B	✓	✗	✗	~33 GB
1.3B	✓	✓	✓	~16 GB
14B	✓	✗	✗	~97 GB
14B	✓	✓	✓	~77 GB

Inference (4-step, prompt-only)

python scripts/inference.py --config configs/inference.yaml \
    --ckpt /path/to/student.pt --prompts prompts.txt --sampler sde

Method in one paragraph

Three networks (decoupled-DMD): a student 4-step generator, a fake score that tracks the student's distribution, and a frozen real score (the pretrained Wan with CFG). Each iteration runs three optimizer sub-steps — (0) the DMD generator loss (a CFG-Augmentation real-score term plus the distribution-matching term -(fake − real), each normalized by a per-sample magnitude weight) and, once enabled, the GAN generator loss; (1) the critic flow-matching loss that keeps the fake score on the student distribution; (2) the LSGAN discriminator loss over real-data vs student latents (features from the frozen real score). The student denoises in 4 steps at timesteps [1000, 934, 822, 587]. See docs/method.md.

References

• Decoupled DMD (CA/DM decomposition) — https://arxiv.org/abs/2511.22677
• DMD2 (separate real/fake networks, GAN without ODE init) — https://arxiv.org/abs/2405.14867
• CausVid (DMD for video, teacher-simulated ODE-pair init) — https://arxiv.org/abs/2412.07772
• Flash Diffusion (GAN discriminator design) — https://arxiv.org/abs/2406.02347
• Wan2.1 — https://github.com/Wan-Video/Wan2.1

License

Apache-2.0. Wan2.1 base weights are subject to their own license.