Skip to content

Conversation

@prishajain1
Copy link
Collaborator

@prishajain1 prishajain1 commented Oct 29, 2025

Add support for WAN 2.2 model

This PR introduces support for the WAN 2.2 model and refactors the existing pipeline to dynamically handle multiple model versions.

New Model Support (WAN 2.2)

The core logic, configuration, and checkpointing for the WAN 2.2 model are added:

base_wan_27b.yml: New configuration file for the WAN 2.2 (27B) model.
wan_pipeline2_2.py: The dedicated processing pipeline for WAN 2.2.
wan_checkpointer2_2.py: Handles checkpoint loading/saving specific to the v2.2 architecture.
wan_checkpointer2_2_test.py: Tests for wan_checkpointer2_2.py

Refactoring for Multi-Model Support

To allow the system to select between WAN 2.1 and 2.2, several core files were modified:

generate_wan.py: Adds a --model_name flag. This flag is now used to dynamically load the correct configuration file (e.g., base_wan_14b.yml or base_wan_27b.yml).
wan_utils.py: Updated all transformer loading functions to include a subfolder parameter.
base_wan_14b.yml: Added the model_name parameter to conform to the new configuration standard.

Note on Future Work

As a follow-up to this PR, the separate wan_pipeline.py/wan_pipeline2_2.py and wan_checkpointer.py/wan_checkpointer2_2.py files will be merged. This current separation is a temporary step to land WAN 2.2 support.

We tested the above pipeline for both WAN 2.1 and WAN 2.2 with exact same configurations, and the numbers for both the cases are similar.

@github-actions
Copy link

@prishajain1 prishajain1 requested a review from Perseus14 October 29, 2025 08:15
@prishajain1 prishajain1 changed the title changes for WAN 2.2 Support for WAN 2.2 Oct 29, 2025
@prishajain1 prishajain1 changed the title Support for WAN 2.2 Add support for WAN 2.2 Oct 29, 2025
@prishajain1 prishajain1 force-pushed the prishajain/wan2.2 branch 3 times, most recently from 6c99a8d to 13e3cbc Compare October 29, 2025 20:39
@prishajain1 prishajain1 requested a review from entrpn October 29, 2025 20:44
@prishajain1 prishajain1 force-pushed the prishajain/wan2.2 branch 2 times, most recently from 714ed15 to f414753 Compare October 30, 2025 09:23
@entrpn entrpn merged commit d843dc0 into main Nov 4, 2025
2 of 3 checks passed
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.

3 participants