Skip to content

Add Wan2.2 T2V-A14B diffusion training presets#861

Draft
zirui wants to merge 4 commits into
dev/zirui/wanfrom
dev/zirui/wan-A2.2-A14B
Draft

Add Wan2.2 T2V-A14B diffusion training presets#861
zirui wants to merge 4 commits into
dev/zirui/wanfrom
dev/zirui/wan-A2.2-A14B

Conversation

@zirui

@zirui zirui commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Adds native diffusion presets for Wan2.2-T2V-A14B high-noise and low-noise SFT stages, aligned with NeMo AutoModel's two-transformer recipe (transformer / transformer_2, boundary ratio 0.875, flow shift 3.0).
  • Teaches the Wan model loader to select Diffusers transformer subfolders and remap WanTransformer3DModel state dict keys into Primus WanDiT keys.
  • Extends scheduler overrides and adds focused tests for scheduler mapping, Diffusers subfolder selection, and full A14B key remapping.

Base / merge plan

This draft PR intentionally targets dev/zirui/wan (PR #779) to keep the review diff small. After #779 lands in main, this branch should be rebased/retargeted to main before final review.

Test plan

  • python -m pytest tests/unit_tests/backends/diffusion
  • pre-commit run --all-files
  • Prepare-hook validation for both A14B examples with local /mnt/shared/zirui/models/Wan2.2-T2V-A14B-Diffusers and Wan2.1-14B encoder assets
  • Full key-level smoke check: Diffusers A14B transformer keys remap exactly to a meta-initialized Primus A14B WanDiT state dict
  • Add benchmark tests and align loss and performance with AutoModel.

Notes

  • The A14B DiT loads from the Diffusers pipeline directory, but the native Primus Wan path still expects frozen UMT5/VAE encoder assets in original Wan .pth format. The presets default those to compatible Wan2.1-14B assets while using the A14B Diffusers tokenizer.

zirui and others added 4 commits July 7, 2026 11:32
Support Diffusers two-transformer Wan2.2 T2V-A14B DiT loading, add AutoModel-aligned high/low-noise SFT configs, and cover scheduler/key mapping with focused tests.

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant