Zengtudor/Wan2.1
1
0
Fork 0
mirror of https://github.com/Wan-Video/Wan2.1.git synced 2025-11-04 06:15:17 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

beta version

Browse Source
This commit is contained in:
DeepBeepMeep 2025-03-02 04:05:49 +01:00
parent 6797c48002
commit 18940291d4
17 changed files with 1964 additions and 729 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

208
LICENSE.txt
View File

@ -1,201 +1,17 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
FREE for Non Commercial USE
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
You are free to:
- Share — copy and redistribute the material in any medium or format
- Adapt — remix, transform, and build upon the material
The licensor cannot revoke these freedoms as long as you follow the license terms.
1. Definitions.
Under the following terms:
- Attribution — You must give appropriate credit , provide a link to the license, and indicate if changes were made . You may do so in any reasonable manner, but not in any way that suggests the licensor endorses you or your use.
NonCommercial — You may not use the material for commercial purposes .
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
- No additional restrictions — You may not apply legal terms or technological measures that legally restrict others from doing anything the license permits.
Notices:
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
- You do not have to comply with the license for elements of the material in the public domain or where your use is permitted by an applicable exception or limitation .
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
No warranties are given. The license may not give you all of the permissions necessary for your intended use. For example, other rights such as publicity, privacy, or moral rights may limit how you use the material.

465
README.md
View File

@ -27,378 +27,159 @@ In this repository, we present **Wan2.1**, a comprehensive and open suite of vid
## 🔥 Latest News!!
* Mar 03, 2025: Wan2.1GP DeepBeepMeep out of this World version ! Reduced memory consumption by 2, with possiblity to generate more than 10s of video at 720p
* Feb 25, 2025: 👋 We've released the inference code and weights of Wan2.1.
* Feb 27, 2025: 👋 Wan2.1 has been integrated into [ComfyUI](https://comfyanonymous.github.io/ComfyUI_examples/wan/). Enjoy!
## 📑 Todo List
- Wan2.1 Text-to-Video
- [x] Multi-GPU Inference code of the 14B and 1.3B models
- [x] Checkpoints of the 14B and 1.3B models
- [x] Gradio demo
- [x] ComfyUI integration
- [ ] Diffusers integration
- Wan2.1 Image-to-Video
- [x] Multi-GPU Inference code of the 14B model
- [x] Checkpoints of the 14B model
- [x] Gradio demo
- [X] ComfyUI integration
- [ ] Diffusers integration
## Features
*GPU Poor version by **DeepBeepMeep**. This great video generator can now run smoothly on any GPU.*
This version has the following improvements over the original Hunyuan Video model:
- Reduce greatly the RAM requirements and VRAM requirements
- Much faster thanks to compilation and fast loading / unloading
- 5 profiles in order to able to run the model at a decent speed on a low end consumer config (32 GB of RAM and 12 VRAM) and to run it at a very good speed on a high end consumer config (48 GB of RAM and 24 GB of VRAM)
- Autodownloading of the needed model files
- Improved gradio interface with progression bar and more options
- Multiples prompts / multiple generations per prompt
- Support multiple pretrained Loras with 32 GB of RAM or less
- Switch easily between Hunyuan and Fast Hunyuan models and quantized / non quantized models
- Much simpler installation
This fork by DeepBeepMeep is an integration of the mmpg module on the gradio_server.py.
It is an illustration on how one can set up on an existing model some fast and properly working CPU offloading with changing only a few lines of code in the core model.
For more information on how to use the mmpg module, please go to: https://github.com/deepbeepmeep/mmgp
You will find the original Hunyuan Video repository here: https://github.com/deepbeepmeep/Wan2GP
## Installation Guide for Linux and Windows
We provide an `environment.yml` file for setting up a Conda environment.
Conda's installation instructions are available [here](https://docs.anaconda.com/free/miniconda/index.html).
This app has been tested on Python 3.10 / 2.6.0 / Cuda 12.4.\
```shell
# 1 - conda. Prepare and activate a conda environment
conda env create -f environment.yml
conda activate Wan2
# OR
# 1 - venv. Alternatively create a python 3.10 venv and then do the following
pip install torch==2.6.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu124
# 2. Install pip dependencies
python -m pip install -r requirements.txt
# 3.1 optional Sage attention support (30% faster, easy to install on Linux but much harder on Windows)
python -m pip install sageattention==1.0.6
# or for Sage Attention 2 (40% faster, sorry only manual compilation for the moment)
git pull https://github.com/thu-ml/SageAttention
cd sageattention
pip install -e .
# 3.2 optional Flash attention support (easy to install on Linux but much harder on Windows)
python -m pip install flash-attn==2.7.2.post1
## Quickstart
#### Installation
Clone the repo:
```
git clone https://github.com/Wan-Video/Wan2.1.git
cd Wan2.1
```
Install dependencies:
Note that *Flash attention* and *Sage attention* are quite complex to install on Windows but offers a better memory management (and consequently longer videos) than the default *sdpa attention*.
Likewise *Pytorch Compilation* will work on Windows only if you manage to install Triton. It is quite a complex process (see below for links).
### Ready to use python wheels for Windows users
I provide here links to simplify the installation for Windows users with Python 3.10 / Pytorch 2.51 / Cuda 12.4. As I am not hosting these files I won't be able to provide support neither guarantee they do what they should do.
- Triton attention (needed for *pytorch compilation* and *Sage attention*)
```
# Ensure torch >= 2.4.0
pip install -r requirements.txt
pip install https://github.com/woct0rdho/triton-windows/releases/download/v3.2.0-windows.post9/triton-3.2.0-cp310-cp310-win_amd64.whl # triton for pytorch 2.6.0
```
- Sage attention
```
pip install https://github.com/deepbeepmeep/SageAttention/raw/refs/heads/main/releases/sageattention-2.1.0-cp310-cp310-win_amd64.whl # for pytorch 2.6.0 (experimental, if it works, otherwise you you will need to install and compile manually, see above)
```
## Run the application
### Run a Gradio Server on port 7860 (recommended)
```bash
python gradio_server.py
```
#### Model Download
### Loras support
| Models | Download Link | Notes |
| --------------|-------------------------------------------------------------------------------|-------------------------------|
| T2V-14B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B) 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B) | Supports both 480P and 720P
| I2V-14B-720P | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P) 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P) | Supports 720P
| I2V-14B-480P | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P) 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P) | Supports 480P
| T2V-1.3B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B) 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B) | Supports 480P
-- Ready to be used but theorical as no lora for Wan have been released as today.
> 💡Note: The 1.3B model is capable of generating videos at 720P resolution. However, due to limited training at this resolution, the results are generally less stable compared to 480P. For optimal performance, we recommend using 480P resolution.
Every lora stored in the subfoler 'loras' will be automatically loaded. You will be then able to activate / desactive any of them when running the application.
For each activated Lora, you may specify a *multiplier* that is one float number that corresponds to its weight (default is 1.0), alternatively you may specify a list of floats multipliers separated by a "," that gives the evolution of this Lora's multiplier over the steps. For instance let's assume there are 30 denoising steps and the multiplier is *0.9,0.8,0.7* then for the steps ranges 0-9, 10-19 and 20-29 the Lora multiplier will be respectively 0.9, 0.8 and 0.7.
Download models using huggingface-cli:
```
pip install "huggingface_hub[cli]"
huggingface-cli download Wan-AI/Wan2.1-T2V-14B --local-dir ./Wan2.1-T2V-14B
You can edit, save or delete Loras presets (combinations of loras with their corresponding multipliers) directly from the gradio interface. Each preset, is a file with ".lset" extension stored in the loras directory and can be shared with other users
Then you can pre activate loras corresponding to a preset when launching the gradio server:
```bash
python gradio_server.py --lora-preset mylorapreset.lset # where 'mylorapreset.lset' is a preset stored in the 'loras' folder
```
Download models using modelscope-cli:
```
pip install modelscope
modelscope download Wan-AI/Wan2.1-T2V-14B --local_dir ./Wan2.1-T2V-14B
```
#### Run Text-to-Video Generation
Please note that command line parameters *--lora-weight* and *--lora-multiplier* have been deprecated since they are redundant with presets.
This repository supports two Text-to-Video models (1.3B and 14B) and two resolutions (480P and 720P). The parameters and configurations for these models are as follows:
You will find prebuilt Loras on https://civitai.com/ or you will be able to build them with tools such as kohya or onetrainer.
<table>
<thead>
<tr>
<th rowspan="2">Task</th>
<th colspan="2">Resolution</th>
<th rowspan="2">Model</th>
</tr>
<tr>
<th>480P</th>
<th>720P</th>
</tr>
</thead>
<tbody>
<tr>
<td>t2v-14B</td>
<td style="color: green;">✔️</td>
<td style="color: green;">✔️</td>
<td>Wan2.1-T2V-14B</td>
</tr>
<tr>
<td>t2v-1.3B</td>
<td style="color: green;">✔️</td>
<td style="color: red;">❌</td>
<td>Wan2.1-T2V-1.3B</td>
</tr>
</tbody>
</table>
### Command line parameters for Gradio Server
--profile no : default (4) : no of profile between 1 and 5\
--quantize-transformer bool: (default True) : enable / disable on the fly transformer quantization\
--lora-dir path : Path of directory that contains Loras in diffusers / safetensor format\
--lora-preset preset : name of preset gile (without the extension) to preload
--verbose level : default (1) : level of information between 0 and 2\
--server-port portno : default (7860) : Gradio port no\
--server-name name : default (0.0.0.0) : Gradio server name\
--open-browser : open automatically Browser when launching Gradio Server\
--compile : turn on pytorch compilation\
--attention mode: force attention mode among, sdpa, flash, sage, sage2\
##### (1) Without Prompt Extention
### Profiles (for power users only)
You can choose between 5 profiles, these will try to leverage the most your hardware, but have little impact for HunyuanVideo GP:
- HighRAM_HighVRAM (1): the fastest well suited for a RTX 3090 / RTX 4090 but consumes much more VRAM, adapted for fast shorter video
- HighRAM_LowVRAM (2): a bit slower, better suited for RTX 3070/3080/4070/4080 or for RTX 3090 / RTX 4090 with large pictures batches or long videos
- LowRAM_HighVRAM (3): adapted for RTX 3090 / RTX 4090 with limited RAM but at the cost of VRAM (shorter videos)
- LowRAM_LowVRAM (4): if you have little VRAM or want to generate longer videos
- VerylowRAM_LowVRAM (5): at least 24 GB of RAM and 10 GB of VRAM : if you don't have much it won't be fast but maybe it will work
To facilitate implementation, we will start with a basic version of the inference process that skips the [prompt extension](#2-using-prompt-extention) step.
Profile 2 (High RAM) and 4 (Low RAM)are the most recommended profiles since they are versatile (support for long videos for a slight performance cost).\
However, a safe approach is to start from profile 5 (default profile) and then go down progressively to profile 4 and then to profile 2 as long as the app remains responsive or doesn't trigger any out of memory error.
- Single-GPU inference
### Other Models for the GPU Poor
```
python generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
```
- HuanyuanVideoGP: https://github.com/deepbeepmeep/HunyuanVideoGP :\
One of the best open source Text to Video generator
If you encounter OOM (Out-of-Memory) issues, you can use the `--offload_model True` and `--t5_cpu` options to reduce GPU memory usage. For example, on an RTX 4090 GPU:
- Hunyuan3D-2GP: https://github.com/deepbeepmeep/Hunyuan3D-2GP :\
A great image to 3D and text to 3D tool by the Tencent team. Thanks to mmgp it can run with less than 6 GB of VRAM
```
python generate.py --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --offload_model True --t5_cpu --sample_shift 8 --sample_guide_scale 6 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
```
- FluxFillGP: https://github.com/deepbeepmeep/FluxFillGP :\
One of the best inpainting / outpainting tools based on Flux that can run with less than 12 GB of VRAM.
> 💡Note: If you are using the `T2V-1.3B` model, we recommend setting the parameter `--sample_guide_scale 6`. The `--sample_shift parameter` can be adjusted within the range of 8 to 12 based on the performance.
- Cosmos1GP: https://github.com/deepbeepmeep/Cosmos1GP :\
This application include two models: a text to world generator and a image / video to world (probably the best open source image to video generator).
- OminiControlGP: https://github.com/deepbeepmeep/OminiControlGP :\
A Flux derived application very powerful that can be used to transfer an object of your choice in a prompted scene. With mmgp you can run it with only 6 GB of VRAM.
- Multi-GPU inference using FSDP + xDiT USP
- YuE GP: https://github.com/deepbeepmeep/YuEGP :\
A great song generator (instruments + singer's voice) based on prompted Lyrics and a genre description. Thanks to mmgp you can run it with less than 10 GB of VRAM without waiting forever.
```
pip install "xfuser>=0.4.1"
torchrun --nproc_per_node=8 generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --dit_fsdp --t5_fsdp --ulysses_size 8 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
```
##### (2) Using Prompt Extention
Extending the prompts can effectively enrich the details in the generated videos, further enhancing the video quality. Therefore, we recommend enabling prompt extension. We provide the following two methods for prompt extension:
- Use the Dashscope API for extension.
- Apply for a `dashscope.api_key` in advance ([EN](https://www.alibabacloud.com/help/en/model-studio/getting-started/first-api-call-to-qwen) | [CN](https://help.aliyun.com/zh/model-studio/getting-started/first-api-call-to-qwen)).
- Configure the environment variable `DASH_API_KEY` to specify the Dashscope API key. For users of Alibaba Cloud's international site, you also need to set the environment variable `DASH_API_URL` to 'https://dashscope-intl.aliyuncs.com/api/v1'. For more detailed instructions, please refer to the [dashscope document](https://www.alibabacloud.com/help/en/model-studio/developer-reference/use-qwen-by-calling-api?spm=a2c63.p38356.0.i1).
- Use the `qwen-plus` model for text-to-video tasks and `qwen-vl-max` for image-to-video tasks.
- You can modify the model used for extension with the parameter `--prompt_extend_model`. For example:
```
DASH_API_KEY=your_key python generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" --use_prompt_extend --prompt_extend_method 'dashscope' --prompt_extend_target_lang 'ch'
```
- Using a local model for extension.
- By default, the Qwen model on HuggingFace is used for this extension. Users can choose Qwen models or other models based on the available GPU memory size.
- For text-to-video tasks, you can use models like `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-7B-Instruct` and `Qwen/Qwen2.5-3B-Instruct`.
- For image-to-video tasks, you can use models like `Qwen/Qwen2.5-VL-7B-Instruct` and `Qwen/Qwen2.5-VL-3B-Instruct`.
- Larger models generally provide better extension results but require more GPU memory.
- You can modify the model used for extension with the parameter `--prompt_extend_model` , allowing you to specify either a local model path or a Hugging Face model. For example:
```
python generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" --use_prompt_extend --prompt_extend_method 'local_qwen' --prompt_extend_target_lang 'ch'
```
##### (3) Runing local gradio
```
cd gradio
# if one uses dashscope’s API for prompt extension
DASH_API_KEY=your_key python t2v_14B_singleGPU.py --prompt_extend_method 'dashscope' --ckpt_dir ./Wan2.1-T2V-14B
# if one uses a local model for prompt extension
python t2v_14B_singleGPU.py --prompt_extend_method 'local_qwen' --ckpt_dir ./Wan2.1-T2V-14B
```
#### Run Image-to-Video Generation
Similar to Text-to-Video, Image-to-Video is also divided into processes with and without the prompt extension step. The specific parameters and their corresponding settings are as follows:
<table>
<thead>
<tr>
<th rowspan="2">Task</th>
<th colspan="2">Resolution</th>
<th rowspan="2">Model</th>
</tr>
<tr>
<th>480P</th>
<th>720P</th>
</tr>
</thead>
<tbody>
<tr>
<td>i2v-14B</td>
<td style="color: green;">❌</td>
<td style="color: green;">✔️</td>
<td>Wan2.1-I2V-14B-720P</td>
</tr>
<tr>
<td>i2v-14B</td>
<td style="color: green;">✔️</td>
<td style="color: red;">❌</td>
<td>Wan2.1-T2V-14B-480P</td>
</tr>
</tbody>
</table>
##### (1) Without Prompt Extention
- Single-GPU inference
```
python generate.py --task i2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-I2V-14B-720P --image examples/i2v_input.JPG --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
```
> 💡For the Image-to-Video task, the `size` parameter represents the area of the generated video, with the aspect ratio following that of the original input image.
- Multi-GPU inference using FSDP + xDiT USP
```
pip install "xfuser>=0.4.1"
torchrun --nproc_per_node=8 generate.py --task i2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-I2V-14B-720P --image examples/i2v_input.JPG --dit_fsdp --t5_fsdp --ulysses_size 8 --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
```
##### (2) Using Prompt Extention
The process of prompt extension can be referenced [here](#2-using-prompt-extention).
Run with local prompt extention using `Qwen/Qwen2.5-VL-7B-Instruct`:
```
python generate.py --task i2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-I2V-14B-720P --image examples/i2v_input.JPG --use_prompt_extend --prompt_extend_model Qwen/Qwen2.5-VL-7B-Instruct --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
```
Run with remote prompt extention using `dashscope`:
```
DASH_API_KEY=your_key python generate.py --task i2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-I2V-14B-720P --image examples/i2v_input.JPG --use_prompt_extend --prompt_extend_method 'dashscope' --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
```
##### (3) Runing local gradio
```
cd gradio
# if one only uses 480P model in gradio
DASH_API_KEY=your_key python i2v_14B_singleGPU.py --prompt_extend_method 'dashscope' --ckpt_dir_480p ./Wan2.1-I2V-14B-480P
# if one only uses 720P model in gradio
DASH_API_KEY=your_key python i2v_14B_singleGPU.py --prompt_extend_method 'dashscope' --ckpt_dir_720p ./Wan2.1-I2V-14B-720P
# if one uses both 480P and 720P models in gradio
DASH_API_KEY=your_key python i2v_14B_singleGPU.py --prompt_extend_method 'dashscope' --ckpt_dir_480p ./Wan2.1-I2V-14B-480P --ckpt_dir_720p ./Wan2.1-I2V-14B-720P
```
#### Run Text-to-Image Generation
Wan2.1 is a unified model for both image and video generation. Since it was trained on both types of data, it can also generate images. The command for generating images is similar to video generation, as follows:
##### (1) Without Prompt Extention
- Single-GPU inference
```
python generate.py --task t2i-14B --size 1024*1024 --ckpt_dir ./Wan2.1-T2V-14B --prompt '一个朴素端庄的美人'
```
- Multi-GPU inference using FSDP + xDiT USP
```
torchrun --nproc_per_node=8 generate.py --dit_fsdp --t5_fsdp --ulysses_size 8 --base_seed 0 --frame_num 1 --task t2i-14B --size 1024*1024 --prompt '一个朴素端庄的美人' --ckpt_dir ./Wan2.1-T2V-14B
```
##### (2) With Prompt Extention
- Single-GPU inference
```
python generate.py --task t2i-14B --size 1024*1024 --ckpt_dir ./Wan2.1-T2V-14B --prompt '一个朴素端庄的美人' --use_prompt_extend
```
- Multi-GPU inference using FSDP + xDiT USP
```
torchrun --nproc_per_node=8 generate.py --dit_fsdp --t5_fsdp --ulysses_size 8 --base_seed 0 --frame_num 1 --task t2i-14B --size 1024*1024 --ckpt_dir ./Wan2.1-T2V-14B --prompt '一个朴素端庄的美人' --use_prompt_extend
```
## Manual Evaluation
##### (1) Text-to-Video Evaluation
Through manual evaluation, the results generated after prompt extension are superior to those from both closed-source and open-source models.
<div align="center">
<img src="assets/t2v_res.jpg" alt="" style="width: 80%;" />
</div>
##### (2) Image-to-Video Evaluation
We also conducted extensive manual evaluations to evaluate the performance of the Image-to-Video model, and the results are presented in the table below. The results clearly indicate that **Wan2.1** outperforms both closed-source and open-source models.
<div align="center">
<img src="assets/i2v_res.png" alt="" style="width: 80%;" />
</div>
## Computational Efficiency on Different GPUs
We test the computational efficiency of different **Wan2.1** models on different GPUs in the following table. The results are presented in the format: **Total time (s) / peak GPU memory (GB)**.
<div align="center">
<img src="assets/comp_effic.png" alt="" style="width: 80%;" />
</div>
> The parameter settings for the tests presented in this table are as follows:
> (1) For the 1.3B model on 8 GPUs, set `--ring_size 8` and `--ulysses_size 1`;
> (2) For the 14B model on 1 GPU, use `--offload_model True`;
> (3) For the 1.3B model on a single 4090 GPU, set `--offload_model True --t5_cpu`;
> (4) For all testings, no prompt extension was applied, meaning `--use_prompt_extend` was not enabled.
> 💡Note: T2V-14B is slower than I2V-14B because the former samples 50 steps while the latter uses 40 steps.
## Community Contributions
- [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio) provides more support for **Wan2.1**, including video-to-video, FP8 quantization, VRAM optimization, LoRA training, and more. Please refer to [their examples](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/wanvideo).
-------
## Introduction of Wan2.1
**Wan2.1** is designed on the mainstream diffusion transformer paradigm, achieving significant advancements in generative capabilities through a series of innovations. These include our novel spatio-temporal variational autoencoder (VAE), scalable training strategies, large-scale data construction, and automated evaluation metrics. Collectively, these contributions enhance the model’s performance and versatility.
##### (1) 3D Variational Autoencoders
We propose a novel 3D causal VAE architecture, termed **Wan-VAE** specifically designed for video generation. By combining multiple strategies, we improve spatio-temporal compression, reduce memory usage, and ensure temporal causality. **Wan-VAE** demonstrates significant advantages in performance efficiency compared to other open-source VAEs. Furthermore, our **Wan-VAE** can encode and decode unlimited-length 1080P videos without losing historical temporal information, making it particularly well-suited for video generation tasks.
<div align="center">
<img src="assets/video_vae_res.jpg" alt="" style="width: 80%;" />
</div>
##### (2) Video Diffusion DiT
**Wan2.1** is designed using the Flow Matching framework within the paradigm of mainstream Diffusion Transformers. Our model's architecture uses the T5 Encoder to encode multilingual text input, with cross-attention in each transformer block embedding the text into the model structure. Additionally, we employ an MLP with a Linear layer and a SiLU layer to process the input time embeddings and predict six modulation parameters individually. This MLP is shared across all transformer blocks, with each block learning a distinct set of biases. Our experimental findings reveal a significant performance improvement with this approach at the same parameter scale.
<div align="center">
<img src="assets/video_dit_arch.jpg" alt="" style="width: 80%;" />
</div>
| Model | Dimension | Input Dimension | Output Dimension | Feedforward Dimension | Frequency Dimension | Number of Heads | Number of Layers |
|--------|-----------|-----------------|------------------|-----------------------|---------------------|-----------------|------------------|
| 1.3B | 1536 | 16 | 16 | 8960 | 256 | 12 | 30 |
| 14B | 5120 | 16 | 16 | 13824 | 256 | 40 | 40 |
##### Data
We curated and deduplicated a candidate dataset comprising a vast amount of image and video data. During the data curation process, we designed a four-step data cleaning process, focusing on fundamental dimensions, visual quality and motion quality. Through the robust data processing pipeline, we can easily obtain high-quality, diverse, and large-scale training sets of images and videos.
![figure1](assets/data_for_diff_stage.jpg "figure1")
##### Comparisons to SOTA
We compared **Wan2.1** with leading open-source and closed-source models to evaluate the performace. Using our carefully designed set of 1,035 internal prompts, we tested across 14 major dimensions and 26 sub-dimensions. We then compute the total score by performing a weighted calculation on the scores of each dimension, utilizing weights derived from human preferences in the matching process. The detailed results are shown in the table below. These results demonstrate our model's superior performance compared to both open-source and closed-source models.
![figure1](assets/vben_vs_sota.png "figure1")
## Citation
If you find our work helpful, please cite us.
```
@article{wan2.1,
title = {Wan: Open and Advanced Large-Scale Video Generative Models},
author = {Wan Team},
journal = {},
year = {2025}
}
```
## License Agreement
The models in this repository are licensed under the Apache 2.0 License. We claim no rights over the your generate contents, granting you the freedom to use them while ensuring that your usage complies with the provisions of this license. You are fully accountable for your use of the models, which must not involve sharing any content that violates applicable laws, causes harm to individuals or groups, disseminates personal information intended for harm, spreads misinformation, or targets vulnerable populations. For a complete list of restrictions and details regarding your rights, please refer to the full text of the [license](LICENSE.txt).
## Acknowledgements
We would like to thank the contributors to the [SD3](https://huggingface.co/stabilityai/stable-diffusion-3-medium), [Qwen](https://huggingface.co/Qwen), [umt5-xxl](https://huggingface.co/google/umt5-xxl), [diffusers](https://github.com/huggingface/diffusers) and [HuggingFace](https://huggingface.co) repositories, for their open research.
## Contact Us
If you would like to leave a message to our research or product teams, feel free to join our [Discord](https://discord.gg/p5XbdQV7) or [WeChat groups](https://gw.alicdn.com/imgextra/i2/O1CN01tqjWFi1ByuyehkTSB_!!6000000000015-0-tps-611-1279.jpg)!

34
gradio/i2v_14B_singleGPU.py
View File

@ -24,8 +24,9 @@ wan_i2v_720P = None
# Button Func
def load_model(value):
def load_i2v_model(value):
global wan_i2v_480P, wan_i2v_720P
from mmgp import offload
if value == '------':
print("No model loaded")
@ -52,8 +53,11 @@ def load_model(value):
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
)
i2v720p= True
)
print("done", flush=True)
pipe = {"transformer": wan_i2v_720P.model, "text_encoder" : wan_i2v_720P.text_encoder.model, "text_encoder_2": wan_i2v_720P.clip.model, "vae": wan_i2v_720P.vae.model } #
offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, compile="transformer", quantizeTransformer = False, pinnedMemory = False)
return '720P'
if value == '480P':
@ -77,11 +81,16 @@ def load_model(value):
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
i2v720p= False
)
print("done", flush=True)
pipe = {"transformer": wan_i2v_480P.model, "text_encoder" : wan_i2v_480P.text_encoder.model, "text_encoder_2": wan_i2v_480P.clip.model, "vae": wan_i2v_480P.vae.model } #
offload.profile(pipe, profile_no=4, budgets = {"model":100, "*":3000}, verboseLevel=2, compile="transformer" )
return '480P'
def prompt_enc(prompt, img, tar_lang):
print('prompt extend...')
if img is None:
@ -96,10 +105,12 @@ def prompt_enc(prompt, img, tar_lang):
return prompt_output.prompt
def i2v_generation(img2vid_prompt, img2vid_image, resolution, sd_steps,
def i2v_generation(img2vid_prompt, img2vid_image, res, sd_steps,
guide_scale, shift_scale, seed, n_prompt):
# print(f"{img2vid_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}")
global resolution
from PIL import Image
img2vid_image = Image.open("d:\mammoth2.jpg")
if resolution == '------':
print(
'Please specify at least one resolution ckpt dir or specify the resolution'
@ -118,19 +129,19 @@ def i2v_generation(img2vid_prompt, img2vid_image, resolution, sd_steps,
guide_scale=guide_scale,
n_prompt=n_prompt,
seed=seed,
offload_model=True)
offload_model=False)
else:
global wan_i2v_480P
video = wan_i2v_480P.generate(
img2vid_prompt,
img2vid_image,
max_area=MAX_AREA_CONFIGS['480*832'],
shift=shift_scale,
shift=3.0, #shift_scale
sampling_steps=sd_steps,
guide_scale=guide_scale,
n_prompt=n_prompt,
seed=seed,
offload_model=True)
offload_model=False)
cache_video(
tensor=video[None],
@ -169,6 +180,7 @@ def gradio_interface():
)
img2vid_prompt = gr.Textbox(
label="Prompt",
value="Several giant wooly mammoths approach treading through a snowy meadow, their long wooly fur lightly blows in the wind as they walk, snow covered trees and dramatic snow capped mountains in the distance, mid afternoon light with wispy clouds and a sun high in the distance creates a warm glow, the low camera view is stunning capturing the large furry mammal with beautiful photography, depth of field.",
placeholder="Describe the video you want to generate",
)
tar_lang = gr.Radio(
@ -262,6 +274,8 @@ def _parse_args():
help="The prompt extend model to use.")
args = parser.parse_args()
args.ckpt_dir_720p = "../ckpts" # os.path.join("ckpt")
args.ckpt_dir_480p = "../ckpts" # os.path.join("ckpt")
assert args.ckpt_dir_720p is not None or args.ckpt_dir_480p is not None, "Please specify at least one checkpoint directory."
return args
@ -269,6 +283,12 @@ def _parse_args():
if __name__ == '__main__':
args = _parse_args()
global resolution
# load_model('720P')
# resolution = '720P'
resolution = '480P'
load_model(resolution)
print("Step1: Init prompt_expander...", end='', flush=True)
if args.prompt_extend_method == "dashscope":

1
gradio/t2i_14B_singleGPU.py
View File

@ -190,6 +190,7 @@ if __name__ == '__main__':
print("Step2: Init 14B t2i model...", end='', flush=True)
cfg = WAN_CONFIGS['t2i-14B']
# cfg = WAN_CONFIGS['t2v-1.3B']
wan_t2i = wan.WanT2V(
config=cfg,
checkpoint_dir=args.ckpt_dir,

35
gradio/t2v_14B_singleGPU.py
View File

@ -46,7 +46,7 @@ def t2v_generation(txt2vid_prompt, resolution, sd_steps, guide_scale,
guide_scale=guide_scale,
n_prompt=n_prompt,
seed=seed,
offload_model=True)
offload_model=False)
cache_video(
tensor=video[None],
@ -177,28 +177,39 @@ if __name__ == '__main__':
args = _parse_args()
print("Step1: Init prompt_expander...", end='', flush=True)
if args.prompt_extend_method == "dashscope":
prompt_expander = DashScopePromptExpander(
model_name=args.prompt_extend_model, is_vl=False)
elif args.prompt_extend_method == "local_qwen":
prompt_expander = QwenPromptExpander(
model_name=args.prompt_extend_model, is_vl=False, device=0)
else:
raise NotImplementedError(
f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
print("done", flush=True)
prompt_expander = None
# if args.prompt_extend_method == "dashscope":
# prompt_expander = DashScopePromptExpander(
# model_name=args.prompt_extend_model, is_vl=False)
# elif args.prompt_extend_method == "local_qwen":
# prompt_expander = QwenPromptExpander(
# model_name=args.prompt_extend_model, is_vl=False, device=0)
# else:
# raise NotImplementedError(
# f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
# print("done", flush=True)
from mmgp import offload
print("Step2: Init 14B t2v model...", end='', flush=True)
cfg = WAN_CONFIGS['t2v-14B']
# cfg = WAN_CONFIGS['t2v-1.3B']
wan_t2v = wan.WanT2V(
config=cfg,
checkpoint_dir=args.ckpt_dir,
checkpoint_dir="../ckpts",
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
)
pipe = {"transformer": wan_t2v.model, "text_encoder" : wan_t2v.text_encoder.model, "vae": wan_t2v.vae.model } #
# offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, quantizeTransformer = False, compile = "transformer") #
offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, quantizeTransformer = False) #
# offload.profile(pipe, profile_no=4, budgets = {"transformer":3000, "*":3000}, verboseLevel=2, quantizeTransformer = False)
print("done", flush=True)
demo = gradio_interface()

1275
gradio_server.py Normal file
View File

File diff suppressed because it is too large Load Diff

1
loras/README.txt Normal file
View File

@ -0,0 +1 @@
Put here Loras

1
loras_i2v/README.txt Normal file
View File

@ -0,0 +1 @@
Put here Loras

5
requirements.txt
View File

@ -11,6 +11,9 @@ easydict
ftfy
dashscope
imageio-ffmpeg
flash_attn
# flash_attn
gradio>=5.0.0
numpy>=1.23.5,<2
einops
moviepy==1.0.3
mmgp==3.2.1

109
wan/image2video.py
View File

@ -39,6 +39,9 @@ class WanI2V:
use_usp=False,
t5_cpu=False,
init_on_cpu=True,
i2v720p= True,
model_filename ="",
text_encoder_filename="",
):
r"""
Initializes the image-to-video generation model components.
@ -77,7 +80,7 @@ class WanI2V:
text_len=config.text_len,
dtype=config.t5_dtype,
device=torch.device('cpu'),
checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
checkpoint_path=text_encoder_filename,
tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
shard_fn=shard_fn if t5_fsdp else None,
)
@ -95,8 +98,10 @@ class WanI2V:
config.clip_checkpoint),
tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
logging.info(f"Creating WanModel from {checkpoint_dir}")
self.model = WanModel.from_pretrained(checkpoint_dir)
logging.info(f"Creating WanModel from {model_filename}")
from mmgp import offload
self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel)
self.model.eval().requires_grad_(False)
if t5_fsdp or dit_fsdp or use_usp:
@ -116,28 +121,30 @@ class WanI2V:
else:
self.sp_size = 1
if dist.is_initialized():
dist.barrier()
if dit_fsdp:
self.model = shard_fn(self.model)
else:
if not init_on_cpu:
self.model.to(self.device)
# if dist.is_initialized():
# dist.barrier()
# if dit_fsdp:
# self.model = shard_fn(self.model)
# else:
# if not init_on_cpu:
# self.model.to(self.device)
self.sample_neg_prompt = config.sample_neg_prompt
def generate(self,
input_prompt,
img,
max_area=720 * 1280,
frame_num=81,
shift=5.0,
sample_solver='unipc',
sampling_steps=40,
guide_scale=5.0,
n_prompt="",
seed=-1,
offload_model=True):
input_prompt,
img,
max_area=720 * 1280,
frame_num=81,
shift=5.0,
sample_solver='unipc',
sampling_steps=40,
guide_scale=5.0,
n_prompt="",
seed=-1,
offload_model=True,
callback = None
):
r"""
Generates video frames from input image and text prompt using diffusion process.
@ -197,14 +204,14 @@ class WanI2V:
seed_g.manual_seed(seed)
noise = torch.randn(
16,
21,
int((frame_num - 1)/4 + 1), #21,
lat_h,
lat_w,
dtype=torch.float32,
generator=seed_g,
device=self.device)
msk = torch.ones(1, 81, lat_h, lat_w, device=self.device)
msk = torch.ones(1, frame_num, lat_h, lat_w, device=self.device)
msk[:, 1:] = 0
msk = torch.concat([
torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]
@ -218,7 +225,7 @@ class WanI2V:
# preprocess
if not self.t5_cpu:
self.text_encoder.model.to(self.device)
# self.text_encoder.model.to(self.device)
context = self.text_encoder([input_prompt], self.device)
context_null = self.text_encoder([n_prompt], self.device)
if offload_model:
@ -229,20 +236,23 @@ class WanI2V:
context = [t.to(self.device) for t in context]
context_null = [t.to(self.device) for t in context_null]
self.clip.model.to(self.device)
# self.clip.model.to(self.device)
clip_context = self.clip.visual([img[:, None, :, :]])
if offload_model:
self.clip.model.cpu()
y = self.vae.encode([
torch.concat([
torch.nn.functional.interpolate(
img[None].cpu(), size=(h, w), mode='bicubic').transpose(
0, 1),
torch.zeros(3, 80, h, w)
],
dim=1).to(self.device)
])[0]
from mmgp import offload
offload.last_offload_obj.unload_all()
enc= torch.concat([
torch.nn.functional.interpolate(
img[None].cpu(), size=(h, w), mode='bicubic').transpose(
0, 1).to(torch.bfloat16),
torch.zeros(3, frame_num-1, h, w, device="cpu", dtype= torch.bfloat16)
], dim=1).to(self.device)
# enc = None
y = self.vae.encode([enc])[0]
y = torch.concat([msk, y])
@contextmanager
@ -283,6 +293,7 @@ class WanI2V:
'clip_fea': clip_context,
'seq_len': max_seq_len,
'y': [y],
'pipeline' : self
}
arg_null = {
@ -290,30 +301,39 @@ class WanI2V:
'clip_fea': clip_context,
'seq_len': max_seq_len,
'y': [y],
'pipeline' : self
}
if offload_model:
torch.cuda.empty_cache()
self.model.to(self.device)
for _, t in enumerate(tqdm(timesteps)):
# self.model.to(self.device)
if callback != None:
callback(-1, None)
self._interrupt = False
for i, t in enumerate(tqdm(timesteps)):
latent_model_input = [latent.to(self.device)]
timestep = [t]
timestep = torch.stack(timestep).to(self.device)
noise_pred_cond = self.model(
latent_model_input, t=timestep, **arg_c)[0].to(
torch.device('cpu') if offload_model else self.device)
latent_model_input, t=timestep, **arg_c)[0]
if self._interrupt:
return None
if offload_model:
torch.cuda.empty_cache()
noise_pred_uncond = self.model(
latent_model_input, t=timestep, **arg_null)[0].to(
torch.device('cpu') if offload_model else self.device)
latent_model_input, t=timestep, **arg_null)[0]
if self._interrupt:
return None
del latent_model_input
if offload_model:
torch.cuda.empty_cache()
noise_pred = noise_pred_uncond + guide_scale * (
noise_pred_cond - noise_pred_uncond)
del noise_pred_uncond
latent = latent.to(
torch.device('cpu') if offload_model else self.device)
@ -325,9 +345,14 @@ class WanI2V:
return_dict=False,
generator=seed_g)[0]
latent = temp_x0.squeeze(0)
del temp_x0
del timestep
x0 = [latent.to(self.device)]
del latent_model_input, timestep
if callback is not None:
callback(i, latent)
x0 = [latent.to(self.device)]
if offload_model:
self.model.cpu()

4
wan/modules/__init__.py
View File

@ -1,4 +1,4 @@
from .attention import flash_attention
from .attention import pay_attention
from .model import WanModel
from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
from .tokenizers import HuggingfaceTokenizer
@ -12,5 +12,5 @@ __all__ = [
'T5Decoder',
'T5EncoderModel',
'HuggingfaceTokenizer',
'flash_attention',
'pay_attention',
]

125
wan/modules/attention.py
View File

@ -1,5 +1,9 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
from importlib.metadata import version
from mmgp import offload
import torch.nn.functional as F
try:
import flash_attn_interface
@ -12,19 +16,99 @@ try:
FLASH_ATTN_2_AVAILABLE = True
except ModuleNotFoundError:
FLASH_ATTN_2_AVAILABLE = False
flash_attn = None
try:
from sageattention import sageattn_varlen
def sageattn_varlen_wrapper(
q,
k,
v,
cu_seqlens_q,
cu_seqlens_kv,
max_seqlen_q,
max_seqlen_kv,
):
return sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
except ImportError:
sageattn_varlen_wrapper = None
import warnings
try:
from sageattention import sageattn
@torch.compiler.disable()
def sageattn_wrapper(
qkv_list,
attention_length
):
q,k, v = qkv_list
padding_length = q.shape[0] -attention_length
q = q[:attention_length, :, : ].unsqueeze(0)
k = k[:attention_length, :, : ].unsqueeze(0)
v = v[:attention_length, :, : ].unsqueeze(0)
o = sageattn(q, k, v, tensor_layout="NHD").squeeze(0)
del q, k ,v
qkv_list.clear()
if padding_length > 0:
o = torch.cat([o, torch.empty( (padding_length, *o.shape[-2:]), dtype= o.dtype, device=o.device ) ], 0)
return o
except ImportError:
sageattn = None
@torch.compiler.disable()
def sdpa_wrapper(
qkv_list,
attention_length
):
q,k, v = qkv_list
padding_length = q.shape[0] -attention_length
q = q[:attention_length, :].transpose(0,1).unsqueeze(0)
k = k[:attention_length, :].transpose(0,1).unsqueeze(0)
v = v[:attention_length, :].transpose(0,1).unsqueeze(0)
o = F.scaled_dot_product_attention(
q, k, v, attn_mask=None, is_causal=False
).squeeze(0).transpose(0,1)
del q, k ,v
qkv_list.clear()
if padding_length > 0:
o = torch.cat([o, torch.empty( (padding_length, *o.shape[-2:]), dtype= o.dtype, device=o.device ) ], 0)
return o
def get_attention_modes():
ret = ["sdpa", "auto"]
if flash_attn != None:
ret.append("flash")
# if memory_efficient_attention != None:
# ret.append("xformers")
if sageattn_varlen_wrapper != None:
ret.append("sage")
if sageattn != None and version("sageattention").startswith("2") :
ret.append("sage2")
return ret
__all__ = [
'flash_attention',
'pay_attention',
'attention',
]
def flash_attention(
q,
k,
v,
def pay_attention(
qkv_list,
# q,
# k,
# v,
q_lens=None,
k_lens=None,
dropout_p=0.,
@ -49,6 +133,10 @@ def flash_attention(
deterministic: bool. If True, slightly slower and uses more memory.
dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
"""
attn = offload.shared_state["_attention"]
q,k,v = qkv_list
qkv_list.clear()
half_dtypes = (torch.float16, torch.bfloat16)
assert dtype in half_dtypes
assert q.device.type == 'cuda' and q.size(-1) <= 256
@ -91,7 +179,27 @@ def flash_attention(
)
# apply attention
if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
if attn=="sage":
x = sageattn_varlen_wrapper(
q=q,
k=k,
v=v,
cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
0, dtype=torch.int32).to(q.device, non_blocking=True),
cu_seqlens_kv=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
0, dtype=torch.int32).to(q.device, non_blocking=True),
max_seqlen_q=lq,
max_seqlen_kv=lk,
).unflatten(0, (b, lq))
elif attn=="sage2":
qkv_list = [q,k,v]
del q,k,v
x = sageattn_wrapper(qkv_list, lq).unsqueeze(0)
elif attn=="sdpa":
qkv_list = [q, k, v]
del q, k , v
x = sdpa_wrapper( qkv_list, lq).unsqueeze(0)
elif attn=="flash" and (version is None or version == 3):
# Note: dropout_p, window_size are not supported in FA3 now.
x = flash_attn_interface.flash_attn_varlen_func(
q=q,
@ -108,8 +216,7 @@ def flash_attention(
softmax_scale=softmax_scale,
causal=causal,
deterministic=deterministic)[0].unflatten(0, (b, lq))
else:
assert FLASH_ATTN_2_AVAILABLE
elif attn=="flash":
x = flash_attn.flash_attn_varlen_func(
q=q,
k=k,
@ -146,7 +253,7 @@ def attention(
fa_version=None,
):
if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
return flash_attention(
return pay_attention(
q=q,
k=k,
v=v,

26
wan/modules/clip.py
View File

@ -8,7 +8,7 @@ import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
from .attention import flash_attention
from .attention import pay_attention
from .tokenizers import HuggingfaceTokenizer
from .xlm_roberta import XLMRoberta
@ -82,7 +82,7 @@ class SelfAttention(nn.Module):
# compute attention
p = self.attn_dropout if self.training else 0.0
x = flash_attention(q, k, v, dropout_p=p, causal=self.causal, version=2)
x = pay_attention([q, k, v], dropout_p=p, causal=self.causal, version=2)
x = x.reshape(b, s, c)
# output
@ -194,7 +194,7 @@ class AttentionPool(nn.Module):
k, v = self.to_kv(x).view(b, s, 2, n, d).unbind(2)
# compute attention
x = flash_attention(q, k, v, version=2)
x = pay_attention(q, k, v, version=2)
x = x.reshape(b, 1, c)
# output
@ -441,11 +441,12 @@ def _clip(pretrained=False,
device='cpu',
**kwargs):
# init a model on device
device ="cpu"
with torch.device(device):
model = model_cls(**kwargs)
# set device
model = model.to(dtype=dtype, device=device)
# model = model.to(dtype=dtype, device=device)
output = (model,)
# init transforms
@ -507,16 +508,19 @@ class CLIPModel:
self.tokenizer_path = tokenizer_path
# init model
self.model, self.transforms = clip_xlm_roberta_vit_h_14(
pretrained=False,
return_transforms=True,
return_tokenizer=False,
dtype=dtype,
device=device)
from accelerate import init_empty_weights
with init_empty_weights():
self.model, self.transforms = clip_xlm_roberta_vit_h_14(
pretrained=False,
return_transforms=True,
return_tokenizer=False,
dtype=dtype,
device=device)
self.model = self.model.eval().requires_grad_(False)
logging.info(f'loading {checkpoint_path}')
self.model.load_state_dict(
torch.load(checkpoint_path, map_location='cpu'))
torch.load(checkpoint_path, map_location='cpu'), assign= True)
# init tokenizer
self.tokenizer = HuggingfaceTokenizer(

270
wan/modules/model.py
View File

@ -7,7 +7,7 @@ import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from .attention import flash_attention
from .attention import pay_attention
__all__ = ['WanModel']
@ -16,7 +16,7 @@ def sinusoidal_embedding_1d(dim, position):
# preprocess
assert dim % 2 == 0
half = dim // 2
position = position.type(torch.float64)
position = position.type(torch.float32)
# calculation
sinusoid = torch.outer(
@ -25,18 +25,47 @@ def sinusoidal_embedding_1d(dim, position):
return x
@amp.autocast(enabled=False)
# @amp.autocast(enabled=False)
def rope_params(max_seq_len, dim, theta=10000):
assert dim % 2 == 0
freqs = torch.outer(
torch.arange(max_seq_len),
1.0 / torch.pow(theta,
torch.arange(0, dim, 2).to(torch.float64).div(dim)))
torch.arange(0, dim, 2).to(torch.float32).div(dim)))
freqs = torch.polar(torch.ones_like(freqs), freqs)
return freqs
@amp.autocast(enabled=False)
def rope_apply_(x, grid_sizes, freqs):
assert x.shape[0]==1
n, c = x.size(2), x.size(3) // 2
# split freqs
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
f, h, w = grid_sizes[0]
seq_len = f * h * w
x_i = x[0, :seq_len, :, :]
x_i = x_i.to(torch.float32)
x_i = x_i.reshape(seq_len, n, -1, 2)
x_i = torch.view_as_complex(x_i)
freqs_i = torch.cat([
freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
], dim=-1)
freqs_i= freqs_i.reshape(seq_len, 1, -1)
# apply rotary embedding
x_i *= freqs_i
x_i = torch.view_as_real(x_i).flatten(2)
x[0, :seq_len, :, :] = x_i.to(torch.bfloat16)
# x_i = torch.cat([x_i, x[0, seq_len:]])
return x
# @amp.autocast(enabled=False)
def rope_apply(x, grid_sizes, freqs):
n, c = x.size(2), x.size(3) // 2
@ -45,12 +74,17 @@ def rope_apply(x, grid_sizes, freqs):
# loop over samples
output = []
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
for i, (f, h, w) in enumerate(grid_sizes):
seq_len = f * h * w
# precompute multipliers
x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
seq_len, n, -1, 2))
# x_i = x[i, :seq_len]
x_i = x[i]
x_i = x_i[:seq_len, :, :]
x_i = x_i.to(torch.float32)
x_i = x_i.reshape(seq_len, n, -1, 2)
x_i = torch.view_as_complex(x_i)
freqs_i = torch.cat([
freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
@ -59,12 +93,14 @@ def rope_apply(x, grid_sizes, freqs):
dim=-1).reshape(seq_len, 1, -1)
# apply rotary embedding
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
x_i *= freqs_i
x_i = torch.view_as_real(x_i).flatten(2)
x_i = x_i.to(torch.bfloat16)
x_i = torch.cat([x_i, x[i, seq_len:]])
# append to collection
output.append(x_i)
return torch.stack(output).float()
return torch.stack(output) #.float()
class WanRMSNorm(nn.Module):
@ -80,11 +116,31 @@ class WanRMSNorm(nn.Module):
Args:
x(Tensor): Shape [B, L, C]
"""
return self._norm(x.float()).type_as(x) * self.weight
y = x.float()
y.pow_(2)
y = y.mean(dim=-1, keepdim=True)
y += self.eps
y.rsqrt_()
x *= y
x *= self.weight
return x
# return self._norm(x).type_as(x) * self.weight
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
def my_LayerNorm(norm, x):
y = x.float()
y_m = y.mean(dim=-1, keepdim=True)
y -= y_m
del y_m
y.pow_(2)
y = y.mean(dim=-1, keepdim=True)
y += norm.eps
y.rsqrt_()
x = x * y
return x
class WanLayerNorm(nn.LayerNorm):
@ -96,7 +152,13 @@ class WanLayerNorm(nn.LayerNorm):
Args:
x(Tensor): Shape [B, L, C]
"""
return super().forward(x.float()).type_as(x)
# return F.layer_norm(
# input, self.normalized_shape, self.weight, self.bias, self.eps
# )
y = super().forward(x)
x = y.type_as(x)
return x
# return super().forward(x).type_as(x)
class WanSelfAttention(nn.Module):
@ -124,7 +186,7 @@ class WanSelfAttention(nn.Module):
self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
def forward(self, x, seq_lens, grid_sizes, freqs):
def forward(self, xlist, seq_lens, grid_sizes, freqs):
r"""
Args:
x(Tensor): Shape [B, L, num_heads, C / num_heads]
@ -132,24 +194,31 @@ class WanSelfAttention(nn.Module):
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
"""
x = xlist[0]
xlist.clear()
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
# query, key, value function
def qkv_fn(x):
q = self.norm_q(self.q(x)).view(b, s, n, d)
k = self.norm_k(self.k(x)).view(b, s, n, d)
v = self.v(x).view(b, s, n, d)
return q, k, v
q, k, v = qkv_fn(x)
x = flash_attention(
q=rope_apply(q, grid_sizes, freqs),
k=rope_apply(k, grid_sizes, freqs),
v=v,
k_lens=seq_lens,
q = self.q(x)
self.norm_q(q)
q = q.view(b, s, n, d) # !!!
k = self.k(x)
self.norm_k(k)
k = k.view(b, s, n, d)
v = self.v(x).view(b, s, n, d)
del x
rope_apply_(q, grid_sizes, freqs)
rope_apply_(k, grid_sizes, freqs)
qkv_list = [q,k,v]
del q,k,v
x = pay_attention(
qkv_list,
# q=q,
# k=k,
# v=v,
# k_lens=seq_lens,
window_size=self.window_size)
# output
x = x.flatten(2)
x = self.o(x)
@ -158,22 +227,31 @@ class WanSelfAttention(nn.Module):
class WanT2VCrossAttention(WanSelfAttention):
def forward(self, x, context, context_lens):
def forward(self, xlist, context, context_lens):
r"""
Args:
x(Tensor): Shape [B, L1, C]
context(Tensor): Shape [B, L2, C]
context_lens(Tensor): Shape [B]
"""
x = xlist[0]
xlist.clear()
b, n, d = x.size(0), self.num_heads, self.head_dim
# compute query, key, value
q = self.norm_q(self.q(x)).view(b, -1, n, d)
k = self.norm_k(self.k(context)).view(b, -1, n, d)
q = self.q(x)
del x
self.norm_q(q)
q= q.view(b, -1, n, d)
k = self.k(context)
self.norm_k(k)
k = k.view(b, -1, n, d)
v = self.v(context).view(b, -1, n, d)
# compute attention
x = flash_attention(q, k, v, k_lens=context_lens)
qvl_list=[q, k, v]
del q, k, v
x = pay_attention(qvl_list, k_lens=context_lens)
# output
x = x.flatten(2)
@ -196,31 +274,54 @@ class WanI2VCrossAttention(WanSelfAttention):
# self.alpha = nn.Parameter(torch.zeros((1, )))
self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
def forward(self, x, context, context_lens):
def forward(self, xlist, context, context_lens):
r"""
Args:
x(Tensor): Shape [B, L1, C]
context(Tensor): Shape [B, L2, C]
context_lens(Tensor): Shape [B]
"""
##### Enjoy this spagheti VRAM optimizations done by DeepBeepMeep !
# I am sure you are a nice person and as you copy this code, you will give me officially proper credits:
# Please link to https://github.com/deepbeepmeep/Wan2GP and @deepbeepmeep on twitter
x = xlist[0]
xlist.clear()
context_img = context[:, :257]
context = context[:, 257:]
b, n, d = x.size(0), self.num_heads, self.head_dim
# compute query, key, value
q = self.norm_q(self.q(x)).view(b, -1, n, d)
k = self.norm_k(self.k(context)).view(b, -1, n, d)
q = self.q(x)
del x
self.norm_q(q)
q= q.view(b, -1, n, d)
k = self.k(context)
self.norm_k(k)
k = k.view(b, -1, n, d)
v = self.v(context).view(b, -1, n, d)
k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)
qkv_list = [q, k, v]
del k,v
x = pay_attention(qkv_list, k_lens=context_lens)
k_img = self.k_img(context_img)
self.norm_k_img(k_img)
k_img = k_img.view(b, -1, n, d)
v_img = self.v_img(context_img).view(b, -1, n, d)
img_x = flash_attention(q, k_img, v_img, k_lens=None)
qkv_list = [q, k_img, v_img]
del q, k_img, v_img
img_x = pay_attention(qkv_list, k_lens=None)
# compute attention
x = flash_attention(q, k, v, k_lens=context_lens)
# output
x = x.flatten(2)
img_x = img_x.flatten(2)
x = x + img_x
x += img_x
del img_x
x = self.o(x)
return x
@ -289,27 +390,46 @@ class WanAttentionBlock(nn.Module):
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
"""
assert e.dtype == torch.float32
with amp.autocast(dtype=torch.float32):
e = (self.modulation + e).chunk(6, dim=1)
assert e[0].dtype == torch.float32
e = (self.modulation + e).chunk(6, dim=1)
# self-attention
y = self.self_attn(
self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, grid_sizes,
freqs)
with amp.autocast(dtype=torch.float32):
x = x + y * e[2]
x_mod = self.norm1(x)
x_mod *= 1 + e[1]
x_mod += e[0]
xlist = [x_mod]
del x_mod
y = self.self_attn( xlist, seq_lens, grid_sizes,freqs)
x.addcmul_(y, e[2])
del y
y = self.norm3(x)
ylist= [y]
del y
x += self.cross_attn(ylist, context, context_lens)
y = self.norm2(x)
# cross-attention & ffn function
def cross_attn_ffn(x, context, context_lens, e):
x = x + self.cross_attn(self.norm3(x), context, context_lens)
y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3])
with amp.autocast(dtype=torch.float32):
x = x + y * e[5]
return x
y *= 1 + e[4]
y += e[3]
ffn = self.ffn[0]
gelu = self.ffn[1]
ffn2= self.ffn[2]
y_shape = y.shape
y = y.view(-1, y_shape[-1])
chunk_size = int(y_shape[1]/2.7)
chunks =torch.split(y, chunk_size)
for y_chunk in chunks:
mlp_chunk = ffn(y_chunk)
mlp_chunk = gelu(mlp_chunk)
y_chunk[...] = ffn2(mlp_chunk)
del mlp_chunk
y = y.view(y_shape)
x.addcmul_(y, e[5])
x = cross_attn_ffn(x, context, context_lens, e)
return x
@ -336,10 +456,13 @@ class Head(nn.Module):
x(Tensor): Shape [B, L1, C]
e(Tensor): Shape [B, C]
"""
assert e.dtype == torch.float32
with amp.autocast(dtype=torch.float32):
e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
# assert e.dtype == torch.float32
e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
x = self.norm(x).to(torch.bfloat16)
x *= (1 + e[1])
x += e[0]
x = self.head(x)
return x
@ -384,7 +507,8 @@ class WanModel(ModelMixin, ConfigMixin):
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=True,
eps=1e-6):
eps=1e-6,
):
r"""
Initialize the diffusion model backbone.
@ -466,7 +590,7 @@ class WanModel(ModelMixin, ConfigMixin):
# buffers (don't use register_buffer otherwise dtype will be changed in to())
assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
d = dim // num_heads
self.freqs = torch.cat([
self.freqs = torch.cat([
rope_params(1024, d - 4 * (d // 6)),
rope_params(1024, 2 * (d // 6)),
rope_params(1024, 2 * (d // 6))
@ -487,6 +611,7 @@ class WanModel(ModelMixin, ConfigMixin):
seq_len,
clip_fea=None,
y=None,
pipeline = None,
):
r"""
Forward pass through the diffusion model
@ -521,8 +646,11 @@ class WanModel(ModelMixin, ConfigMixin):
# embeddings
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
grid_sizes = torch.stack(
[torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
# grid_sizes = torch.stack(
# [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
grid_sizes = [ list(u.shape[2:]) for u in x]
x = [u.flatten(2).transpose(1, 2) for u in x]
seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
assert seq_lens.max() <= seq_len
@ -532,11 +660,10 @@ class WanModel(ModelMixin, ConfigMixin):
])
# time embeddings
with amp.autocast(dtype=torch.float32):
e = self.time_embedding(
sinusoidal_embedding_1d(self.freq_dim, t).float())
e0 = self.time_projection(e).unflatten(1, (6, self.dim))
assert e.dtype == torch.float32 and e0.dtype == torch.float32
e = self.time_embedding(
sinusoidal_embedding_1d(self.freq_dim, t))
e0 = self.time_projection(e).unflatten(1, (6, self.dim)).to(torch.bfloat16)
# assert e.dtype == torch.float32 and e0.dtype == torch.float32
# context
context_lens = None
@ -561,6 +688,9 @@ class WanModel(ModelMixin, ConfigMixin):
context_lens=context_lens)
for block in self.blocks:
if pipeline._interrupt:
return [None]
x = block(x, **kwargs)
# head
@ -588,7 +718,7 @@ class WanModel(ModelMixin, ConfigMixin):
c = self.out_dim
out = []
for u, v in zip(x, grid_sizes.tolist()):
for u, v in zip(x, grid_sizes):
u = u[:math.prod(v)].view(*v, *self.patch_size, c)
u = torch.einsum('fhwpqrc->cfphqwr', u)
u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])

19
wan/modules/t5.py
View File

@ -442,7 +442,7 @@ def _t5(name,
model = model_cls(**kwargs)
# set device
model = model.to(dtype=dtype, device=device)
# model = model.to(dtype=dtype, device=device)
# init tokenizer
if return_tokenizer:
@ -486,20 +486,25 @@ class T5EncoderModel:
self.checkpoint_path = checkpoint_path
self.tokenizer_path = tokenizer_path
from accelerate import init_empty_weights
# init model
model = umt5_xxl(
encoder_only=True,
return_tokenizer=False,
dtype=dtype,
device=device).eval().requires_grad_(False)
with init_empty_weights():
model = umt5_xxl(
encoder_only=True,
return_tokenizer=False,
dtype=dtype,
device=device).eval().requires_grad_(False)
logging.info(f'loading {checkpoint_path}')
model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
from mmgp import offload
offload.load_model_data(model,checkpoint_path )
self.model = model
if shard_fn is not None:
self.model = shard_fn(self.model, sync_module_states=False)
else:
self.model.to(self.device)
# init tokenizer
tokenizer_path= "google/umt5-xxl"
self.tokenizer = HuggingfaceTokenizer(
name=tokenizer_path, seq_len=text_len, clean='whitespace')

59
wan/modules/vae.py
View File

@ -1,6 +1,6 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import logging
from mmgp import offload
import torch
import torch.cuda.amp as amp
import torch.nn as nn
@ -31,9 +31,16 @@ class CausalConv3d(nn.Conv3d):
cache_x = cache_x.to(x.device)
x = torch.cat([cache_x, x], dim=2)
padding[4] -= cache_x.shape[2]
cache_x = None
x = F.pad(x, padding)
x = super().forward(x)
return super().forward(x)
mem_threshold = offload.shared_state.get("_vae_threshold",0)
vae_config = offload.shared_state.get("_vae",1)
if vae_config == 0 and torch.cuda.memory_reserved() > mem_threshold or vae_config == 2:
torch.cuda.empty_cache()
return x
class RMS_norm(nn.Module):
@ -49,10 +56,11 @@ class RMS_norm(nn.Module):
self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
def forward(self, x):
return F.normalize(
x = F.normalize(
x, dim=(1 if self.channel_first else
-1)) * self.scale * self.gamma + self.bias
x = x.to(torch.bfloat16)
return x
class Upsample(nn.Upsample):
@ -107,11 +115,12 @@ class Resample(nn.Module):
feat_cache[idx] = 'Rep'
feat_idx[0] += 1
else:
cache_x = x[:, :, -CACHE_T:, :, :].clone()
clone = True
cache_x = x[:, :, -CACHE_T:, :, :]#.clone()
if cache_x.shape[2] < 2 and feat_cache[
idx] is not None and feat_cache[idx] != 'Rep':
# cache last frame of last two chunk
clone = False
cache_x = torch.cat([
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
cache_x.device), cache_x
@ -119,11 +128,14 @@ class Resample(nn.Module):
dim=2)
if cache_x.shape[2] < 2 and feat_cache[
idx] is not None and feat_cache[idx] == 'Rep':
clone = False
cache_x = torch.cat([
torch.zeros_like(cache_x).to(cache_x.device),
cache_x
],
dim=2)
if clone:
cache_x = cache_x.clone()
if feat_cache[idx] == 'Rep':
x = self.time_conv(x)
else:
@ -144,7 +156,7 @@ class Resample(nn.Module):
if feat_cache is not None:
idx = feat_idx[0]
if feat_cache[idx] is None:
feat_cache[idx] = x.clone()
feat_cache[idx] = x #.to("cpu") #x.clone() yyyy
feat_idx[0] += 1
else:
@ -155,7 +167,7 @@ class Resample(nn.Module):
x = self.time_conv(
torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
feat_cache[idx] = cache_x
feat_cache[idx] = cache_x#.to("cpu") #yyyyy
feat_idx[0] += 1
return x
@ -212,11 +224,11 @@ class ResidualBlock(nn.Module):
cache_x.device), cache_x
],
dim=2)
x = layer(x, feat_cache[idx])
feat_cache[idx] = cache_x
x = layer(x, feat_cache[idx]).to(torch.bfloat16)
feat_cache[idx] = cache_x#.to("cpu")
feat_idx[0] += 1
else:
x = layer(x)
x = layer(x).to(torch.bfloat16)
return x + h
@ -326,12 +338,16 @@ class Encoder3d(nn.Module):
cache_x.device), cache_x
],
dim=2)
x = self.conv1(x, feat_cache[idx])
x = self.conv1(x, feat_cache[idx]).to(torch.bfloat16)
feat_cache[idx] = cache_x
del cache_x
feat_idx[0] += 1
else:
x = self.conv1(x)
# torch.cuda.empty_cache()
## downsamples
for layer in self.downsamples:
if feat_cache is not None:
@ -339,6 +355,8 @@ class Encoder3d(nn.Module):
else:
x = layer(x)
# torch.cuda.empty_cache()
## middle
for layer in self.middle:
if isinstance(layer, ResidualBlock) and feat_cache is not None:
@ -346,6 +364,8 @@ class Encoder3d(nn.Module):
else:
x = layer(x)
# torch.cuda.empty_cache()
## head
for layer in self.head:
if isinstance(layer, CausalConv3d) and feat_cache is not None:
@ -360,9 +380,13 @@ class Encoder3d(nn.Module):
dim=2)
x = layer(x, feat_cache[idx])
feat_cache[idx] = cache_x
del cache_x
feat_idx[0] += 1
else:
x = layer(x)
# torch.cuda.empty_cache()
return x
@ -433,10 +457,12 @@ class Decoder3d(nn.Module):
],
dim=2)
x = self.conv1(x, feat_cache[idx])
feat_cache[idx] = cache_x
feat_cache[idx] = cache_x#.to("cpu")
del cache_x
feat_idx[0] += 1
else:
x = self.conv1(x)
cache_x = None
## middle
for layer in self.middle:
@ -456,7 +482,7 @@ class Decoder3d(nn.Module):
for layer in self.head:
if isinstance(layer, CausalConv3d) and feat_cache is not None:
idx = feat_idx[0]
cache_x = x[:, :, -CACHE_T:, :, :].clone()
cache_x = x[:, :, -CACHE_T:, :, :] .clone()
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
# cache last frame of last two chunk
cache_x = torch.cat([
@ -465,7 +491,8 @@ class Decoder3d(nn.Module):
],
dim=2)
x = layer(x, feat_cache[idx])
feat_cache[idx] = cache_x
feat_cache[idx] = cache_x#.to("cpu")
del cache_x
feat_idx[0] += 1
else:
x = layer(x)
@ -532,6 +559,8 @@ class WanVAE_(nn.Module):
feat_cache=self._enc_feat_map,
feat_idx=self._enc_conv_idx)
out = torch.cat([out, out_], 2)
mu, log_var = self.conv1(out).chunk(2, dim=1)
if isinstance(scale[0], torch.Tensor):
mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(

56
wan/text2video.py
View File

@ -35,6 +35,8 @@ class WanT2V:
dit_fsdp=False,
use_usp=False,
t5_cpu=False,
model_filename = None,
text_encoder_filename = None
):
r"""
Initializes the Wan text-to-video generation model components.
@ -70,18 +72,26 @@ class WanT2V:
text_len=config.text_len,
dtype=config.t5_dtype,
device=torch.device('cpu'),
checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
checkpoint_path=text_encoder_filename,
tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
shard_fn=shard_fn if t5_fsdp else None)
self.vae_stride = config.vae_stride
self.patch_size = config.patch_size
self.vae = WanVAE(
vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
device=self.device)
logging.info(f"Creating WanModel from {checkpoint_dir}")
self.model = WanModel.from_pretrained(checkpoint_dir)
logging.info(f"Creating WanModel from {model_filename}")
from mmgp import offload
self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel)
self.model.eval().requires_grad_(False)
if use_usp:
@ -98,12 +108,12 @@ class WanT2V:
else:
self.sp_size = 1
if dist.is_initialized():
dist.barrier()
if dit_fsdp:
self.model = shard_fn(self.model)
else:
self.model.to(self.device)
# if dist.is_initialized():
# dist.barrier()
# if dit_fsdp:
# self.model = shard_fn(self.model)
# else:
# self.model.to(self.device)
self.sample_neg_prompt = config.sample_neg_prompt
@ -117,7 +127,9 @@ class WanT2V:
guide_scale=5.0,
n_prompt="",
seed=-1,
offload_model=True):
offload_model=True,
callback = None
):
r"""
Generates video frames from text prompt using diffusion process.
@ -168,7 +180,7 @@ class WanT2V:
seed_g.manual_seed(seed)
if not self.t5_cpu:
self.text_encoder.model.to(self.device)
# self.text_encoder.model.to(self.device)
context = self.text_encoder([input_prompt], self.device)
context_null = self.text_encoder([n_prompt], self.device)
if offload_model:
@ -223,23 +235,32 @@ class WanT2V:
# sample videos
latents = noise
arg_c = {'context': context, 'seq_len': seq_len}
arg_null = {'context': context_null, 'seq_len': seq_len}
arg_c = {'context': context, 'seq_len': seq_len, 'pipeline': self}
arg_null = {'context': context_null, 'seq_len': seq_len, 'pipeline': self}
for _, t in enumerate(tqdm(timesteps)):
if callback != None:
callback(-1, None)
self._interrupt = False
for i, t in enumerate(tqdm(timesteps)):
latent_model_input = latents
timestep = [t]
timestep = torch.stack(timestep)
self.model.to(self.device)
# self.model.to(self.device)
noise_pred_cond = self.model(
latent_model_input, t=timestep, **arg_c)[0]
if self._interrupt:
return None
noise_pred_uncond = self.model(
latent_model_input, t=timestep, **arg_null)[0]
if self._interrupt:
return None
del latent_model_input
noise_pred = noise_pred_uncond + guide_scale * (
noise_pred_cond - noise_pred_uncond)
del noise_pred_uncond
temp_x0 = sample_scheduler.step(
noise_pred.unsqueeze(0),
@ -248,6 +269,10 @@ class WanT2V:
return_dict=False,
generator=seed_g)[0]
latents = [temp_x0.squeeze(0)]
del temp_x0
if callback is not None:
callback(i, latents)
x0 = latents
if offload_model:
@ -256,6 +281,7 @@ class WanT2V:
if self.rank == 0:
videos = self.vae.decode(x0)
del noise, latents
del sample_scheduler
if offload_model: