```
├── .gitignore
├── INSTALL.md (200 tokens)
├── LICENSE.txt (2.3k tokens)
├── Makefile
├── README.md (7k tokens)
├── assets/
├── comp_effic.png
├── logo.png
├── moe_2.png
├── moe_arch.png
├── performance.png
├── vae.png
├── examples/
├── Five Hundred Miles.MP3
├── Five Hundred Miles.png
├── i2v_input.JPG
├── pose.mp4
├── pose.png
├── sing.MP3
├── talk.wav
├── wan_animate/
├── animate/
├── image.jpeg
├── video.mp4
├── replace/
├── image.jpeg
├── video.mp4
├── zero_shot_prompt.wav
├── generate.py (4.1k tokens)
├── pyproject.toml (300 tokens)
├── requirements.txt
├── requirements_animate.txt
├── requirements_s2v.txt
├── tests/
├── README.md
├── test.sh (1200 tokens)
├── wan/
├── __init__.py (100 tokens)
├── animate.py (5.5k tokens)
├── configs/
├── __init__.py (300 tokens)
├── shared_config.py (100 tokens)
├── wan_animate_14B.py (300 tokens)
├── wan_i2v_A14B.py (200 tokens)
├── wan_s2v_14B.py (400 tokens)
├── wan_t2v_A14B.py (200 tokens)
├── wan_ti2v_5B.py (200 tokens)
├── distributed/
├── __init__.py
├── fsdp.py (300 tokens)
├── sequence_parallel.py (1000 tokens)
├── ulysses.py (300 tokens)
├── util.py (300 tokens)
├── image2video.py (3.3k tokens)
├── modules/
├── __init__.py (100 tokens)
├── animate/
├── __init__.py
├── animate_utils.py (900 tokens)
├── clip.py (3.4k tokens)
├── face_blocks.py (2.4k tokens)
├── model_animate.py (3.2k tokens)
├── motion_encoder.py (1600 tokens)
├── preprocess/
├── UserGuider.md (1000 tokens)
├── __init__.py
├── human_visualization.py (8.8k tokens)
├── pose2d.py (3.5k tokens)
├── pose2d_utils.py (8.4k tokens)
├── preprocess_data.py (900 tokens)
├── process_pipepline.py (3.6k tokens)
├── retarget_pose.py (7.4k tokens)
├── sam_utils.py (1200 tokens)
├── utils.py (1500 tokens)
├── video_predictor.py (1700 tokens)
├── xlm_roberta.py (1000 tokens)
├── attention.py (1100 tokens)
├── model.py (3.6k tokens)
├── s2v/
├── __init__.py
├── audio_encoder.py (1500 tokens)
├── audio_utils.py (700 tokens)
├── auxi_blocks.py (1600 tokens)
├── model_s2v.py (6.9k tokens)
├── motioner.py (5.7k tokens)
├── s2v_utils.py (600 tokens)
├── t5.py (3.4k tokens)
├── tokenizers.py (500 tokens)
├── vae2_1.py (4.6k tokens)
├── vae2_2.py (6.4k tokens)
├── speech2video.py (5.8k tokens)
├── text2video.py (2.9k tokens)
├── textimage2video.py (4.8k tokens)
├── utils/
├── __init__.py (100 tokens)
├── fm_solvers.py (8k tokens)
├── fm_solvers_unipc.py (6.5k tokens)
├── prompt_extend.py (3.9k tokens)
├── qwen_vl_utils.py (2.6k tokens)
├── system_prompt.py (2.5k tokens)
├── utils.py (1500 tokens)
```
## /.gitignore
```gitignore path="/.gitignore"
__pycache__/
.DS_Store
.vscode*
tmp_examples*
new_checkpoint*
batch_test*
nohup*
```
## /INSTALL.md
# Installation Guide
## Install with pip
```bash
pip install .
pip install .[dev] # Installe aussi les outils de dev
```
## Install with Poetry
Ensure you have [Poetry](https://python-poetry.org/docs/#installation) installed on your system.
To install all dependencies:
```bash
poetry install
```
### Handling `flash-attn` Installation Issues
If `flash-attn` fails due to **PEP 517 build issues**, you can try one of the following fixes.
#### No-Build-Isolation Installation (Recommended)
```bash
poetry run pip install --upgrade pip setuptools wheel
poetry run pip install flash-attn --no-build-isolation
poetry install
```
#### Install from Git (Alternative)
```bash
poetry run pip install git+https://github.com/Dao-AILab/flash-attention.git
```
---
### Running the Model
Once the installation is complete, you can run **Wan2.2** using:
```bash
poetry run python generate.py --task t2v-A14B --size '1280*720' --ckpt_dir ./Wan2.2-T2V-A14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
```
#### Test
```bash
bash tests/test.sh
```
#### Format
```bash
black .
isort .
```
## /LICENSE.txt
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"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.
## /Makefile
``` path="/Makefile"
.PHONY: format
format:
isort generate.py wan
yapf -i -r *.py generate.py wan
```
## /README.md
# Wan2.2
<p align="center">
<img src="assets/logo.png" width="400"/>
<p>
<p align="center">
💜 <a href="https://wan.video"><b>Wan</b></a>    ｜    🖥️ <a href="https://github.com/Wan-Video/Wan2.2">GitHub</a>    |   🤗 <a href="https://huggingface.co/Wan-AI/">Hugging Face</a>   |   🤖 <a href="https://modelscope.cn/organization/Wan-AI">ModelScope</a>   |    📑 <a href="https://arxiv.org/abs/2503.20314">Paper</a>    |    📑 <a href="https://wan.video/welcome?spm=a2ty_o02.30011076.0.0.6c9ee41eCcluqg">Blog</a>    |    💬 <a href="https://discord.gg/AKNgpMK4Yj">Discord</a>  
<br>
📕 <a href="https://alidocs.dingtalk.com/i/nodes/jb9Y4gmKWrx9eo4dCql9LlbYJGXn6lpz">使用指南(中文)</a>   |    📘 <a href="https://alidocs.dingtalk.com/i/nodes/EpGBa2Lm8aZxe5myC99MelA2WgN7R35y">User Guide(English)</a>   |   💬 <a href="https://gw.alicdn.com/imgextra/i2/O1CN01tqjWFi1ByuyehkTSB_!!6000000000015-0-tps-611-1279.jpg">WeChat(微信)</a>  
<br>
-----
[**Wan: Open and Advanced Large-Scale Video Generative Models**](https://arxiv.org/abs/2503.20314) <be>
We are excited to introduce **Wan2.2**, a major upgrade to our foundational video models. With **Wan2.2**, we have focused on incorporating the following innovations:
- 👍 **Effective MoE Architecture**: Wan2.2 introduces a Mixture-of-Experts (MoE) architecture into video diffusion models. By separating the denoising process cross timesteps with specialized powerful expert models, this enlarges the overall model capacity while maintaining the same computational cost.
- 👍 **Cinematic-level Aesthetics**: Wan2.2 incorporates meticulously curated aesthetic data, complete with detailed labels for lighting, composition, contrast, color tone, and more. This allows for more precise and controllable cinematic style generation, facilitating the creation of videos with customizable aesthetic preferences.
- 👍 **Complex Motion Generation**: Compared to Wan2.1, Wan2.2 is trained on a significantly larger data, with +65.6% more images and +83.2% more videos. This expansion notably enhances the model's generalization across multiple dimensions such as motions, semantics, and aesthetics, achieving TOP performance among all open-sourced and closed-sourced models.
- 👍 **Efficient High-Definition Hybrid TI2V**: Wan2.2 open-sources a 5B model built with our advanced Wan2.2-VAE that achieves a compression ratio of **16×16×4**. This model supports both text-to-video and image-to-video generation at 720P resolution with 24fps and can also run on consumer-grade graphics cards like 4090. It is one of the fastest **720P@24fps** models currently available, capable of serving both the industrial and academic sectors simultaneously.
## Video Demos
<div align="center">
<video src="https://github.com/user-attachments/assets/b63bfa58-d5d7-4de6-a1a2-98970b06d9a7" width="70%" poster=""> </video>
</div>
## 🔥 Latest News!!
* Nov 13, 2025: 👋 Wan2.2-Animate-14B has been integrated into Diffusers ([PR](https://github.com/huggingface/diffusers/pull/12526),[Weights](https://huggingface.co/Wan-AI/Wan2.2-Animate-14B-Diffusers)). Thanks to all community contributors. Enjoy!
* Sep 19, 2025: 💃 We introduct **[Wan2.2-Animate-14B](https://humanaigc.github.io/wan-animate)**, an unified model for character animation and replacement with holistic movement and expression replication. We released the [model weights](#model-download) and [inference code](#run-wan-animate). And you can try it on [wan.video](https://wan.video/), [ModelScope Studio](https://www.modelscope.cn/studios/Wan-AI/Wan2.2-Animate) or [HuggingFace Space](https://huggingface.co/spaces/Wan-AI/Wan2.2-Animate)!
* Aug 26, 2025: 🎵 We introduce **[Wan2.2-S2V-14B](https://humanaigc.github.io/wan-s2v-webpage)**, an audio-driven cinematic video generation model, including [inference code](#run-speech-to-video-generation), [model weights](#model-download), and [technical report](https://humanaigc.github.io/wan-s2v-webpage/content/wan-s2v.pdf)! Now you can try it on [wan.video](https://wan.video/), [ModelScope Gradio](https://www.modelscope.cn/studios/Wan-AI/Wan2.2-S2V) or [HuggingFace Gradio](https://huggingface.co/spaces/Wan-AI/Wan2.2-S2V)!
* Jul 28, 2025: 👋 We have open a [HF space](https://huggingface.co/spaces/Wan-AI/Wan-2.2-5B) using the TI2V-5B model. Enjoy!
* Jul 28, 2025: 👋 Wan2.2 has been integrated into ComfyUI ([CN](https://docs.comfy.org/zh-CN/tutorials/video/wan/wan2_2) | [EN](https://docs.comfy.org/tutorials/video/wan/wan2_2)). Enjoy!
* Jul 28, 2025: 👋 Wan2.2's T2V, I2V and TI2V have been integrated into Diffusers ([T2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers) | [I2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers) | [TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers)). Feel free to give it a try!
* Jul 28, 2025: 👋 We've released the inference code and model weights of **Wan2.2**.
* Sep 5, 2025: 👋 We add text-to-speech synthesis support with [CosyVoice](https://github.com/FunAudioLLM/CosyVoice) for Speech-to-Video generation task.
## Community Works
If your research or project builds upon [**Wan2.1**](https://github.com/Wan-Video/Wan2.1) or [**Wan2.2**](https://github.com/Wan-Video/Wan2.2), and you would like more people to see it, please inform us.
- [Prompt Relay](https://github.com/GordonChen19/Prompt-Relay), a plug-and-play, inference-time method for temporal control in video generation. Prompt Relay improves video quality and gives users precise control over what happens at each moment in the video. Visit their [webpage](https://gordonchen19.github.io/Prompt-Relay/) for more details.
- [Helios](https://github.com/PKU-YuanGroup/Helios), a breakthrough video generation model base on **Wan2.1** that achieves minute-scale, high-quality video synthesis at 19.5 FPS on a single H100 GPU (about 10 FPS on a single Ascend NPU) —without relying on conventional long video anti-drifting strategies or standard video acceleration techniques. Visit their [webpage](https://pku-yuangroup.github.io/Helios-Page/) for more details.
- [LightX2V](https://github.com/ModelTC/LightX2V), a lightweight and efficient video generation framework that integrates **Wan2.1** and **Wan2.2**, supporting multiple engineering acceleration techniques for fast inference. [LightX2V-HuggingFace](https://huggingface.co/lightx2v), offers a variety of Wan-based step-distillation models, quantized models, and lightweight VAE models.
- [HuMo](https://github.com/Phantom-video/HuMo) proposed a unified, human-centric framework based on **Wan** to produce high-quality, fine-grained, and controllable human videos from multimodal inputs—including text, images, and audio. Visit their [webpage](https://phantom-video.github.io/HuMo/) for more details.
- [FastVideo](https://github.com/hao-ai-lab/FastVideo) includes distilled **Wan** models with sparse attention that significanly speed up the inference time.
- [Cache-dit](https://github.com/vipshop/cache-dit) offers Fully Cache Acceleration support for **Wan2.2** MoE with DBCache, TaylorSeer and Cache CFG. Visit their [example](https://github.com/vipshop/cache-dit/blob/main/examples/pipeline/run_wan_2.2.py) for more details.
- [Kijai's ComfyUI WanVideoWrapper](https://github.com/kijai/ComfyUI-WanVideoWrapper) is an alternative implementation of **Wan** models for ComfyUI. Thanks to its Wan-only focus, it's on the frontline of getting cutting edge optimizations and hot research features, which are often hard to integrate into ComfyUI quickly due to its more rigid structure.
- [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio) provides comprehensive support for **Wan 2.2**, including low-GPU-memory layer-by-layer offload, FP8 quantization, sequence parallelism, LoRA training, full training.
## 📑 Todo List
- Wan2.2 Text-to-Video
- [x] Multi-GPU Inference code of the A14B and 14B models
- [x] Checkpoints of the A14B and 14B models
- [x] ComfyUI integration
- [x] Diffusers integration
- Wan2.2 Image-to-Video
- [x] Multi-GPU Inference code of the A14B model
- [x] Checkpoints of the A14B model
- [x] ComfyUI integration
- [x] Diffusers integration
- Wan2.2 Text-Image-to-Video
- [x] Multi-GPU Inference code of the 5B model
- [x] Checkpoints of the 5B model
- [x] ComfyUI integration
- [x] Diffusers integration
- Wan2.2-S2V Speech-to-Video
- [x] Inference code of Wan2.2-S2V
- [x] Checkpoints of Wan2.2-S2V-14B
- [x] ComfyUI integration
- [x] Diffusers integration
- Wan2.2-Animate Character Animation and Replacement
- [x] Inference code of Wan2.2-Animate
- [x] Checkpoints of Wan2.2-Animate
- [x] ComfyUI integration
- [x] Diffusers integration
## Run Wan2.2
#### Installation
Clone the repo:
```sh
git clone https://github.com/Wan-Video/Wan2.2.git
cd Wan2.2
```
Install dependencies:
```sh
# Ensure torch >= 2.4.0
# If the installation of `flash_attn` fails, try installing the other packages first and install `flash_attn` last
pip install -r requirements.txt
# If you want to use CosyVoice to synthesize speech for Speech-to-Video Generation, please install requirements_s2v.txt additionally
pip install -r requirements_s2v.txt
```
#### Model Download
| Models | Download Links | Description |
|--------------------|---------------------------------------------------------------------------------------------------------------------------------------------|-------------|
| T2V-A14B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B) 🤖 [ModelScope](https://modelscope.cn/models/Wan-AI/Wan2.2-T2V-A14B) | Text-to-Video MoE model, supports 480P & 720P |
| I2V-A14B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B) 🤖 [ModelScope](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B) | Image-to-Video MoE model, supports 480P & 720P |
| TI2V-5B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) 🤖 [ModelScope](https://modelscope.cn/models/Wan-AI/Wan2.2-TI2V-5B) | High-compression VAE, T2V+I2V, supports 720P |
| S2V-14B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-S2V-14B) 🤖 [ModelScope](https://modelscope.cn/models/Wan-AI/Wan2.2-S2V-14B) | Speech-to-Video model, supports 480P & 720P |
| Animate-14B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-Animate-14B) 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.2-Animate-14B) | Character animation and replacement | |
> 💡Note:
> The TI2V-5B model supports 720P video generation at **24 FPS**.
Download models using huggingface-cli:
``` sh
pip install "huggingface_hub[cli]"
huggingface-cli download Wan-AI/Wan2.2-T2V-A14B --local-dir ./Wan2.2-T2V-A14B
```
Download models using modelscope-cli:
``` sh
pip install modelscope
modelscope download Wan-AI/Wan2.2-T2V-A14B --local_dir ./Wan2.2-T2V-A14B
```
#### Run Text-to-Video Generation
This repository supports the `Wan2.2-T2V-A14B` Text-to-Video model and can simultaneously support video generation at 480P and 720P resolutions.
##### (1) Without Prompt Extension
To facilitate implementation, we will start with a basic version of the inference process that skips the [prompt extension](#2-using-prompt-extention) step.
- Single-GPU inference
``` sh
python generate.py --task t2v-A14B --size 1280*720 --ckpt_dir ./Wan2.2-T2V-A14B --offload_model True --convert_model_dtype --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
```
> 💡 This command can run on a GPU with at least 80GB VRAM.
> 💡If you encounter OOM (Out-of-Memory) issues, you can use the `--offload_model True`, `--convert_model_dtype` and `--t5_cpu` options to reduce GPU memory usage.
- Multi-GPU inference using FSDP + DeepSpeed Ulysses
We use [PyTorch FSDP](https://docs.pytorch.org/docs/stable/fsdp.html) and [DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509) to accelerate inference.
``` sh
torchrun --nproc_per_node=8 generate.py --task t2v-A14B --size 1280*720 --ckpt_dir ./Wan2.2-T2V-A14B --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 Extension
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:
```sh
DASH_API_KEY=your_key torchrun --nproc_per_node=8 generate.py --task t2v-A14B --size 1280*720 --ckpt_dir ./Wan2.2-T2V-A14B --dit_fsdp --t5_fsdp --ulysses_size 8 --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 'zh'
```
- 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:
``` sh
torchrun --nproc_per_node=8 generate.py --task t2v-A14B --size 1280*720 --ckpt_dir ./Wan2.2-T2V-A14B --dit_fsdp --t5_fsdp --ulysses_size 8 --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 'zh'
```
#### Run Image-to-Video Generation
This repository supports the `Wan2.2-I2V-A14B` Image-to-Video model and can simultaneously support video generation at 480P and 720P resolutions.
- Single-GPU inference
```sh
python generate.py --task i2v-A14B --size 1280*720 --ckpt_dir ./Wan2.2-I2V-A14B --offload_model True --convert_model_dtype --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."
```
> This command can run on a GPU with at least 80GB VRAM.
> 💡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 + DeepSpeed Ulysses
```sh
torchrun --nproc_per_node=8 generate.py --task i2v-A14B --size 1280*720 --ckpt_dir ./Wan2.2-I2V-A14B --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."
```
- Image-to-Video Generation without prompt
```sh
DASH_API_KEY=your_key torchrun --nproc_per_node=8 generate.py --task i2v-A14B --size 1280*720 --ckpt_dir ./Wan2.2-I2V-A14B --prompt '' --image examples/i2v_input.JPG --dit_fsdp --t5_fsdp --ulysses_size 8 --use_prompt_extend --prompt_extend_method 'dashscope'
```
> 💡The model can generate videos solely from the input image. You can use prompt extension to generate prompt from the image.
> The process of prompt extension can be referenced [here](#2-using-prompt-extention).
#### Run Text-Image-to-Video Generation
This repository supports the `Wan2.2-TI2V-5B` Text-Image-to-Video model and can support video generation at 720P resolutions.
- Single-GPU Text-to-Video inference
```sh
python generate.py --task ti2v-5B --size 1280*704 --ckpt_dir ./Wan2.2-TI2V-5B --offload_model True --convert_model_dtype --t5_cpu --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage"
```
> 💡Unlike other tasks, the 720P resolution of the Text-Image-to-Video task is `1280*704` or `704*1280`.
> This command can run on a GPU with at least 24GB VRAM (e.g, RTX 4090 GPU).
> 💡If you are running on a GPU with at least 80GB VRAM, you can remove the `--offload_model True`, `--convert_model_dtype` and `--t5_cpu` options to speed up execution.
- Single-GPU Image-to-Video inference
```sh
python generate.py --task ti2v-5B --size 1280*704 --ckpt_dir ./Wan2.2-TI2V-5B --offload_model True --convert_model_dtype --t5_cpu --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."
```
> 💡If the image parameter is configured, it is an Image-to-Video generation; otherwise, it defaults to a Text-to-Video generation.
> 💡Similar to Image-to-Video, 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 + DeepSpeed Ulysses
```sh
torchrun --nproc_per_node=8 generate.py --task ti2v-5B --size 1280*704 --ckpt_dir ./Wan2.2-TI2V-5B --dit_fsdp --t5_fsdp --ulysses_size 8 --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."
```
> The process of prompt extension can be referenced [here](#2-using-prompt-extention).
#### Run Speech-to-Video Generation
This repository supports the `Wan2.2-S2V-14B` Speech-to-Video model and can simultaneously support video generation at 480P and 720P resolutions.
- Single-GPU Speech-to-Video inference
```sh
python generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./Wan2.2-S2V-14B/ --offload_model True --convert_model_dtype --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard." --image "examples/i2v_input.JPG" --audio "examples/talk.wav"
# Without setting --num_clip, the generated video length will automatically adjust based on the input audio length
# You can use CosyVoice to generate audio with --enable_tts
python generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./Wan2.2-S2V-14B/ --offload_model True --convert_model_dtype --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard." --image "examples/i2v_input.JPG" --enable_tts --tts_prompt_audio "examples/zero_shot_prompt.wav" --tts_prompt_text "希望你以后能够做的比我还好呦。" --tts_text "收到好友从远方寄来的生日礼物，那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐，笑容如花儿般绽放。"
```
> 💡 This command can run on a GPU with at least 80GB VRAM.
- Multi-GPU inference using FSDP + DeepSpeed Ulysses
```sh
torchrun --nproc_per_node=8 generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./Wan2.2-S2V-14B/ --dit_fsdp --t5_fsdp --ulysses_size 8 --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard." --image "examples/i2v_input.JPG" --audio "examples/talk.wav"
```
- Pose + Audio driven generation
```sh
torchrun --nproc_per_node=8 generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./Wan2.2-S2V-14B/ --dit_fsdp --t5_fsdp --ulysses_size 8 --prompt "a person is singing" --image "examples/pose.png" --audio "examples/sing.MP3" --pose_video "./examples/pose.mp4"
```
> 💡For the Speech-to-Video task, the `size` parameter represents the area of the generated video, with the aspect ratio following that of the original input image.
> 💡The model can generate videos from audio input combined with reference image and optional text prompt.
> 💡The `--pose_video` parameter enables pose-driven generation, allowing the model to follow specific pose sequences while generating videos synchronized with audio input.
> 💡The `--num_clip` parameter controls the number of video clips generated, useful for quick preview with shorter generation time.
Please visit our project page to see more examples and learn about the scenarios suitable for this model.
#### Run Wan-Animate
Wan-Animate takes a video and a character image as input, and generates a video in either "animation" or "replacement" mode.
1. animation mode： The model generates a video of the character image that mimics the human motion in the input video.
2. replacement mode: The model replaces the character image with the input video.
Please visit our [project page](https://humanaigc.github.io/wan-animate) to see more examples and learn about the scenarios suitable for this model.
##### (1) Preprocessing
The input video should be preprocessed into several materials before be feed into the inference process. Please refer to the following processing flow, and more details about preprocessing can be found in [UserGuider](https://github.com/Wan-Video/Wan2.2/blob/main/wan/modules/animate/preprocess/UserGuider.md).
* For animation
```bash
python ./wan/modules/animate/preprocess/preprocess_data.py \
--ckpt_path ./Wan2.2-Animate-14B/process_checkpoint \
--video_path ./examples/wan_animate/animate/video.mp4 \
--refer_path ./examples/wan_animate/animate/image.jpeg \
--save_path ./examples/wan_animate/animate/process_results \
--resolution_area 1280 720 \
--retarget_flag \
--use_flux
```
* For replacement
```bash
python ./wan/modules/animate/preprocess/preprocess_data.py \
--ckpt_path ./Wan2.2-Animate-14B/process_checkpoint \
--video_path ./examples/wan_animate/replace/video.mp4 \
--refer_path ./examples/wan_animate/replace/image.jpeg \
--save_path ./examples/wan_animate/replace/process_results \
--resolution_area 1280 720 \
--iterations 3 \
--k 7 \
--w_len 1 \
--h_len 1 \
--replace_flag
```
##### (2) Run in animation mode
* Single-GPU inference
```bash
python generate.py --task animate-14B --ckpt_dir ./Wan2.2-Animate-14B/ --src_root_path ./examples/wan_animate/animate/process_results/ --refert_num 1
```
* Multi-GPU inference using FSDP + DeepSpeed Ulysses
```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 8 generate.py --task animate-14B --ckpt_dir ./Wan2.2-Animate-14B/ --src_root_path ./examples/wan_animate/animate/process_results/ --refert_num 1 --dit_fsdp --t5_fsdp --ulysses_size 8
```
* Diffusers Pipeline
```python
from diffusers import WanAnimatePipeline
from diffusers.utils import export_to_video, load_image, load_video
device = "cuda:0"
dtype = torch.bfloat16
model_id = "Wan-AI/Wan2.2-Animate-14B-Diffusers"
pipe = WanAnimatePipeline.from_pretrained(model_id torch_dtype=dtype)
pipe.to(device)
seed = 42
prompt = "People in the video are doing actions."
# Animation
image = load_image("/path/to/animate/reference/image/src_ref.png")
pose_video = load_video("/path/to/animate/pose/video/src_pose.mp4")
face_video = load_video("/path/to/animate/face/video/src_face.mp4")
animate_video = pipe(
image=image,
pose_video=pose_video,
face_video=face_video,
prompt=prompt,
mode="animate",
segment_frame_length=77, # clip_len in original code
prev_segment_conditioning_frames=1, # refert_num in original code
guidance_scale=1.0,
num_inference_steps=20,
generator=torch.Generator(device=device).manual_seed(seed),
).frames[0]
export_to_video(animate_video, "diffusers_animate.mp4", fps=30)
```
##### (3) Run in replacement mode
* Single-GPU inference
```bash
python generate.py --task animate-14B --ckpt_dir ./Wan2.2-Animate-14B/ --src_root_path ./examples/wan_animate/replace/process_results/ --refert_num 1 --replace_flag --use_relighting_lora
```
* Multi-GPU inference using FSDP + DeepSpeed Ulysses
```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 8 generate.py --task animate-14B --ckpt_dir ./Wan2.2-Animate-14B/ --src_root_path ./examples/wan_animate/replace/process_results/src_pose.mp4 --refert_num 1 --replace_flag --use_relighting_lora --dit_fsdp --t5_fsdp --ulysses_size 8
```
* Diffusers Pipeline
```python
# create pipeline as in the Animation code ☝️
# Replacement
image = load_image("/path/to/replace/reference/image/src_ref.png")
pose_video = load_video("/path/to/replace/pose/video/src_pose.mp4")
face_video = load_video("/path/to/replace/face/video/src_face.mp4")
background_video = load_video("/path/to/replace/background/video/src_bg.mp4")
mask_video = load_video("/path/to/replace/mask/video/src_mask.mp4")
replace_video = pipe(
image=image,
pose_video=pose_video,
face_video=face_video,
background_video=background_video,
mask_video=mask_video,
prompt=prompt,
mode="replace",
segment_frame_length=77, # clip_len in original code
prev_segment_conditioning_frames=1, # refert_num in original code
guidance_scale=1.0,
num_inference_steps=20,
generator=torch.Generator(device=device).manual_seed(seed),
).frames[0]
export_to_video(replace_video, "diffusers_replace.mp4", fps=30)
```
> 💡 If you're using **Wan-Animate**, we do not recommend using LoRA models trained on `Wan2.2`, since weight changes during training may lead to unexpected behavior.
## Computational Efficiency on Different GPUs
We test the computational efficiency of different **Wan2.2** 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) Multi-GPU: 14B: `--ulysses_size 4/8 --dit_fsdp --t5_fsdp`, 5B: `--ulysses_size 4/8 --offload_model True --convert_model_dtype --t5_cpu`; Single-GPU: 14B: `--offload_model True --convert_model_dtype`, 5B: `--offload_model True --convert_model_dtype --t5_cpu`
(--convert_model_dtype converts model parameter types to config.param_dtype);
> (2) The distributed testing utilizes the built-in FSDP and Ulysses implementations, with FlashAttention3 deployed on Hopper architecture GPUs;
> (3) Tests were run without the `--use_prompt_extend` flag;
> (4) Reported results are the average of multiple samples taken after the warm-up phase.
-------
## Introduction of Wan2.2
**Wan2.2** builds on the foundation of Wan2.1 with notable improvements in generation quality and model capability. This upgrade is driven by a series of key technical innovations, mainly including the Mixture-of-Experts (MoE) architecture, upgraded training data, and high-compression video generation.
##### (1) Mixture-of-Experts (MoE) Architecture
Wan2.2 introduces Mixture-of-Experts (MoE) architecture into the video generation diffusion model. MoE has been widely validated in large language models as an efficient approach to increase total model parameters while keeping inference cost nearly unchanged. In Wan2.2, the A14B model series adopts a two-expert design tailored to the denoising process of diffusion models: a high-noise expert for the early stages, focusing on overall layout; and a low-noise expert for the later stages, refining video details. Each expert model has about 14B parameters, resulting in a total of 27B parameters but only 14B active parameters per step, keeping inference computation and GPU memory nearly unchanged.
<div align="center">
<img src="assets/moe_arch.png" alt="" style="width: 90%;" />
</div>
The transition point between the two experts is determined by the signal-to-noise ratio (SNR), a metric that decreases monotonically as the denoising step $t$ increases. At the beginning of the denoising process, $t$ is large and the noise level is high, so the SNR is at its minimum, denoted as ${SNR}_{min}$. In this stage, the high-noise expert is activated. We define a threshold step ${t}_{moe}$ corresponding to half of the ${SNR}_{min}$, and switch to the low-noise expert when $t<{t}_{moe}$.
<div align="center">
<img src="assets/moe_2.png" alt="" style="width: 90%;" />
</div>
To validate the effectiveness of the MoE architecture, four settings are compared based on their validation loss curves. The baseline **Wan2.1** model does not employ the MoE architecture. Among the MoE-based variants, the **Wan2.1 & High-Noise Expert** reuses the Wan2.1 model as the low-noise expert while uses the Wan2.2's high-noise expert, while the **Wan2.1 & Low-Noise Expert** uses Wan2.1 as the high-noise expert and employ the Wan2.2's low-noise expert. The **Wan2.2 (MoE)** (our final version) achieves the lowest validation loss, indicating that its generated video distribution is closest to ground-truth and exhibits superior convergence.
##### (2) Efficient High-Definition Hybrid TI2V
To enable more efficient deployment, Wan2.2 also explores a high-compression design. In addition to the 27B MoE models, a 5B dense model, i.e., TI2V-5B, is released. It is supported by a high-compression Wan2.2-VAE, which achieves a $T\times H\times W$ compression ratio of $4\times16\times16$, increasing the overall compression rate to 64 while maintaining high-quality video reconstruction. With an additional patchification layer, the total compression ratio of TI2V-5B reaches $4\times32\times32$. Without specific optimization, TI2V-5B can generate a 5-second 720P video in under 9 minutes on a single consumer-grade GPU, ranking among the fastest 720P@24fps video generation models. This model also natively supports both text-to-video and image-to-video tasks within a single unified framework, covering both academic research and practical applications.
<div align="center">
<img src="assets/vae.png" alt="" style="width: 80%;" />
</div>
##### Comparisons to SOTAs
We compared Wan2.2 with leading closed-source commercial models on our new Wan-Bench 2.0, evaluating performance across multiple crucial dimensions. The results demonstrate that Wan2.2 achieves superior performance compared to these leading models.
<div align="center">
<img src="assets/performance.png" alt="" style="width: 90%;" />
</div>
## Citation
If you find our work helpful, please cite us.
```
@article{wan2025,
title={Wan: Open and Advanced Large-Scale Video Generative Models},
author={Team Wan and Ang Wang and Baole Ai and Bin Wen and Chaojie Mao and Chen-Wei Xie and Di Chen and Feiwu Yu and Haiming Zhao and Jianxiao Yang and Jianyuan Zeng and Jiayu Wang and Jingfeng Zhang and Jingren Zhou and Jinkai Wang and Jixuan Chen and Kai Zhu and Kang Zhao and Keyu Yan and Lianghua Huang and Mengyang Feng and Ningyi Zhang and Pandeng Li and Pingyu Wu and Ruihang Chu and Ruili Feng and Shiwei Zhang and Siyang Sun and Tao Fang and Tianxing Wang and Tianyi Gui and Tingyu Weng and Tong Shen and Wei Lin and Wei Wang and Wei Wang and Wenmeng Zhou and Wente Wang and Wenting Shen and Wenyuan Yu and Xianzhong Shi and Xiaoming Huang and Xin Xu and Yan Kou and Yangyu Lv and Yifei Li and Yijing Liu and Yiming Wang and Yingya Zhang and Yitong Huang and Yong Li and You Wu and Yu Liu and Yulin Pan and Yun Zheng and Yuntao Hong and Yupeng Shi and Yutong Feng and Zeyinzi Jiang and Zhen Han and Zhi-Fan Wu and Ziyu Liu},
journal = {arXiv preprint arXiv:2503.20314},
year={2025}
}
```
## License Agreement
The models in this repository are licensed under the Apache 2.0 License. We claim no rights over the your generated 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/AKNgpMK4Yj) or [WeChat groups](https://gw.alicdn.com/imgextra/i2/O1CN01tqjWFi1ByuyehkTSB_!!6000000000015-0-tps-611-1279.jpg)!
## /assets/comp_effic.png
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/assets/comp_effic.png
## /assets/logo.png
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/assets/logo.png
## /assets/moe_2.png
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/assets/moe_2.png
## /assets/moe_arch.png
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/assets/moe_arch.png
## /assets/performance.png
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/assets/performance.png
## /assets/vae.png
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/assets/vae.png
## /examples/Five Hundred Miles.MP3
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/Five Hundred Miles.MP3
## /examples/Five Hundred Miles.png
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/Five Hundred Miles.png
## /examples/i2v_input.JPG
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/i2v_input.JPG
## /examples/pose.mp4
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/pose.mp4
## /examples/pose.png
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/pose.png
## /examples/sing.MP3
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/sing.MP3
## /examples/talk.wav
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/talk.wav
## /examples/wan_animate/animate/image.jpeg
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/wan_animate/animate/image.jpeg
## /examples/wan_animate/animate/video.mp4
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/wan_animate/animate/video.mp4
## /examples/wan_animate/replace/image.jpeg
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/wan_animate/replace/image.jpeg
## /examples/wan_animate/replace/video.mp4
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/wan_animate/replace/video.mp4
## /examples/zero_shot_prompt.wav
Binary file available at https://raw.githubusercontent.com/Wan-Video/Wan2.2/refs/heads/main/examples/zero_shot_prompt.wav
## /generate.py
```py path="/generate.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import argparse
import logging
import os
import sys
import warnings
from datetime import datetime
warnings.filterwarnings('ignore')
import random
import torch
import torch.distributed as dist
from PIL import Image
import wan
from wan.configs import MAX_AREA_CONFIGS, SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
from wan.distributed.util import init_distributed_group
from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
from wan.utils.utils import merge_video_audio, save_video, str2bool
EXAMPLE_PROMPT = {
"t2v-A14B": {
"prompt":
"Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
},
"i2v-A14B": {
"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.",
"image":
"examples/i2v_input.JPG",
},
"ti2v-5B": {
"prompt":
"Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
},
"animate-14B": {
"prompt": "视频中的人在做动作",
"video": "",
"pose": "",
"mask": "",
},
"s2v-14B": {
"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.",
"image":
"examples/i2v_input.JPG",
"audio":
"examples/talk.wav",
"tts_prompt_audio":
"examples/zero_shot_prompt.wav",
"tts_prompt_text":
"希望你以后能够做的比我还好呦。",
"tts_text":
"收到好友从远方寄来的生日礼物，那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐，笑容如花儿般绽放。"
},
}
def _validate_args(args):
# Basic check
assert args.ckpt_dir is not None, "Please specify the checkpoint directory."
assert args.task in WAN_CONFIGS, f"Unsupport task: {args.task}"
assert args.task in EXAMPLE_PROMPT, f"Unsupport task: {args.task}"
if args.prompt is None:
args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
if args.image is None and "image" in EXAMPLE_PROMPT[args.task]:
args.image = EXAMPLE_PROMPT[args.task]["image"]
if args.audio is None and args.enable_tts is False and "audio" in EXAMPLE_PROMPT[args.task]:
args.audio = EXAMPLE_PROMPT[args.task]["audio"]
if (args.tts_prompt_audio is None or args.tts_text is None) and args.enable_tts is True and "audio" in EXAMPLE_PROMPT[args.task]:
args.tts_prompt_audio = EXAMPLE_PROMPT[args.task]["tts_prompt_audio"]
args.tts_prompt_text = EXAMPLE_PROMPT[args.task]["tts_prompt_text"]
args.tts_text = EXAMPLE_PROMPT[args.task]["tts_text"]
if args.task == "i2v-A14B":
assert args.image is not None, "Please specify the image path for i2v."
cfg = WAN_CONFIGS[args.task]
if args.sample_steps is None:
args.sample_steps = cfg.sample_steps
if args.sample_shift is None:
args.sample_shift = cfg.sample_shift
if args.sample_guide_scale is None:
args.sample_guide_scale = cfg.sample_guide_scale
if args.frame_num is None:
args.frame_num = cfg.frame_num
args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
0, sys.maxsize)
# Size check
if not 's2v' in args.task:
assert args.size in SUPPORTED_SIZES[
args.
task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"
def _parse_args():
parser = argparse.ArgumentParser(
description="Generate a image or video from a text prompt or image using Wan"
)
parser.add_argument(
"--task",
type=str,
default="t2v-A14B",
choices=list(WAN_CONFIGS.keys()),
help="The task to run.")
parser.add_argument(
"--size",
type=str,
default="1280*720",
choices=list(SIZE_CONFIGS.keys()),
help="The area (width*height) of the generated video. For the I2V task, the aspect ratio of the output video will follow that of the input image."
)
parser.add_argument(
"--frame_num",
type=int,
default=None,
help="How many frames of video are generated. The number should be 4n+1"
)
parser.add_argument(
"--ckpt_dir",
type=str,
default=None,
help="The path to the checkpoint directory.")
parser.add_argument(
"--offload_model",
type=str2bool,
default=None,
help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage."
)
parser.add_argument(
"--ulysses_size",
type=int,
default=1,
help="The size of the ulysses parallelism in DiT.")
parser.add_argument(
"--t5_fsdp",
action="store_true",
default=False,
help="Whether to use FSDP for T5.")
parser.add_argument(
"--t5_cpu",
action="store_true",
default=False,
help="Whether to place T5 model on CPU.")
parser.add_argument(
"--dit_fsdp",
action="store_true",
default=False,
help="Whether to use FSDP for DiT.")
parser.add_argument(
"--save_file",
type=str,
default=None,
help="The file to save the generated video to.")
parser.add_argument(
"--prompt",
type=str,
default=None,
help="The prompt to generate the video from.")
parser.add_argument(
"--use_prompt_extend",
action="store_true",
default=False,
help="Whether to use prompt extend.")
parser.add_argument(
"--prompt_extend_method",
type=str,
default="local_qwen",
choices=["dashscope", "local_qwen"],
help="The prompt extend method to use.")
parser.add_argument(
"--prompt_extend_model",
type=str,
default=None,
help="The prompt extend model to use.")
parser.add_argument(
"--prompt_extend_target_lang",
type=str,
default="zh",
choices=["zh", "en"],
help="The target language of prompt extend.")
parser.add_argument(
"--base_seed",
type=int,
default=-1,
help="The seed to use for generating the video.")
parser.add_argument(
"--image",
type=str,
default=None,
help="The image to generate the video from.")
parser.add_argument(
"--sample_solver",
type=str,
default='unipc',
choices=['unipc', 'dpm++'],
help="The solver used to sample.")
parser.add_argument(
"--sample_steps", type=int, default=None, help="The sampling steps.")
parser.add_argument(
"--sample_shift",
type=float,
default=None,
help="Sampling shift factor for flow matching schedulers.")
parser.add_argument(
"--sample_guide_scale",
type=float,
default=None,
help="Classifier free guidance scale.")
parser.add_argument(
"--convert_model_dtype",
action="store_true",
default=False,
help="Whether to convert model paramerters dtype.")
# animate
parser.add_argument(
"--src_root_path",
type=str,
default=None,
help="The file of the process output path. Default None.")
parser.add_argument(
"--refert_num",
type=int,
default=77,
help="How many frames used for temporal guidance. Recommended to be 1 or 5."
)
parser.add_argument(
"--replace_flag",
action="store_true",
default=False,
help="Whether to use replace.")
parser.add_argument(
"--use_relighting_lora",
action="store_true",
default=False,
help="Whether to use relighting lora.")
# following args only works for s2v
parser.add_argument(
"--num_clip",
type=int,
default=None,
help="Number of video clips to generate, the whole video will not exceed the length of audio."
)
parser.add_argument(
"--audio",
type=str,
default=None,
help="Path to the audio file, e.g. wav, mp3")
parser.add_argument(
"--enable_tts",
action="store_true",
default=False,
help="Use CosyVoice to synthesis audio")
parser.add_argument(
"--tts_prompt_audio",
type=str,
default=None,
help="Path to the tts prompt audio file, e.g. wav, mp3. Must be greater than 16khz, and between 5s to 15s.")
parser.add_argument(
"--tts_prompt_text",
type=str,
default=None,
help="Content to the tts prompt audio. If provided, must exactly match tts_prompt_audio")
parser.add_argument(
"--tts_text",
type=str,
default=None,
help="Text wish to synthesize")
parser.add_argument(
"--pose_video",
type=str,
default=None,
help="Provide Dw-pose sequence to do Pose Driven")
parser.add_argument(
"--start_from_ref",
action="store_true",
default=False,
help="whether set the reference image as the starting point for generation"
)
parser.add_argument(
"--infer_frames",
type=int,
default=80,
help="Number of frames per clip, 48 or 80 or others (must be multiple of 4) for 14B s2v"
)
args = parser.parse_args()
_validate_args(args)
return args
def _init_logging(rank):
# logging
if rank == 0:
# set format
logging.basicConfig(
level=logging.INFO,
format="[%(asctime)s] %(levelname)s: %(message)s",
handlers=[logging.StreamHandler(stream=sys.stdout)])
else:
logging.basicConfig(level=logging.ERROR)
def generate(args):
rank = int(os.getenv("RANK", 0))
world_size = int(os.getenv("WORLD_SIZE", 1))
local_rank = int(os.getenv("LOCAL_RANK", 0))
device = local_rank
_init_logging(rank)
if args.offload_model is None:
args.offload_model = False if world_size > 1 else True
logging.info(
f"offload_model is not specified, set to {args.offload_model}.")
if world_size > 1:
torch.cuda.set_device(local_rank)
dist.init_process_group(
backend="nccl",
init_method="env://",
rank=rank,
world_size=world_size)
else:
assert not (
args.t5_fsdp or args.dit_fsdp
), f"t5_fsdp and dit_fsdp are not supported in non-distributed environments."
assert not (
args.ulysses_size > 1
), f"sequence parallel are not supported in non-distributed environments."
if args.ulysses_size > 1:
assert args.ulysses_size == world_size, f"The number of ulysses_size should be equal to the world size."
init_distributed_group()
if args.use_prompt_extend:
if args.prompt_extend_method == "dashscope":
prompt_expander = DashScopePromptExpander(
model_name=args.prompt_extend_model,
task=args.task,
is_vl=args.image is not None)
elif args.prompt_extend_method == "local_qwen":
prompt_expander = QwenPromptExpander(
model_name=args.prompt_extend_model,
task=args.task,
is_vl=args.image is not None,
device=rank)
else:
raise NotImplementedError(
f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
cfg = WAN_CONFIGS[args.task]
if args.ulysses_size > 1:
assert cfg.num_heads % args.ulysses_size == 0, f"`{cfg.num_heads=}` cannot be divided evenly by `{args.ulysses_size=}`."
logging.info(f"Generation job args: {args}")
logging.info(f"Generation model config: {cfg}")
if dist.is_initialized():
base_seed = [args.base_seed] if rank == 0 else [None]
dist.broadcast_object_list(base_seed, src=0)
args.base_seed = base_seed[0]
logging.info(f"Input prompt: {args.prompt}")
img = None
if args.image is not None:
img = Image.open(args.image).convert("RGB")
logging.info(f"Input image: {args.image}")
# prompt extend
if args.use_prompt_extend:
logging.info("Extending prompt ...")
if rank == 0:
prompt_output = prompt_expander(
args.prompt,
image=img,
tar_lang=args.prompt_extend_target_lang,
seed=args.base_seed)
if prompt_output.status == False:
logging.info(
f"Extending prompt failed: {prompt_output.message}")
logging.info("Falling back to original prompt.")
input_prompt = args.prompt
else:
input_prompt = prompt_output.prompt
input_prompt = [input_prompt]
else:
input_prompt = [None]
if dist.is_initialized():
dist.broadcast_object_list(input_prompt, src=0)
args.prompt = input_prompt[0]
logging.info(f"Extended prompt: {args.prompt}")
if "t2v" in args.task:
logging.info("Creating WanT2V pipeline.")
wan_t2v = wan.WanT2V(
config=cfg,
checkpoint_dir=args.ckpt_dir,
device_id=device,
rank=rank,
t5_fsdp=args.t5_fsdp,
dit_fsdp=args.dit_fsdp,
use_sp=(args.ulysses_size > 1),
t5_cpu=args.t5_cpu,
convert_model_dtype=args.convert_model_dtype,
)
logging.info(f"Generating video ...")
video = wan_t2v.generate(
args.prompt,
size=SIZE_CONFIGS[args.size],
frame_num=args.frame_num,
shift=args.sample_shift,
sample_solver=args.sample_solver,
sampling_steps=args.sample_steps,
guide_scale=args.sample_guide_scale,
seed=args.base_seed,
offload_model=args.offload_model)
elif "ti2v" in args.task:
logging.info("Creating WanTI2V pipeline.")
wan_ti2v = wan.WanTI2V(
config=cfg,
checkpoint_dir=args.ckpt_dir,
device_id=device,
rank=rank,
t5_fsdp=args.t5_fsdp,
dit_fsdp=args.dit_fsdp,
use_sp=(args.ulysses_size > 1),
t5_cpu=args.t5_cpu,
convert_model_dtype=args.convert_model_dtype,
)
logging.info(f"Generating video ...")
video = wan_ti2v.generate(
args.prompt,
img=img,
size=SIZE_CONFIGS[args.size],
max_area=MAX_AREA_CONFIGS[args.size],
frame_num=args.frame_num,
shift=args.sample_shift,
sample_solver=args.sample_solver,
sampling_steps=args.sample_steps,
guide_scale=args.sample_guide_scale,
seed=args.base_seed,
offload_model=args.offload_model)
elif "animate" in args.task:
logging.info("Creating Wan-Animate pipeline.")
wan_animate = wan.WanAnimate(
config=cfg,
checkpoint_dir=args.ckpt_dir,
device_id=device,
rank=rank,
t5_fsdp=args.t5_fsdp,
dit_fsdp=args.dit_fsdp,
use_sp=(args.ulysses_size > 1),
t5_cpu=args.t5_cpu,
convert_model_dtype=args.convert_model_dtype,
use_relighting_lora=args.use_relighting_lora
)
logging.info(f"Generating video ...")
video = wan_animate.generate(
src_root_path=args.src_root_path,
replace_flag=args.replace_flag,
refert_num = args.refert_num,
clip_len=args.frame_num,
shift=args.sample_shift,
sample_solver=args.sample_solver,
sampling_steps=args.sample_steps,
guide_scale=args.sample_guide_scale,
seed=args.base_seed,
offload_model=args.offload_model)
elif "s2v" in args.task:
logging.info("Creating WanS2V pipeline.")
wan_s2v = wan.WanS2V(
config=cfg,
checkpoint_dir=args.ckpt_dir,
device_id=device,
rank=rank,
t5_fsdp=args.t5_fsdp,
dit_fsdp=args.dit_fsdp,
use_sp=(args.ulysses_size > 1),
t5_cpu=args.t5_cpu,
convert_model_dtype=args.convert_model_dtype,
)
logging.info(f"Generating video ...")
video = wan_s2v.generate(
input_prompt=args.prompt,
ref_image_path=args.image,
audio_path=args.audio,
enable_tts=args.enable_tts,
tts_prompt_audio=args.tts_prompt_audio,
tts_prompt_text=args.tts_prompt_text,
tts_text=args.tts_text,
num_repeat=args.num_clip,
pose_video=args.pose_video,
max_area=MAX_AREA_CONFIGS[args.size],
infer_frames=args.infer_frames,
shift=args.sample_shift,
sample_solver=args.sample_solver,
sampling_steps=args.sample_steps,
guide_scale=args.sample_guide_scale,
seed=args.base_seed,
offload_model=args.offload_model,
init_first_frame=args.start_from_ref,
)
else:
logging.info("Creating WanI2V pipeline.")
wan_i2v = wan.WanI2V(
config=cfg,
checkpoint_dir=args.ckpt_dir,
device_id=device,
rank=rank,
t5_fsdp=args.t5_fsdp,
dit_fsdp=args.dit_fsdp,
use_sp=(args.ulysses_size > 1),
t5_cpu=args.t5_cpu,
convert_model_dtype=args.convert_model_dtype,
)
logging.info("Generating video ...")
video = wan_i2v.generate(
args.prompt,
img,
max_area=MAX_AREA_CONFIGS[args.size],
frame_num=args.frame_num,
shift=args.sample_shift,
sample_solver=args.sample_solver,
sampling_steps=args.sample_steps,
guide_scale=args.sample_guide_scale,
seed=args.base_seed,
offload_model=args.offload_model)
if rank == 0:
if args.save_file is None:
formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
formatted_prompt = args.prompt.replace(" ", "_").replace("/",
"_")[:50]
suffix = '.mp4'
args.save_file = f"{args.task}_{args.size.replace('*','x') if sys.platform=='win32' else args.size}_{args.ulysses_size}_{formatted_prompt}_{formatted_time}" + suffix
logging.info(f"Saving generated video to {args.save_file}")
save_video(
tensor=video[None],
save_file=args.save_file,
fps=cfg.sample_fps,
nrow=1,
normalize=True,
value_range=(-1, 1))
if "s2v" in args.task:
if args.enable_tts is False:
merge_video_audio(video_path=args.save_file, audio_path=args.audio)
else:
merge_video_audio(video_path=args.save_file, audio_path="tts.wav")
del video
torch.cuda.synchronize()
if dist.is_initialized():
dist.barrier()
dist.destroy_process_group()
logging.info("Finished.")
if __name__ == "__main__":
args = _parse_args()
generate(args)
```
## /pyproject.toml
```toml path="/pyproject.toml"
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[project]
name = "wan"
version = "2.2.0"
description = "Wan: Open and Advanced Large-Scale Video Generative Models"
authors = [
{ name = "Wan Team", email = "wan.ai@alibabacloud.com" }
]
license = { file = "LICENSE.txt" }
readme = "README.md"
requires-python = ">=3.10,<4.0"
dependencies = [
"torch>=2.4.0",
"torchvision>=0.19.0",
"opencv-python>=4.9.0.80",
"diffusers>=0.31.0",
"transformers>=4.49.0",
"tokenizers>=0.20.3",
"accelerate>=1.1.1",
"tqdm",
"imageio",
"easydict",
"ftfy",
"dashscope",
"imageio-ffmpeg",
"flash_attn",
"numpy>=1.23.5,<2"
]
[project.optional-dependencies]
dev = [
"pytest",
"black",
"flake8",
"isort",
"mypy",
"huggingface-hub[cli]"
]
[project.urls]
homepage = "https://wanxai.com"
documentation = "https://github.com/Wan-Video/Wan2.2"
repository = "https://github.com/Wan-Video/Wan2.2"
huggingface = "https://huggingface.co/Wan-AI/"
modelscope = "https://modelscope.cn/organization/Wan-AI"
discord = "https://discord.gg/p5XbdQV7"
[tool.setuptools]
packages = ["wan"]
[tool.setuptools.package-data]
"wan" = ["**/*.py"]
[tool.black]
line-length = 88
[tool.isort]
profile = "black"
[tool.mypy]
strict = true
```
## /requirements.txt
torch>=2.4.0
torchvision>=0.19.0
torchaudio
opencv-python>=4.9.0.80
diffusers>=0.31.0
transformers>=4.49.0,<=4.51.3
tokenizers>=0.20.3
accelerate>=1.1.1
tqdm
imageio[ffmpeg]
easydict
ftfy
dashscope
imageio-ffmpeg
flash_attn
numpy>=1.23.5,<2
## /requirements_animate.txt
decord
peft
onnxruntime
pandas
matplotlib
-e git+https://github.com/facebookresearch/sam2.git@0e78a118995e66bb27d78518c4bd9a3e95b4e266#egg=SAM-2
loguru
sentencepiece
## /requirements_s2v.txt
openai-whisper
HyperPyYAML
onnxruntime
inflect
wetext
omegaconf
conformer
hydra-core
lightning
rich
gdown
matplotlib
wget
pyarrow
pyworld
librosa
decord
modelscope
GitPython
## /tests/README.md
Put all your models (Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B) in a folder and specify the max GPU number you want to use.
```bash
bash ./tests/test.sh <local model dir> <gpu number>
```
## /tests/test.sh
```sh path="/tests/test.sh"
#!/bin/bash
set -x
unset NCCL_DEBUG
if [ "$#" -eq 2 ]; then
MODEL_DIR=$(realpath "$1")
GPUS=$2
else
echo "Usage: $0 <local model dir> <gpu number>"
exit 1
fi
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
REPO_ROOT="$(dirname "$SCRIPT_DIR")"
cd "$REPO_ROOT" || exit 1
PY_FILE=./generate.py
function t2v_A14B() {
CKPT_DIR="$MODEL_DIR/Wan2.2-T2V-A14B"
# # 1-GPU Test
# echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B 1-GPU Test: "
# python $PY_FILE --task t2v-A14B --size 480*832 --ckpt_dir $CKPT_DIR
# Multiple GPU Test
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B Multiple GPU Test: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task t2v-A14B --ckpt_dir $CKPT_DIR --size 832*480 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B Multiple GPU Test: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task t2v-A14B --ckpt_dir $CKPT_DIR --size 720*1280 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B Multiple GPU Test: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task t2v-A14B --ckpt_dir $CKPT_DIR --size 1280*720 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B Multiple GPU, prompt extend local_qwen: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task t2v-A14B --ckpt_dir $CKPT_DIR --size 480*832 --dit_fsdp --t5_fsdp --ulysses_size $GPUS --use_prompt_extend --prompt_extend_model "Qwen/Qwen2.5-3B-Instruct" --prompt_extend_target_lang "en"
}
function i2v_A14B() {
CKPT_DIR="$MODEL_DIR/Wan2.2-I2V-A14B"
# echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B 1-GPU Test: "
# python $PY_FILE --task i2v-A14B --size 832*480 --ckpt_dir $CKPT_DIR
# Multiple GPU Test
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B Multiple GPU Test: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task i2v-A14B --ckpt_dir $CKPT_DIR --size 832*480 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B Multiple GPU, prompt extend local_qwen: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task i2v-A14B --ckpt_dir $CKPT_DIR --size 720*1280 --dit_fsdp --t5_fsdp --ulysses_size $GPUS --use_prompt_extend --prompt_extend_model "Qwen/Qwen2.5-VL-3B-Instruct" --prompt_extend_target_lang "en"
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B Multiple GPU, prompt extend local_qwen: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task i2v-A14B --ckpt_dir $CKPT_DIR --size 1280*720 --dit_fsdp --t5_fsdp --ulysses_size $GPUS --use_prompt_extend --prompt_extend_model "Qwen/Qwen2.5-VL-3B-Instruct" --prompt_extend_target_lang "en"
if [ -n "${DASH_API_KEY+x}" ]; then
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B Multiple GPU, prompt extend dashscope: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task i2v-A14B --ckpt_dir $CKPT_DIR --size 480*832 --dit_fsdp --t5_fsdp --ulysses_size $GPUS --use_prompt_extend --prompt_extend_method "dashscope"
else
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> No DASH_API_KEY found, skip the dashscope extend test."
fi
}
function ti2v_5B() {
CKPT_DIR="$MODEL_DIR/Wan2.2-TI2V-5B"
# # 1-GPU Test
# echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B t2v 1-GPU Test: "
# python $PY_FILE --task ti2v-5B --size 1280*704 --ckpt_dir $CKPT_DIR
# Multiple GPU Test
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B t2v Multiple GPU Test: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task ti2v-5B --ckpt_dir $CKPT_DIR --size 1280*704 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B t2v Multiple GPU, prompt extend local_qwen: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task ti2v-5B --ckpt_dir $CKPT_DIR --size 704*1280 --dit_fsdp --t5_fsdp --ulysses_size $GPUS --use_prompt_extend --prompt_extend_model "Qwen/Qwen2.5-3B-Instruct" --prompt_extend_target_lang "en"
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B i2v Multiple GPU Test: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task ti2v-5B --ckpt_dir $CKPT_DIR --size 704*1280 --dit_fsdp --t5_fsdp --ulysses_size $GPUS --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." --image "examples/i2v_input.JPG"
echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B i2v Multiple GPU, prompt extend local_qwen: "
torchrun --nproc_per_node=$GPUS $PY_FILE --task ti2v-5B --ckpt_dir $CKPT_DIR --size 1280*704 --dit_fsdp --t5_fsdp --ulysses_size $GPUS --use_prompt_extend --prompt_extend_model "Qwen/Qwen2.5-3B-Instruct" --prompt_extend_target_lang 'en' --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." --image "examples/i2v_input.JPG"
}
t2v_A14B
i2v_A14B
ti2v_5B
```
## /wan/__init__.py
```py path="/wan/__init__.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from . import configs, distributed, modules
from .image2video import WanI2V
from .speech2video import WanS2V
from .text2video import WanT2V
from .textimage2video import WanTI2V
from .animate import WanAnimate
```
## /wan/animate.py
```py path="/wan/animate.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import logging
import math
import os
import cv2
import types
from copy import deepcopy
from functools import partial
from einops import rearrange
import numpy as np
import torch
import torch.distributed as dist
from peft import set_peft_model_state_dict
from decord import VideoReader
from tqdm import tqdm
import torch.nn.functional as F
from .distributed.fsdp import shard_model
from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward
from .distributed.util import get_world_size
from .modules.animate import WanAnimateModel
from .modules.animate import CLIPModel
from .modules.t5 import T5EncoderModel
from .modules.vae2_1 import Wan2_1_VAE
from .modules.animate.animate_utils import TensorList, get_loraconfig
from .utils.fm_solvers import (
FlowDPMSolverMultistepScheduler,
get_sampling_sigmas,
retrieve_timesteps,
)
from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
class WanAnimate:
def __init__(
self,
config,
checkpoint_dir,
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_sp=False,
t5_cpu=False,
init_on_cpu=True,
convert_model_dtype=False,
use_relighting_lora=False
):
r"""
Initializes the generation model components.
Args:
config (EasyDict):
Object containing model parameters initialized from config.py
checkpoint_dir (`str`):
Path to directory containing model checkpoints
device_id (`int`, *optional*, defaults to 0):
Id of target GPU device
rank (`int`, *optional*, defaults to 0):
Process rank for distributed training
t5_fsdp (`bool`, *optional*, defaults to False):
Enable FSDP sharding for T5 model
dit_fsdp (`bool`, *optional*, defaults to False):
Enable FSDP sharding for DiT model
use_sp (`bool`, *optional*, defaults to False):
Enable distribution strategy of sequence parallel.
t5_cpu (`bool`, *optional*, defaults to False):
Whether to place T5 model on CPU. Only works without t5_fsdp.
init_on_cpu (`bool`, *optional*, defaults to True):
Enable initializing Transformer Model on CPU. Only works without FSDP or USP.
convert_model_dtype (`bool`, *optional*, defaults to False):
Convert DiT model parameters dtype to 'config.param_dtype'.
Only works without FSDP.
use_relighting_lora (`bool`, *optional*, defaults to False):
Whether to use relighting lora for character replacement.
"""
self.device = torch.device(f"cuda:{device_id}")
self.config = config
self.rank = rank
self.t5_cpu = t5_cpu
self.init_on_cpu = init_on_cpu
self.num_train_timesteps = config.num_train_timesteps
self.param_dtype = config.param_dtype
if t5_fsdp or dit_fsdp or use_sp:
self.init_on_cpu = False
shard_fn = partial(shard_model, device_id=device_id)
self.text_encoder = T5EncoderModel(
text_len=config.text_len,
dtype=config.t5_dtype,
device=torch.device('cpu'),
checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
shard_fn=shard_fn if t5_fsdp else None,
)
self.clip = CLIPModel(
dtype=torch.float16,
device=self.device,
checkpoint_path=os.path.join(checkpoint_dir,
config.clip_checkpoint),
tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
self.vae = Wan2_1_VAE(
vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
device=self.device)
logging.info(f"Creating WanAnimate from {checkpoint_dir}")
if not dit_fsdp:
self.noise_model = WanAnimateModel.from_pretrained(
checkpoint_dir,
torch_dtype=self.param_dtype,
device_map=self.device)
else:
self.noise_model = WanAnimateModel.from_pretrained(
checkpoint_dir, torch_dtype=self.param_dtype)
self.noise_model = self._configure_model(
model=self.noise_model,
use_sp=use_sp,
dit_fsdp=dit_fsdp,
shard_fn=shard_fn,
convert_model_dtype=convert_model_dtype,
use_lora=use_relighting_lora,
checkpoint_dir=checkpoint_dir,
config=config
)
if use_sp:
self.sp_size = get_world_size()
else:
self.sp_size = 1
self.sample_neg_prompt = config.sample_neg_prompt
self.sample_prompt = config.prompt
def _configure_model(self, model, use_sp, dit_fsdp, shard_fn,
convert_model_dtype, use_lora, checkpoint_dir, config):
"""
Configures a model object. This includes setting evaluation modes,
applying distributed parallel strategy, and handling device placement.
Args:
model (torch.nn.Module):
The model instance to configure.
use_sp (`bool`):
Enable distribution strategy of sequence parallel.
dit_fsdp (`bool`):
Enable FSDP sharding for DiT model.
shard_fn (callable):
The function to apply FSDP sharding.
convert_model_dtype (`bool`):
Convert DiT model parameters dtype to 'config.param_dtype'.
Only works without FSDP.
Returns:
torch.nn.Module:
The configured model.
"""
model.eval().requires_grad_(False)
if use_sp:
for block in model.blocks:
block.self_attn.forward = types.MethodType(
sp_attn_forward, block.self_attn)
model.use_context_parallel = True
if dist.is_initialized():
dist.barrier()
if use_lora:
logging.info("Loading Relighting Lora. ")
lora_config = get_loraconfig(
transformer=model,
rank=128,
alpha=128
)
model.add_adapter(lora_config)
lora_path = os.path.join(checkpoint_dir, config.lora_checkpoint)
peft_state_dict = torch.load(lora_path)["state_dict"]
set_peft_model_state_dict(model, peft_state_dict)
if dit_fsdp:
model = shard_fn(model, use_lora=use_lora)
else:
if convert_model_dtype:
model.to(self.param_dtype)
if not self.init_on_cpu:
model.to(self.device)
return model
def inputs_padding(self, array, target_len):
idx = 0
flip = False
target_array = []
while len(target_array) < target_len:
target_array.append(deepcopy(array[idx]))
if flip:
idx -= 1
else:
idx += 1
if idx == 0 or idx == len(array) - 1:
flip = not flip
return target_array[:target_len]
def get_valid_len(self, real_len, clip_len=81, overlap=1):
real_clip_len = clip_len - overlap
last_clip_num = (real_len - overlap) % real_clip_len
if last_clip_num == 0:
extra = 0
else:
extra = real_clip_len - last_clip_num
target_len = real_len + extra
return target_len
def get_i2v_mask(self, lat_t, lat_h, lat_w, mask_len=1, mask_pixel_values=None, device="cuda"):
if mask_pixel_values is None:
msk = torch.zeros(1, (lat_t-1) * 4 + 1, lat_h, lat_w, device=device)
else:
msk = mask_pixel_values.clone()
msk[:, :mask_len] = 1
msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1)
msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
msk = msk.transpose(1, 2)[0]
return msk
def padding_resize(self, img_ori, height=512, width=512, padding_color=(0, 0, 0), interpolation=cv2.INTER_LINEAR):
ori_height = img_ori.shape[0]
ori_width = img_ori.shape[1]
channel = img_ori.shape[2]
img_pad = np.zeros((height, width, channel))
if channel == 1:
img_pad[:, :, 0] = padding_color[0]
else:
img_pad[:, :, 0] = padding_color[0]
img_pad[:, :, 1] = padding_color[1]
img_pad[:, :, 2] = padding_color[2]
if (ori_height / ori_width) > (height / width):
new_width = int(height / ori_height * ori_width)
img = cv2.resize(img_ori, (new_width, height), interpolation=interpolation)
padding = int((width - new_width) / 2)
if len(img.shape) == 2:
img = img[:, :, np.newaxis]
img_pad[:, padding: padding + new_width, :] = img
else:
new_height = int(width / ori_width * ori_height)
img = cv2.resize(img_ori, (width, new_height), interpolation=interpolation)
padding = int((height - new_height) / 2)
if len(img.shape) == 2:
img = img[:, :, np.newaxis]
img_pad[padding: padding + new_height, :, :] = img
img_pad = np.uint8(img_pad)
return img_pad
def prepare_source(self, src_pose_path, src_face_path, src_ref_path):
pose_video_reader = VideoReader(src_pose_path)
pose_len = len(pose_video_reader)
pose_idxs = list(range(pose_len))
cond_images = pose_video_reader.get_batch(pose_idxs).asnumpy()
face_video_reader = VideoReader(src_face_path)
face_len = len(face_video_reader)
face_idxs = list(range(face_len))
face_images = face_video_reader.get_batch(face_idxs).asnumpy()
height, width = cond_images[0].shape[:2]
refer_images = cv2.imread(src_ref_path)[..., ::-1]
refer_images = self.padding_resize(refer_images, height=height, width=width)
return cond_images, face_images, refer_images
def prepare_source_for_replace(self, src_bg_path, src_mask_path):
bg_video_reader = VideoReader(src_bg_path)
bg_len = len(bg_video_reader)
bg_idxs = list(range(bg_len))
bg_images = bg_video_reader.get_batch(bg_idxs).asnumpy()
mask_video_reader = VideoReader(src_mask_path)
mask_len = len(mask_video_reader)
mask_idxs = list(range(mask_len))
mask_images = mask_video_reader.get_batch(mask_idxs).asnumpy()
mask_images = mask_images[:, :, :, 0] / 255
return bg_images, mask_images
def generate(
self,
src_root_path,
replace_flag=False,
clip_len=77,
refert_num=1,
shift=5.0,
sample_solver='dpm++',
sampling_steps=20,
guide_scale=1,
input_prompt="",
n_prompt="",
seed=-1,
offload_model=True,
):
r"""
Generates video frames from input image using diffusion process.
Args:
src_root_path ('str'):
Process output path
replace_flag (`bool`, *optional*, defaults to False):
Whether to use character replace.
clip_len (`int`, *optional*, defaults to 77):
How many frames to generate per clips. The number should be 4n+1
refert_num (`int`, *optional*, defaults to 1):
How many frames used for temporal guidance. Recommended to be 1 or 5.
shift (`float`, *optional*, defaults to 5.0):
Noise schedule shift parameter.
sample_solver (`str`, *optional*, defaults to 'dpm++'):
Solver used to sample the video.
sampling_steps (`int`, *optional*, defaults to 20):
Number of diffusion sampling steps. Higher values improve quality but slow generation
guide_scale (`float` or tuple[`float`], *optional*, defaults 1.0):
Classifier-free guidance scale. We only use it for expression control.
In most cases, it's not necessary and faster generation can be achieved without it.
When expression adjustments are needed, you may consider using this feature.
input_prompt (`str`):
Text prompt for content generation. We don't recommend custom prompts (although they work)
n_prompt (`str`, *optional*, defaults to ""):
Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
seed (`int`, *optional*, defaults to -1):
Random seed for noise generation. If -1, use random seed
offload_model (`bool`, *optional*, defaults to True):
If True, offloads models to CPU during generation to save VRAM
Returns:
torch.Tensor:
Generated video frames tensor. Dimensions: (C, N, H, W) where:
- C: Color channels (3 for RGB)
- N: Number of frames
- H: Frame height
- W: Frame width
"""
assert refert_num == 1 or refert_num == 5, "refert_num should be 1 or 5."
seed_g = torch.Generator(device=self.device)
seed_g.manual_seed(seed)
if n_prompt == "":
n_prompt = self.sample_neg_prompt
if input_prompt == "":
input_prompt = self.sample_prompt
src_pose_path = os.path.join(src_root_path, "src_pose.mp4")
src_face_path = os.path.join(src_root_path, "src_face.mp4")
src_ref_path = os.path.join(src_root_path, "src_ref.png")
cond_images, face_images, refer_images = self.prepare_source(src_pose_path=src_pose_path, src_face_path=src_face_path, src_ref_path=src_ref_path)
if not self.t5_cpu:
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:
self.text_encoder.model.cpu()
else:
context = self.text_encoder([input_prompt], torch.device('cpu'))
context_null = self.text_encoder([n_prompt], torch.device('cpu'))
context = [t.to(self.device) for t in context]
context_null = [t.to(self.device) for t in context_null]
real_frame_len = len(cond_images)
target_len = self.get_valid_len(real_frame_len, clip_len, overlap=refert_num)
logging.info('real frames: {} target frames: {}'.format(real_frame_len, target_len))
cond_images = self.inputs_padding(cond_images, target_len)
face_images = self.inputs_padding(face_images, target_len)
if replace_flag:
src_bg_path = os.path.join(src_root_path, "src_bg.mp4")
src_mask_path = os.path.join(src_root_path, "src_mask.mp4")
bg_images, mask_images = self.prepare_source_for_replace(src_bg_path, src_mask_path)
bg_images = self.inputs_padding(bg_images, target_len)
mask_images = self.inputs_padding(mask_images, target_len)
height, width = refer_images.shape[:2]
start = 0
end = clip_len
all_out_frames = []
while True:
if start + refert_num >= len(cond_images):
break
if start == 0:
mask_reft_len = 0
else:
mask_reft_len = refert_num
batch = {
"conditioning_pixel_values": torch.zeros(1, 3, clip_len, height, width),
"bg_pixel_values": torch.zeros(1, 3, clip_len, height, width),
"mask_pixel_values": torch.zeros(1, 1, clip_len, height, width),
"face_pixel_values": torch.zeros(1, 3, clip_len, 512, 512),
"refer_pixel_values": torch.zeros(1, 3, height, width),
"refer_t_pixel_values": torch.zeros(refert_num, 3, height, width)
}
batch["conditioning_pixel_values"] = rearrange(
torch.tensor(np.stack(cond_images[start:end]) / 127.5 - 1),
"t h w c -> 1 c t h w",
)
batch["face_pixel_values"] = rearrange(
torch.tensor(np.stack(face_images[start:end]) / 127.5 - 1),
"t h w c -> 1 c t h w",
)
batch["refer_pixel_values"] = rearrange(
torch.tensor(refer_images / 127.5 - 1), "h w c -> 1 c h w"
)
if start > 0:
batch["refer_t_pixel_values"] = rearrange(
out_frames[0, :, -refert_num:].clone().detach(),
"c t h w -> t c h w",
)
batch["refer_t_pixel_values"] = rearrange(batch["refer_t_pixel_values"],
"t c h w -> 1 c t h w",
)
if replace_flag:
batch["bg_pixel_values"] = rearrange(
torch.tensor(np.stack(bg_images[start:end]) / 127.5 - 1),
"t h w c -> 1 c t h w",
)
batch["mask_pixel_values"] = rearrange(
torch.tensor(np.stack(mask_images[start:end])[:, :, :, None]),
"t h w c -> 1 t c h w",
)
for key, value in batch.items():
if isinstance(value, torch.Tensor):
batch[key] = value.to(device=self.device, dtype=torch.bfloat16)
ref_pixel_values = batch["refer_pixel_values"]
refer_t_pixel_values = batch["refer_t_pixel_values"]
conditioning_pixel_values = batch["conditioning_pixel_values"]
face_pixel_values = batch["face_pixel_values"]
B, _, H, W = ref_pixel_values.shape
T = clip_len
lat_h = H // 8
lat_w = W // 8
lat_t = T // 4 + 1
target_shape = [lat_t + 1, lat_h, lat_w]
noise = [
torch.randn(
16,
target_shape[0],
target_shape[1],
target_shape[2],
dtype=torch.float32,
device=self.device,
generator=seed_g,
)
]
max_seq_len = int(math.ceil(np.prod(target_shape) // 4 / self.sp_size)) * self.sp_size
if max_seq_len % self.sp_size != 0:
raise ValueError(f"max_seq_len {max_seq_len} is not divisible by sp_size {self.sp_size}")
with (
torch.autocast(device_type=str(self.device), dtype=torch.bfloat16, enabled=True),
torch.no_grad()
):
if sample_solver == 'unipc':
sample_scheduler = FlowUniPCMultistepScheduler(
num_train_timesteps=self.num_train_timesteps,
shift=1,
use_dynamic_shifting=False)
sample_scheduler.set_timesteps(
sampling_steps, device=self.device, shift=shift)
timesteps = sample_scheduler.timesteps
elif sample_solver == 'dpm++':
sample_scheduler = FlowDPMSolverMultistepScheduler(
num_train_timesteps=self.num_train_timesteps,
shift=1,
use_dynamic_shifting=False)
sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
timesteps, _ = retrieve_timesteps(
sample_scheduler,
device=self.device,
sigmas=sampling_sigmas)
else:
raise NotImplementedError("Unsupported solver.")
latents = noise
pose_latents_no_ref = self.vae.encode(conditioning_pixel_values.to(torch.bfloat16))
pose_latents_no_ref = torch.stack(pose_latents_no_ref)
pose_latents = torch.cat([pose_latents_no_ref], dim=2)
ref_pixel_values = rearrange(ref_pixel_values, "t c h w -> 1 c t h w")
ref_latents = self.vae.encode(ref_pixel_values.to(torch.bfloat16))
ref_latents = torch.stack(ref_latents)
mask_ref = self.get_i2v_mask(1, lat_h, lat_w, 1, device=self.device)
y_ref = torch.concat([mask_ref, ref_latents[0]]).to(dtype=torch.bfloat16, device=self.device)
img = ref_pixel_values[0, :, 0]
clip_context = self.clip.visual([img[:, None, :, :]]).to(dtype=torch.bfloat16, device=self.device)
if mask_reft_len > 0:
if replace_flag:
bg_pixel_values = batch["bg_pixel_values"]
y_reft = self.vae.encode(
[
torch.concat([refer_t_pixel_values[0, :, :mask_reft_len], bg_pixel_values[0, :, mask_reft_len:]], dim=1).to(self.device)
]
)[0]
mask_pixel_values = 1 - batch["mask_pixel_values"]
mask_pixel_values = rearrange(mask_pixel_values, "b t c h w -> (b t) c h w")
mask_pixel_values = F.interpolate(mask_pixel_values, size=(H//8, W//8), mode='nearest')
mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b t c h w", b=1)[:,:,0]
msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, mask_pixel_values=mask_pixel_values, device=self.device)
else:
y_reft = self.vae.encode(
[
torch.concat(
[
torch.nn.functional.interpolate(refer_t_pixel_values[0, :, :mask_reft_len].cpu(),
size=(H, W), mode="bicubic"),
torch.zeros(3, T - mask_reft_len, H, W),
],
dim=1,
).to(self.device)
]
)[0]
msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, device=self.device)
else:
if replace_flag:
bg_pixel_values = batch["bg_pixel_values"]
mask_pixel_values = 1 - batch["mask_pixel_values"]
mask_pixel_values = rearrange(mask_pixel_values, "b t c h w -> (b t) c h w")
mask_pixel_values = F.interpolate(mask_pixel_values, size=(H//8, W//8), mode='nearest')
mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b t c h w", b=1)[:,:,0]
y_reft = self.vae.encode(
[
torch.concat(
[
bg_pixel_values[0],
],
dim=1,
).to(self.device)
]
)[0]
msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, mask_pixel_values=mask_pixel_values, device=self.device)
else:
y_reft = self.vae.encode(
[
torch.concat(
[
torch.zeros(3, T - mask_reft_len, H, W),
],
dim=1,
).to(self.device)
]
)[0]
msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, device=self.device)
y_reft = torch.concat([msk_reft, y_reft]).to(dtype=torch.bfloat16, device=self.device)
y = torch.concat([y_ref, y_reft], dim=1)
arg_c = {
"context": context,
"seq_len": max_seq_len,
"clip_fea": clip_context.to(dtype=torch.bfloat16, device=self.device),
"y": [y],
"pose_latents": pose_latents,
"face_pixel_values": face_pixel_values,
}
if guide_scale > 1:
face_pixel_values_uncond = face_pixel_values * 0 - 1
arg_null = {
"context": context_null,
"seq_len": max_seq_len,
"clip_fea": clip_context.to(dtype=torch.bfloat16, device=self.device),
"y": [y],
"pose_latents": pose_latents,
"face_pixel_values": face_pixel_values_uncond,
}
for i, t in enumerate(tqdm(timesteps)):
latent_model_input = latents
timestep = [t]
timestep = torch.stack(timestep)
noise_pred_cond = TensorList(
self.noise_model(TensorList(latent_model_input), t=timestep, **arg_c)
)
if guide_scale > 1:
noise_pred_uncond = TensorList(
self.noise_model(
TensorList(latent_model_input), t=timestep, **arg_null
)
)
noise_pred = noise_pred_uncond + guide_scale * (
noise_pred_cond - noise_pred_uncond
)
else:
noise_pred = noise_pred_cond
temp_x0 = sample_scheduler.step(
noise_pred[0].unsqueeze(0),
t,
latents[0].unsqueeze(0),
return_dict=False,
generator=seed_g,
)[0]
latents[0] = temp_x0.squeeze(0)
x0 = latents
x0 = [x.to(dtype=torch.float32) for x in x0]
out_frames = torch.stack(self.vae.decode([x0[0][:, 1:]]))
if start != 0:
out_frames = out_frames[:, :, refert_num:]
all_out_frames.append(out_frames.cpu())
start += clip_len - refert_num
end += clip_len - refert_num
videos = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len]
return videos[0] if self.rank == 0 else None
```
## /wan/configs/__init__.py
```py path="/wan/configs/__init__.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import copy
import os
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
from .wan_i2v_A14B import i2v_A14B
from .wan_s2v_14B import s2v_14B
from .wan_t2v_A14B import t2v_A14B
from .wan_ti2v_5B import ti2v_5B
from .wan_animate_14B import animate_14B
WAN_CONFIGS = {
't2v-A14B': t2v_A14B,
'i2v-A14B': i2v_A14B,
'ti2v-5B': ti2v_5B,
'animate-14B': animate_14B,
's2v-14B': s2v_14B,
}
SIZE_CONFIGS = {
'720*1280': (720, 1280),
'1280*720': (1280, 720),
'480*832': (480, 832),
'832*480': (832, 480),
'704*1280': (704, 1280),
'1280*704': (1280, 704),
'1024*704': (1024, 704),
'704*1024': (704, 1024),
}
MAX_AREA_CONFIGS = {
'720*1280': 720 * 1280,
'1280*720': 1280 * 720,
'480*832': 480 * 832,
'832*480': 832 * 480,
'704*1280': 704 * 1280,
'1280*704': 1280 * 704,
'1024*704': 1024 * 704,
'704*1024': 704 * 1024,
}
SUPPORTED_SIZES = {
't2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'),
'i2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'),
'ti2v-5B': ('704*1280', '1280*704'),
's2v-14B': ('720*1280', '1280*720', '480*832', '832*480', '1024*704',
'704*1024', '704*1280', '1280*704'),
'animate-14B': ('720*1280', '1280*720')
}
```
## /wan/configs/shared_config.py
```py path="/wan/configs/shared_config.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
from easydict import EasyDict
#------------------------ Wan shared config ------------------------#
wan_shared_cfg = EasyDict()
# t5
wan_shared_cfg.t5_model = 'umt5_xxl'
wan_shared_cfg.t5_dtype = torch.bfloat16
wan_shared_cfg.text_len = 512
# transformer
wan_shared_cfg.param_dtype = torch.bfloat16
# inference
wan_shared_cfg.num_train_timesteps = 1000
wan_shared_cfg.sample_fps = 16
wan_shared_cfg.sample_neg_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
wan_shared_cfg.frame_num = 81
```
## /wan/configs/wan_animate_14B.py
```py path="/wan/configs/wan_animate_14B.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from easydict import EasyDict
from .shared_config import wan_shared_cfg
#------------------------ Wan animate 14B ------------------------#
animate_14B = EasyDict(__name__='Config: Wan animate 14B')
animate_14B.update(wan_shared_cfg)
animate_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
animate_14B.t5_tokenizer = 'google/umt5-xxl'
animate_14B.clip_checkpoint = 'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
animate_14B.clip_tokenizer = 'xlm-roberta-large'
animate_14B.lora_checkpoint = 'relighting_lora.ckpt'
# vae
animate_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
animate_14B.vae_stride = (4, 8, 8)
# transformer
animate_14B.patch_size = (1, 2, 2)
animate_14B.dim = 5120
animate_14B.ffn_dim = 13824
animate_14B.freq_dim = 256
animate_14B.num_heads = 40
animate_14B.num_layers = 40
animate_14B.window_size = (-1, -1)
animate_14B.qk_norm = True
animate_14B.cross_attn_norm = True
animate_14B.eps = 1e-6
animate_14B.use_face_encoder = True
animate_14B.motion_encoder_dim = 512
# inference
animate_14B.sample_shift = 5.0
animate_14B.sample_steps = 20
animate_14B.sample_guide_scale = 1.0
animate_14B.frame_num = 77
animate_14B.sample_fps = 30
animate_14B.prompt = '视频中的人在做动作'
```
## /wan/configs/wan_i2v_A14B.py
```py path="/wan/configs/wan_i2v_A14B.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
from easydict import EasyDict
from .shared_config import wan_shared_cfg
#------------------------ Wan I2V A14B ------------------------#
i2v_A14B = EasyDict(__name__='Config: Wan I2V A14B')
i2v_A14B.update(wan_shared_cfg)
i2v_A14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
i2v_A14B.t5_tokenizer = 'google/umt5-xxl'
# vae
i2v_A14B.vae_checkpoint = 'Wan2.1_VAE.pth'
i2v_A14B.vae_stride = (4, 8, 8)
# transformer
i2v_A14B.patch_size = (1, 2, 2)
i2v_A14B.dim = 5120
i2v_A14B.ffn_dim = 13824
i2v_A14B.freq_dim = 256
i2v_A14B.num_heads = 40
i2v_A14B.num_layers = 40
i2v_A14B.window_size = (-1, -1)
i2v_A14B.qk_norm = True
i2v_A14B.cross_attn_norm = True
i2v_A14B.eps = 1e-6
i2v_A14B.low_noise_checkpoint = 'low_noise_model'
i2v_A14B.high_noise_checkpoint = 'high_noise_model'
# inference
i2v_A14B.sample_shift = 5.0
i2v_A14B.sample_steps = 40
i2v_A14B.boundary = 0.900
i2v_A14B.sample_guide_scale = (3.5, 3.5) # low noise, high noise
```
## /wan/configs/wan_s2v_14B.py
```py path="/wan/configs/wan_s2v_14B.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from easydict import EasyDict
from .shared_config import wan_shared_cfg
#------------------------ Wan S2V 14B ------------------------#
s2v_14B = EasyDict(__name__='Config: Wan S2V 14B')
s2v_14B.update(wan_shared_cfg)
# t5
s2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
s2v_14B.t5_tokenizer = 'google/umt5-xxl'
# vae
s2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
s2v_14B.vae_stride = (4, 8, 8)
# wav2vec
s2v_14B.wav2vec = "wav2vec2-large-xlsr-53-english"
s2v_14B.num_heads = 40
# transformer
s2v_14B.transformer = EasyDict(
__name__="Config: Transformer config for WanModel_S2V")
s2v_14B.transformer.patch_size = (1, 2, 2)
s2v_14B.transformer.dim = 5120
s2v_14B.transformer.ffn_dim = 13824
s2v_14B.transformer.freq_dim = 256
s2v_14B.transformer.num_heads = 40
s2v_14B.transformer.num_layers = 40
s2v_14B.transformer.window_size = (-1, -1)
s2v_14B.transformer.qk_norm = True
s2v_14B.transformer.cross_attn_norm = True
s2v_14B.transformer.eps = 1e-6
s2v_14B.transformer.enable_adain = True
s2v_14B.transformer.adain_mode = "attn_norm"
s2v_14B.transformer.audio_inject_layers = [
0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39
]
s2v_14B.transformer.zero_init = True
s2v_14B.transformer.zero_timestep = True
s2v_14B.transformer.enable_motioner = False
s2v_14B.transformer.add_last_motion = True
s2v_14B.transformer.trainable_token = False
s2v_14B.transformer.enable_tsm = False
s2v_14B.transformer.enable_framepack = True
s2v_14B.transformer.framepack_drop_mode = 'padd'
s2v_14B.transformer.audio_dim = 1024
s2v_14B.transformer.motion_frames = 73
s2v_14B.transformer.cond_dim = 16
# inference
s2v_14B.sample_neg_prompt = "画面模糊，最差质量，画面模糊，细节模糊不清，情绪激动剧烈，手快速抖动，字幕，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
s2v_14B.drop_first_motion = True
s2v_14B.sample_shift = 3
s2v_14B.sample_steps = 40
s2v_14B.sample_guide_scale = 4.5
```
## /wan/configs/wan_t2v_A14B.py
```py path="/wan/configs/wan_t2v_A14B.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from easydict import EasyDict
from .shared_config import wan_shared_cfg
#------------------------ Wan T2V A14B ------------------------#
t2v_A14B = EasyDict(__name__='Config: Wan T2V A14B')
t2v_A14B.update(wan_shared_cfg)
# t5
t2v_A14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
t2v_A14B.t5_tokenizer = 'google/umt5-xxl'
# vae
t2v_A14B.vae_checkpoint = 'Wan2.1_VAE.pth'
t2v_A14B.vae_stride = (4, 8, 8)
# transformer
t2v_A14B.patch_size = (1, 2, 2)
t2v_A14B.dim = 5120
t2v_A14B.ffn_dim = 13824
t2v_A14B.freq_dim = 256
t2v_A14B.num_heads = 40
t2v_A14B.num_layers = 40
t2v_A14B.window_size = (-1, -1)
t2v_A14B.qk_norm = True
t2v_A14B.cross_attn_norm = True
t2v_A14B.eps = 1e-6
t2v_A14B.low_noise_checkpoint = 'low_noise_model'
t2v_A14B.high_noise_checkpoint = 'high_noise_model'
# inference
t2v_A14B.sample_shift = 12.0
t2v_A14B.sample_steps = 40
t2v_A14B.boundary = 0.875
t2v_A14B.sample_guide_scale = (3.0, 4.0) # low noise, high noise
```
## /wan/configs/wan_ti2v_5B.py
```py path="/wan/configs/wan_ti2v_5B.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from easydict import EasyDict
from .shared_config import wan_shared_cfg
#------------------------ Wan TI2V 5B ------------------------#
ti2v_5B = EasyDict(__name__='Config: Wan TI2V 5B')
ti2v_5B.update(wan_shared_cfg)
# t5
ti2v_5B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
ti2v_5B.t5_tokenizer = 'google/umt5-xxl'
# vae
ti2v_5B.vae_checkpoint = 'Wan2.2_VAE.pth'
ti2v_5B.vae_stride = (4, 16, 16)
# transformer
ti2v_5B.patch_size = (1, 2, 2)
ti2v_5B.dim = 3072
ti2v_5B.ffn_dim = 14336
ti2v_5B.freq_dim = 256
ti2v_5B.num_heads = 24
ti2v_5B.num_layers = 30
ti2v_5B.window_size = (-1, -1)
ti2v_5B.qk_norm = True
ti2v_5B.cross_attn_norm = True
ti2v_5B.eps = 1e-6
# inference
ti2v_5B.sample_fps = 24
ti2v_5B.sample_shift = 5.0
ti2v_5B.sample_steps = 50
ti2v_5B.sample_guide_scale = 5.0
ti2v_5B.frame_num = 121
```
## /wan/distributed/__init__.py
```py path="/wan/distributed/__init__.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
```
## /wan/distributed/fsdp.py
```py path="/wan/distributed/fsdp.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import gc
from functools import partial
import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
from torch.distributed.utils import _free_storage
def shard_model(
model,
device_id,
param_dtype=torch.bfloat16,
reduce_dtype=torch.float32,
buffer_dtype=torch.float32,
process_group=None,
sharding_strategy=ShardingStrategy.FULL_SHARD,
sync_module_states=True,
use_lora=False
):
model = FSDP(
module=model,
process_group=process_group,
sharding_strategy=sharding_strategy,
auto_wrap_policy=partial(
lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks),
mixed_precision=MixedPrecision(
param_dtype=param_dtype,
reduce_dtype=reduce_dtype,
buffer_dtype=buffer_dtype),
device_id=device_id,
sync_module_states=sync_module_states,
use_orig_params=True if use_lora else False)
return model
def free_model(model):
for m in model.modules():
if isinstance(m, FSDP):
_free_storage(m._handle.flat_param.data)
del model
gc.collect()
torch.cuda.empty_cache()
```
## /wan/distributed/sequence_parallel.py
```py path="/wan/distributed/sequence_parallel.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
import torch.cuda.amp as amp
from ..modules.model import sinusoidal_embedding_1d
from .ulysses import distributed_attention
from .util import gather_forward, get_rank, get_world_size
def pad_freqs(original_tensor, target_len):
seq_len, s1, s2 = original_tensor.shape
pad_size = target_len - seq_len
padding_tensor = torch.ones(
pad_size,
s1,
s2,
dtype=original_tensor.dtype,
device=original_tensor.device)
padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
return padded_tensor
@torch.amp.autocast('cuda', enabled=False)
def rope_apply(x, grid_sizes, freqs):
"""
x: [B, L, N, C].
grid_sizes: [B, 3].
freqs: [M, C // 2].
"""
s, n, c = x.size(1), x.size(2), x.size(3) // 2
# split freqs
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
# loop over samples
output = []
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
seq_len = f * h * w
# precompute multipliers
x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(
s, n, -1, 2))
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).reshape(seq_len, 1, -1)
# apply rotary embedding
sp_size = get_world_size()
sp_rank = get_rank()
freqs_i = pad_freqs(freqs_i, s * sp_size)
s_per_rank = s
freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *
s_per_rank), :, :]
x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
x_i = torch.cat([x_i, x[i, s:]])
# append to collection
output.append(x_i)
return torch.stack(output).float()
def sp_dit_forward(
self,
x,
t,
context,
seq_len,
y=None,
):
"""
x: A list of videos each with shape [C, T, H, W].
t: [B].
context: A list of text embeddings each with shape [L, C].
"""
if self.model_type == 'i2v':
assert y is not None
# params
device = self.patch_embedding.weight.device
if self.freqs.device != device:
self.freqs = self.freqs.to(device)
if y is not None:
x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
# 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])
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
x = torch.cat([
torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
for u in x
])
# time embeddings
if t.dim() == 1:
t = t.expand(t.size(0), seq_len)
with torch.amp.autocast('cuda', dtype=torch.float32):
bt = t.size(0)
t = t.flatten()
e = self.time_embedding(
sinusoidal_embedding_1d(self.freq_dim,
t).unflatten(0, (bt, seq_len)).float())
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
assert e.dtype == torch.float32 and e0.dtype == torch.float32
# context
context_lens = None
context = self.text_embedding(
torch.stack([
torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
for u in context
]))
# Context Parallel
x = torch.chunk(x, get_world_size(), dim=1)[get_rank()]
e = torch.chunk(e, get_world_size(), dim=1)[get_rank()]
e0 = torch.chunk(e0, get_world_size(), dim=1)[get_rank()]
# arguments
kwargs = dict(
e=e0,
seq_lens=seq_lens,
grid_sizes=grid_sizes,
freqs=self.freqs,
context=context,
context_lens=context_lens)
for block in self.blocks:
x = block(x, **kwargs)
# head
x = self.head(x, e)
# Context Parallel
x = gather_forward(x, dim=1)
# unpatchify
x = self.unpatchify(x, grid_sizes)
return [u.float() for u in x]
def sp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=torch.bfloat16):
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
half_dtypes = (torch.float16, torch.bfloat16)
def half(x):
return x if x.dtype in half_dtypes else x.to(dtype)
# 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)
q = rope_apply(q, grid_sizes, freqs)
k = rope_apply(k, grid_sizes, freqs)
x = distributed_attention(
half(q),
half(k),
half(v),
seq_lens,
window_size=self.window_size,
)
# output
x = x.flatten(2)
x = self.o(x)
return x
```
## /wan/distributed/ulysses.py
```py path="/wan/distributed/ulysses.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
import torch.distributed as dist
from ..modules.attention import flash_attention
from .util import all_to_all
def distributed_attention(
q,
k,
v,
seq_lens,
window_size=(-1, -1),
):
"""
Performs distributed attention based on DeepSpeed Ulysses attention mechanism.
please refer to https://arxiv.org/pdf/2309.14509
Args:
q: [B, Lq // p, Nq, C1].
k: [B, Lk // p, Nk, C1].
v: [B, Lk // p, Nk, C2]. Nq must be divisible by Nk.
seq_lens: [B], length of each sequence in batch
window_size: (left right). If not (-1, -1), apply sliding window local attention.
"""
if not dist.is_initialized():
raise ValueError("distributed group should be initialized.")
b = q.shape[0]
# gather q/k/v sequence
q = all_to_all(q, scatter_dim=2, gather_dim=1)
k = all_to_all(k, scatter_dim=2, gather_dim=1)
v = all_to_all(v, scatter_dim=2, gather_dim=1)
# apply attention
x = flash_attention(
q,
k,
v,
k_lens=seq_lens,
window_size=window_size,
)
# scatter q/k/v sequence
x = all_to_all(x, scatter_dim=1, gather_dim=2)
return x
```
## /wan/distributed/util.py
```py path="/wan/distributed/util.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
import torch.distributed as dist
def init_distributed_group():
"""r initialize sequence parallel group.
"""
if not dist.is_initialized():
dist.init_process_group(backend='nccl')
def get_rank():
return dist.get_rank()
def get_world_size():
return dist.get_world_size()
def all_to_all(x, scatter_dim, gather_dim, group=None, **kwargs):
"""
`scatter` along one dimension and `gather` along another.
"""
world_size = get_world_size()
if world_size > 1:
inputs = [u.contiguous() for u in x.chunk(world_size, dim=scatter_dim)]
outputs = [torch.empty_like(u) for u in inputs]
dist.all_to_all(outputs, inputs, group=group, **kwargs)
x = torch.cat(outputs, dim=gather_dim).contiguous()
return x
def all_gather(tensor):
world_size = dist.get_world_size()
if world_size == 1:
return [tensor]
tensor_list = [torch.empty_like(tensor) for _ in range(world_size)]
torch.distributed.all_gather(tensor_list, tensor)
return tensor_list
def gather_forward(input, dim):
# skip if world_size == 1
world_size = dist.get_world_size()
if world_size == 1:
return input
# gather sequence
output = all_gather(input)
return torch.cat(output, dim=dim).contiguous()
```
## /wan/image2video.py
```py path="/wan/image2video.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import gc
import logging
import math
import os
import random
import sys
import types
from contextlib import contextmanager
from functools import partial
import numpy as np
import torch
import torch.cuda.amp as amp
import torch.distributed as dist
import torchvision.transforms.functional as TF
from tqdm import tqdm
from .distributed.fsdp import shard_model
from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward
from .distributed.util import get_world_size
from .modules.model import WanModel
from .modules.t5 import T5EncoderModel
from .modules.vae2_1 import Wan2_1_VAE
from .utils.fm_solvers import (
FlowDPMSolverMultistepScheduler,
get_sampling_sigmas,
retrieve_timesteps,
)
from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
class WanI2V:
def __init__(
self,
config,
checkpoint_dir,
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_sp=False,
t5_cpu=False,
init_on_cpu=True,
convert_model_dtype=False,
):
r"""
Initializes the image-to-video generation model components.
Args:
config (EasyDict):
Object containing model parameters initialized from config.py
checkpoint_dir (`str`):
Path to directory containing model checkpoints
device_id (`int`, *optional*, defaults to 0):
Id of target GPU device
rank (`int`, *optional*, defaults to 0):
Process rank for distributed training
t5_fsdp (`bool`, *optional*, defaults to False):
Enable FSDP sharding for T5 model
dit_fsdp (`bool`, *optional*, defaults to False):
Enable FSDP sharding for DiT model
use_sp (`bool`, *optional*, defaults to False):
Enable distribution strategy of sequence parallel.
t5_cpu (`bool`, *optional*, defaults to False):
Whether to place T5 model on CPU. Only works without t5_fsdp.
init_on_cpu (`bool`, *optional*, defaults to True):
Enable initializing Transformer Model on CPU. Only works without FSDP or USP.
convert_model_dtype (`bool`, *optional*, defaults to False):
Convert DiT model parameters dtype to 'config.param_dtype'.
Only works without FSDP.
"""
self.device = torch.device(f"cuda:{device_id}")
self.config = config
self.rank = rank
self.t5_cpu = t5_cpu
self.init_on_cpu = init_on_cpu
self.num_train_timesteps = config.num_train_timesteps
self.boundary = config.boundary
self.param_dtype = config.param_dtype
if t5_fsdp or dit_fsdp or use_sp:
self.init_on_cpu = False
shard_fn = partial(shard_model, device_id=device_id)
self.text_encoder = T5EncoderModel(
text_len=config.text_len,
dtype=config.t5_dtype,
device=torch.device('cpu'),
checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
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 = Wan2_1_VAE(
vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
device=self.device)
logging.info(f"Creating WanModel from {checkpoint_dir}")
self.low_noise_model = WanModel.from_pretrained(
checkpoint_dir, subfolder=config.low_noise_checkpoint)
self.low_noise_model = self._configure_model(
model=self.low_noise_model,
use_sp=use_sp,
dit_fsdp=dit_fsdp,
shard_fn=shard_fn,
convert_model_dtype=convert_model_dtype)
self.high_noise_model = WanModel.from_pretrained(
checkpoint_dir, subfolder=config.high_noise_checkpoint)
self.high_noise_model = self._configure_model(
model=self.high_noise_model,
use_sp=use_sp,
dit_fsdp=dit_fsdp,
shard_fn=shard_fn,
convert_model_dtype=convert_model_dtype)
if use_sp:
self.sp_size = get_world_size()
else:
self.sp_size = 1
self.sample_neg_prompt = config.sample_neg_prompt
def _configure_model(self, model, use_sp, dit_fsdp, shard_fn,
convert_model_dtype):
"""
Configures a model object. This includes setting evaluation modes,
applying distributed parallel strategy, and handling device placement.
Args:
model (torch.nn.Module):
The model instance to configure.
use_sp (`bool`):
Enable distribution strategy of sequence parallel.
dit_fsdp (`bool`):
Enable FSDP sharding for DiT model.
shard_fn (callable):
The function to apply FSDP sharding.
convert_model_dtype (`bool`):
Convert DiT model parameters dtype to 'config.param_dtype'.
Only works without FSDP.
Returns:
torch.nn.Module:
The configured model.
"""
model.eval().requires_grad_(False)
if use_sp:
for block in model.blocks:
block.self_attn.forward = types.MethodType(
sp_attn_forward, block.self_attn)
model.forward = types.MethodType(sp_dit_forward, model)
if dist.is_initialized():
dist.barrier()
if dit_fsdp:
model = shard_fn(model)
else:
if convert_model_dtype:
model.to(self.param_dtype)
if not self.init_on_cpu:
model.to(self.device)
return model
def _prepare_model_for_timestep(self, t, boundary, offload_model):
r"""
Prepares and returns the required model for the current timestep.
Args:
t (torch.Tensor):
current timestep.
boundary (`int`):
The timestep threshold. If `t` is at or above this value,
the `high_noise_model` is considered as the required model.
offload_model (`bool`):
A flag intended to control the offloading behavior.
Returns:
torch.nn.Module:
The active model on the target device for the current timestep.
"""
if t.item() >= boundary:
required_model_name = 'high_noise_model'
offload_model_name = 'low_noise_model'
else:
required_model_name = 'low_noise_model'
offload_model_name = 'high_noise_model'
if offload_model or self.init_on_cpu:
if next(getattr(
self,
offload_model_name).parameters()).device.type == 'cuda':
getattr(self, offload_model_name).to('cpu')
if next(getattr(
self,
required_model_name).parameters()).device.type == 'cpu':
getattr(self, required_model_name).to(self.device)
return getattr(self, required_model_name)
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):
r"""
Generates video frames from input image and text prompt using diffusion process.
Args:
input_prompt (`str`):
Text prompt for content generation.
img (PIL.Image.Image):
Input image tensor. Shape: [3, H, W]
max_area (`int`, *optional*, defaults to 720*1280):
Maximum pixel area for latent space calculation. Controls video resolution scaling
frame_num (`int`, *optional*, defaults to 81):
How many frames to sample from a video. The number should be 4n+1
shift (`float`, *optional*, defaults to 5.0):
Noise schedule shift parameter. Affects temporal dynamics
[NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0.
sample_solver (`str`, *optional*, defaults to 'unipc'):
Solver used to sample the video.
sampling_steps (`int`, *optional*, defaults to 40):
Number of diffusion sampling steps. Higher values improve quality but slow generation
guide_scale (`float` or tuple[`float`], *optional*, defaults 5.0):
Classifier-free guidance scale. Controls prompt adherence vs. creativity.
If tuple, the first guide_scale will be used for low noise model and
the second guide_scale will be used for high noise model.
n_prompt (`str`, *optional*, defaults to ""):
Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
seed (`int`, *optional*, defaults to -1):
Random seed for noise generation. If -1, use random seed
offload_model (`bool`, *optional*, defaults to True):
If True, offloads models to CPU during generation to save VRAM
Returns:
torch.Tensor:
Generated video frames tensor. Dimensions: (C, N H, W) where:
- C: Color channels (3 for RGB)
- N: Number of frames (81)
- H: Frame height (from max_area)
- W: Frame width from max_area)
"""
# preprocess
guide_scale = (guide_scale, guide_scale) if isinstance(
guide_scale, float) else guide_scale
img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device)
F = frame_num
h, w = img.shape[1:]
aspect_ratio = h / w
lat_h = round(
np.sqrt(max_area * aspect_ratio) // self.vae_stride[1] //
self.patch_size[1] * self.patch_size[1])
lat_w = round(
np.sqrt(max_area / aspect_ratio) // self.vae_stride[2] //
self.patch_size[2] * self.patch_size[2])
h = lat_h * self.vae_stride[1]
w = lat_w * self.vae_stride[2]
max_seq_len = ((F - 1) // self.vae_stride[0] + 1) * lat_h * lat_w // (
self.patch_size[1] * self.patch_size[2])
max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size
seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
seed_g = torch.Generator(device=self.device)
seed_g.manual_seed(seed)
noise = torch.randn(
16,
(F - 1) // self.vae_stride[0] + 1,
lat_h,
lat_w,
dtype=torch.float32,
generator=seed_g,
device=self.device)
msk = torch.ones(1, F, 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:]
],
dim=1)
msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
msk = msk.transpose(1, 2)[0]
if n_prompt == "":
n_prompt = self.sample_neg_prompt
# preprocess
if not self.t5_cpu:
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:
self.text_encoder.model.cpu()
else:
context = self.text_encoder([input_prompt], torch.device('cpu'))
context_null = self.text_encoder([n_prompt], torch.device('cpu'))
context = [t.to(self.device) for t in context]
context_null = [t.to(self.device) for t in context_null]
y = self.vae.encode([
torch.concat([
torch.nn.functional.interpolate(
img[None].cpu(), size=(h, w), mode='bicubic').transpose(
0, 1),
torch.zeros(3, F - 1, h, w)
],
dim=1).to(self.device)
])[0]
y = torch.concat([msk, y])
@contextmanager
def noop_no_sync():
yield
no_sync_low_noise = getattr(self.low_noise_model, 'no_sync',
noop_no_sync)
no_sync_high_noise = getattr(self.high_noise_model, 'no_sync',
noop_no_sync)
# evaluation mode
with (
torch.amp.autocast('cuda', dtype=self.param_dtype),
torch.no_grad(),
no_sync_low_noise(),
no_sync_high_noise(),
):
boundary = self.boundary * self.num_train_timesteps
if sample_solver == 'unipc':
sample_scheduler = FlowUniPCMultistepScheduler(
num_train_timesteps=self.num_train_timesteps,
shift=1,
use_dynamic_shifting=False)
sample_scheduler.set_timesteps(
sampling_steps, device=self.device, shift=shift)
timesteps = sample_scheduler.timesteps
elif sample_solver == 'dpm++':
sample_scheduler = FlowDPMSolverMultistepScheduler(
num_train_timesteps=self.num_train_timesteps,
shift=1,
use_dynamic_shifting=False)
sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
timesteps, _ = retrieve_timesteps(
sample_scheduler,
device=self.device,
sigmas=sampling_sigmas)
else:
raise NotImplementedError("Unsupported solver.")
# sample videos
latent = noise
arg_c = {
'context': [context[0]],
'seq_len': max_seq_len,
'y': [y],
}
arg_null = {
'context': context_null,
'seq_len': max_seq_len,
'y': [y],
}
if offload_model:
torch.cuda.empty_cache()
for _, t in enumerate(tqdm(timesteps)):
latent_model_input = [latent.to(self.device)]
timestep = [t]
timestep = torch.stack(timestep).to(self.device)
model = self._prepare_model_for_timestep(
t, boundary, offload_model)
sample_guide_scale = guide_scale[1] if t.item(
) >= boundary else guide_scale[0]
noise_pred_cond = model(
latent_model_input, t=timestep, **arg_c)[0]
if offload_model:
torch.cuda.empty_cache()
noise_pred_uncond = model(
latent_model_input, t=timestep, **arg_null)[0]
if offload_model:
torch.cuda.empty_cache()
noise_pred = noise_pred_uncond + sample_guide_scale * (
noise_pred_cond - noise_pred_uncond)
temp_x0 = sample_scheduler.step(
noise_pred.unsqueeze(0),
t,
latent.unsqueeze(0),
return_dict=False,
generator=seed_g)[0]
latent = temp_x0.squeeze(0)
x0 = [latent]
del latent_model_input, timestep
if offload_model:
self.low_noise_model.cpu()
self.high_noise_model.cpu()
torch.cuda.empty_cache()
if self.rank == 0:
videos = self.vae.decode(x0)
del noise, latent, x0
del sample_scheduler
if offload_model:
gc.collect()
torch.cuda.synchronize()
if dist.is_initialized():
dist.barrier()
return videos[0] if self.rank == 0 else None
```
## /wan/modules/__init__.py
```py path="/wan/modules/__init__.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from .attention import flash_attention
from .model import WanModel
from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
from .tokenizers import HuggingfaceTokenizer
from .vae2_1 import Wan2_1_VAE
from .vae2_2 import Wan2_2_VAE
__all__ = [
'Wan2_1_VAE',
'Wan2_2_VAE',
'WanModel',
'T5Model',
'T5Encoder',
'T5Decoder',
'T5EncoderModel',
'HuggingfaceTokenizer',
'flash_attention',
]
```
## /wan/modules/animate/__init__.py
```py path="/wan/modules/animate/__init__.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from .model_animate import WanAnimateModel
from .clip import CLIPModel
__all__ = ['WanAnimateModel', 'CLIPModel']
```
## /wan/modules/animate/animate_utils.py
```py path="/wan/modules/animate/animate_utils.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
import numbers
from peft import LoraConfig
def get_loraconfig(transformer, rank=128, alpha=128, init_lora_weights="gaussian"):
target_modules = []
for name, module in transformer.named_modules():
if "blocks" in name and "face" not in name and "modulation" not in name and isinstance(module, torch.nn.Linear):
target_modules.append(name)
transformer_lora_config = LoraConfig(
r=rank,
lora_alpha=alpha,
init_lora_weights=init_lora_weights,
target_modules=target_modules,
)
return transformer_lora_config
class TensorList(object):
def __init__(self, tensors):
"""
tensors: a list of torch.Tensor objects. No need to have uniform shape.
"""
assert isinstance(tensors, (list, tuple))
assert all(isinstance(u, torch.Tensor) for u in tensors)
assert len(set([u.ndim for u in tensors])) == 1
assert len(set([u.dtype for u in tensors])) == 1
assert len(set([u.device for u in tensors])) == 1
self.tensors = tensors
def to(self, *args, **kwargs):
return TensorList([u.to(*args, **kwargs) for u in self.tensors])
def size(self, dim):
assert dim == 0, 'only support get the 0th size'
return len(self.tensors)
def pow(self, *args, **kwargs):
return TensorList([u.pow(*args, **kwargs) for u in self.tensors])
def squeeze(self, dim):
assert dim != 0
if dim > 0:
dim -= 1
return TensorList([u.squeeze(dim) for u in self.tensors])
def type(self, *args, **kwargs):
return TensorList([u.type(*args, **kwargs) for u in self.tensors])
def type_as(self, other):
assert isinstance(other, (torch.Tensor, TensorList))
if isinstance(other, torch.Tensor):
return TensorList([u.type_as(other) for u in self.tensors])
else:
return TensorList([u.type(other.dtype) for u in self.tensors])
@property
def dtype(self):
return self.tensors[0].dtype
@property
def device(self):
return self.tensors[0].device
@property
def ndim(self):
return 1 + self.tensors[0].ndim
def __getitem__(self, index):
return self.tensors[index]
def __len__(self):
return len(self.tensors)
def __add__(self, other):
return self._apply(other, lambda u, v: u + v)
def __radd__(self, other):
return self._apply(other, lambda u, v: v + u)
def __sub__(self, other):
return self._apply(other, lambda u, v: u - v)
def __rsub__(self, other):
return self._apply(other, lambda u, v: v - u)
def __mul__(self, other):
return self._apply(other, lambda u, v: u * v)
def __rmul__(self, other):
return self._apply(other, lambda u, v: v * u)
def __floordiv__(self, other):
return self._apply(other, lambda u, v: u // v)
def __truediv__(self, other):
return self._apply(other, lambda u, v: u / v)
def __rfloordiv__(self, other):
return self._apply(other, lambda u, v: v // u)
def __rtruediv__(self, other):
return self._apply(other, lambda u, v: v / u)
def __pow__(self, other):
return self._apply(other, lambda u, v: u ** v)
def __rpow__(self, other):
return self._apply(other, lambda u, v: v ** u)
def __neg__(self):
return TensorList([-u for u in self.tensors])
def __iter__(self):
for tensor in self.tensors:
yield tensor
def __repr__(self):
return 'TensorList: \n' + repr(self.tensors)
def _apply(self, other, op):
if isinstance(other, (list, tuple, TensorList)) or (
isinstance(other, torch.Tensor) and (
other.numel() > 1 or other.ndim > 1
)
):
assert len(other) == len(self.tensors)
return TensorList([op(u, v) for u, v in zip(self.tensors, other)])
elif isinstance(other, numbers.Number) or (
isinstance(other, torch.Tensor) and (
other.numel() == 1 and other.ndim <= 1
)
):
return TensorList([op(u, other) for u in self.tensors])
else:
raise TypeError(
f'unsupported operand for *: "TensorList" and "{type(other)}"'
)
```
## /wan/modules/animate/clip.py
```py path="/wan/modules/animate/clip.py"
# Modified from ``https://github.com/openai/CLIP'' and ``https://github.com/mlfoundations/open_clip''
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import logging
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
from ..attention import flash_attention
from ..tokenizers import HuggingfaceTokenizer
from .xlm_roberta import XLMRoberta
__all__ = [
'XLMRobertaCLIP',
'clip_xlm_roberta_vit_h_14',
'CLIPModel',
]
def pos_interpolate(pos, seq_len):
if pos.size(1) == seq_len:
return pos
else:
src_grid = int(math.sqrt(pos.size(1)))
tar_grid = int(math.sqrt(seq_len))
n = pos.size(1) - src_grid * src_grid
return torch.cat([
pos[:, :n],
F.interpolate(
pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute(
0, 3, 1, 2),
size=(tar_grid, tar_grid),
mode='bicubic',
align_corners=False).flatten(2).transpose(1, 2)
],
dim=1)
class QuickGELU(nn.Module):
def forward(self, x):
return x * torch.sigmoid(1.702 * x)
class LayerNorm(nn.LayerNorm):
def forward(self, x):
return super().forward(x.float()).type_as(x)
class SelfAttention(nn.Module):
def __init__(self,
dim,
num_heads,
causal=False,
attn_dropout=0.0,
proj_dropout=0.0):
assert dim % num_heads == 0
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.causal = causal
self.attn_dropout = attn_dropout
self.proj_dropout = proj_dropout
# layers
self.to_qkv = nn.Linear(dim, dim * 3)
self.proj = nn.Linear(dim, dim)
def forward(self, x):
"""
x: [B, L, C].
"""
b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
# compute query, key, value
q, k, v = self.to_qkv(x).view(b, s, 3, n, d).unbind(2)
# 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 = x.reshape(b, s, c)
# output
x = self.proj(x)
x = F.dropout(x, self.proj_dropout, self.training)
return x
class SwiGLU(nn.Module):
def __init__(self, dim, mid_dim):
super().__init__()
self.dim = dim
self.mid_dim = mid_dim
# layers
self.fc1 = nn.Linear(dim, mid_dim)
self.fc2 = nn.Linear(dim, mid_dim)
self.fc3 = nn.Linear(mid_dim, dim)
def forward(self, x):
x = F.silu(self.fc1(x)) * self.fc2(x)
x = self.fc3(x)
return x
class AttentionBlock(nn.Module):
def __init__(self,
dim,
mlp_ratio,
num_heads,
post_norm=False,
causal=False,
activation='quick_gelu',
attn_dropout=0.0,
proj_dropout=0.0,
norm_eps=1e-5):
assert activation in ['quick_gelu', 'gelu', 'swi_glu']
super().__init__()
self.dim = dim
self.mlp_ratio = mlp_ratio
self.num_heads = num_heads
self.post_norm = post_norm
self.causal = causal
self.norm_eps = norm_eps
# layers
self.norm1 = LayerNorm(dim, eps=norm_eps)
self.attn = SelfAttention(dim, num_heads, causal, attn_dropout,
proj_dropout)
self.norm2 = LayerNorm(dim, eps=norm_eps)
if activation == 'swi_glu':
self.mlp = SwiGLU(dim, int(dim * mlp_ratio))
else:
self.mlp = nn.Sequential(
nn.Linear(dim, int(dim * mlp_ratio)),
QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
def forward(self, x):
if self.post_norm:
x = x + self.norm1(self.attn(x))
x = x + self.norm2(self.mlp(x))
else:
x = x + self.attn(self.norm1(x))
x = x + self.mlp(self.norm2(x))
return x
class AttentionPool(nn.Module):
def __init__(self,
dim,
mlp_ratio,
num_heads,
activation='gelu',
proj_dropout=0.0,
norm_eps=1e-5):
assert dim % num_heads == 0
super().__init__()
self.dim = dim
self.mlp_ratio = mlp_ratio
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.proj_dropout = proj_dropout
self.norm_eps = norm_eps
# layers
gain = 1.0 / math.sqrt(dim)
self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
self.to_q = nn.Linear(dim, dim)
self.to_kv = nn.Linear(dim, dim * 2)
self.proj = nn.Linear(dim, dim)
self.norm = LayerNorm(dim, eps=norm_eps)
self.mlp = nn.Sequential(
nn.Linear(dim, int(dim * mlp_ratio)),
QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
def forward(self, x):
"""
x: [B, L, C].
"""
b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
# compute query, key, value
q = self.to_q(self.cls_embedding).view(1, 1, n, d).expand(b, -1, -1, -1)
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 = x.reshape(b, 1, c)
# output
x = self.proj(x)
x = F.dropout(x, self.proj_dropout, self.training)
# mlp
x = x + self.mlp(self.norm(x))
return x[:, 0]
class VisionTransformer(nn.Module):
def __init__(self,
image_size=224,
patch_size=16,
dim=768,
mlp_ratio=4,
out_dim=512,
num_heads=12,
num_layers=12,
pool_type='token',
pre_norm=True,
post_norm=False,
activation='quick_gelu',
attn_dropout=0.0,
proj_dropout=0.0,
embedding_dropout=0.0,
norm_eps=1e-5):
if image_size % patch_size != 0:
print(
'[WARNING] image_size is not divisible by patch_size',
flush=True)
assert pool_type in ('token', 'token_fc', 'attn_pool')
out_dim = out_dim or dim
super().__init__()
self.image_size = image_size
self.patch_size = patch_size
self.num_patches = (image_size // patch_size)**2
self.dim = dim
self.mlp_ratio = mlp_ratio
self.out_dim = out_dim
self.num_heads = num_heads
self.num_layers = num_layers
self.pool_type = pool_type
self.post_norm = post_norm
self.norm_eps = norm_eps
# embeddings
gain = 1.0 / math.sqrt(dim)
self.patch_embedding = nn.Conv2d(
3,
dim,
kernel_size=patch_size,
stride=patch_size,
bias=not pre_norm)
if pool_type in ('token', 'token_fc'):
self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
self.pos_embedding = nn.Parameter(gain * torch.randn(
1, self.num_patches +
(1 if pool_type in ('token', 'token_fc') else 0), dim))
self.dropout = nn.Dropout(embedding_dropout)
# transformer
self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None
self.transformer = nn.Sequential(*[
AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False,
activation, attn_dropout, proj_dropout, norm_eps)
for _ in range(num_layers)
])
self.post_norm = LayerNorm(dim, eps=norm_eps)
# head
if pool_type == 'token':
self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
elif pool_type == 'token_fc':
self.head = nn.Linear(dim, out_dim)
elif pool_type == 'attn_pool':
self.head = AttentionPool(dim, mlp_ratio, num_heads, activation,
proj_dropout, norm_eps)
def forward(self, x, interpolation=False, use_31_block=False):
b = x.size(0)
# embeddings
x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
if self.pool_type in ('token', 'token_fc'):
x = torch.cat([self.cls_embedding.expand(b, -1, -1), x], dim=1)
if interpolation:
e = pos_interpolate(self.pos_embedding, x.size(1))
else:
e = self.pos_embedding
x = self.dropout(x + e)
if self.pre_norm is not None:
x = self.pre_norm(x)
# transformer
if use_31_block:
x = self.transformer[:-1](x)
return x
else:
x = self.transformer(x)
return x
class XLMRobertaWithHead(XLMRoberta):
def __init__(self, **kwargs):
self.out_dim = kwargs.pop('out_dim')
super().__init__(**kwargs)
# head
mid_dim = (self.dim + self.out_dim) // 2
self.head = nn.Sequential(
nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(),
nn.Linear(mid_dim, self.out_dim, bias=False))
def forward(self, ids):
# xlm-roberta
x = super().forward(ids)
# average pooling
mask = ids.ne(self.pad_id).unsqueeze(-1).to(x)
x = (x * mask).sum(dim=1) / mask.sum(dim=1)
# head
x = self.head(x)
return x
class XLMRobertaCLIP(nn.Module):
def __init__(self,
embed_dim=1024,
image_size=224,
patch_size=14,
vision_dim=1280,
vision_mlp_ratio=4,
vision_heads=16,
vision_layers=32,
vision_pool='token',
vision_pre_norm=True,
vision_post_norm=False,
activation='gelu',
vocab_size=250002,
max_text_len=514,
type_size=1,
pad_id=1,
text_dim=1024,
text_heads=16,
text_layers=24,
text_post_norm=True,
text_dropout=0.1,
attn_dropout=0.0,
proj_dropout=0.0,
embedding_dropout=0.0,
norm_eps=1e-5):
super().__init__()
self.embed_dim = embed_dim
self.image_size = image_size
self.patch_size = patch_size
self.vision_dim = vision_dim
self.vision_mlp_ratio = vision_mlp_ratio
self.vision_heads = vision_heads
self.vision_layers = vision_layers
self.vision_pre_norm = vision_pre_norm
self.vision_post_norm = vision_post_norm
self.activation = activation
self.vocab_size = vocab_size
self.max_text_len = max_text_len
self.type_size = type_size
self.pad_id = pad_id
self.text_dim = text_dim
self.text_heads = text_heads
self.text_layers = text_layers
self.text_post_norm = text_post_norm
self.norm_eps = norm_eps
# models
self.visual = VisionTransformer(
image_size=image_size,
patch_size=patch_size,
dim=vision_dim,
mlp_ratio=vision_mlp_ratio,
out_dim=embed_dim,
num_heads=vision_heads,
num_layers=vision_layers,
pool_type=vision_pool,
pre_norm=vision_pre_norm,
post_norm=vision_post_norm,
activation=activation,
attn_dropout=attn_dropout,
proj_dropout=proj_dropout,
embedding_dropout=embedding_dropout,
norm_eps=norm_eps)
self.textual = XLMRobertaWithHead(
vocab_size=vocab_size,
max_seq_len=max_text_len,
type_size=type_size,
pad_id=pad_id,
dim=text_dim,
out_dim=embed_dim,
num_heads=text_heads,
num_layers=text_layers,
post_norm=text_post_norm,
dropout=text_dropout)
self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([]))
def forward(self, imgs, txt_ids):
"""
imgs: [B, 3, H, W] of torch.float32.
- mean: [0.48145466, 0.4578275, 0.40821073]
- std: [0.26862954, 0.26130258, 0.27577711]
txt_ids: [B, L] of torch.long.
Encoded by data.CLIPTokenizer.
"""
xi = self.visual(imgs)
xt = self.textual(txt_ids)
return xi, xt
def param_groups(self):
groups = [{
'params': [
p for n, p in self.named_parameters()
if 'norm' in n or n.endswith('bias')
],
'weight_decay': 0.0
}, {
'params': [
p for n, p in self.named_parameters()
if not ('norm' in n or n.endswith('bias'))
]
}]
return groups
def _clip(pretrained=False,
pretrained_name=None,
model_cls=XLMRobertaCLIP,
return_transforms=False,
return_tokenizer=False,
tokenizer_padding='eos',
dtype=torch.float32,
device='cpu',
**kwargs):
# init a model on device
with torch.device(device):
model = model_cls(**kwargs)
# set device
model = model.to(dtype=dtype, device=device)
output = (model,)
# init transforms
if return_transforms:
# mean and std
if 'siglip' in pretrained_name.lower():
mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
else:
mean = [0.48145466, 0.4578275, 0.40821073]
std = [0.26862954, 0.26130258, 0.27577711]
# transforms
transforms = T.Compose([
T.Resize((model.image_size, model.image_size),
interpolation=T.InterpolationMode.BICUBIC),
T.ToTensor(),
T.Normalize(mean=mean, std=std)
])
output += (transforms,)
return output[0] if len(output) == 1 else output
def clip_xlm_roberta_vit_h_14(
pretrained=False,
pretrained_name='open-clip-xlm-roberta-large-vit-huge-14',
**kwargs):
cfg = dict(
embed_dim=1024,
image_size=224,
patch_size=14,
vision_dim=1280,
vision_mlp_ratio=4,
vision_heads=16,
vision_layers=32,
vision_pool='token',
activation='gelu',
vocab_size=250002,
max_text_len=514,
type_size=1,
pad_id=1,
text_dim=1024,
text_heads=16,
text_layers=24,
text_post_norm=True,
text_dropout=0.1,
attn_dropout=0.0,
proj_dropout=0.0,
embedding_dropout=0.0)
cfg.update(**kwargs)
return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg)
class CLIPModel:
def __init__(self, dtype, device, checkpoint_path, tokenizer_path):
self.dtype = dtype
self.device = device
self.checkpoint_path = checkpoint_path
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)
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'))
# init tokenizer
self.tokenizer = HuggingfaceTokenizer(
name=tokenizer_path,
seq_len=self.model.max_text_len - 2,
clean='whitespace')
def visual(self, videos):
# preprocess
size = (self.model.image_size,) * 2
videos = torch.cat([
F.interpolate(
u.transpose(0, 1),
size=size,
mode='bicubic',
align_corners=False) for u in videos
])
videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))
# forward
with torch.cuda.amp.autocast(dtype=self.dtype):
out = self.model.visual(videos, use_31_block=True)
return out
```
## /wan/modules/animate/face_blocks.py
```py path="/wan/modules/animate/face_blocks.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from torch import nn
import torch
from typing import Tuple, Optional
from einops import rearrange
import torch.nn.functional as F
import math
from ...distributed.util import gather_forward, get_rank, get_world_size
try:
from flash_attn import flash_attn_qkvpacked_func, flash_attn_func
except ImportError:
flash_attn_func = None
MEMORY_LAYOUT = {
"flash": (
lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]),
lambda x: x,
),
"torch": (
lambda x: x.transpose(1, 2),
lambda x: x.transpose(1, 2),
),
"vanilla": (
lambda x: x.transpose(1, 2),
lambda x: x.transpose(1, 2),
),
}
def attention(
q,
k,
v,
mode="flash",
drop_rate=0,
attn_mask=None,
causal=False,
max_seqlen_q=None,
batch_size=1,
):
"""
Perform QKV self attention.
Args:
q (torch.Tensor): Query tensor with shape [b, s, a, d], where a is the number of heads.
k (torch.Tensor): Key tensor with shape [b, s1, a, d]
v (torch.Tensor): Value tensor with shape [b, s1, a, d]
mode (str): Attention mode. Choose from 'self_flash', 'cross_flash', 'torch', and 'vanilla'.
drop_rate (float): Dropout rate in attention map. (default: 0)
attn_mask (torch.Tensor): Attention mask with shape [b, s1] (cross_attn), or [b, a, s, s1] (torch or vanilla).
(default: None)
causal (bool): Whether to use causal attention. (default: False)
cu_seqlens_q (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
used to index into q.
cu_seqlens_kv (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
used to index into kv.
max_seqlen_q (int): The maximum sequence length in the batch of q.
max_seqlen_kv (int): The maximum sequence length in the batch of k and v.
Returns:
torch.Tensor: Output tensor after self attention with shape [b, s, ad]
"""
pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode]
if mode == "torch":
if attn_mask is not None and attn_mask.dtype != torch.bool:
attn_mask = attn_mask.to(q.dtype)
x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal)
elif mode == "flash":
x = flash_attn_func(
q,
k,
v,
)
x = x.view(batch_size, max_seqlen_q, x.shape[-2], x.shape[-1]) # reshape x to [b, s, a, d]
elif mode == "vanilla":
scale_factor = 1 / math.sqrt(q.size(-1))
b, a, s, _ = q.shape
s1 = k.size(2)
attn_bias = torch.zeros(b, a, s, s1, dtype=q.dtype, device=q.device)
if causal:
# Only applied to self attention
assert attn_mask is None, "Causal mask and attn_mask cannot be used together"
temp_mask = torch.ones(b, a, s, s, dtype=torch.bool, device=q.device).tril(diagonal=0)
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
attn_bias.to(q.dtype)
if attn_mask is not None:
if attn_mask.dtype == torch.bool:
attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
else:
attn_bias += attn_mask
attn = (q @ k.transpose(-2, -1)) * scale_factor
attn += attn_bias
attn = attn.softmax(dim=-1)
attn = torch.dropout(attn, p=drop_rate, train=True)
x = attn @ v
else:
raise NotImplementedError(f"Unsupported attention mode: {mode}")
x = post_attn_layout(x)
b, s, a, d = x.shape
out = x.reshape(b, s, -1)
return out
class CausalConv1d(nn.Module):
def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode="replicate", **kwargs):
super().__init__()
self.pad_mode = pad_mode
padding = (kernel_size - 1, 0) # T
self.time_causal_padding = padding
self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
def forward(self, x):
x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
return self.conv(x)
class FaceEncoder(nn.Module):
def __init__(self, in_dim: int, hidden_dim: int, num_heads=int, dtype=None, device=None):
factory_kwargs = {"dtype": dtype, "device": device}
super().__init__()
self.num_heads = num_heads
self.conv1_local = CausalConv1d(in_dim, 1024 * num_heads, 3, stride=1)
self.norm1 = nn.LayerNorm(hidden_dim // 8, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.act = nn.SiLU()
self.conv2 = CausalConv1d(1024, 1024, 3, stride=2)
self.conv3 = CausalConv1d(1024, 1024, 3, stride=2)
self.out_proj = nn.Linear(1024, hidden_dim)
self.norm1 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.norm2 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.norm3 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, hidden_dim))
def forward(self, x):
x = rearrange(x, "b t c -> b c t")
b, c, t = x.shape
x = self.conv1_local(x)
x = rearrange(x, "b (n c) t -> (b n) t c", n=self.num_heads)
x = self.norm1(x)
x = self.act(x)
x = rearrange(x, "b t c -> b c t")
x = self.conv2(x)
x = rearrange(x, "b c t -> b t c")
x = self.norm2(x)
x = self.act(x)
x = rearrange(x, "b t c -> b c t")
x = self.conv3(x)
x = rearrange(x, "b c t -> b t c")
x = self.norm3(x)
x = self.act(x)
x = self.out_proj(x)
x = rearrange(x, "(b n) t c -> b t n c", b=b)
padding = self.padding_tokens.repeat(b, x.shape[1], 1, 1)
x = torch.cat([x, padding], dim=-2)
x_local = x.clone()
return x_local
class RMSNorm(nn.Module):
def __init__(
self,
dim: int,
elementwise_affine=True,
eps: float = 1e-6,
device=None,
dtype=None,
):
"""
Initialize the RMSNorm normalization layer.
Args:
dim (int): The dimension of the input tensor.
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
Attributes:
eps (float): A small value added to the denominator for numerical stability.
weight (nn.Parameter): Learnable scaling parameter.
"""
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__()
self.eps = eps
if elementwise_affine:
self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
def _norm(self, x):
"""
Apply the RMSNorm normalization to the input tensor.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The normalized tensor.
"""
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
"""
Forward pass through the RMSNorm layer.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The output tensor after applying RMSNorm.
"""
output = self._norm(x.float()).type_as(x)
if hasattr(self, "weight"):
output = output * self.weight
return output
def get_norm_layer(norm_layer):
"""
Get the normalization layer.
Args:
norm_layer (str): The type of normalization layer.
Returns:
norm_layer (nn.Module): The normalization layer.
"""
if norm_layer == "layer":
return nn.LayerNorm
elif norm_layer == "rms":
return RMSNorm
else:
raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")
class FaceAdapter(nn.Module):
def __init__(
self,
hidden_dim: int,
heads_num: int,
qk_norm: bool = True,
qk_norm_type: str = "rms",
num_adapter_layers: int = 1,
dtype=None,
device=None,
):
factory_kwargs = {"dtype": dtype, "device": device}
super().__init__()
self.hidden_size = hidden_dim
self.heads_num = heads_num
self.fuser_blocks = nn.ModuleList(
[
FaceBlock(
self.hidden_size,
self.heads_num,
qk_norm=qk_norm,
qk_norm_type=qk_norm_type,
**factory_kwargs,
)
for _ in range(num_adapter_layers)
]
)
def forward(
self,
x: torch.Tensor,
motion_embed: torch.Tensor,
idx: int,
freqs_cis_q: Tuple[torch.Tensor, torch.Tensor] = None,
freqs_cis_k: Tuple[torch.Tensor, torch.Tensor] = None,
) -> torch.Tensor:
return self.fuser_blocks[idx](x, motion_embed, freqs_cis_q, freqs_cis_k)
class FaceBlock(nn.Module):
def __init__(
self,
hidden_size: int,
heads_num: int,
qk_norm: bool = True,
qk_norm_type: str = "rms",
qk_scale: float = None,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
):
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__()
self.deterministic = False
self.hidden_size = hidden_size
self.heads_num = heads_num
head_dim = hidden_size // heads_num
self.scale = qk_scale or head_dim**-0.5
self.linear1_kv = nn.Linear(hidden_size, hidden_size * 2, **factory_kwargs)
self.linear1_q = nn.Linear(hidden_size, hidden_size, **factory_kwargs)
self.linear2 = nn.Linear(hidden_size, hidden_size, **factory_kwargs)
qk_norm_layer = get_norm_layer(qk_norm_type)
self.q_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
)
self.k_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
)
self.pre_norm_feat = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.pre_norm_motion = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
def forward(
self,
x: torch.Tensor,
motion_vec: torch.Tensor,
motion_mask: Optional[torch.Tensor] = None,
use_context_parallel=False,
) -> torch.Tensor:
B, T, N, C = motion_vec.shape
T_comp = T
x_motion = self.pre_norm_motion(motion_vec)
x_feat = self.pre_norm_feat(x)
kv = self.linear1_kv(x_motion)
q = self.linear1_q(x_feat)
k, v = rearrange(kv, "B L N (K H D) -> K B L N H D", K=2, H=self.heads_num)
q = rearrange(q, "B S (H D) -> B S H D", H=self.heads_num)
# Apply QK-Norm if needed.
q = self.q_norm(q).to(v)
k = self.k_norm(k).to(v)
k = rearrange(k, "B L N H D -> (B L) N H D")
v = rearrange(v, "B L N H D -> (B L) N H D")
if use_context_parallel:
q = gather_forward(q, dim=1)
q = rearrange(q, "B (L S) H D -> (B L) S H D", L=T_comp)
# Compute attention.
attn = attention(
q,
k,
v,
max_seqlen_q=q.shape[1],
batch_size=q.shape[0],
)
attn = rearrange(attn, "(B L) S C -> B (L S) C", L=T_comp)
if use_context_parallel:
attn = torch.chunk(attn, get_world_size(), dim=1)[get_rank()]
output = self.linear2(attn)
if motion_mask is not None:
output = output * rearrange(motion_mask, "B T H W -> B (T H W)").unsqueeze(-1)
return output
```
## /wan/modules/animate/model_animate.py
```py path="/wan/modules/animate/model_animate.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import math
import types
from copy import deepcopy
from einops import rearrange
from typing import List
import numpy as np
import torch
import torch.cuda.amp as amp
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from diffusers.loaders import PeftAdapterMixin
from ...distributed.sequence_parallel import (
distributed_attention,
gather_forward,
get_rank,
get_world_size,
)
from ..model import (
Head,
WanAttentionBlock,
WanLayerNorm,
WanRMSNorm,
WanModel,
WanSelfAttention,
flash_attention,
rope_params,
sinusoidal_embedding_1d,
rope_apply
)
from .face_blocks import FaceEncoder, FaceAdapter
from .motion_encoder import Generator
class HeadAnimate(Head):
def forward(self, x, e):
"""
Args:
x(Tensor): Shape [B, L1, C]
e(Tensor): Shape [B, L1, 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]))
return x
class WanAnimateSelfAttention(WanSelfAttention):
def forward(self, x, seq_lens, grid_sizes, freqs):
"""
Args:
x(Tensor): Shape [B, L, num_heads, C / num_heads]
seq_lens(Tensor): Shape [B]
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
"""
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,
window_size=self.window_size)
# output
x = x.flatten(2)
x = self.o(x)
return x
class WanAnimateCrossAttention(WanSelfAttention):
def __init__(
self,
dim,
num_heads,
window_size=(-1, -1),
qk_norm=True,
eps=1e-6,
use_img_emb=True
):
super().__init__(
dim,
num_heads,
window_size,
qk_norm,
eps
)
self.use_img_emb = use_img_emb
if use_img_emb:
self.k_img = nn.Linear(dim, dim)
self.v_img = nn.Linear(dim, dim)
self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
def forward(self, x, context, context_lens):
"""
x: [B, L1, C].
context: [B, L2, C].
context_lens: [B].
"""
if self.use_img_emb:
context_img = context[:, :257]
context = context[:, 257:]
else:
context = context
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)
v = self.v(context).view(b, -1, n, d)
if self.use_img_emb:
k_img = self.norm_k_img(self.k_img(context_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)
# compute attention
x = flash_attention(q, k, v, k_lens=context_lens)
# output
x = x.flatten(2)
if self.use_img_emb:
img_x = img_x.flatten(2)
x = x + img_x
x = self.o(x)
return x
class WanAnimateAttentionBlock(nn.Module):
def __init__(self,
dim,
ffn_dim,
num_heads,
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=True,
eps=1e-6,
use_img_emb=True):
super().__init__()
self.dim = dim
self.ffn_dim = ffn_dim
self.num_heads = num_heads
self.window_size = window_size
self.qk_norm = qk_norm
self.cross_attn_norm = cross_attn_norm
self.eps = eps
# layers
self.norm1 = WanLayerNorm(dim, eps)
self.self_attn = WanAnimateSelfAttention(dim, num_heads, window_size, qk_norm, eps)
self.norm3 = WanLayerNorm(
dim, eps, elementwise_affine=True
) if cross_attn_norm else nn.Identity()
self.cross_attn = WanAnimateCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, use_img_emb=use_img_emb)
self.norm2 = WanLayerNorm(dim, eps)
self.ffn = nn.Sequential(
nn.Linear(dim, ffn_dim),
nn.GELU(approximate='tanh'),
nn.Linear(ffn_dim, dim)
)
# modulation
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim ** 0.5)
def forward(
self,
x,
e,
seq_lens,
grid_sizes,
freqs,
context,
context_lens,
):
"""
Args:
x(Tensor): Shape [B, L, C]
e(Tensor): Shape [B, L1, 6, C]
seq_lens(Tensor): Shape [B], length of each sequence in batch
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
# 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]
# 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
x = cross_attn_ffn(x, context, context_lens, e)
return x
class MLPProj(torch.nn.Module):
def __init__(self, in_dim, out_dim):
super().__init__()
self.proj = torch.nn.Sequential(
torch.nn.LayerNorm(in_dim),
torch.nn.Linear(in_dim, in_dim),
torch.nn.GELU(),
torch.nn.Linear(in_dim, out_dim),
torch.nn.LayerNorm(out_dim),
)
def forward(self, image_embeds):
clip_extra_context_tokens = self.proj(image_embeds)
return clip_extra_context_tokens
class WanAnimateModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
_no_split_modules = ['WanAttentionBlock']
@register_to_config
def __init__(self,
patch_size=(1, 2, 2),
text_len=512,
in_dim=36,
dim=5120,
ffn_dim=13824,
freq_dim=256,
text_dim=4096,
out_dim=16,
num_heads=40,
num_layers=40,
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=True,
eps=1e-6,
motion_encoder_dim=512,
use_context_parallel=False,
use_img_emb=True):
super().__init__()
self.patch_size = patch_size
self.text_len = text_len
self.in_dim = in_dim
self.dim = dim
self.ffn_dim = ffn_dim
self.freq_dim = freq_dim
self.text_dim = text_dim
self.out_dim = out_dim
self.num_heads = num_heads
self.num_layers = num_layers
self.window_size = window_size
self.qk_norm = qk_norm
self.cross_attn_norm = cross_attn_norm
self.eps = eps
self.motion_encoder_dim = motion_encoder_dim
self.use_context_parallel = use_context_parallel
self.use_img_emb = use_img_emb
# embeddings
self.patch_embedding = nn.Conv3d(
in_dim, dim, kernel_size=patch_size, stride=patch_size)
self.pose_patch_embedding = nn.Conv3d(
16, dim, kernel_size=patch_size, stride=patch_size
)
self.text_embedding = nn.Sequential(
nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
nn.Linear(dim, dim))
self.time_embedding = nn.Sequential(
nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
# blocks
self.blocks = nn.ModuleList([
WanAnimateAttentionBlock(dim, ffn_dim, num_heads, window_size, qk_norm,
cross_attn_norm, eps, use_img_emb) for _ in range(num_layers)
])
# head
self.head = HeadAnimate(dim, out_dim, patch_size, eps)
# 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([
rope_params(1024, d - 4 * (d // 6)),
rope_params(1024, 2 * (d // 6)),
rope_params(1024, 2 * (d // 6))
], dim=1)
self.img_emb = MLPProj(1280, dim)
# initialize weights
self.init_weights()
self.motion_encoder = Generator(size=512, style_dim=512, motion_dim=20)
self.face_adapter = FaceAdapter(
heads_num=self.num_heads,
hidden_dim=self.dim,
num_adapter_layers=self.num_layers // 5,
)
self.face_encoder = FaceEncoder(
in_dim=motion_encoder_dim,
hidden_dim=self.dim,
num_heads=4,
)
def after_patch_embedding(self, x: List[torch.Tensor], pose_latents, face_pixel_values):
pose_latents = [self.pose_patch_embedding(u.unsqueeze(0)) for u in pose_latents]
for x_, pose_latents_ in zip(x, pose_latents):
x_[:, :, 1:] += pose_latents_
b,c,T,h,w = face_pixel_values.shape
face_pixel_values = rearrange(face_pixel_values, "b c t h w -> (b t) c h w")
encode_bs = 8
face_pixel_values_tmp = []
for i in range(math.ceil(face_pixel_values.shape[0]/encode_bs)):
face_pixel_values_tmp.append(self.motion_encoder.get_motion(face_pixel_values[i*encode_bs:(i+1)*encode_bs]))
motion_vec = torch.cat(face_pixel_values_tmp)
motion_vec = rearrange(motion_vec, "(b t) c -> b t c", t=T)
motion_vec = self.face_encoder(motion_vec)
B, L, H, C = motion_vec.shape
pad_face = torch.zeros(B, 1, H, C).type_as(motion_vec)
motion_vec = torch.cat([pad_face, motion_vec], dim=1)
return x, motion_vec
def after_transformer_block(self, block_idx, x, motion_vec, motion_masks=None):
if block_idx % 5 == 0:
adapter_args = [x, motion_vec, motion_masks, self.use_context_parallel]
residual_out = self.face_adapter.fuser_blocks[block_idx // 5](*adapter_args)
x = residual_out + x
return x
def forward(
self,
x,
t,
clip_fea,
context,
seq_len,
y=None,
pose_latents=None,
face_pixel_values=None
):
# params
device = self.patch_embedding.weight.device
if self.freqs.device != device:
self.freqs = self.freqs.to(device)
if y is not None:
x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
# embeddings
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
x, motion_vec = self.after_patch_embedding(x, pose_latents, face_pixel_values)
grid_sizes = torch.stack(
[torch.tensor(u.shape[2:], dtype=torch.long) 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
x = torch.cat([
torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
dim=1) for u in x
])
# 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
# context
context_lens = None
context = self.text_embedding(
torch.stack([
torch.cat(
[u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
for u in context
]))
if self.use_img_emb:
context_clip = self.img_emb(clip_fea) # bs x 257 x dim
context = torch.concat([context_clip, context], dim=1)
# arguments
kwargs = dict(
e=e0,
seq_lens=seq_lens,
grid_sizes=grid_sizes,
freqs=self.freqs,
context=context,
context_lens=context_lens)
if self.use_context_parallel:
x = torch.chunk(x, get_world_size(), dim=1)[get_rank()]
for idx, block in enumerate(self.blocks):
x = block(x, **kwargs)
x = self.after_transformer_block(idx, x, motion_vec)
# head
x = self.head(x, e)
if self.use_context_parallel:
x = gather_forward(x, dim=1)
# unpatchify
x = self.unpatchify(x, grid_sizes)
return [u.float() for u in x]
def unpatchify(self, x, grid_sizes):
r"""
Reconstruct video tensors from patch embeddings.
Args:
x (List[Tensor]):
List of patchified features, each with shape [L, C_out * prod(patch_size)]
grid_sizes (Tensor):
Original spatial-temporal grid dimensions before patching,
shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
Returns:
List[Tensor]:
Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
"""
c = self.out_dim
out = []
for u, v in zip(x, grid_sizes.tolist()):
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)])
out.append(u)
return out
def init_weights(self):
r"""
Initialize model parameters using Xavier initialization.
"""
# basic init
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
if m.bias is not None:
nn.init.zeros_(m.bias)
# init embeddings
nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
for m in self.text_embedding.modules():
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, std=.02)
for m in self.time_embedding.modules():
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, std=.02)
# init output layer
nn.init.zeros_(self.head.head.weight)
```
## /wan/modules/animate/motion_encoder.py
```py path="/wan/modules/animate/motion_encoder.py"
# Modified from ``https://github.com/wyhsirius/LIA``
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
import torch.nn as nn
from torch.nn import functional as F
import math
def custom_qr(input_tensor):
original_dtype = input_tensor.dtype
if original_dtype == torch.bfloat16:
q, r = torch.linalg.qr(input_tensor.to(torch.float32))
return q.to(original_dtype), r.to(original_dtype)
return torch.linalg.qr(input_tensor)
def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
return F.leaky_relu(input + bias, negative_slope) * scale
def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
_, minor, in_h, in_w = input.shape
kernel_h, kernel_w = kernel.shape
out = input.view(-1, minor, in_h, 1, in_w, 1)
out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0])
out = out.view(-1, minor, in_h * up_y, in_w * up_x)
out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0),
max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0), ]
out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
out = F.conv2d(out, w)
out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, )
return out[:, :, ::down_y, ::down_x]
def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
def make_kernel(k):
k = torch.tensor(k, dtype=torch.float32)
if k.ndim == 1:
k = k[None, :] * k[:, None]
k /= k.sum()
return k
class FusedLeakyReLU(nn.Module):
def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
super().__init__()
self.bias = nn.Parameter(torch.zeros(1, channel, 1, 1))
self.negative_slope = negative_slope
self.scale = scale
def forward(self, input):
out = fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
return out
class Blur(nn.Module):
def __init__(self, kernel, pad, upsample_factor=1):
super().__init__()
kernel = make_kernel(kernel)
if upsample_factor > 1:
kernel = kernel * (upsample_factor ** 2)
self.register_buffer('kernel', kernel)
self.pad = pad
def forward(self, input):
return upfirdn2d(input, self.kernel, pad=self.pad)
class ScaledLeakyReLU(nn.Module):
def __init__(self, negative_slope=0.2):
super().__init__()
self.negative_slope = negative_slope
def forward(self, input):
return F.leaky_relu(input, negative_slope=self.negative_slope)
class EqualConv2d(nn.Module):
def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
super().__init__()
self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size))
self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
self.stride = stride
self.padding = padding
if bias:
self.bias = nn.Parameter(torch.zeros(out_channel))
else:
self.bias = None
def forward(self, input):
return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding)
def __repr__(self):
return (
f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
)
class EqualLinear(nn.Module):
def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):
super().__init__()
self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
if bias:
self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
else:
self.bias = None
self.activation = activation
self.scale = (1 / math.sqrt(in_dim)) * lr_mul
self.lr_mul = lr_mul
def forward(self, input):
if self.activation:
out = F.linear(input, self.weight * self.scale)
out = fused_leaky_relu(out, self.bias * self.lr_mul)
else:
out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul)
return out
def __repr__(self):
return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})')
class ConvLayer(nn.Sequential):
def __init__(
self,
in_channel,
out_channel,
kernel_size,
downsample=False,
blur_kernel=[1, 3, 3, 1],
bias=True,
activate=True,
):
layers = []
if downsample:
factor = 2
p = (len(blur_kernel) - factor) + (kernel_size - 1)
pad0 = (p + 1) // 2
pad1 = p // 2
layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
stride = 2
self.padding = 0
else:
stride = 1
self.padding = kernel_size // 2
layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=self.padding, stride=stride,
bias=bias and not activate))
if activate:
if bias:
layers.append(FusedLeakyReLU(out_channel))
else:
layers.append(ScaledLeakyReLU(0.2))
super().__init__(*layers)
class ResBlock(nn.Module):
def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
super().__init__()
self.conv1 = ConvLayer(in_channel, in_channel, 3)
self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)
def forward(self, input):
out = self.conv1(input)
out = self.conv2(out)
skip = self.skip(input)
out = (out + skip) / math.sqrt(2)
return out
class EncoderApp(nn.Module):
def __init__(self, size, w_dim=512):
super(EncoderApp, self).__init__()
channels = {
4: 512,
8: 512,
16: 512,
32: 512,
64: 256,
128: 128,
256: 64,
512: 32,
1024: 16
}
self.w_dim = w_dim
log_size = int(math.log(size, 2))
self.convs = nn.ModuleList()
self.convs.append(ConvLayer(3, channels[size], 1))
in_channel = channels[size]
for i in range(log_size, 2, -1):
out_channel = channels[2 ** (i - 1)]
self.convs.append(ResBlock(in_channel, out_channel))
in_channel = out_channel
self.convs.append(EqualConv2d(in_channel, self.w_dim, 4, padding=0, bias=False))
def forward(self, x):
res = []
h = x
for conv in self.convs:
h = conv(h)
res.append(h)
return res[-1].squeeze(-1).squeeze(-1), res[::-1][2:]
class Encoder(nn.Module):
def __init__(self, size, dim=512, dim_motion=20):
super(Encoder, self).__init__()
# appearance netmork
self.net_app = EncoderApp(size, dim)
# motion network
fc = [EqualLinear(dim, dim)]
for i in range(3):
fc.append(EqualLinear(dim, dim))
fc.append(EqualLinear(dim, dim_motion))
self.fc = nn.Sequential(*fc)
def enc_app(self, x):
h_source = self.net_app(x)
return h_source
def enc_motion(self, x):
h, _ = self.net_app(x)
h_motion = self.fc(h)
return h_motion
class Direction(nn.Module):
def __init__(self, motion_dim):
super(Direction, self).__init__()
self.weight = nn.Parameter(torch.randn(512, motion_dim))
def forward(self, input):
weight = self.weight + 1e-8
Q, R = custom_qr(weight)
if input is None:
return Q
else:
input_diag = torch.diag_embed(input) # alpha, diagonal matrix
out = torch.matmul(input_diag, Q.T)
out = torch.sum(out, dim=1)
return out
class Synthesis(nn.Module):
def __init__(self, motion_dim):
super(Synthesis, self).__init__()
self.direction = Direction(motion_dim)
class Generator(nn.Module):
def __init__(self, size, style_dim=512, motion_dim=20):
super().__init__()
self.enc = Encoder(size, style_dim, motion_dim)
self.dec = Synthesis(motion_dim)
def get_motion(self, img):
#motion_feat = self.enc.enc_motion(img)
motion_feat = torch.utils.checkpoint.checkpoint((self.enc.enc_motion), img, use_reentrant=True)
with torch.cuda.amp.autocast(dtype=torch.float32):
motion = self.dec.direction(motion_feat)
return motion
```
## /wan/modules/animate/preprocess/UserGuider.md
# Wan-animate Preprocessing User Guider
## 1. Introductions
Wan-animate offers two generation modes: `animation` and `replacement`. While both modes extract the skeleton from the reference video, they each have a distinct preprocessing pipeline.
### 1.1 Animation Mode
In this mode, it is highly recommended to enable pose retargeting, especially if the body proportions of the reference and driving characters are dissimilar.
- A simplified version of pose retargeting pipeline is provided to help developers quickly implement this functionality.
- **NOTE:** Due to the potential complexity of input data, the results from this simplified retargeting version are NOT guaranteed to be perfect. It is strongly advised to verify the preprocessing results before proceeding.
- Community contributions to improve on this feature are welcome.
### 1.2 Replacement Mode
- Pose retargeting is DISABLED by default in this mode. This is a deliberate choice to account for potential spatial interactions between the character and the environment.
- **WARNING**: If there is a significant mismatch in body proportions between the reference and driving characters, artifacts or deformations may appear in the final output.
- A simplified version for extracting the character's mask is also provided.
- **WARNING:** This mask extraction process is designed for **single-person videos ONLY** and may produce incorrect results or fail in multi-person videos (incorrect pose tracking). For multi-person video, users are required to either develop their own solution or integrate a suitable open-source tool.
---
## 2. Preprocessing Instructions and Recommendations
### 2.1 Basic Usage
- The preprocessing process requires some additional models, including pose detection (mandatory), and mask extraction and image editing models (optional, as needed). Place them according to the following directory structure:
```
/path/to/your/ckpt_path/
├── det/
│ └── yolov10m.onnx
├── pose2d/
│ └── vitpose_h_wholebody.onnx
├── sam2/
│ └── sam2_hiera_large.pt
└── FLUX.1-Kontext-dev/
```
- `video_path`, `refer_path`, and `save_path` correspond to the paths for the input driving video, the character image, and the preprocessed results.
- When using `animation` mode, two videos, `src_face.mp4` and `src_pose.mp4`, will be generated in `save_path`. When using `replacement` mode, two additional videos, `src_bg.mp4` and `src_mask.mp4`, will also be generated.
- The `resolution_area` parameter determines the resolution for both preprocessing and the generation model. Its size is determined by pixel area.
- The `fps` parameter can specify the frame rate for video processing. A lower frame rate can improve generation efficiency, but may cause stuttering or choppiness.
---
### 2.2 Animation Mode
- We support three forms: not using pose retargeting, using basic pose retargeting, and using enhanced pose retargeting based on the `FLUX.1-Kontext-dev` image editing model. These are specified via the `retarget_flag` and `use_flux` parameters.
- Specifying `retarget_flag` to use basic pose retargeting requires ensuring that both the reference character and the character in the first frame of the driving video are in a front-facing, stretched pose.
- Other than that, we recommend using enhanced pose retargeting by specifying both `retarget_flag` and `use_flux`. **NOTE:** Due to the limited capabilities of `FLUX.1-Kontext-dev`, it is NOT guaranteed to produce the expected results (e.g., consistency is not maintained, the pose is incorrect, etc.). It is recommended to check the intermediate results as well as the finally generated pose video; both are stored in `save_path`. Of course, users can also use a better image editing model, or explore the prompts for Flux on their own.
---
### 2.3 Replacement Mode
- Specifying `replace_flag` to enable data preprocessing for this mode. The preprocessing will additionally process a mask for the character in the video, and its size and shape can be adjusted by specifying some parameters.
- `iterations` and `k` can make the mask larger, covering more area.
- `w_len` and `h_len` can adjust the mask's shape. Smaller values will make the outline coarser, while larger values will make it finer.
- A smaller, finer-contoured mask can allow for more of the original background to be preserved, but may potentially limit the character's generation area (considering potential appearance differences, this can lead to some shape leakage). A larger, coarser mask can allow the character generation to be more flexible and consistent, but because it includes more of the background, it might affect the background's consistency. We recommend users to adjust the relevant parameters based on their specific input data.
## /wan/modules/animate/preprocess/__init__.py
```py path="/wan/modules/animate/preprocess/__init__.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from .process_pipepline import ProcessPipeline
from .video_predictor import SAM2VideoPredictor
```
## /wan/modules/animate/preprocess/pose2d.py
```py path="/wan/modules/animate/preprocess/pose2d.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import os
import cv2
from typing import Union, List
import numpy as np
import torch
import onnxruntime
from pose2d_utils import (
read_img,
box_convert_simple,
bbox_from_detector,
crop,
keypoints_from_heatmaps,
load_pose_metas_from_kp2ds_seq
)
class SimpleOnnxInference(object):
def __init__(self, checkpoint, device='cuda', reverse_input=False, **kwargs):
if isinstance(device, str):
device = torch.device(device)
if device.type == 'cuda':
device = '{}:{}'.format(device.type, device.index)
providers = [("CUDAExecutionProvider", {"device_id": device[-1:] if device[-1] in [str(_i) for _i in range(10)] else "0"}), "CPUExecutionProvider"]
else:
providers = ["CPUExecutionProvider"]
self.device = device
if not os.path.exists(checkpoint):
raise RuntimeError("{} is not existed!".format(checkpoint))
if os.path.isdir(checkpoint):
checkpoint = os.path.join(checkpoint, 'end2end.onnx')
self.session = onnxruntime.InferenceSession(checkpoint,
providers=providers
)
self.input_name = self.session.get_inputs()[0].name
self.output_name = self.session.get_outputs()[0].name
self.input_resolution = self.session.get_inputs()[0].shape[2:] if not reverse_input else self.session.get_inputs()[0].shape[2:][::-1]
self.input_resolution = np.array(self.input_resolution)
def __call__(self, *args, **kwargs):
return self.forward(*args, **kwargs)
def get_output_names(self):
output_names = []
for node in self.session.get_outputs():
output_names.append(node.name)
return output_names
def set_device(self, device):
if isinstance(device, str):
device = torch.device(device)
if device.type == 'cuda':
device = '{}:{}'.format(device.type, device.index)
providers = [("CUDAExecutionProvider", {"device_id": device[-1:] if device[-1] in [str(_i) for _i in range(10)] else "0"}), "CPUExecutionProvider"]
else:
providers = ["CPUExecutionProvider"]
self.session.set_providers(providers)
self.device = device
class Yolo(SimpleOnnxInference):
def __init__(self, checkpoint, device='cuda', threshold_conf=0.05, threshold_multi_persons=0.1, input_resolution=(640, 640), threshold_iou=0.5, threshold_bbox_shape_ratio=0.4, cat_id=[1], select_type='max', strict=True, sorted_func=None, **kwargs):
super(Yolo, self).__init__(checkpoint, device=device, **kwargs)
model_inputs = self.session.get_inputs()
input_shape = model_inputs[0].shape
self.input_width = 640
self.input_height = 640
self.threshold_multi_persons = threshold_multi_persons
self.threshold_conf = threshold_conf
self.threshold_iou = threshold_iou
self.threshold_bbox_shape_ratio = threshold_bbox_shape_ratio
self.input_resolution = input_resolution
self.cat_id = cat_id
self.select_type = select_type
self.strict = strict
self.sorted_func = sorted_func
def preprocess(self, input_image):
"""
Preprocesses the input image before performing inference.
Returns:
image_data: Preprocessed image data ready for inference.
"""
img = read_img(input_image)
# Get the height and width of the input image
img_height, img_width = img.shape[:2]
# Resize the image to match the input shape
img = cv2.resize(img, (self.input_resolution[1], self.input_resolution[0]))
# Normalize the image data by dividing it by 255.0
image_data = np.array(img) / 255.0
# Transpose the image to have the channel dimension as the first dimension
image_data = np.transpose(image_data, (2, 0, 1)) # Channel first
# Expand the dimensions of the image data to match the expected input shape
# image_data = np.expand_dims(image_data, axis=0).astype(np.float32)
image_data = image_data.astype(np.float32)
# Return the preprocessed image data
return image_data, np.array([img_height, img_width])
def postprocess(self, output, shape_raw, cat_id=[1]):
"""
Performs post-processing on the model's output to extract bounding boxes, scores, and class IDs.
Args:
input_image (numpy.ndarray): The input image.
output (numpy.ndarray): The output of the model.
Returns:
numpy.ndarray: The input image with detections drawn on it.
"""
# Transpose and squeeze the output to match the expected shape
outputs = np.squeeze(output)
if len(outputs.shape) == 1:
outputs = outputs[None]
if output.shape[-1] != 6 and output.shape[1] == 84:
outputs = np.transpose(outputs)
# Get the number of rows in the outputs array
rows = outputs.shape[0]
# Calculate the scaling factors for the bounding box coordinates
x_factor = shape_raw[1] / self.input_width
y_factor = shape_raw[0] / self.input_height
# Lists to store the bounding boxes, scores, and class IDs of the detections
boxes = []
scores = []
class_ids = []
if outputs.shape[-1] == 6:
max_scores = outputs[:, 4]
classid = outputs[:, -1]
threshold_conf_masks = max_scores >= self.threshold_conf
classid_masks = classid[threshold_conf_masks] != 3.14159
max_scores = max_scores[threshold_conf_masks][classid_masks]
classid = classid[threshold_conf_masks][classid_masks]
boxes = outputs[:, :4][threshold_conf_masks][classid_masks]
boxes[:, [0, 2]] *= x_factor
boxes[:, [1, 3]] *= y_factor
boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
boxes = boxes.astype(np.int32)
else:
classes_scores = outputs[:, 4:]
max_scores = np.amax(classes_scores, -1)
threshold_conf_masks = max_scores >= self.threshold_conf
classid = np.argmax(classes_scores[threshold_conf_masks], -1)
classid_masks = classid!=3.14159
classes_scores = classes_scores[threshold_conf_masks][classid_masks]
max_scores = max_scores[threshold_conf_masks][classid_masks]
classid = classid[classid_masks]
xywh = outputs[:, :4][threshold_conf_masks][classid_masks]
x = xywh[:, 0:1]
y = xywh[:, 1:2]
w = xywh[:, 2:3]
h = xywh[:, 3:4]
left = ((x - w / 2) * x_factor)
top = ((y - h / 2) * y_factor)
width = (w * x_factor)
height = (h * y_factor)
boxes = np.concatenate([left, top, width, height], axis=-1).astype(np.int32)
boxes = boxes.tolist()
scores = max_scores.tolist()
class_ids = classid.tolist()
# Apply non-maximum suppression to filter out overlapping bounding boxes
indices = cv2.dnn.NMSBoxes(boxes, scores, self.threshold_conf, self.threshold_iou)
# Iterate over the selected indices after non-maximum suppression
results = []
for i in indices:
# Get the box, score, and class ID corresponding to the index
box = box_convert_simple(boxes[i], 'xywh2xyxy')
score = scores[i]
class_id = class_ids[i]
results.append(box + [score] + [class_id])
# # Draw the detection on the input image
# Return the modified input image
return np.array(results)
def process_results(self, results, shape_raw, cat_id=[1], single_person=True):
if isinstance(results, tuple):
det_results = results[0]
else:
det_results = results
person_results = []
person_count = 0
if len(results):
max_idx = -1
max_bbox_size = shape_raw[0] * shape_raw[1] * -10
max_bbox_shape = -1
bboxes = []
idx_list = []
for i in range(results.shape[0]):
bbox = results[i]
if (bbox[-1] + 1 in cat_id) and (bbox[-2] > self.threshold_conf):
idx_list.append(i)
bbox_shape = max((bbox[2] - bbox[0]), ((bbox[3] - bbox[1])))
if bbox_shape > max_bbox_shape:
max_bbox_shape = bbox_shape
results = results[idx_list]
for i in range(results.shape[0]):
bbox = results[i]
bboxes.append(bbox)
if self.select_type == 'max':
bbox_size = (bbox[2] - bbox[0]) * ((bbox[3] - bbox[1]))
elif self.select_type == 'center':
bbox_size = (abs((bbox[2] + bbox[0]) / 2 - shape_raw[1]/2)) * -1
bbox_shape = max((bbox[2] - bbox[0]), ((bbox[3] - bbox[1])))
if bbox_size > max_bbox_size:
if (self.strict or max_idx != -1) and bbox_shape < max_bbox_shape * self.threshold_bbox_shape_ratio:
continue
max_bbox_size = bbox_size
max_bbox_shape = bbox_shape
max_idx = i
if self.sorted_func is not None and len(bboxes) > 0:
max_idx = self.sorted_func(bboxes, shape_raw)
bbox = bboxes[max_idx]
if self.select_type == 'max':
max_bbox_size = (bbox[2] - bbox[0]) * ((bbox[3] - bbox[1]))
elif self.select_type == 'center':
max_bbox_size = (abs((bbox[2] + bbox[0]) / 2 - shape_raw[1]/2)) * -1
if max_idx != -1:
person_count = 1
if max_idx != -1:
person = {}
person['bbox'] = results[max_idx, :5]
person['track_id'] = int(0)
person_results.append(person)
for i in range(results.shape[0]):
bbox = results[i]
if (bbox[-1] + 1 in cat_id) and (bbox[-2] > self.threshold_conf):
if self.select_type == 'max':
bbox_size = (bbox[2] - bbox[0]) * ((bbox[3] - bbox[1]))
elif self.select_type == 'center':
bbox_size = (abs((bbox[2] + bbox[0]) / 2 - shape_raw[1]/2)) * -1
if i != max_idx and bbox_size > max_bbox_size * self.threshold_multi_persons and bbox_size < max_bbox_size:
person_count += 1
if not single_person:
person = {}
person['bbox'] = results[i, :5]
person['track_id'] = int(person_count - 1)
person_results.append(person)
return person_results
else:
return None
def postprocess_threading(self, outputs, shape_raw, person_results, i, single_person=True, **kwargs):
result = self.postprocess(outputs[i], shape_raw[i], cat_id=self.cat_id)
result = self.process_results(result, shape_raw[i], cat_id=self.cat_id, single_person=single_person)
if result is not None and len(result) != 0:
person_results[i] = result
def forward(self, img, shape_raw, **kwargs):
"""
Performs inference using an ONNX model and returns the output image with drawn detections.
Returns:
output_img: The output image with drawn detections.
"""
if isinstance(img, torch.Tensor):
img = img.cpu().numpy()
shape_raw = shape_raw.cpu().numpy()
outputs = self.session.run(None, {self.session.get_inputs()[0].name: img})[0]
person_results = [[{'bbox': np.array([0., 0., 1.*shape_raw[i][1], 1.*shape_raw[i][0], -1]), 'track_id': -1}] for i in range(len(outputs))]
for i in range(len(outputs)):
self.postprocess_threading(outputs, shape_raw, person_results, i, **kwargs)
return person_results
class ViTPose(SimpleOnnxInference):
def __init__(self, checkpoint, device='cuda', **kwargs):
super(ViTPose, self).__init__(checkpoint, device=device)
def forward(self, img, center, scale, **kwargs):
heatmaps = self.session.run([], {self.session.get_inputs()[0].name: img})[0]
points, prob = keypoints_from_heatmaps(heatmaps=heatmaps,
center=center,
scale=scale*200,
unbiased=True,
use_udp=False)
return np.concatenate([points, prob], axis=2)
@staticmethod
def preprocess(img, bbox=None, input_resolution=(256, 192), rescale=1.25, mask=None, **kwargs):
if bbox is None or bbox[-1] <= 0 or (bbox[2] - bbox[0]) < 10 or (bbox[3] - bbox[1]) < 10:
bbox = np.array([0, 0, img.shape[1], img.shape[0]])
bbox_xywh = bbox
if mask is not None:
img = np.where(mask>128, img, mask)
if isinstance(input_resolution, int):
center, scale = bbox_from_detector(bbox_xywh, (input_resolution, input_resolution), rescale=rescale)
img, new_shape, old_xy, new_xy = crop(img, center, scale, (input_resolution, input_resolution))
else:
center, scale = bbox_from_detector(bbox_xywh, input_resolution, rescale=rescale)
img, new_shape, old_xy, new_xy = crop(img, center, scale, (input_resolution[0], input_resolution[1]))
IMG_NORM_MEAN = np.array([0.485, 0.456, 0.406])
IMG_NORM_STD = np.array([0.229, 0.224, 0.225])
img_norm = (img / 255. - IMG_NORM_MEAN) / IMG_NORM_STD
img_norm = img_norm.transpose(2, 0, 1).astype(np.float32)
return img_norm, np.array(center), np.array(scale)
class Pose2d:
def __init__(self, checkpoint, detector_checkpoint=None, device='cuda', **kwargs):
if detector_checkpoint is not None:
self.detector = Yolo(detector_checkpoint, device)
else:
self.detector = None
self.model = ViTPose(checkpoint, device)
self.device = device
def load_images(self, inputs):
"""
Load images from various input types.
Args:
inputs (Union[str, np.ndarray, List[np.ndarray]]): Input can be file path,
single image array, or list of image arrays
Returns:
List[np.ndarray]: List of RGB image arrays
Raises:
ValueError: If file format is unsupported or image cannot be read
"""
if isinstance(inputs, str):
if inputs.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
cap = cv2.VideoCapture(inputs)
frames = []
while True:
ret, frame = cap.read()
if not ret:
break
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
cap.release()
images = frames
elif inputs.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
img = cv2.cvtColor(cv2.imread(inputs), cv2.COLOR_BGR2RGB)
if img is None:
raise ValueError(f"Cannot read image: {inputs}")
images = [img]
else:
raise ValueError(f"Unsupported file format: {inputs}")
elif isinstance(inputs, np.ndarray):
images = [cv2.cvtColor(image, cv2.COLOR_BGR2RGB) for image in inputs]
elif isinstance(inputs, list):
images = [cv2.cvtColor(image, cv2.COLOR_BGR2RGB) for image in inputs]
return images
def __call__(
self,
inputs: Union[str, np.ndarray, List[np.ndarray]],
return_image: bool = False,
**kwargs
):
"""
Process input and estimate 2D keypoints.
Args:
inputs (Union[str, np.ndarray, List[np.ndarray]]): Input can be file path,
single image array, or list of image arrays
**kwargs: Additional arguments for processing
Returns:
np.ndarray: Array of detected 2D keypoints for all input images
"""
images = self.load_images(inputs)
H, W = images[0].shape[:2]
if self.detector is not None:
bboxes = []
for _image in images:
img, shape = self.detector.preprocess(_image)
bboxes.append(self.detector(img[None], shape[None])[0][0]["bbox"])
else:
bboxes = [None] * len(images)
kp2ds = []
for _image, _bbox in zip(images, bboxes):
img, center, scale = self.model.preprocess(_image, _bbox)
kp2ds.append(self.model(img[None], center[None], scale[None]))
kp2ds = np.concatenate(kp2ds, 0)
metas = load_pose_metas_from_kp2ds_seq(kp2ds, width=W, height=H)
return metas
```
## /wan/modules/animate/preprocess/preprocess_data.py
```py path="/wan/modules/animate/preprocess/preprocess_data.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import os
import argparse
from process_pipepline import ProcessPipeline
def _parse_args():
parser = argparse.ArgumentParser(
description="The preprocessing pipeline for Wan-animate."
)
parser.add_argument(
"--ckpt_path",
type=str,
default=None,
help="The path to the preprocessing model's checkpoint directory. ")
parser.add_argument(
"--video_path",
type=str,
default=None,
help="The path to the driving video.")
parser.add_argument(
"--refer_path",
type=str,
default=None,
help="The path to the refererence image.")
parser.add_argument(
"--save_path",
type=str,
default=None,
help="The path to save the processed results.")
parser.add_argument(
"--resolution_area",
type=int,
nargs=2,
default=[1280, 720],
help="The target resolution for processing, specified as [width, height]. To handle different aspect ratios, the video is resized to have a total area equivalent to width * height, while preserving the original aspect ratio."
)
parser.add_argument(
"--fps",
type=int,
default=30,
help="The target FPS for processing the driving video. Set to -1 to use the video's original FPS."
)
parser.add_argument(
"--replace_flag",
action="store_true",
default=False,
help="Whether to use replacement mode.")
parser.add_argument(
"--retarget_flag",
action="store_true",
default=False,
help="Whether to use pose retargeting. Currently only supported in animation mode")
parser.add_argument(
"--use_flux",
action="store_true",
default=False,
help="Whether to use image editing in pose retargeting. Recommended if the character in the reference image or the first frame of the driving video is not in a standard, front-facing pose")
# Parameters for the mask strategy in replacement mode. These control the mask's size and shape. Refer to https://arxiv.org/pdf/2502.06145
parser.add_argument(
"--iterations",
type=int,
default=3,
help="Number of iterations for mask dilation."
)
parser.add_argument(
"--k",
type=int,
default=7,
help="Number of kernel size for mask dilation."
)
parser.add_argument(
"--w_len",
type=int,
default=1,
help="The number of subdivisions for the grid along the 'w' dimension. A higher value results in a more detailed contour. A value of 1 means no subdivision is performed."
)
parser.add_argument(
"--h_len",
type=int,
default=1,
help="The number of subdivisions for the grid along the 'h' dimension. A higher value results in a more detailed contour. A value of 1 means no subdivision is performed."
)
args = parser.parse_args()
return args
if __name__ == '__main__':
args = _parse_args()
args_dict = vars(args)
print(args_dict)
assert len(args.resolution_area) == 2, "resolution_area should be a list of two integers [width, height]"
assert not args.use_flux or args.retarget_flag, "Image editing with FLUX can only be used when pose retargeting is enabled."
pose2d_checkpoint_path = os.path.join(args.ckpt_path, 'pose2d/vitpose_h_wholebody.onnx')
det_checkpoint_path = os.path.join(args.ckpt_path, 'det/yolov10m.onnx')
sam2_checkpoint_path = os.path.join(args.ckpt_path, 'sam2/sam2_hiera_large.pt') if args.replace_flag else None
flux_kontext_path = os.path.join(args.ckpt_path, 'FLUX.1-Kontext-dev') if args.use_flux else None
process_pipeline = ProcessPipeline(det_checkpoint_path=det_checkpoint_path, pose2d_checkpoint_path=pose2d_checkpoint_path, sam_checkpoint_path=sam2_checkpoint_path, flux_kontext_path=flux_kontext_path)
os.makedirs(args.save_path, exist_ok=True)
process_pipeline(video_path=args.video_path,
refer_image_path=args.refer_path,
output_path=args.save_path,
resolution_area=args.resolution_area,
fps=args.fps,
iterations=args.iterations,
k=args.k,
w_len=args.w_len,
h_len=args.h_len,
retarget_flag=args.retarget_flag,
use_flux=args.use_flux,
replace_flag=args.replace_flag)
```
## /wan/modules/animate/preprocess/process_pipepline.py
```py path="/wan/modules/animate/preprocess/process_pipepline.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import os
import numpy as np
import shutil
import torch
from diffusers import FluxKontextPipeline
import cv2
from loguru import logger
from PIL import Image
try:
import moviepy.editor as mpy
except:
import moviepy as mpy
from decord import VideoReader
from pose2d import Pose2d
from pose2d_utils import AAPoseMeta
from utils import resize_by_area, get_frame_indices, padding_resize, get_face_bboxes, get_aug_mask, get_mask_body_img
from human_visualization import draw_aapose_by_meta_new
from retarget_pose import get_retarget_pose
import sam2.modeling.sam.transformer as transformer
transformer.USE_FLASH_ATTN = False
transformer.MATH_KERNEL_ON = True
transformer.OLD_GPU = True
from sam_utils import build_sam2_video_predictor
class ProcessPipeline():
def __init__(self, det_checkpoint_path, pose2d_checkpoint_path, sam_checkpoint_path, flux_kontext_path):
self.pose2d = Pose2d(checkpoint=pose2d_checkpoint_path, detector_checkpoint=det_checkpoint_path)
model_cfg = "sam2_hiera_l.yaml"
if sam_checkpoint_path is not None:
self.predictor = build_sam2_video_predictor(model_cfg, sam_checkpoint_path)
if flux_kontext_path is not None:
self.flux_kontext = FluxKontextPipeline.from_pretrained(flux_kontext_path, torch_dtype=torch.bfloat16).to("cuda")
def __call__(self, video_path, refer_image_path, output_path, resolution_area=[1280, 720], fps=30, iterations=3, k=7, w_len=1, h_len=1, retarget_flag=False, use_flux=False, replace_flag=False):
if replace_flag:
video_reader = VideoReader(video_path)
frame_num = len(video_reader)
print('frame_num: {}'.format(frame_num))
video_fps = video_reader.get_avg_fps()
print('video_fps: {}'.format(video_fps))
print('fps: {}'.format(fps))
# TODO: Maybe we can switch to PyAV later, which can get accurate frame num
duration = video_reader.get_frame_timestamp(-1)[-1]
expected_frame_num = int(duration * video_fps + 0.5)
ratio = abs((frame_num - expected_frame_num)/frame_num)
if ratio > 0.1:
print("Warning: The difference between the actual number of frames and the expected number of frames is two large")
frame_num = expected_frame_num
if fps == -1:
fps = video_fps
target_num = int(frame_num / video_fps * fps)
print('target_num: {}'.format(target_num))
idxs = get_frame_indices(frame_num, video_fps, target_num, fps)
frames = video_reader.get_batch(idxs).asnumpy()
frames = [resize_by_area(frame, resolution_area[0] * resolution_area[1], divisor=16) for frame in frames]
height, width = frames[0].shape[:2]
logger.info(f"Processing pose meta")
tpl_pose_metas = self.pose2d(frames)
face_images = []
for idx, meta in enumerate(tpl_pose_metas):
face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3,
image_shape=(frames[0].shape[0], frames[0].shape[1]))
x1, x2, y1, y2 = face_bbox_for_image
face_image = frames[idx][y1:y2, x1:x2]
face_image = cv2.resize(face_image, (512, 512))
face_images.append(face_image)
logger.info(f"Processing reference image: {refer_image_path}")
refer_img = cv2.imread(refer_image_path)
src_ref_path = os.path.join(output_path, 'src_ref.png')
shutil.copy(refer_image_path, src_ref_path)
refer_img = refer_img[..., ::-1]
refer_img = padding_resize(refer_img, height, width)
logger.info(f"Processing template video: {video_path}")
tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas]
cond_images = []
for idx, meta in enumerate(tpl_retarget_pose_metas):
canvas = np.zeros_like(refer_img)
conditioning_image = draw_aapose_by_meta_new(canvas, meta)
cond_images.append(conditioning_image)
masks = self.get_mask(frames, 400, tpl_pose_metas)
bg_images = []
aug_masks = []
for frame, mask in zip(frames, masks):
if iterations > 0:
_, each_mask = get_mask_body_img(frame, mask, iterations=iterations, k=k)
each_aug_mask = get_aug_mask(each_mask, w_len=w_len, h_len=h_len)
else:
each_aug_mask = mask
each_bg_image = frame * (1 - each_aug_mask[:, :, None])
bg_images.append(each_bg_image)
aug_masks.append(each_aug_mask)
src_face_path = os.path.join(output_path, 'src_face.mp4')
mpy.ImageSequenceClip(face_images, fps=fps).write_videofile(src_face_path)
src_pose_path = os.path.join(output_path, 'src_pose.mp4')
mpy.ImageSequenceClip(cond_images, fps=fps).write_videofile(src_pose_path)
src_bg_path = os.path.join(output_path, 'src_bg.mp4')
mpy.ImageSequenceClip(bg_images, fps=fps).write_videofile(src_bg_path)
aug_masks_new = [np.stack([mask * 255, mask * 255, mask * 255], axis=2) for mask in aug_masks]
src_mask_path = os.path.join(output_path, 'src_mask.mp4')
mpy.ImageSequenceClip(aug_masks_new, fps=fps).write_videofile(src_mask_path)
return True
else:
logger.info(f"Processing reference image: {refer_image_path}")
refer_img = cv2.imread(refer_image_path)
src_ref_path = os.path.join(output_path, 'src_ref.png')
shutil.copy(refer_image_path, src_ref_path)
refer_img = refer_img[..., ::-1]
refer_img = resize_by_area(refer_img, resolution_area[0] * resolution_area[1], divisor=16)
refer_pose_meta = self.pose2d([refer_img])[0]
logger.info(f"Processing template video: {video_path}")
video_reader = VideoReader(video_path)
frame_num = len(video_reader)
print('frame_num: {}'.format(frame_num))
video_fps = video_reader.get_avg_fps()
print('video_fps: {}'.format(video_fps))
print('fps: {}'.format(fps))
# TODO: Maybe we can switch to PyAV later, which can get accurate frame num
duration = video_reader.get_frame_timestamp(-1)[-1]
expected_frame_num = int(duration * video_fps + 0.5)
ratio = abs((frame_num - expected_frame_num)/frame_num)
if ratio > 0.1:
print("Warning: The difference between the actual number of frames and the expected number of frames is two large")
frame_num = expected_frame_num
if fps == -1:
fps = video_fps
target_num = int(frame_num / video_fps * fps)
print('target_num: {}'.format(target_num))
idxs = get_frame_indices(frame_num, video_fps, target_num, fps)
frames = video_reader.get_batch(idxs).asnumpy()
logger.info(f"Processing pose meta")
tpl_pose_meta0 = self.pose2d(frames[:1])[0]
tpl_pose_metas = self.pose2d(frames)
face_images = []
for idx, meta in enumerate(tpl_pose_metas):
face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3,
image_shape=(frames[0].shape[0], frames[0].shape[1]))
x1, x2, y1, y2 = face_bbox_for_image
face_image = frames[idx][y1:y2, x1:x2]
face_image = cv2.resize(face_image, (512, 512))
face_images.append(face_image)
if retarget_flag:
if use_flux:
tpl_prompt, refer_prompt = self.get_editing_prompts(tpl_pose_metas, refer_pose_meta)
refer_input = Image.fromarray(refer_img)
refer_edit = self.flux_kontext(
image=refer_input,
height=refer_img.shape[0],
width=refer_img.shape[1],
prompt=refer_prompt,
guidance_scale=2.5,
num_inference_steps=28,
).images[0]
refer_edit = Image.fromarray(padding_resize(np.array(refer_edit), refer_img.shape[0], refer_img.shape[1]))
refer_edit_path = os.path.join(output_path, 'refer_edit.png')
refer_edit.save(refer_edit_path)
refer_edit_pose_meta = self.pose2d([np.array(refer_edit)])[0]
tpl_img = frames[1]
tpl_input = Image.fromarray(tpl_img)
tpl_edit = self.flux_kontext(
image=tpl_input,
height=tpl_img.shape[0],
width=tpl_img.shape[1],
prompt=tpl_prompt,
guidance_scale=2.5,
num_inference_steps=28,
).images[0]
tpl_edit = Image.fromarray(padding_resize(np.array(tpl_edit), tpl_img.shape[0], tpl_img.shape[1]))
tpl_edit_path = os.path.join(output_path, 'tpl_edit.png')
tpl_edit.save(tpl_edit_path)
tpl_edit_pose_meta0 = self.pose2d([np.array(tpl_edit)])[0]
tpl_retarget_pose_metas = get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, tpl_edit_pose_meta0, refer_edit_pose_meta)
else:
tpl_retarget_pose_metas = get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, None, None)
else:
tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas]
cond_images = []
for idx, meta in enumerate(tpl_retarget_pose_metas):
if retarget_flag:
canvas = np.zeros_like(refer_img)
conditioning_image = draw_aapose_by_meta_new(canvas, meta)
else:
canvas = np.zeros_like(frames[0])
conditioning_image = draw_aapose_by_meta_new(canvas, meta)
conditioning_image = padding_resize(conditioning_image, refer_img.shape[0], refer_img.shape[1])
cond_images.append(conditioning_image)
src_face_path = os.path.join(output_path, 'src_face.mp4')
mpy.ImageSequenceClip(face_images, fps=fps).write_videofile(src_face_path)
src_pose_path = os.path.join(output_path, 'src_pose.mp4')
mpy.ImageSequenceClip(cond_images, fps=fps).write_videofile(src_pose_path)
return True
def get_editing_prompts(self, tpl_pose_metas, refer_pose_meta):
arm_visible = False
leg_visible = False
for tpl_pose_meta in tpl_pose_metas:
tpl_keypoints = tpl_pose_meta['keypoints_body']
if tpl_keypoints[3].all() != 0 or tpl_keypoints[4].all() != 0 or tpl_keypoints[6].all() != 0 or tpl_keypoints[7].all() != 0:
if (tpl_keypoints[3][0] <= 1 and tpl_keypoints[3][1] <= 1 and tpl_keypoints[3][2] >= 0.75) or (tpl_keypoints[4][0] <= 1 and tpl_keypoints[4][1] <= 1 and tpl_keypoints[4][2] >= 0.75) or \
(tpl_keypoints[6][0] <= 1 and tpl_keypoints[6][1] <= 1 and tpl_keypoints[6][2] >= 0.75) or (tpl_keypoints[7][0] <= 1 and tpl_keypoints[7][1] <= 1 and tpl_keypoints[7][2] >= 0.75):
arm_visible = True
if tpl_keypoints[9].all() != 0 or tpl_keypoints[12].all() != 0 or tpl_keypoints[10].all() != 0 or tpl_keypoints[13].all() != 0:
if (tpl_keypoints[9][0] <= 1 and tpl_keypoints[9][1] <= 1 and tpl_keypoints[9][2] >= 0.75) or (tpl_keypoints[12][0] <= 1 and tpl_keypoints[12][1] <= 1 and tpl_keypoints[12][2] >= 0.75) or \
(tpl_keypoints[10][0] <= 1 and tpl_keypoints[10][1] <= 1 and tpl_keypoints[10][2] >= 0.75) or (tpl_keypoints[13][0] <= 1 and tpl_keypoints[13][1] <= 1 and tpl_keypoints[13][2] >= 0.75):
leg_visible = True
if arm_visible and leg_visible:
break
if leg_visible:
if tpl_pose_meta['width'] > tpl_pose_meta['height']:
tpl_prompt = "Change the person to a standard T-pose (facing forward with arms extended). The person is standing. Feet and Hands are visible in the image."
else:
tpl_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. The person is standing. Feet and Hands are visible in the image."
if refer_pose_meta['width'] > refer_pose_meta['height']:
refer_prompt = "Change the person to a standard T-pose (facing forward with arms extended). The person is standing. Feet and Hands are visible in the image."
else:
refer_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. The person is standing. Feet and Hands are visible in the image."
elif arm_visible:
if tpl_pose_meta['width'] > tpl_pose_meta['height']:
tpl_prompt = "Change the person to a standard T-pose (facing forward with arms extended). Hands are visible in the image."
else:
tpl_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. Hands are visible in the image."
if refer_pose_meta['width'] > refer_pose_meta['height']:
refer_prompt = "Change the person to a standard T-pose (facing forward with arms extended). Hands are visible in the image."
else:
refer_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. Hands are visible in the image."
else:
tpl_prompt = "Change the person to face forward."
refer_prompt = "Change the person to face forward."
return tpl_prompt, refer_prompt
def get_mask(self, frames, th_step, kp2ds_all):
frame_num = len(frames)
if frame_num < th_step:
num_step = 1
else:
num_step = (frame_num + th_step) // th_step
all_mask = []
for index in range(num_step):
each_frames = frames[index * th_step:(index + 1) * th_step]
kp2ds = kp2ds_all[index * th_step:(index + 1) * th_step]
if len(each_frames) > 4:
key_frame_num = 4
elif 4 >= len(each_frames) > 0:
key_frame_num = 1
else:
continue
key_frame_step = len(kp2ds) // key_frame_num
key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))
key_points_index = [0, 1, 2, 5, 8, 11, 10, 13]
key_frame_body_points_list = []
for key_frame_index in key_frame_index_list:
keypoints_body_list = []
body_key_points = kp2ds[key_frame_index]['keypoints_body']
for each_index in key_points_index:
each_keypoint = body_key_points[each_index]
if None is each_keypoint:
continue
keypoints_body_list.append(each_keypoint)
keypoints_body = np.array(keypoints_body_list)[:, :2]
wh = np.array([[kp2ds[0]['width'], kp2ds[0]['height']]])
points = (keypoints_body * wh).astype(np.int32)
key_frame_body_points_list.append(points)
inference_state = self.predictor.init_state_v2(frames=each_frames)
self.predictor.reset_state(inference_state)
ann_obj_id = 1
for ann_frame_idx, points in zip(key_frame_index_list, key_frame_body_points_list):
labels = np.array([1] * points.shape[0], np.int32)
_, out_obj_ids, out_mask_logits = self.predictor.add_new_points(
inference_state=inference_state,
frame_idx=ann_frame_idx,
obj_id=ann_obj_id,
points=points,
labels=labels,
)
video_segments = {}
for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video(inference_state):
video_segments[out_frame_idx] = {
out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
for i, out_obj_id in enumerate(out_obj_ids)
}
for out_frame_idx in range(len(video_segments)):
for out_obj_id, out_mask in video_segments[out_frame_idx].items():
out_mask = out_mask[0].astype(np.uint8)
all_mask.append(out_mask)
return all_mask
def convert_list_to_array(self, metas):
metas_list = []
for meta in metas:
for key, value in meta.items():
if type(value) is list:
value = np.array(value)
meta[key] = value
metas_list.append(meta)
return metas_list
```
## /wan/modules/s2v/__init__.py
```py path="/wan/modules/s2v/__init__.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from .audio_encoder import AudioEncoder
from .model_s2v import WanModel_S2V
__all__ = ['WanModel_S2V', 'AudioEncoder']
```
## /wan/modules/tokenizers.py
```py path="/wan/modules/tokenizers.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import html
import string
import ftfy
import regex as re
from transformers import AutoTokenizer
__all__ = ['HuggingfaceTokenizer']
def basic_clean(text):
text = ftfy.fix_text(text)
text = html.unescape(html.unescape(text))
return text.strip()
def whitespace_clean(text):
text = re.sub(r'\s+', ' ', text)
text = text.strip()
return text
def canonicalize(text, keep_punctuation_exact_string=None):
text = text.replace('_', ' ')
if keep_punctuation_exact_string:
text = keep_punctuation_exact_string.join(
part.translate(str.maketrans('', '', string.punctuation))
for part in text.split(keep_punctuation_exact_string))
else:
text = text.translate(str.maketrans('', '', string.punctuation))
text = text.lower()
text = re.sub(r'\s+', ' ', text)
return text.strip()
class HuggingfaceTokenizer:
def __init__(self, name, seq_len=None, clean=None, **kwargs):
assert clean in (None, 'whitespace', 'lower', 'canonicalize')
self.name = name
self.seq_len = seq_len
self.clean = clean
# init tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)
self.vocab_size = self.tokenizer.vocab_size
def __call__(self, sequence, **kwargs):
return_mask = kwargs.pop('return_mask', False)
# arguments
_kwargs = {'return_tensors': 'pt'}
if self.seq_len is not None:
_kwargs.update({
'padding': 'max_length',
'truncation': True,
'max_length': self.seq_len
})
_kwargs.update(**kwargs)
# tokenization
if isinstance(sequence, str):
sequence = [sequence]
if self.clean:
sequence = [self._clean(u) for u in sequence]
ids = self.tokenizer(sequence, **_kwargs)
# output
if return_mask:
return ids.input_ids, ids.attention_mask
else:
return ids.input_ids
def _clean(self, text):
if self.clean == 'whitespace':
text = whitespace_clean(basic_clean(text))
elif self.clean == 'lower':
text = whitespace_clean(basic_clean(text)).lower()
elif self.clean == 'canonicalize':
text = canonicalize(basic_clean(text))
return text
```
## /wan/utils/__init__.py
```py path="/wan/utils/__init__.py"
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from .fm_solvers import (
FlowDPMSolverMultistepScheduler,
get_sampling_sigmas,
retrieve_timesteps,
)
from .fm_solvers_unipc import FlowUniPCMultistepScheduler
__all__ = [
'HuggingfaceTokenizer', 'get_sampling_sigmas', 'retrieve_timesteps',
'FlowDPMSolverMultistepScheduler', 'FlowUniPCMultistepScheduler'
]
```
The content has been capped at 50000 tokens. The user could consider applying other filters to refine the result. The better and more specific the context, the better the LLM can follow instructions. If the context seems verbose, the user can refine the filter using uithub. Thank you for using https://uithub.com - Perfect LLM context for any GitHub repo.