-
Notifications
You must be signed in to change notification settings - Fork 43
Added support for WAN 2.2 model #281
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Added support for WAN 2.2 model #281
Conversation
0837512 to
56b1761
Compare
9c812c4 to
11d30fc
Compare
entrpn
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you also make sure Wan2.1 training and inference still runs without issues before this PR goes into main?
sample logs: |
Added 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 configuration for the WAN 2.2 model are added:
base_wan_27b.yml: New configuration file for the WAN 2.2 (27B) model.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: calls the correct pipeline (WAN 2.1 / WAN 2.2) depending on the model_name parameter in config files.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.wan_pipeline.py: The processing pipeline which supports both WAN 2.1 and WAN 2.2 by using separate base classeswan_checkpointer.py: Handles checkpoint loading/saving for both WAN 2.1 and WAN 2.2.wan_checkpointer_test.py: Tests for WAN 2.1 and WAN 2.2 checkpointerWe tested the above pipeline for both WAN 2.1 and WAN 2.2 and the benchmarking numbers can be found here: https://docs.google.com/spreadsheets/d/1GJqOoTUikAqLWedQG0Qu4YQbdQ5IeveHMqnOgi3d1tY/edit?pli=1&resourcekey=0-ywo0BVMZ-W1Vf7kCJhekXQ&gid=1962169097#gid=1962169097