ForgeSpecter AI commited on
Commit
a115ebd
·
verified ·
1 Parent(s): 1fcf4ea

Upload folder using huggingface_hub

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +17 -0
  2. .gitignore +7 -0
  3. INSTALL.md +55 -0
  4. LICENSE.txt +201 -0
  5. Makefile +5 -0
  6. README.md +512 -3
  7. Tugra_Wan2.2_Colab.ipynb +158 -0
  8. assets/comp_effic.png +3 -0
  9. assets/logo.png +3 -0
  10. assets/moe_2.png +3 -0
  11. assets/moe_arch.png +0 -0
  12. assets/performance.png +3 -0
  13. assets/vae.png +3 -0
  14. examples/Five Hundred Miles.MP3 +3 -0
  15. examples/Five Hundred Miles.png +3 -0
  16. examples/i2v_input.JPG +3 -0
  17. examples/pose.mp4 +3 -0
  18. examples/pose.png +3 -0
  19. examples/sing.MP3 +3 -0
  20. examples/talk.wav +3 -0
  21. examples/wan_animate/animate/image.jpeg +3 -0
  22. examples/wan_animate/animate/video.mp4 +3 -0
  23. examples/wan_animate/replace/image.jpeg +3 -0
  24. examples/wan_animate/replace/video.mp4 +3 -0
  25. examples/zero_shot_prompt.wav +3 -0
  26. generate.py +575 -0
  27. pyproject.toml +66 -0
  28. requirements.txt +16 -0
  29. requirements_animate.txt +8 -0
  30. requirements_s2v.txt +19 -0
  31. test.html +46 -0
  32. tests/README.md +6 -0
  33. tests/test.sh +91 -0
  34. tugra/__init__.py +7 -0
  35. tugra/__pycache__/__init__.cpython-314.pyc +0 -0
  36. tugra/animate.py +648 -0
  37. tugra/configs/__init__.py +51 -0
  38. tugra/configs/__pycache__/__init__.cpython-314.pyc +0 -0
  39. tugra/configs/__pycache__/shared_config.cpython-314.pyc +0 -0
  40. tugra/configs/__pycache__/tugra_animate_14B.cpython-314.pyc +0 -0
  41. tugra/configs/__pycache__/tugra_i2v_A14B.cpython-314.pyc +0 -0
  42. tugra/configs/__pycache__/tugra_s2v_14B.cpython-314.pyc +0 -0
  43. tugra/configs/__pycache__/tugra_t2v_A14B.cpython-314.pyc +0 -0
  44. tugra/configs/__pycache__/tugra_ti2v_5B.cpython-314.pyc +0 -0
  45. tugra/configs/__pycache__/wan_i2v_A14B.cpython-314.pyc +0 -0
  46. tugra/configs/shared_config.py +20 -0
  47. tugra/configs/tugra_animate_14B.py +41 -0
  48. tugra/configs/tugra_i2v_A14B.py +37 -0
  49. tugra/configs/tugra_s2v_14B.py +59 -0
  50. tugra/configs/tugra_t2v_A14B.py +37 -0
.gitattributes CHANGED
@@ -33,3 +33,20 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ assets/comp_effic.png filter=lfs diff=lfs merge=lfs -text
37
+ assets/logo.png filter=lfs diff=lfs merge=lfs -text
38
+ assets/moe_2.png filter=lfs diff=lfs merge=lfs -text
39
+ assets/performance.png filter=lfs diff=lfs merge=lfs -text
40
+ assets/vae.png filter=lfs diff=lfs merge=lfs -text
41
+ examples/Five[[:space:]]Hundred[[:space:]]Miles.MP3 filter=lfs diff=lfs merge=lfs -text
42
+ examples/Five[[:space:]]Hundred[[:space:]]Miles.png filter=lfs diff=lfs merge=lfs -text
43
+ examples/i2v_input.JPG filter=lfs diff=lfs merge=lfs -text
44
+ examples/pose.mp4 filter=lfs diff=lfs merge=lfs -text
45
+ examples/pose.png filter=lfs diff=lfs merge=lfs -text
46
+ examples/sing.MP3 filter=lfs diff=lfs merge=lfs -text
47
+ examples/talk.wav filter=lfs diff=lfs merge=lfs -text
48
+ examples/wan_animate/animate/image.jpeg filter=lfs diff=lfs merge=lfs -text
49
+ examples/wan_animate/animate/video.mp4 filter=lfs diff=lfs merge=lfs -text
50
+ examples/wan_animate/replace/image.jpeg filter=lfs diff=lfs merge=lfs -text
51
+ examples/wan_animate/replace/video.mp4 filter=lfs diff=lfs merge=lfs -text
52
+ examples/zero_shot_prompt.wav filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ __pycache__/
2
+ .DS_Store
3
+ .vscode*
4
+ tmp_examples*
5
+ new_checkpoint*
6
+ batch_test*
7
+ nohup*
INSTALL.md ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Installation Guide
2
+
3
+ ## Install with pip
4
+
5
+ ```bash
6
+ pip install .
7
+ pip install .[dev] # Installe aussi les outils de dev
8
+ ```
9
+
10
+ ## Install with Poetry
11
+
12
+ Ensure you have [Poetry](https://python-poetry.org/docs/#installation) installed on your system.
13
+
14
+ To install all dependencies:
15
+
16
+ ```bash
17
+ poetry install
18
+ ```
19
+
20
+ ### Handling `flash-attn` Installation Issues
21
+
22
+ If `flash-attn` fails due to **PEP 517 build issues**, you can try one of the following fixes.
23
+
24
+ #### No-Build-Isolation Installation (Recommended)
25
+ ```bash
26
+ poetry run pip install --upgrade pip setuptools wheel
27
+ poetry run pip install flash-attn --no-build-isolation
28
+ poetry install
29
+ ```
30
+
31
+ #### Install from Git (Alternative)
32
+ ```bash
33
+ poetry run pip install git+https://github.com/Dao-AILab/flash-attention.git
34
+ ```
35
+
36
+ ---
37
+
38
+ ### Running the Model
39
+
40
+ Once the installation is complete, you can run **TUĞRA** using:
41
+
42
+ ```bash
43
+ poetry run python generate.py --task t2v-A14B --size '1280*720' --ckpt_dir ./TUĞRA-T2V-A14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
44
+ ```
45
+
46
+ #### Test
47
+ ```bash
48
+ bash tests/test.sh
49
+ ```
50
+
51
+ #### Format
52
+ ```bash
53
+ black .
54
+ isort .
55
+ ```
LICENSE.txt ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright [yyyy] [name of copyright owner]
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
Makefile ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ .PHONY: format
2
+
3
+ format:
4
+ isort generate.py tugra
5
+ yapf -i -r *.py generate.py tugra
README.md CHANGED
@@ -1,3 +1,512 @@
1
- ---
2
- license: apache-2.0
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TUĞRA
2
+
3
+ **TUĞRA: Teknolojik Ulusal Grafik ve Render Algoritması**
4
+
5
+ <p align="center">
6
+ <img src="assets/logo.png" width="400"/>
7
+ <p>
8
+
9
+ <p align="center">
10
+ 💜 <a href="https://modelscope.cn/models?page=1&page_type=2&owner=Tugra-AI"><b>WAN</b></a> &nbsp&nbsp | &nbsp&nbsp 🖥️ <a href="https://github.com/Tugra-Video/TUĞRA">GitHub</a> &nbsp&nbsp | &nbsp&nbsp🤗 <a href="https://huggingface.co/Tugra-AI/">Hugging Face</a>&nbsp&nbsp | &nbsp&nbsp🤖 <a href="https://modelscope.cn/organization/Tugra-AI">ModelScope</a>&nbsp&nbsp | &nbsp&nbsp 📑 <a href="https://arxiv.org/abs/2503.20314">Paper</a> &nbsp&nbsp | &nbsp&nbsp 📑 <a href="https://wan.video/welcome?spm=a2ty_o02.30011076.0.0.6c9ee41eCcluqg">Blog</a> &nbsp&nbsp | &nbsp&nbsp 💬 <a href="https://discord.gg/AKNgpMK4Yj">Discord</a>&nbsp&nbsp
11
+ <br>
12
+ 📕 <a href="https://alidocs.dingtalk.com/i/nodes/jb9Y4gmKWrx9eo4dCql9LlbYJGXn6lpz">使用指南(中文)</a>&nbsp&nbsp | &nbsp&nbsp 📘 <a href="https://alidocs.dingtalk.com/i/nodes/EpGBa2Lm8aZxe5myC99MelA2WgN7R35y">User Guide(English)</a>&nbsp&nbsp | &nbsp&nbsp💬 <a href="https://gw.alicdn.com/imgextra/i2/O1CN01tqjWFi1ByuyehkTSB_!!6000000000015-0-tps-611-1279.jpg">WeChat(微信)</a>&nbsp&nbsp
13
+ <br>
14
+
15
+ -----
16
+
17
+ [**TUĞRA: Open and Advanced Large-Scale Video Generative Models**](https://arxiv.org/abs/2503.20314) <be>
18
+
19
+
20
+ We are excited to introduce **TUĞRA**, a major upgrade to our foundational video models. With **TUĞRA**, we have focused on incorporating the following innovations:
21
+
22
+ - 👍 **Effective MoE Architecture**: TUĞRA 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.
23
+
24
+ - 👍 **Cinematic-level Aesthetics**: TUĞRA 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.
25
+
26
+ - 👍 **Complex Motion Generation**: Compared to previous versions, TUĞRA 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.
27
+
28
+ - 👍 **Efficient High-Definition Hybrid TI2V**: TUĞRA open-sources a 5B model built with our advanced TUĞRA-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.
29
+
30
+
31
+ ## Video Demos
32
+
33
+ <div align="center">
34
+ <video src="https://github.com/user-attachments/assets/b63bfa58-d5d7-4de6-a1a2-98970b06d9a7" width="70%" poster=""> </video>
35
+ </div>
36
+
37
+ ## 🔥 Latest News!!
38
+ * Nov 13, 2025: 👋 TUĞRA-Animate-14B has been integrated into Diffusers ([PR](https://github.com/huggingface/diffusers/pull/12526),[Weights](https://huggingface.co/Tugra-AI/TUĞRA-Animate-14B-Diffusers)). Thanks to all community contributors. Enjoy!
39
+
40
+ * Sep 19, 2025: 💃 We introduct **[TUĞRA-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/Tugra-AI/TUĞRA-Animate) or [HuggingFace Space](https://huggingface.co/spaces/Tugra-AI/TUĞRA-Animate)!
41
+ * Aug 26, 2025: 🎵 We introduce **[TUĞRA-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/Tugra-AI/TUĞRA-S2V) or [HuggingFace Gradio](https://huggingface.co/spaces/Tugra-AI/TUĞRA-S2V)!
42
+ * Jul 28, 2025: 👋 We have open a [HF space](https://huggingface.co/spaces/Tugra-AI/Wan-2.2-5B) using the TI2V-5B model. Enjoy!
43
+ * Jul 28, 2025: 👋 TUĞRA 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!
44
+ * Jul 28, 2025: 👋 TUĞRA's T2V, I2V and TI2V have been integrated into Diffusers ([T2V-A14B](https://huggingface.co/Tugra-AI/TUĞRA-T2V-A14B-Diffusers) | [I2V-A14B](https://huggingface.co/Tugra-AI/TUĞRA-I2V-A14B-Diffusers) | [TI2V-5B](https://huggingface.co/Tugra-AI/TUĞRA-TI2V-5B-Diffusers)). Feel free to give it a try!
45
+ * Jul 28, 2025: 👋 We've released the inference code and model weights of **TUĞRA**.
46
+ * Sep 5, 2025: 👋 We add text-to-speech synthesis support with [CosyVoice](https://github.com/FunAudioLLM/CosyVoice) for Speech-to-Video generation task.
47
+
48
+
49
+ ## Community Works
50
+ If your research or project builds upon [**Tugra v2.1**](https://github.com/Tugra-Video/Tugra v2.1) or [**TUĞRA**](https://github.com/Tugra-Video/TUĞRA), and you would like more people to see it, please inform us.
51
+
52
+ - [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.
53
+ - [Helios](https://github.com/PKU-YuanGroup/Helios), a breakthrough video generation model base on **Tugra v2.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.
54
+ - [LightX2V](https://github.com/ModelTC/LightX2V), a lightweight and efficient video generation framework that integrates **Tugra v2.1** and **TUĞRA**, 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.
55
+ - [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.
56
+ - [FastVideo](https://github.com/hao-ai-lab/FastVideo) includes distilled **Wan** models with sparse attention that significanly speed up the inference time.
57
+ - [Cache-dit](https://github.com/vipshop/cache-dit) offers Fully Cache Acceleration support for **TUĞRA** 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.
58
+ - [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.
59
+ - [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.
60
+
61
+
62
+ ## 📑 Todo List
63
+ - TUĞRA Text-to-Video
64
+ - [x] Multi-GPU Inference code of the A14B and 14B models
65
+ - [x] Checkpoints of the A14B and 14B models
66
+ - [x] ComfyUI integration
67
+ - [x] Diffusers integration
68
+ - TUĞRA Image-to-Video
69
+ - [x] Multi-GPU Inference code of the A14B model
70
+ - [x] Checkpoints of the A14B model
71
+ - [x] ComfyUI integration
72
+ - [x] Diffusers integration
73
+ - TUĞRA Text-Image-to-Video
74
+ - [x] Multi-GPU Inference code of the 5B model
75
+ - [x] Checkpoints of the 5B model
76
+ - [x] ComfyUI integration
77
+ - [x] Diffusers integration
78
+ - TUĞRA-S2V Speech-to-Video
79
+ - [x] Inference code of TUĞRA-S2V
80
+ - [x] Checkpoints of TUĞRA-S2V-14B
81
+ - [x] ComfyUI integration
82
+ - [x] Diffusers integration
83
+ - TUĞRA-Animate Character Animation and Replacement
84
+ - [x] Inference code of TUĞRA-Animate
85
+ - [x] Checkpoints of TUĞRA-Animate
86
+ - [x] ComfyUI integration
87
+ - [x] Diffusers integration
88
+
89
+ ## Run TUĞRA
90
+
91
+ #### Installation
92
+ Clone the repo:
93
+ ```sh
94
+ git clone https://github.com/Tugra-Video/TUĞRA.git
95
+ cd TUĞRA
96
+ ```
97
+
98
+ Install dependencies:
99
+ ```sh
100
+ # Ensure torch >= 2.4.0
101
+ # If the installation of `flash_attn` fails, try installing the other packages first and install `flash_attn` last
102
+ pip install -r requirements.txt
103
+ # If you want to use CosyVoice to synthesize speech for Speech-to-Video Generation, please install requirements_s2v.txt additionally
104
+ pip install -r requirements_s2v.txt
105
+ ```
106
+
107
+
108
+ #### Model Download
109
+
110
+ | Models | Download Links | Description |
111
+ |--------------------|---------------------------------------------------------------------------------------------------------------------------------------------|-------------|
112
+ | T2V-A14B | 🤗 [Huggingface](https://huggingface.co/Tugra-AI/TUĞRA-T2V-A14B) �� [ModelScope](https://modelscope.cn/models/Tugra-AI/TUĞRA-T2V-A14B) | Text-to-Video MoE model, supports 480P & 720P |
113
+ | I2V-A14B | 🤗 [Huggingface](https://huggingface.co/Tugra-AI/TUĞRA-I2V-A14B) 🤖 [ModelScope](https://modelscope.cn/models/Tugra-AI/TUĞRA-I2V-A14B) | Image-to-Video MoE model, supports 480P & 720P |
114
+ | TI2V-5B | 🤗 [Huggingface](https://huggingface.co/Tugra-AI/TUĞRA-TI2V-5B) 🤖 [ModelScope](https://modelscope.cn/models/Tugra-AI/TUĞRA-TI2V-5B) | High-compression VAE, T2V+I2V, supports 720P |
115
+ | S2V-14B | 🤗 [Huggingface](https://huggingface.co/Tugra-AI/TUĞRA-S2V-14B) 🤖 [ModelScope](https://modelscope.cn/models/Tugra-AI/TUĞRA-S2V-14B) | Speech-to-Video model, supports 480P & 720P |
116
+ | Animate-14B | 🤗 [Huggingface](https://huggingface.co/Tugra-AI/TUĞRA-Animate-14B) 🤖 [ModelScope](https://www.modelscope.cn/models/Tugra-AI/TUĞRA-Animate-14B) | Character animation and replacement | |
117
+
118
+
119
+
120
+ > 💡Note:
121
+ > The TI2V-5B model supports 720P video generation at **24 FPS**.
122
+
123
+
124
+ Download models using huggingface-cli:
125
+ ``` sh
126
+ pip install "huggingface_hub[cli]"
127
+ huggingface-cli download Tugra-AI/TUĞRA-T2V-A14B --local-dir ./TUĞRA-T2V-A14B
128
+ ```
129
+
130
+ Download models using modelscope-cli:
131
+ ``` sh
132
+ pip install modelscope
133
+ modelscope download Tugra-AI/TUĞRA-T2V-A14B --local_dir ./TUĞRA-T2V-A14B
134
+ ```
135
+
136
+ #### Run Text-to-Video Generation
137
+
138
+ This repository supports the `TUĞRA-T2V-A14B` Text-to-Video model and can simultaneously support video generation at 480P and 720P resolutions.
139
+
140
+
141
+ ##### (1) Without Prompt Extension
142
+
143
+ To facilitate implementation, we will start with a basic version of the inference process that skips the [prompt extension](#2-using-prompt-extention) step.
144
+
145
+ - Single-GPU inference
146
+
147
+ ``` sh
148
+ python generate.py --task t2v-A14B --size 1280*720 --ckpt_dir ./TUĞRA-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."
149
+ ```
150
+
151
+ > 💡 This command can run on a GPU with at least 80GB VRAM.
152
+
153
+ > 💡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.
154
+
155
+
156
+ - Multi-GPU inference using FSDP + DeepSpeed Ulysses
157
+
158
+ We use [PyTorch FSDP](https://docs.pytorch.org/docs/stable/fsdp.html) and [DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509) to accelerate inference.
159
+
160
+
161
+ ``` sh
162
+ torchrun --nproc_per_node=8 generate.py --task t2v-A14B --size 1280*720 --ckpt_dir ./TUĞRA-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."
163
+ ```
164
+
165
+
166
+ ##### (2) Using Prompt Extension
167
+
168
+ 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:
169
+
170
+ - Use the Dashscope API for extension.
171
+ - 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)).
172
+ - 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).
173
+ - Use the `qwen-plus` model for text-to-video tasks and `qwen-vl-max` for image-to-video tasks.
174
+ - You can modify the model used for extension with the parameter `--prompt_extend_model`. For example:
175
+ ```sh
176
+ DASH_API_KEY=your_key torchrun --nproc_per_node=8 generate.py --task t2v-A14B --size 1280*720 --ckpt_dir ./TUĞRA-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'
177
+ ```
178
+
179
+ - Using a local model for extension.
180
+
181
+ - 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.
182
+ - 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`.
183
+ - For image-to-video tasks, you can use models like `Qwen/Qwen2.5-VL-7B-Instruct` and `Qwen/Qwen2.5-VL-3B-Instruct`.
184
+ - Larger models generally provide better extension results but require more GPU memory.
185
+ - 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:
186
+
187
+ ``` sh
188
+ torchrun --nproc_per_node=8 generate.py --task t2v-A14B --size 1280*720 --ckpt_dir ./TUĞRA-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'
189
+ ```
190
+
191
+
192
+ #### Run Image-to-Video Generation
193
+
194
+ This repository supports the `TUĞRA-I2V-A14B` Image-to-Video model and can simultaneously support video generation at 480P and 720P resolutions.
195
+
196
+
197
+ - Single-GPU inference
198
+ ```sh
199
+ python generate.py --task i2v-A14B --size 1280*720 --ckpt_dir ./TUĞRA-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."
200
+ ```
201
+
202
+ > This command can run on a GPU with at least 80GB VRAM.
203
+
204
+ > 💡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.
205
+
206
+
207
+ - Multi-GPU inference using FSDP + DeepSpeed Ulysses
208
+
209
+ ```sh
210
+ torchrun --nproc_per_node=8 generate.py --task i2v-A14B --size 1280*720 --ckpt_dir ./TUĞRA-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."
211
+ ```
212
+
213
+ - Image-to-Video Generation without prompt
214
+
215
+ ```sh
216
+ DASH_API_KEY=your_key torchrun --nproc_per_node=8 generate.py --task i2v-A14B --size 1280*720 --ckpt_dir ./TUĞRA-I2V-A14B --prompt '' --image examples/i2v_input.JPG --dit_fsdp --t5_fsdp --ulysses_size 8 --use_prompt_extend --prompt_extend_method 'dashscope'
217
+ ```
218
+
219
+ > 💡The model can generate videos solely from the input image. You can use prompt extension to generate prompt from the image.
220
+
221
+ > The process of prompt extension can be referenced [here](#2-using-prompt-extention).
222
+
223
+ #### Run Text-Image-to-Video Generation
224
+
225
+ This repository supports the `TUĞRA-TI2V-5B` Text-Image-to-Video model and can support video generation at 720P resolutions.
226
+
227
+
228
+ - Single-GPU Text-to-Video inference
229
+ ```sh
230
+ python generate.py --task ti2v-5B --size 1280*704 --ckpt_dir ./TUĞRA-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"
231
+ ```
232
+
233
+ > 💡Unlike other tasks, the 720P resolution of the Text-Image-to-Video task is `1280*704` or `704*1280`.
234
+
235
+ > This command can run on a GPU with at least 24GB VRAM (e.g, RTX 4090 GPU).
236
+
237
+ > 💡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.
238
+
239
+
240
+ - Single-GPU Image-to-Video inference
241
+ ```sh
242
+ python generate.py --task ti2v-5B --size 1280*704 --ckpt_dir ./TUĞRA-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."
243
+ ```
244
+
245
+ > 💡If the image parameter is configured, it is an Image-to-Video generation; otherwise, it defaults to a Text-to-Video generation.
246
+
247
+ > 💡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.
248
+
249
+
250
+ - Multi-GPU inference using FSDP + DeepSpeed Ulysses
251
+
252
+ ```sh
253
+ torchrun --nproc_per_node=8 generate.py --task ti2v-5B --size 1280*704 --ckpt_dir ./TUĞRA-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."
254
+ ```
255
+
256
+ > The process of prompt extension can be referenced [here](#2-using-prompt-extention).
257
+
258
+ #### Run Speech-to-Video Generation
259
+
260
+ This repository supports the `TUĞRA-S2V-14B` Speech-to-Video model and can simultaneously support video generation at 480P and 720P resolutions.
261
+
262
+ - Single-GPU Speech-to-Video inference
263
+
264
+ ```sh
265
+ python generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./TUĞRA-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"
266
+ # Without setting --num_clip, the generated video length will automatically adjust based on the input audio length
267
+
268
+ # You can use CosyVoice to generate audio with --enable_tts
269
+ python generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./TUĞRA-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 "收到好友从远方寄来的生日礼物,那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐,笑容如花儿般绽放。"
270
+ ```
271
+
272
+ > 💡 This command can run on a GPU with at least 80GB VRAM.
273
+
274
+ - Multi-GPU inference using FSDP + DeepSpeed Ulysses
275
+
276
+ ```sh
277
+ torchrun --nproc_per_node=8 generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./TUĞRA-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"
278
+ ```
279
+
280
+ - Pose + Audio driven generation
281
+
282
+ ```sh
283
+ torchrun --nproc_per_node=8 generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./TUĞRA-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"
284
+ ```
285
+
286
+ > 💡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.
287
+
288
+ > 💡The model can generate videos from audio input combined with reference image and optional text prompt.
289
+
290
+ > 💡The `--pose_video` parameter enables pose-driven generation, allowing the model to follow specific pose sequences while generating videos synchronized with audio input.
291
+
292
+ > 💡The `--num_clip` parameter controls the number of video clips generated, useful for quick preview with shorter generation time.
293
+
294
+ Please visit our project page to see more examples and learn about the scenarios suitable for this model.
295
+
296
+ #### Run Tugra-Animate
297
+
298
+ Tugra-Animate takes a video and a character image as input, and generates a video in either "animation" or "replacement" mode.
299
+
300
+ 1. animation mode: The model generates a video of the character image that mimics the human motion in the input video.
301
+ 2. replacement mode: The model replaces the character image with the input video.
302
+
303
+ Please visit our [project page](https://humanaigc.github.io/wan-animate) to see more examples and learn about the scenarios suitable for this model.
304
+
305
+ ##### (1) Preprocessing
306
+ 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/Tugra-Video/TUĞRA/blob/main/wan/modules/animate/preprocess/UserGuider.md).
307
+
308
+ * For animation
309
+ ```bash
310
+ python ./wan/modules/animate/preprocess/preprocess_data.py \
311
+ --ckpt_path ./TUĞRA-Animate-14B/process_checkpoint \
312
+ --video_path ./examples/wan_animate/animate/video.mp4 \
313
+ --refer_path ./examples/wan_animate/animate/image.jpeg \
314
+ --save_path ./examples/wan_animate/animate/process_results \
315
+ --resolution_area 1280 720 \
316
+ --retarget_flag \
317
+ --use_flux
318
+ ```
319
+ * For replacement
320
+ ```bash
321
+ python ./wan/modules/animate/preprocess/preprocess_data.py \
322
+ --ckpt_path ./TUĞRA-Animate-14B/process_checkpoint \
323
+ --video_path ./examples/wan_animate/replace/video.mp4 \
324
+ --refer_path ./examples/wan_animate/replace/image.jpeg \
325
+ --save_path ./examples/wan_animate/replace/process_results \
326
+ --resolution_area 1280 720 \
327
+ --iterations 3 \
328
+ --k 7 \
329
+ --w_len 1 \
330
+ --h_len 1 \
331
+ --replace_flag
332
+ ```
333
+ ##### (2) Run in animation mode
334
+
335
+ * Single-GPU inference
336
+
337
+ ```bash
338
+ python generate.py --task animate-14B --ckpt_dir ./TUĞRA-Animate-14B/ --src_root_path ./examples/wan_animate/animate/process_results/ --refert_num 1
339
+ ```
340
+
341
+ * Multi-GPU inference using FSDP + DeepSpeed Ulysses
342
+
343
+ ```bash
344
+ python -m torch.distributed.run --nnodes 1 --nproc_per_node 8 generate.py --task animate-14B --ckpt_dir ./TUĞRA-Animate-14B/ --src_root_path ./examples/wan_animate/animate/process_results/ --refert_num 1 --dit_fsdp --t5_fsdp --ulysses_size 8
345
+ ```
346
+
347
+ * Diffusers Pipeline
348
+
349
+ ```python
350
+ from diffusers import WanAnimatePipeline
351
+ from diffusers.utils import export_to_video, load_image, load_video
352
+
353
+ device = "cuda:0"
354
+ dtype = torch.bfloat16
355
+ model_id = "Tugra-AI/TUĞRA-Animate-14B-Diffusers"
356
+ pipe = WanAnimatePipeline.from_pretrained(model_id torch_dtype=dtype)
357
+ pipe.to(device)
358
+
359
+ seed = 42
360
+ prompt = "People in the video are doing actions."
361
+
362
+ # Animation
363
+ image = load_image("/path/to/animate/reference/image/src_ref.png")
364
+ pose_video = load_video("/path/to/animate/pose/video/src_pose.mp4")
365
+ face_video = load_video("/path/to/animate/face/video/src_face.mp4")
366
+
367
+ animate_video = pipe(
368
+ image=image,
369
+ pose_video=pose_video,
370
+ face_video=face_video,
371
+ prompt=prompt,
372
+ mode="animate",
373
+ segment_frame_length=77, # clip_len in original code
374
+ prev_segment_conditioning_frames=1, # refert_num in original code
375
+ guidance_scale=1.0,
376
+ num_inference_steps=20,
377
+ generator=torch.Generator(device=device).manual_seed(seed),
378
+ ).frames[0]
379
+ export_to_video(animate_video, "diffusers_animate.mp4", fps=30)
380
+ ```
381
+
382
+ ##### (3) Run in replacement mode
383
+
384
+ * Single-GPU inference
385
+
386
+ ```bash
387
+ python generate.py --task animate-14B --ckpt_dir ./TUĞRA-Animate-14B/ --src_root_path ./examples/wan_animate/replace/process_results/ --refert_num 1 --replace_flag --use_relighting_lora
388
+ ```
389
+
390
+ * Multi-GPU inference using FSDP + DeepSpeed Ulysses
391
+
392
+ ```bash
393
+ python -m torch.distributed.run --nnodes 1 --nproc_per_node 8 generate.py --task animate-14B --ckpt_dir ./TUĞRA-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
394
+ ```
395
+
396
+ * Diffusers Pipeline
397
+
398
+ ```python
399
+ # create pipeline as in the Animation code ☝️
400
+
401
+ # Replacement
402
+ image = load_image("/path/to/replace/reference/image/src_ref.png")
403
+ pose_video = load_video("/path/to/replace/pose/video/src_pose.mp4")
404
+ face_video = load_video("/path/to/replace/face/video/src_face.mp4")
405
+ background_video = load_video("/path/to/replace/background/video/src_bg.mp4")
406
+ mask_video = load_video("/path/to/replace/mask/video/src_mask.mp4")
407
+
408
+ replace_video = pipe(
409
+ image=image,
410
+ pose_video=pose_video,
411
+ face_video=face_video,
412
+ background_video=background_video,
413
+ mask_video=mask_video,
414
+ prompt=prompt,
415
+ mode="replace",
416
+ segment_frame_length=77, # clip_len in original code
417
+ prev_segment_conditioning_frames=1, # refert_num in original code
418
+ guidance_scale=1.0,
419
+ num_inference_steps=20,
420
+ generator=torch.Generator(device=device).manual_seed(seed),
421
+ ).frames[0]
422
+ export_to_video(replace_video, "diffusers_replace.mp4", fps=30)
423
+ ```
424
+
425
+ > 💡 If you're using **Tugra-Animate**, we do not recommend using LoRA models trained on `TUĞRA`, since weight changes during training may lead to unexpected behavior.
426
+
427
+ ## Computational Efficiency on Different GPUs
428
+
429
+ We test the computational efficiency of different **TUĞRA** models on different GPUs in the following table. The results are presented in the format: **Total time (s) / peak GPU memory (GB)**.
430
+
431
+
432
+ <div align="center">
433
+ <img src="assets/comp_effic.png" alt="" style="width: 80%;" />
434
+ </div>
435
+
436
+ > The parameter settings for the tests presented in this table are as follows:
437
+ > (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`
438
+ (--convert_model_dtype converts model parameter types to config.param_dtype);
439
+ > (2) The distributed testing utilizes the built-in FSDP and Ulysses implementations, with FlashAttention3 deployed on Hopper architecture GPUs;
440
+ > (3) Tests were run without the `--use_prompt_extend` flag;
441
+ > (4) Reported results are the average of multiple samples taken after the warm-up phase.
442
+
443
+
444
+ -------
445
+
446
+ ## Introduction of TUĞRA
447
+
448
+ **TUĞRA** builds on the foundation of Tugra v2.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.
449
+
450
+ ##### (1) Mixture-of-Experts (MoE) Architecture
451
+
452
+ TUĞRA 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 TUĞRA, 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.
453
+
454
+ <div align="center">
455
+ <img src="assets/moe_arch.png" alt="" style="width: 90%;" />
456
+ </div>
457
+
458
+ 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}$.
459
+
460
+ <div align="center">
461
+ <img src="assets/moe_2.png" alt="" style="width: 90%;" />
462
+ </div>
463
+
464
+ To validate the effectiveness of the MoE architecture, four settings are compared based on their validation loss curves. The baseline **Tugra v2.1** model does not employ the MoE architecture. Among the MoE-based variants, the **Tugra v2.1 & High-Noise Expert** reuses the Tugra v2.1 model as the low-noise expert while uses the TUĞRA's high-noise expert, while the **Tugra v2.1 & Low-Noise Expert** uses Tugra v2.1 as the high-noise expert and employ the TUĞRA's low-noise expert. The **TUĞRA (MoE)** (our final version) achieves the lowest validation loss, indicating that its generated video distribution is closest to ground-truth and exhibits superior convergence.
465
+
466
+
467
+ ##### (2) Efficient High-Definition Hybrid TI2V
468
+ To enable more efficient deployment, TUĞRA 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 TUĞRA-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.
469
+
470
+
471
+ <div align="center">
472
+ <img src="assets/vae.png" alt="" style="width: 80%;" />
473
+ </div>
474
+
475
+
476
+
477
+ ##### Comparisons to SOTAs
478
+ We compared TUĞRA with leading closed-source commercial models on our new Wan-Bench 2.0, evaluating performance across multiple crucial dimensions. The results demonstrate that TUĞRA achieves superior performance compared to these leading models.
479
+
480
+
481
+ <div align="center">
482
+ <img src="assets/performance.png" alt="" style="width: 90%;" />
483
+ </div>
484
+
485
+ ## Citation
486
+ If you find our work helpful, please cite us.
487
+
488
+ ```
489
+ @article{wan2025,
490
+ title={Wan: Open and Advanced Large-Scale Video Generative Models},
491
+ 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},
492
+ journal = {arXiv preprint arXiv:2503.20314},
493
+ year={2025}
494
+ }
495
+ ```
496
+
497
+ ## License Agreement
498
+ 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).
499
+
500
+
501
+ ## Acknowledgements
502
+
503
+ 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.
504
+
505
+
506
+
507
+ ## Contact Us
508
+ 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)!
509
+
510
+ # TUGRA
511
+
512
+
Tugra_Wan2.2_Colab.ipynb ADDED
@@ -0,0 +1,158 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "markdown",
5
+ "metadata": {},
6
+ "source": [
7
+ "# 💜 TUĞRA (Wan 2.2) Colab Inference Notebook\n",
8
+ "\n",
9
+ "This notebook enables you to run the state-of-the-art **TUĞRA (Wan 2.2)** video generation models in Google Colab.\n",
10
+ "\n",
11
+ "### Recommended GPU:\n",
12
+ "- **T4 GPU** (Free tier): Compatible with the **TI2V-5B** model (using CPU offloading and low VRAM configurations).\n",
13
+ "- **A100 or L4 GPU** (Colab Pro): Best for faster generation and larger models like **T2V-A14B**."
14
+ ]
15
+ },
16
+ {
17
+ "cell_type": "markdown",
18
+ "metadata": {},
19
+ "source": [
20
+ "## 🛠️ Step 1: Clone Repository & Install Dependencies"
21
+ ]
22
+ },
23
+ {
24
+ "cell_type": "code",
25
+ "execution_count": null,
26
+ "metadata": {},
27
+ "outputs": [],
28
+ "source": [
29
+ "# Clone the TUĞRA repository\n",
30
+ "!git clone https://github.com/Tugra-Video/TU\u011eRA.git\n",
31
+ "%cd TU\u011eRA\n",
32
+ "\n",
33
+ "# Install dependencies (Ensure we have CUDA-enabled torch & diffusers)\n",
34
+ "!pip install -r requirements.txt\n",
35
+ "!pip install huggingface_hub[cli] modelscope"
36
+ ]
37
+ },
38
+ {
39
+ "cell_type": "markdown",
40
+ "metadata": {},
41
+ "source": [
42
+ "## 📥 Step 2: Download Model Weights\n",
43
+ "\n",
44
+ "We will download the **TI2V-5B** model (Text-to-Video & Image-to-Video, 720P resolution) which runs efficiently on single GPU setups."
45
+ ]
46
+ },
47
+ {
48
+ "cell_type": "code",
49
+ "execution_count": null,
50
+ "metadata": {},
51
+ "outputs": [],
52
+ "source": [
53
+ "# Download the TU\u011eRA-TI2V-5B model weights\n",
54
+ "!huggingface-cli download Tugra-AI/TU\u011eRA-TI2V-5B --local-dir ./TU\u011eRA-TI2V-5B"
55
+ ]
56
+ },
57
+ {
58
+ "cell_type": "markdown",
59
+ "metadata": {},
60
+ "source": [
61
+ "## 🎬 Step 3: Run Video Generation\n",
62
+ "\n",
63
+ "### Option A: Text-to-Video (T2V) Generation"
64
+ ]
65
+ },
66
+ {
67
+ "cell_type": "code",
68
+ "execution_count": null,
69
+ "metadata": {},
70
+ "outputs": [],
71
+ "source": [
72
+ "PROMPT = \"A majestic purple dragon soaring high above neon-lit futuristic skyscrapers, cyberpunk aesthetic, cinematic shot, 8k resolution.\"\n",
73
+ "\n",
74
+ "!python generate.py \\\n",
75
+ " --task ti2v-5B \\\n",
76
+ " --size 1280*704 \\\n",
77
+ " --ckpt_dir ./TU\u011eRA-TI2V-5B \\\n",
78
+ " --offload_model True \\\n",
79
+ " --convert_model_dtype \\\n",
80
+ " --t5_cpu \\\n",
81
+ " --prompt \"{PROMPT}\""
82
+ ]
83
+ },
84
+ {
85
+ "cell_type": "markdown",
86
+ "metadata": {},
87
+ "source": [
88
+ "### Option B: Image-to-Video (I2V) Generation"
89
+ ]
90
+ },
91
+ {
92
+ "cell_type": "code",
93
+ "execution_count": null,
94
+ "metadata": {},
95
+ "outputs": [],
96
+ "source": [
97
+ "# 1. Upload your reference image using the Colab file explorer on the left.\n",
98
+ "# 2. Update the path below to point to your uploaded image.\n",
99
+ "IMAGE_PATH = \"examples/i2v_input.JPG\"\n",
100
+ "PROMPT = \"Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard, waves gently crashing around.\"\n",
101
+ "\n",
102
+ "!python generate.py \\\n",
103
+ " --task ti2v-5B \\\n",
104
+ " --size 1280*704 \\\n",
105
+ " --ckpt_dir ./TU\u011eRA-TI2V-5B \\\n",
106
+ " --offload_model True \\\n",
107
+ " --convert_model_dtype \\\n",
108
+ " --t5_cpu \\\n",
109
+ " --image \"{IMAGE_PATH}\" \\\n",
110
+ " --prompt \"{PROMPT}\""
111
+ ]
112
+ },
113
+ {
114
+ "cell_type": "markdown",
115
+ "metadata": {},
116
+ "source": [
117
+ "## 📺 Step 4: Display Generated Video\n",
118
+ "\n",
119
+ "Find your output `.mp4` file in the directory and run this cell to play it directly in Colab!"
120
+ ]
121
+ },
122
+ {
123
+ "cell_type": "code",
124
+ "execution_count": null,
125
+ "metadata": {},
126
+ "outputs": [],
127
+ "source": [
128
+ "import glob\n",
129
+ "import os\n",
130
+ "from IPython.display import HTML\n",
131
+ "from base64 import b64encode\n",
132
+ "\n",
133
+ "# Get the latest generated mp4 file\n",
134
+ "videos = glob.glob(\"*.mp4\")\n",
135
+ "if videos:\n",
136
+ " latest_video = max(videos, key=os.path.getctime)\n",
137
+ " print(f\"Displaying: {latest_video}\")\n",
138
+ " \n",
139
+ " mp4 = open(latest_video, 'rb').read()\n",
140
+ " data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
141
+ " display(HTML(f\"\"\"\n",
142
+ " <video width=640 controls>\n",
143
+ " <source src=\"{data_url}\" type=\"video/mp4\">\n",
144
+ " </video>\n",
145
+ " \"\"\"))\n",
146
+ "else:\n",
147
+ " print(\"No generated video found yet!\")"
148
+ ]
149
+ }
150
+ ],
151
+ "metadata": {
152
+ "language_info": {
153
+ "name": "python"
154
+ }
155
+ },
156
+ "nbformat": 4,
157
+ "nbformat_minor": 2
158
+ }
assets/comp_effic.png ADDED

Git LFS Details

  • SHA256: 75ee012dcfb08365bec67a3ec7afc126fc2817f79b9f80e38711792d4770e32b
  • Pointer size: 131 Bytes
  • Size of remote file: 202 kB
assets/logo.png ADDED

Git LFS Details

  • SHA256: b593caf9fb24b69d3df8414ff219f368c944edb3eae8faa8c9fe85f5e6656279
  • Pointer size: 131 Bytes
  • Size of remote file: 374 kB
assets/moe_2.png ADDED

Git LFS Details

  • SHA256: 4ea471ccb64349bd08bc9a78f336ae000e9ca3b40da9a652b8028b214a8c6093
  • Pointer size: 131 Bytes
  • Size of remote file: 528 kB
assets/moe_arch.png ADDED
assets/performance.png ADDED

Git LFS Details

  • SHA256: 97ef99c13c8ae717a8a11c8d8ec927b69077c647cc6689755d08fc38e7fbb830
  • Pointer size: 131 Bytes
  • Size of remote file: 307 kB
assets/vae.png ADDED

Git LFS Details

  • SHA256: 4aaea5e187f1c5908e15ade5bef24c9fb59882986bc3d2ad75f7fe820f3d772f
  • Pointer size: 131 Bytes
  • Size of remote file: 165 kB
examples/Five Hundred Miles.MP3 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5d969412bfe4d5e4b328d3ff92c1307fc39f7988ceba66860b5c1a17e40502d6
3
+ size 121043
examples/Five Hundred Miles.png ADDED

Git LFS Details

  • SHA256: 5775ff7fbb162b937ec7ea2ff028cf4207d77a156720f9be822004b549ab4e98
  • Pointer size: 131 Bytes
  • Size of remote file: 878 kB
examples/i2v_input.JPG ADDED

Git LFS Details

  • SHA256: 077e3d965090c9028c69c00931675f42e1acc815c6eb450ab291b3b72d211a8e
  • Pointer size: 131 Bytes
  • Size of remote file: 251 kB
examples/pose.mp4 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:467b541d19625fdc42bf47f9d4db2d02cf0579f4e3a9233c543b2117dbed8a8e
3
+ size 2192979
examples/pose.png ADDED

Git LFS Details

  • SHA256: 53a5d9b435adaf15dd8ffcab1a833b61da4a63079200fb9cec33127ee10f733b
  • Pointer size: 131 Bytes
  • Size of remote file: 823 kB
examples/sing.MP3 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:520217a826cd078a61ff1eac7a3f8dfa55ade170d07a977d86d9bcb049d7fa59
3
+ size 300144
examples/talk.wav ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c8b0b80ae25baaa402853f34b24f5ba64decd67bcf9a512640d1e8b1d040824f
3
+ size 884814
examples/wan_animate/animate/image.jpeg ADDED

Git LFS Details

  • SHA256: 8123db8e5c47c3a229c288b4c5245e8ee2ce4378b1c09e92873b75939812eb7b
  • Pointer size: 131 Bytes
  • Size of remote file: 123 kB
examples/wan_animate/animate/video.mp4 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:80f3cfe3786a7f8a94844476448fb45e7e115216ddcdaad14b0b88223be597e7
3
+ size 903201
examples/wan_animate/replace/image.jpeg ADDED

Git LFS Details

  • SHA256: 412591418fbb133bd46c41b3376b810bd7e3eb59b916bf9693da337a08ca1b0d
  • Pointer size: 131 Bytes
  • Size of remote file: 143 kB
examples/wan_animate/replace/video.mp4 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:db6da60e5fcb0fda0bff151bfbdbb7085d5a86a78508743cce2a25709de86a19
3
+ size 754294
examples/zero_shot_prompt.wav ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bd199eb7109fd6ce9943cb297e3cf350c1073af014063dfadbdc100230526243
3
+ size 111496
generate.py ADDED
@@ -0,0 +1,575 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 KAYRA. All rights reserved.d.
2
+ import argparse
3
+ import logging
4
+ import os
5
+ import sys
6
+ import warnings
7
+ from datetime import datetime
8
+
9
+ warnings.filterwarnings('ignore')
10
+
11
+ import random
12
+
13
+ import torch
14
+ import torch.distributed as dist
15
+ from PIL import Image
16
+
17
+ import tugra
18
+ from tugra.configs import MAX_AREA_CONFIGS, SIZE_CONFIGS, SUPPORTED_SIZES, TUGRA_CONFIGS
19
+ from tugra.distributed.util import init_distributed_group
20
+ from tugra.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
21
+ from tugra.utils.utils import merge_video_audio, save_video, str2bool
22
+
23
+
24
+ EXAMPLE_PROMPT = {
25
+ "t2v-A14B": {
26
+ "prompt":
27
+ "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
28
+ },
29
+ "i2v-A14B": {
30
+ "prompt":
31
+ "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.",
32
+ "image":
33
+ "examples/i2v_input.JPG",
34
+ },
35
+ "ti2v-5B": {
36
+ "prompt":
37
+ "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
38
+ },
39
+ "animate-14B": {
40
+ "prompt": "视频中的人在做动作",
41
+ "video": "",
42
+ "pose": "",
43
+ "mask": "",
44
+ },
45
+ "s2v-14B": {
46
+ "prompt":
47
+ "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.",
48
+ "image":
49
+ "examples/i2v_input.JPG",
50
+ "audio":
51
+ "examples/talk.wav",
52
+ "tts_prompt_audio":
53
+ "examples/zero_shot_prompt.wav",
54
+ "tts_prompt_text":
55
+ "希望你以后能够做的比我还好呦。",
56
+ "tts_text":
57
+ "收到好友从远方寄来的生日礼物,那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐,笑容如花儿般绽放。"
58
+ },
59
+ }
60
+
61
+
62
+ def _validate_args(args):
63
+ # Basic check
64
+ assert args.ckpt_dir is not None, "Please specify the checkpoint directory."
65
+ assert args.task in TUGRA_CONFIGS, f"Unsupport task: {args.task}"
66
+ assert args.task in EXAMPLE_PROMPT, f"Unsupport task: {args.task}"
67
+
68
+ if args.prompt is None:
69
+ args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
70
+ if args.image is None and "image" in EXAMPLE_PROMPT[args.task]:
71
+ args.image = EXAMPLE_PROMPT[args.task]["image"]
72
+ if args.audio is None and args.enable_tts is False and "audio" in EXAMPLE_PROMPT[args.task]:
73
+ args.audio = EXAMPLE_PROMPT[args.task]["audio"]
74
+ 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]:
75
+ args.tts_prompt_audio = EXAMPLE_PROMPT[args.task]["tts_prompt_audio"]
76
+ args.tts_prompt_text = EXAMPLE_PROMPT[args.task]["tts_prompt_text"]
77
+ args.tts_text = EXAMPLE_PROMPT[args.task]["tts_text"]
78
+
79
+ if args.task == "i2v-A14B":
80
+ assert args.image is not None, "Please specify the image path for i2v."
81
+
82
+ cfg = TUGRA_CONFIGS[args.task]
83
+
84
+ if args.sample_steps is None:
85
+ args.sample_steps = cfg.sample_steps
86
+
87
+ if args.sample_shift is None:
88
+ args.sample_shift = cfg.sample_shift
89
+
90
+ if args.sample_guide_scale is None:
91
+ args.sample_guide_scale = cfg.sample_guide_scale
92
+
93
+ if args.frame_num is None:
94
+ args.frame_num = cfg.frame_num
95
+
96
+ args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
97
+ 0, sys.maxsize)
98
+ # Size check
99
+ if not 's2v' in args.task:
100
+ assert args.size in SUPPORTED_SIZES[
101
+ args.
102
+ task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"
103
+
104
+
105
+ def _parse_args():
106
+ parser = argparse.ArgumentParser(
107
+ description="Generate a image or video from a text prompt or image using TUĞRA"
108
+ )
109
+ parser.add_argument(
110
+ "--task",
111
+ type=str,
112
+ default="t2v-A14B",
113
+ choices=list(TUGRA_CONFIGS.keys()),
114
+ help="The task to run.")
115
+ parser.add_argument(
116
+ "--size",
117
+ type=str,
118
+ default="1280*720",
119
+ choices=list(SIZE_CONFIGS.keys()),
120
+ 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."
121
+ )
122
+ parser.add_argument(
123
+ "--frame_num",
124
+ type=int,
125
+ default=None,
126
+ help="How many frames of video are generated. The number should be 4n+1"
127
+ )
128
+ parser.add_argument(
129
+ "--ckpt_dir",
130
+ type=str,
131
+ default=None,
132
+ help="The path to the checkpoint directory.")
133
+ parser.add_argument(
134
+ "--offload_model",
135
+ type=str2bool,
136
+ default=None,
137
+ help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage."
138
+ )
139
+ parser.add_argument(
140
+ "--ulysses_size",
141
+ type=int,
142
+ default=1,
143
+ help="The size of the ulysses parallelism in DiT.")
144
+ parser.add_argument(
145
+ "--t5_fsdp",
146
+ action="store_true",
147
+ default=False,
148
+ help="Whether to use FSDP for T5.")
149
+ parser.add_argument(
150
+ "--t5_cpu",
151
+ action="store_true",
152
+ default=False,
153
+ help="Whether to place T5 model on CPU.")
154
+ parser.add_argument(
155
+ "--dit_fsdp",
156
+ action="store_true",
157
+ default=False,
158
+ help="Whether to use FSDP for DiT.")
159
+ parser.add_argument(
160
+ "--save_file",
161
+ type=str,
162
+ default=None,
163
+ help="The file to save the generated video to.")
164
+ parser.add_argument(
165
+ "--prompt",
166
+ type=str,
167
+ default=None,
168
+ help="The prompt to generate the video from.")
169
+ parser.add_argument(
170
+ "--use_prompt_extend",
171
+ action="store_true",
172
+ default=False,
173
+ help="Whether to use prompt extend.")
174
+ parser.add_argument(
175
+ "--prompt_extend_method",
176
+ type=str,
177
+ default="local_qwen",
178
+ choices=["dashscope", "local_qwen"],
179
+ help="The prompt extend method to use.")
180
+ parser.add_argument(
181
+ "--prompt_extend_model",
182
+ type=str,
183
+ default=None,
184
+ help="The prompt extend model to use.")
185
+ parser.add_argument(
186
+ "--prompt_extend_target_lang",
187
+ type=str,
188
+ default="zh",
189
+ choices=["zh", "en"],
190
+ help="The target language of prompt extend.")
191
+ parser.add_argument(
192
+ "--base_seed",
193
+ type=int,
194
+ default=-1,
195
+ help="The seed to use for generating the video.")
196
+ parser.add_argument(
197
+ "--image",
198
+ type=str,
199
+ default=None,
200
+ help="The image to generate the video from.")
201
+ parser.add_argument(
202
+ "--sample_solver",
203
+ type=str,
204
+ default='unipc',
205
+ choices=['unipc', 'dpm++'],
206
+ help="The solver used to sample.")
207
+ parser.add_argument(
208
+ "--sample_steps", type=int, default=None, help="The sampling steps.")
209
+ parser.add_argument(
210
+ "--sample_shift",
211
+ type=float,
212
+ default=None,
213
+ help="Sampling shift factor for flow matching schedulers.")
214
+ parser.add_argument(
215
+ "--sample_guide_scale",
216
+ type=float,
217
+ default=None,
218
+ help="Classifier free guidance scale.")
219
+ parser.add_argument(
220
+ "--convert_model_dtype",
221
+ action="store_true",
222
+ default=False,
223
+ help="Whether to convert model paramerters dtype.")
224
+
225
+ # animate
226
+ parser.add_argument(
227
+ "--src_root_path",
228
+ type=str,
229
+ default=None,
230
+ help="The file of the process output path. Default None.")
231
+ parser.add_argument(
232
+ "--refert_num",
233
+ type=int,
234
+ default=77,
235
+ help="How many frames used for temporal guidance. Recommended to be 1 or 5."
236
+ )
237
+ parser.add_argument(
238
+ "--replace_flag",
239
+ action="store_true",
240
+ default=False,
241
+ help="Whether to use replace.")
242
+ parser.add_argument(
243
+ "--use_relighting_lora",
244
+ action="store_true",
245
+ default=False,
246
+ help="Whether to use relighting lora.")
247
+
248
+ # following args only works for s2v
249
+ parser.add_argument(
250
+ "--num_clip",
251
+ type=int,
252
+ default=None,
253
+ help="Number of video clips to generate, the whole video will not exceed the length of audio."
254
+ )
255
+ parser.add_argument(
256
+ "--audio",
257
+ type=str,
258
+ default=None,
259
+ help="Path to the audio file, e.g. wav, mp3")
260
+ parser.add_argument(
261
+ "--enable_tts",
262
+ action="store_true",
263
+ default=False,
264
+ help="Use CosyVoice to synthesis audio")
265
+ parser.add_argument(
266
+ "--tts_prompt_audio",
267
+ type=str,
268
+ default=None,
269
+ help="Path to the tts prompt audio file, e.g. wav, mp3. Must be greater than 16khz, and between 5s to 15s.")
270
+ parser.add_argument(
271
+ "--tts_prompt_text",
272
+ type=str,
273
+ default=None,
274
+ help="Content to the tts prompt audio. If provided, must exactly match tts_prompt_audio")
275
+ parser.add_argument(
276
+ "--tts_text",
277
+ type=str,
278
+ default=None,
279
+ help="Text wish to synthesize")
280
+ parser.add_argument(
281
+ "--pose_video",
282
+ type=str,
283
+ default=None,
284
+ help="Provide Dw-pose sequence to do Pose Driven")
285
+ parser.add_argument(
286
+ "--start_from_ref",
287
+ action="store_true",
288
+ default=False,
289
+ help="whether set the reference image as the starting point for generation"
290
+ )
291
+ parser.add_argument(
292
+ "--infer_frames",
293
+ type=int,
294
+ default=80,
295
+ help="Number of frames per clip, 48 or 80 or others (must be multiple of 4) for 14B s2v"
296
+ )
297
+ args = parser.parse_args()
298
+ _validate_args(args)
299
+
300
+ return args
301
+
302
+
303
+ def _init_logging(rank):
304
+ # logging
305
+ if rank == 0:
306
+ # set format
307
+ logging.basicConfig(
308
+ level=logging.INFO,
309
+ format="[%(asctime)s] %(levelname)s: %(message)s",
310
+ handlers=[logging.StreamHandler(stream=sys.stdout)])
311
+ else:
312
+ logging.basicConfig(level=logging.ERROR)
313
+
314
+
315
+ def generate(args):
316
+ rank = int(os.getenv("RANK", 0))
317
+ world_size = int(os.getenv("WORLD_SIZE", 1))
318
+ local_rank = int(os.getenv("LOCAL_RANK", 0))
319
+ device = local_rank
320
+ _init_logging(rank)
321
+
322
+ if args.offload_model is None:
323
+ args.offload_model = False if world_size > 1 else True
324
+ logging.info(
325
+ f"offload_model is not specified, set to {args.offload_model}.")
326
+ if world_size > 1:
327
+ torch.cuda.set_device(local_rank)
328
+ dist.init_process_group(
329
+ backend="nccl",
330
+ init_method="env://",
331
+ rank=rank,
332
+ world_size=world_size)
333
+ else:
334
+ assert not (
335
+ args.t5_fsdp or args.dit_fsdp
336
+ ), f"t5_fsdp and dit_fsdp are not supported in non-distributed environments."
337
+ assert not (
338
+ args.ulysses_size > 1
339
+ ), f"sequence parallel are not supported in non-distributed environments."
340
+
341
+ if args.ulysses_size > 1:
342
+ assert args.ulysses_size == world_size, f"The number of ulysses_size should be equal to the world size."
343
+ init_distributed_group()
344
+
345
+ if args.use_prompt_extend:
346
+ if args.prompt_extend_method == "dashscope":
347
+ prompt_expander = DashScopePromptExpander(
348
+ model_name=args.prompt_extend_model,
349
+ task=args.task,
350
+ is_vl=args.image is not None)
351
+ elif args.prompt_extend_method == "local_qwen":
352
+ prompt_expander = QwenPromptExpander(
353
+ model_name=args.prompt_extend_model,
354
+ task=args.task,
355
+ is_vl=args.image is not None,
356
+ device=rank)
357
+ else:
358
+ raise NotImplementedError(
359
+ f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
360
+
361
+ cfg = TUGRA_CONFIGS[args.task]
362
+ if args.ulysses_size > 1:
363
+ assert cfg.num_heads % args.ulysses_size == 0, f"`{cfg.num_heads=}` cannot be divided evenly by `{args.ulysses_size=}`."
364
+
365
+ logging.info(f"Generation job args: {args}")
366
+ logging.info(f"Generation model config: {cfg}")
367
+
368
+ if dist.is_initialized():
369
+ base_seed = [args.base_seed] if rank == 0 else [None]
370
+ dist.broadcast_object_list(base_seed, src=0)
371
+ args.base_seed = base_seed[0]
372
+
373
+ logging.info(f"Input prompt: {args.prompt}")
374
+ img = None
375
+ if args.image is not None:
376
+ img = Image.open(args.image).convert("RGB")
377
+ logging.info(f"Input image: {args.image}")
378
+
379
+ # prompt extend
380
+ if args.use_prompt_extend:
381
+ logging.info("Extending prompt ...")
382
+ if rank == 0:
383
+ prompt_output = prompt_expander(
384
+ args.prompt,
385
+ image=img,
386
+ tar_lang=args.prompt_extend_target_lang,
387
+ seed=args.base_seed)
388
+ if prompt_output.status == False:
389
+ logging.info(
390
+ f"Extending prompt failed: {prompt_output.message}")
391
+ logging.info("Falling back to original prompt.")
392
+ input_prompt = args.prompt
393
+ else:
394
+ input_prompt = prompt_output.prompt
395
+ input_prompt = [input_prompt]
396
+ else:
397
+ input_prompt = [None]
398
+ if dist.is_initialized():
399
+ dist.broadcast_object_list(input_prompt, src=0)
400
+ args.prompt = input_prompt[0]
401
+ logging.info(f"Extended prompt: {args.prompt}")
402
+
403
+ if "t2v" in args.task:
404
+ logging.info("Creating TugraT2V pipeline.")
405
+ tugra_t2v = tugra.TugraT2V(
406
+ config=cfg,
407
+ checkpoint_dir=args.ckpt_dir,
408
+ device_id=device,
409
+ rank=rank,
410
+ t5_fsdp=args.t5_fsdp,
411
+ dit_fsdp=args.dit_fsdp,
412
+ use_sp=(args.ulysses_size > 1),
413
+ t5_cpu=args.t5_cpu,
414
+ convert_model_dtype=args.convert_model_dtype,
415
+ )
416
+
417
+ logging.info(f"Generating video ...")
418
+ video = tugra_t2v.generate(
419
+ args.prompt,
420
+ size=SIZE_CONFIGS[args.size],
421
+ frame_num=args.frame_num,
422
+ shift=args.sample_shift,
423
+ sample_solver=args.sample_solver,
424
+ sampling_steps=args.sample_steps,
425
+ guide_scale=args.sample_guide_scale,
426
+ seed=args.base_seed,
427
+ offload_model=args.offload_model)
428
+ elif "ti2v" in args.task:
429
+ logging.info("Creating TugraTI2V pipeline.")
430
+ tugra_ti2v = tugra.TugraTI2V(
431
+ config=cfg,
432
+ checkpoint_dir=args.ckpt_dir,
433
+ device_id=device,
434
+ rank=rank,
435
+ t5_fsdp=args.t5_fsdp,
436
+ dit_fsdp=args.dit_fsdp,
437
+ use_sp=(args.ulysses_size > 1),
438
+ t5_cpu=args.t5_cpu,
439
+ convert_model_dtype=args.convert_model_dtype,
440
+ )
441
+
442
+ logging.info(f"Generating video ...")
443
+ video = tugra_ti2v.generate(
444
+ args.prompt,
445
+ img=img,
446
+ size=SIZE_CONFIGS[args.size],
447
+ max_area=MAX_AREA_CONFIGS[args.size],
448
+ frame_num=args.frame_num,
449
+ shift=args.sample_shift,
450
+ sample_solver=args.sample_solver,
451
+ sampling_steps=args.sample_steps,
452
+ guide_scale=args.sample_guide_scale,
453
+ seed=args.base_seed,
454
+ offload_model=args.offload_model)
455
+ elif "animate" in args.task:
456
+ logging.info("Creating Tugra-Animate pipeline.")
457
+ tugra_animate = tugra.TugraAnimate(
458
+ config=cfg,
459
+ checkpoint_dir=args.ckpt_dir,
460
+ device_id=device,
461
+ rank=rank,
462
+ t5_fsdp=args.t5_fsdp,
463
+ dit_fsdp=args.dit_fsdp,
464
+ use_sp=(args.ulysses_size > 1),
465
+ t5_cpu=args.t5_cpu,
466
+ convert_model_dtype=args.convert_model_dtype,
467
+ use_relighting_lora=args.use_relighting_lora
468
+ )
469
+
470
+ logging.info(f"Generating video ...")
471
+ video = tugra_animate.generate(
472
+ src_root_path=args.src_root_path,
473
+ replace_flag=args.replace_flag,
474
+ refert_num = args.refert_num,
475
+ clip_len=args.frame_num,
476
+ shift=args.sample_shift,
477
+ sample_solver=args.sample_solver,
478
+ sampling_steps=args.sample_steps,
479
+ guide_scale=args.sample_guide_scale,
480
+ seed=args.base_seed,
481
+ offload_model=args.offload_model)
482
+ elif "s2v" in args.task:
483
+ logging.info("Creating TugraS2V pipeline.")
484
+ tugra_s2v = tugra.TugraS2V(
485
+ config=cfg,
486
+ checkpoint_dir=args.ckpt_dir,
487
+ device_id=device,
488
+ rank=rank,
489
+ t5_fsdp=args.t5_fsdp,
490
+ dit_fsdp=args.dit_fsdp,
491
+ use_sp=(args.ulysses_size > 1),
492
+ t5_cpu=args.t5_cpu,
493
+ convert_model_dtype=args.convert_model_dtype,
494
+ )
495
+ logging.info(f"Generating video ...")
496
+ video = tugra_s2v.generate(
497
+ input_prompt=args.prompt,
498
+ ref_image_path=args.image,
499
+ audio_path=args.audio,
500
+ enable_tts=args.enable_tts,
501
+ tts_prompt_audio=args.tts_prompt_audio,
502
+ tts_prompt_text=args.tts_prompt_text,
503
+ tts_text=args.tts_text,
504
+ num_repeat=args.num_clip,
505
+ pose_video=args.pose_video,
506
+ max_area=MAX_AREA_CONFIGS[args.size],
507
+ infer_frames=args.infer_frames,
508
+ shift=args.sample_shift,
509
+ sample_solver=args.sample_solver,
510
+ sampling_steps=args.sample_steps,
511
+ guide_scale=args.sample_guide_scale,
512
+ seed=args.base_seed,
513
+ offload_model=args.offload_model,
514
+ init_first_frame=args.start_from_ref,
515
+ )
516
+ else:
517
+ logging.info("Creating TugraI2V pipeline.")
518
+ tugra_i2v = tugra.TugraI2V(
519
+ config=cfg,
520
+ checkpoint_dir=args.ckpt_dir,
521
+ device_id=device,
522
+ rank=rank,
523
+ t5_fsdp=args.t5_fsdp,
524
+ dit_fsdp=args.dit_fsdp,
525
+ use_sp=(args.ulysses_size > 1),
526
+ t5_cpu=args.t5_cpu,
527
+ convert_model_dtype=args.convert_model_dtype,
528
+ )
529
+ logging.info("Generating video ...")
530
+ video = tugra_i2v.generate(
531
+ args.prompt,
532
+ img,
533
+ max_area=MAX_AREA_CONFIGS[args.size],
534
+ frame_num=args.frame_num,
535
+ shift=args.sample_shift,
536
+ sample_solver=args.sample_solver,
537
+ sampling_steps=args.sample_steps,
538
+ guide_scale=args.sample_guide_scale,
539
+ seed=args.base_seed,
540
+ offload_model=args.offload_model)
541
+
542
+ if rank == 0:
543
+ if args.save_file is None:
544
+ formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
545
+ formatted_prompt = args.prompt.replace(" ", "_").replace("/",
546
+ "_")[:50]
547
+ suffix = '.mp4'
548
+ 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
549
+
550
+ logging.info(f"Saving generated video to {args.save_file}")
551
+ save_video(
552
+ tensor=video[None],
553
+ save_file=args.save_file,
554
+ fps=cfg.sample_fps,
555
+ nrow=1,
556
+ normalize=True,
557
+ value_range=(-1, 1))
558
+ if "s2v" in args.task:
559
+ if args.enable_tts is False:
560
+ merge_video_audio(video_path=args.save_file, audio_path=args.audio)
561
+ else:
562
+ merge_video_audio(video_path=args.save_file, audio_path="tts.wav")
563
+ del video
564
+
565
+ torch.cuda.synchronize()
566
+ if dist.is_initialized():
567
+ dist.barrier()
568
+ dist.destroy_process_group()
569
+
570
+ logging.info("Finished.")
571
+
572
+
573
+ if __name__ == "__main__":
574
+ args = _parse_args()
575
+ generate(args)
pyproject.toml ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["setuptools>=61.0"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "tugra"
7
+ version = "2.2.0"
8
+ description = "TUĞRA: Teknolojik Ulusal Grafik ve Render Algoritması"
9
+ authors = [
10
+ { name = "TUĞRA Ekibi", email = "wan.ai@alibabacloud.com" }
11
+ ]
12
+ license = { file = "LICENSE.txt" }
13
+ readme = "README.md"
14
+ requires-python = ">=3.10,<4.0"
15
+ dependencies = [
16
+ "torch>=2.4.0",
17
+ "torchvision>=0.19.0",
18
+ "opencv-python>=4.9.0.80",
19
+ "diffusers>=0.31.0",
20
+ "transformers>=4.49.0",
21
+ "tokenizers>=0.20.3",
22
+ "accelerate>=1.1.1",
23
+ "tqdm",
24
+ "imageio",
25
+ "easydict",
26
+ "ftfy",
27
+ "dashscope",
28
+ "imageio-ffmpeg",
29
+ "flash_attn",
30
+ "numpy>=1.23.5,<2"
31
+ ]
32
+
33
+ [project.optional-dependencies]
34
+ dev = [
35
+ "pytest",
36
+ "black",
37
+ "flake8",
38
+ "isort",
39
+ "mypy",
40
+ "huggingface-hub[cli]"
41
+ ]
42
+
43
+ [project.urls]
44
+ homepage = "https://tugra.ai"
45
+ documentation = "https://github.com/Tugra-Video/TUĞRA"
46
+ repository = "https://github.com/Tugra-Video/TUĞRA"
47
+ huggingface = "https://huggingface.co/Tugra-AI/"
48
+ modelscope = "https://modelscope.cn/organization/Tugra-AI"
49
+ discord = "https://discord.gg/p5XbdQV7"
50
+
51
+ [tool.setuptools]
52
+ packages = ["tugra"]
53
+
54
+ [tool.setuptools.package-data]
55
+ "tugra" = ["**/*.py"]
56
+
57
+ [tool.black]
58
+ line-length = 88
59
+
60
+ [tool.isort]
61
+ profile = "black"
62
+
63
+ [tool.mypy]
64
+ strict = true
65
+
66
+
requirements.txt ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch>=2.4.0
2
+ torchvision>=0.19.0
3
+ torchaudio
4
+ opencv-python>=4.9.0.80
5
+ diffusers>=0.31.0
6
+ transformers>=4.49.0,<=4.51.3
7
+ tokenizers>=0.20.3
8
+ accelerate>=1.1.1
9
+ tqdm
10
+ imageio[ffmpeg]
11
+ easydict
12
+ ftfy
13
+ dashscope
14
+ imageio-ffmpeg
15
+ #flash_attn
16
+ numpy>=1.23.5,<2
requirements_animate.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ decord
2
+ peft
3
+ onnxruntime
4
+ pandas
5
+ matplotlib
6
+ -e git+https://github.com/facebookresearch/sam2.git@0e78a118995e66bb27d78518c4bd9a3e95b4e266#egg=SAM-2
7
+ loguru
8
+ sentencepiece
requirements_s2v.txt ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ openai-whisper
2
+ HyperPyYAML
3
+ onnxruntime
4
+ inflect
5
+ wetext
6
+ omegaconf
7
+ conformer
8
+ hydra-core
9
+ lightning
10
+ rich
11
+ gdown
12
+ matplotlib
13
+ wget
14
+ pyarrow
15
+ pyworld
16
+ librosa
17
+ decord
18
+ modelscope
19
+ GitPython
test.html ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <!DOCTYPE html>
2
+ <html lang="tr">
3
+
4
+ <head>
5
+ <meta charset="UTF-8">
6
+ <title>README Görüntüleyici</title>
7
+ <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
8
+ <style>
9
+ body {
10
+ font-family: sans-serif;
11
+ line-height: 1.6;
12
+ padding: 40px;
13
+ max-width: 800px;
14
+ margin: auto;
15
+ }
16
+
17
+ #content {
18
+ border: 1px solid #ddd;
19
+ padding: 20px;
20
+ border-radius: 8px;
21
+ }
22
+ </style>
23
+ </head>
24
+
25
+ <body>
26
+
27
+ <div id="content">Yükleniyor...</div>
28
+
29
+ <script>
30
+ // README.md dosyasının yolu
31
+ const readmePath = 'README.md';
32
+
33
+ fetch(readmePath)
34
+ .then(response => response.text())
35
+ .then(markdownText => {
36
+ // Markdown'ı HTML'e çevir ve ekrana bas
37
+ document.getElementById('content').innerHTML = marked.parse(markdownText);
38
+ })
39
+ .catch(err => {
40
+ document.getElementById('content').innerHTML = "Dosya yüklenirken bir hata oluştu.";
41
+ console.error(err);
42
+ });
43
+ </script>
44
+ </body>
45
+
46
+ </html>
tests/README.md ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+
2
+ 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.
3
+
4
+ ```bash
5
+ bash ./tests/test.sh <local model dir> <gpu number>
6
+ ```
tests/test.sh ADDED
@@ -0,0 +1,91 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ set -x
3
+
4
+ unset NCCL_DEBUG
5
+
6
+ if [ "$#" -eq 2 ]; then
7
+ MODEL_DIR=$(realpath "$1")
8
+ GPUS=$2
9
+ else
10
+ echo "Usage: $0 <local model dir> <gpu number>"
11
+ exit 1
12
+ fi
13
+
14
+ SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
15
+ REPO_ROOT="$(dirname "$SCRIPT_DIR")"
16
+ cd "$REPO_ROOT" || exit 1
17
+
18
+ PY_FILE=./generate.py
19
+
20
+
21
+ function t2v_A14B() {
22
+ CKPT_DIR="$MODEL_DIR/Wan2.2-T2V-A14B"
23
+
24
+ # # 1-GPU Test
25
+ # echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B 1-GPU Test: "
26
+ # python $PY_FILE --task t2v-A14B --size 480*832 --ckpt_dir $CKPT_DIR
27
+
28
+ # Multiple GPU Test
29
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B Multiple GPU Test: "
30
+ torchrun --nproc_per_node=$GPUS $PY_FILE --task t2v-A14B --ckpt_dir $CKPT_DIR --size 832*480 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
31
+
32
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B Multiple GPU Test: "
33
+ torchrun --nproc_per_node=$GPUS $PY_FILE --task t2v-A14B --ckpt_dir $CKPT_DIR --size 720*1280 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
34
+
35
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B Multiple GPU Test: "
36
+ torchrun --nproc_per_node=$GPUS $PY_FILE --task t2v-A14B --ckpt_dir $CKPT_DIR --size 1280*720 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
37
+
38
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> t2v_A14B Multiple GPU, prompt extend local_qwen: "
39
+ 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"
40
+ }
41
+
42
+
43
+ function i2v_A14B() {
44
+ CKPT_DIR="$MODEL_DIR/Wan2.2-I2V-A14B"
45
+
46
+ # echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B 1-GPU Test: "
47
+ # python $PY_FILE --task i2v-A14B --size 832*480 --ckpt_dir $CKPT_DIR
48
+
49
+ # Multiple GPU Test
50
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B Multiple GPU Test: "
51
+ torchrun --nproc_per_node=$GPUS $PY_FILE --task i2v-A14B --ckpt_dir $CKPT_DIR --size 832*480 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
52
+
53
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B Multiple GPU, prompt extend local_qwen: "
54
+ 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"
55
+
56
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B Multiple GPU, prompt extend local_qwen: "
57
+ 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"
58
+
59
+ if [ -n "${DASH_API_KEY+x}" ]; then
60
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> i2v_14B Multiple GPU, prompt extend dashscope: "
61
+ 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"
62
+ else
63
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> No DASH_API_KEY found, skip the dashscope extend test."
64
+ fi
65
+ }
66
+
67
+ function ti2v_5B() {
68
+ CKPT_DIR="$MODEL_DIR/Wan2.2-TI2V-5B"
69
+
70
+ # # 1-GPU Test
71
+ # echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B t2v 1-GPU Test: "
72
+ # python $PY_FILE --task ti2v-5B --size 1280*704 --ckpt_dir $CKPT_DIR
73
+
74
+ # Multiple GPU Test
75
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B t2v Multiple GPU Test: "
76
+ torchrun --nproc_per_node=$GPUS $PY_FILE --task ti2v-5B --ckpt_dir $CKPT_DIR --size 1280*704 --dit_fsdp --t5_fsdp --ulysses_size $GPUS
77
+
78
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B t2v Multiple GPU, prompt extend local_qwen: "
79
+ 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"
80
+
81
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B i2v Multiple GPU Test: "
82
+ 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"
83
+
84
+ echo -e "\n\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ti2v_5B i2v Multiple GPU, prompt extend local_qwen: "
85
+ 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"
86
+
87
+ }
88
+
89
+ t2v_A14B
90
+ i2v_A14B
91
+ ti2v_5B
tugra/__init__.py ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 KAYRA. All rights reserved.d.
2
+ from . import configs, distributed, modules
3
+ from .image2video import TugraI2V
4
+ from .speech2video import TugraS2V
5
+ from .text2video import TugraT2V
6
+ from .textimage2video import TugraTI2V
7
+ from .animate import TugraAnimate
tugra/__pycache__/__init__.cpython-314.pyc ADDED
Binary file (462 Bytes). View file
 
tugra/animate.py ADDED
@@ -0,0 +1,648 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 KAYRA. All rights reserved.
2
+ import logging
3
+ import math
4
+ import os
5
+ import cv2
6
+ import types
7
+ from copy import deepcopy
8
+ from functools import partial
9
+ from einops import rearrange
10
+ import numpy as np
11
+ import torch
12
+
13
+ import torch.distributed as dist
14
+ from peft import set_peft_model_state_dict
15
+ from decord import VideoReader
16
+ from tqdm import tqdm
17
+ import torch.nn.functional as F
18
+ from .distributed.fsdp import shard_model
19
+ from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward
20
+ from .distributed.util import get_world_size
21
+
22
+ from .modules.animate import WanAnimateModel
23
+ from .modules.animate import CLIPModel
24
+ from .modules.t5 import T5EncoderModel
25
+ from .modules.vae2_1 import Wan2_1_VAE
26
+ from .modules.animate.animate_utils import TensorList, get_loraconfig
27
+ from .utils.fm_solvers import (
28
+ FlowDPMSolverMultistepScheduler,
29
+ get_sampling_sigmas,
30
+ retrieve_timesteps,
31
+ )
32
+ from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
33
+
34
+
35
+
36
+ class WanAnimate:
37
+
38
+ def __init__(
39
+ self,
40
+ config,
41
+ checkpoint_dir,
42
+ device_id=0,
43
+ rank=0,
44
+ t5_fsdp=False,
45
+ dit_fsdp=False,
46
+ use_sp=False,
47
+ t5_cpu=False,
48
+ init_on_cpu=True,
49
+ convert_model_dtype=False,
50
+ use_relighting_lora=False
51
+ ):
52
+ r"""
53
+ Initializes the generation model components.
54
+
55
+ Args:
56
+ config (EasyDict):
57
+ Object containing model parameters initialized from config.py
58
+ checkpoint_dir (`str`):
59
+ Path to directory containing model checkpoints
60
+ device_id (`int`, *optional*, defaults to 0):
61
+ Id of target GPU device
62
+ rank (`int`, *optional*, defaults to 0):
63
+ Process rank for distributed training
64
+ t5_fsdp (`bool`, *optional*, defaults to False):
65
+ Enable FSDP sharding for T5 model
66
+ dit_fsdp (`bool`, *optional*, defaults to False):
67
+ Enable FSDP sharding for DiT model
68
+ use_sp (`bool`, *optional*, defaults to False):
69
+ Enable distribution strategy of sequence parallel.
70
+ t5_cpu (`bool`, *optional*, defaults to False):
71
+ Whether to place T5 model on CPU. Only works without t5_fsdp.
72
+ init_on_cpu (`bool`, *optional*, defaults to True):
73
+ Enable initializing Transformer Model on CPU. Only works without FSDP or USP.
74
+ convert_model_dtype (`bool`, *optional*, defaults to False):
75
+ Convert DiT model parameters dtype to 'config.param_dtype'.
76
+ Only works without FSDP.
77
+ use_relighting_lora (`bool`, *optional*, defaults to False):
78
+ Whether to use relighting lora for character replacement.
79
+ """
80
+ self.device = torch.device(f"cuda:{device_id}")
81
+ self.config = config
82
+ self.rank = rank
83
+ self.t5_cpu = t5_cpu
84
+ self.init_on_cpu = init_on_cpu
85
+
86
+ self.num_train_timesteps = config.num_train_timesteps
87
+ self.param_dtype = config.param_dtype
88
+
89
+ if t5_fsdp or dit_fsdp or use_sp:
90
+ self.init_on_cpu = False
91
+
92
+ shard_fn = partial(shard_model, device_id=device_id)
93
+ self.text_encoder = T5EncoderModel(
94
+ text_len=config.text_len,
95
+ dtype=config.t5_dtype,
96
+ device=torch.device('cpu'),
97
+ checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
98
+ tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
99
+ shard_fn=shard_fn if t5_fsdp else None,
100
+ )
101
+
102
+ self.clip = CLIPModel(
103
+ dtype=torch.float16,
104
+ device=self.device,
105
+ checkpoint_path=os.path.join(checkpoint_dir,
106
+ config.clip_checkpoint),
107
+ tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
108
+
109
+ self.vae = Wan2_1_VAE(
110
+ vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
111
+ device=self.device)
112
+
113
+ logging.info(f"Creating WanAnimate from {checkpoint_dir}")
114
+
115
+ if not dit_fsdp:
116
+ self.noise_model = WanAnimateModel.from_pretrained(
117
+ checkpoint_dir,
118
+ torch_dtype=self.param_dtype,
119
+ device_map=self.device)
120
+ else:
121
+ self.noise_model = WanAnimateModel.from_pretrained(
122
+ checkpoint_dir, torch_dtype=self.param_dtype)
123
+
124
+ self.noise_model = self._configure_model(
125
+ model=self.noise_model,
126
+ use_sp=use_sp,
127
+ dit_fsdp=dit_fsdp,
128
+ shard_fn=shard_fn,
129
+ convert_model_dtype=convert_model_dtype,
130
+ use_lora=use_relighting_lora,
131
+ checkpoint_dir=checkpoint_dir,
132
+ config=config
133
+ )
134
+
135
+ if use_sp:
136
+ self.sp_size = get_world_size()
137
+ else:
138
+ self.sp_size = 1
139
+
140
+ self.sample_neg_prompt = config.sample_neg_prompt
141
+ self.sample_prompt = config.prompt
142
+
143
+
144
+ def _configure_model(self, model, use_sp, dit_fsdp, shard_fn,
145
+ convert_model_dtype, use_lora, checkpoint_dir, config):
146
+ """
147
+ Configures a model object. This includes setting evaluation modes,
148
+ applying distributed parallel strategy, and handling device placement.
149
+
150
+ Args:
151
+ model (torch.nn.Module):
152
+ The model instance to configure.
153
+ use_sp (`bool`):
154
+ Enable distribution strategy of sequence parallel.
155
+ dit_fsdp (`bool`):
156
+ Enable FSDP sharding for DiT model.
157
+ shard_fn (callable):
158
+ The function to apply FSDP sharding.
159
+ convert_model_dtype (`bool`):
160
+ Convert DiT model parameters dtype to 'config.param_dtype'.
161
+ Only works without FSDP.
162
+
163
+ Returns:
164
+ torch.nn.Module:
165
+ The configured model.
166
+ """
167
+ model.eval().requires_grad_(False)
168
+
169
+ if use_sp:
170
+ for block in model.blocks:
171
+ block.self_attn.forward = types.MethodType(
172
+ sp_attn_forward, block.self_attn)
173
+
174
+ model.use_context_parallel = True
175
+
176
+ if dist.is_initialized():
177
+ dist.barrier()
178
+
179
+ if use_lora:
180
+ logging.info("Loading Relighting Lora. ")
181
+ lora_config = get_loraconfig(
182
+ transformer=model,
183
+ rank=128,
184
+ alpha=128
185
+ )
186
+ model.add_adapter(lora_config)
187
+ lora_path = os.path.join(checkpoint_dir, config.lora_checkpoint)
188
+ peft_state_dict = torch.load(lora_path)["state_dict"]
189
+ set_peft_model_state_dict(model, peft_state_dict)
190
+
191
+ if dit_fsdp:
192
+ model = shard_fn(model, use_lora=use_lora)
193
+ else:
194
+ if convert_model_dtype:
195
+ model.to(self.param_dtype)
196
+ if not self.init_on_cpu:
197
+ model.to(self.device)
198
+
199
+ return model
200
+
201
+ def inputs_padding(self, array, target_len):
202
+ idx = 0
203
+ flip = False
204
+ target_array = []
205
+ while len(target_array) < target_len:
206
+ target_array.append(deepcopy(array[idx]))
207
+ if flip:
208
+ idx -= 1
209
+ else:
210
+ idx += 1
211
+ if idx == 0 or idx == len(array) - 1:
212
+ flip = not flip
213
+ return target_array[:target_len]
214
+
215
+ def get_valid_len(self, real_len, clip_len=81, overlap=1):
216
+ real_clip_len = clip_len - overlap
217
+ last_clip_num = (real_len - overlap) % real_clip_len
218
+ if last_clip_num == 0:
219
+ extra = 0
220
+ else:
221
+ extra = real_clip_len - last_clip_num
222
+ target_len = real_len + extra
223
+ return target_len
224
+
225
+
226
+ def get_i2v_mask(self, lat_t, lat_h, lat_w, mask_len=1, mask_pixel_values=None, device="cuda"):
227
+ if mask_pixel_values is None:
228
+ msk = torch.zeros(1, (lat_t-1) * 4 + 1, lat_h, lat_w, device=device)
229
+ else:
230
+ msk = mask_pixel_values.clone()
231
+ msk[:, :mask_len] = 1
232
+ msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1)
233
+ msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
234
+ msk = msk.transpose(1, 2)[0]
235
+ return msk
236
+
237
+ def padding_resize(self, img_ori, height=512, width=512, padding_color=(0, 0, 0), interpolation=cv2.INTER_LINEAR):
238
+ ori_height = img_ori.shape[0]
239
+ ori_width = img_ori.shape[1]
240
+ channel = img_ori.shape[2]
241
+
242
+ img_pad = np.zeros((height, width, channel))
243
+ if channel == 1:
244
+ img_pad[:, :, 0] = padding_color[0]
245
+ else:
246
+ img_pad[:, :, 0] = padding_color[0]
247
+ img_pad[:, :, 1] = padding_color[1]
248
+ img_pad[:, :, 2] = padding_color[2]
249
+
250
+ if (ori_height / ori_width) > (height / width):
251
+ new_width = int(height / ori_height * ori_width)
252
+ img = cv2.resize(img_ori, (new_width, height), interpolation=interpolation)
253
+ padding = int((width - new_width) / 2)
254
+ if len(img.shape) == 2:
255
+ img = img[:, :, np.newaxis]
256
+ img_pad[:, padding: padding + new_width, :] = img
257
+ else:
258
+ new_height = int(width / ori_width * ori_height)
259
+ img = cv2.resize(img_ori, (width, new_height), interpolation=interpolation)
260
+ padding = int((height - new_height) / 2)
261
+ if len(img.shape) == 2:
262
+ img = img[:, :, np.newaxis]
263
+ img_pad[padding: padding + new_height, :, :] = img
264
+
265
+ img_pad = np.uint8(img_pad)
266
+
267
+ return img_pad
268
+
269
+ def prepare_source(self, src_pose_path, src_face_path, src_ref_path):
270
+ pose_video_reader = VideoReader(src_pose_path)
271
+ pose_len = len(pose_video_reader)
272
+ pose_idxs = list(range(pose_len))
273
+ cond_images = pose_video_reader.get_batch(pose_idxs).asnumpy()
274
+
275
+ face_video_reader = VideoReader(src_face_path)
276
+ face_len = len(face_video_reader)
277
+ face_idxs = list(range(face_len))
278
+ face_images = face_video_reader.get_batch(face_idxs).asnumpy()
279
+ height, width = cond_images[0].shape[:2]
280
+ refer_images = cv2.imread(src_ref_path)[..., ::-1]
281
+ refer_images = self.padding_resize(refer_images, height=height, width=width)
282
+ return cond_images, face_images, refer_images
283
+
284
+ def prepare_source_for_replace(self, src_bg_path, src_mask_path):
285
+ bg_video_reader = VideoReader(src_bg_path)
286
+ bg_len = len(bg_video_reader)
287
+ bg_idxs = list(range(bg_len))
288
+ bg_images = bg_video_reader.get_batch(bg_idxs).asnumpy()
289
+
290
+ mask_video_reader = VideoReader(src_mask_path)
291
+ mask_len = len(mask_video_reader)
292
+ mask_idxs = list(range(mask_len))
293
+ mask_images = mask_video_reader.get_batch(mask_idxs).asnumpy()
294
+ mask_images = mask_images[:, :, :, 0] / 255
295
+ return bg_images, mask_images
296
+
297
+ def generate(
298
+ self,
299
+ src_root_path,
300
+ replace_flag=False,
301
+ clip_len=77,
302
+ refert_num=1,
303
+ shift=5.0,
304
+ sample_solver='dpm++',
305
+ sampling_steps=20,
306
+ guide_scale=1,
307
+ input_prompt="",
308
+ n_prompt="",
309
+ seed=-1,
310
+ offload_model=True,
311
+ ):
312
+ r"""
313
+ Generates video frames from input image using diffusion process.
314
+
315
+ Args:
316
+ src_root_path ('str'):
317
+ Process output path
318
+ replace_flag (`bool`, *optional*, defaults to False):
319
+ Whether to use character replace.
320
+ clip_len (`int`, *optional*, defaults to 77):
321
+ How many frames to generate per clips. The number should be 4n+1
322
+ refert_num (`int`, *optional*, defaults to 1):
323
+ How many frames used for temporal guidance. Recommended to be 1 or 5.
324
+ shift (`float`, *optional*, defaults to 5.0):
325
+ Noise schedule shift parameter.
326
+ sample_solver (`str`, *optional*, defaults to 'dpm++'):
327
+ Solver used to sample the video.
328
+ sampling_steps (`int`, *optional*, defaults to 20):
329
+ Number of diffusion sampling steps. Higher values improve quality but slow generation
330
+ guide_scale (`float` or tuple[`float`], *optional*, defaults 1.0):
331
+ Classifier-free guidance scale. We only use it for expression control.
332
+ In most cases, it's not necessary and faster generation can be achieved without it.
333
+ When expression adjustments are needed, you may consider using this feature.
334
+ input_prompt (`str`):
335
+ Text prompt for content generation. We don't recommend custom prompts (although they work)
336
+ n_prompt (`str`, *optional*, defaults to ""):
337
+ Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
338
+ seed (`int`, *optional*, defaults to -1):
339
+ Random seed for noise generation. If -1, use random seed
340
+ offload_model (`bool`, *optional*, defaults to True):
341
+ If True, offloads models to CPU during generation to save VRAM
342
+
343
+ Returns:
344
+ torch.Tensor:
345
+ Generated video frames tensor. Dimensions: (C, N, H, W) where:
346
+ - C: Color channels (3 for RGB)
347
+ - N: Number of frames
348
+ - H: Frame height
349
+ - W: Frame width
350
+ """
351
+ assert refert_num == 1 or refert_num == 5, "refert_num should be 1 or 5."
352
+
353
+ seed_g = torch.Generator(device=self.device)
354
+ seed_g.manual_seed(seed)
355
+
356
+ if n_prompt == "":
357
+ n_prompt = self.sample_neg_prompt
358
+
359
+ if input_prompt == "":
360
+ input_prompt = self.sample_prompt
361
+
362
+ src_pose_path = os.path.join(src_root_path, "src_pose.mp4")
363
+ src_face_path = os.path.join(src_root_path, "src_face.mp4")
364
+ src_ref_path = os.path.join(src_root_path, "src_ref.png")
365
+
366
+ 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)
367
+
368
+ if not self.t5_cpu:
369
+ self.text_encoder.model.to(self.device)
370
+ context = self.text_encoder([input_prompt], self.device)
371
+ context_null = self.text_encoder([n_prompt], self.device)
372
+ if offload_model:
373
+ self.text_encoder.model.cpu()
374
+ else:
375
+ context = self.text_encoder([input_prompt], torch.device('cpu'))
376
+ context_null = self.text_encoder([n_prompt], torch.device('cpu'))
377
+ context = [t.to(self.device) for t in context]
378
+ context_null = [t.to(self.device) for t in context_null]
379
+
380
+ real_frame_len = len(cond_images)
381
+ target_len = self.get_valid_len(real_frame_len, clip_len, overlap=refert_num)
382
+ logging.info('real frames: {} target frames: {}'.format(real_frame_len, target_len))
383
+ cond_images = self.inputs_padding(cond_images, target_len)
384
+ face_images = self.inputs_padding(face_images, target_len)
385
+
386
+ if replace_flag:
387
+ src_bg_path = os.path.join(src_root_path, "src_bg.mp4")
388
+ src_mask_path = os.path.join(src_root_path, "src_mask.mp4")
389
+ bg_images, mask_images = self.prepare_source_for_replace(src_bg_path, src_mask_path)
390
+ bg_images = self.inputs_padding(bg_images, target_len)
391
+ mask_images = self.inputs_padding(mask_images, target_len)
392
+
393
+ height, width = refer_images.shape[:2]
394
+ start = 0
395
+ end = clip_len
396
+ all_out_frames = []
397
+ while True:
398
+ if start + refert_num >= len(cond_images):
399
+ break
400
+
401
+ if start == 0:
402
+ mask_reft_len = 0
403
+ else:
404
+ mask_reft_len = refert_num
405
+
406
+ batch = {
407
+ "conditioning_pixel_values": torch.zeros(1, 3, clip_len, height, width),
408
+ "bg_pixel_values": torch.zeros(1, 3, clip_len, height, width),
409
+ "mask_pixel_values": torch.zeros(1, 1, clip_len, height, width),
410
+ "face_pixel_values": torch.zeros(1, 3, clip_len, 512, 512),
411
+ "refer_pixel_values": torch.zeros(1, 3, height, width),
412
+ "refer_t_pixel_values": torch.zeros(refert_num, 3, height, width)
413
+ }
414
+
415
+ batch["conditioning_pixel_values"] = rearrange(
416
+ torch.tensor(np.stack(cond_images[start:end]) / 127.5 - 1),
417
+ "t h w c -> 1 c t h w",
418
+ )
419
+ batch["face_pixel_values"] = rearrange(
420
+ torch.tensor(np.stack(face_images[start:end]) / 127.5 - 1),
421
+ "t h w c -> 1 c t h w",
422
+ )
423
+
424
+ batch["refer_pixel_values"] = rearrange(
425
+ torch.tensor(refer_images / 127.5 - 1), "h w c -> 1 c h w"
426
+ )
427
+
428
+ if start > 0:
429
+ batch["refer_t_pixel_values"] = rearrange(
430
+ out_frames[0, :, -refert_num:].clone().detach(),
431
+ "c t h w -> t c h w",
432
+ )
433
+
434
+ batch["refer_t_pixel_values"] = rearrange(batch["refer_t_pixel_values"],
435
+ "t c h w -> 1 c t h w",
436
+ )
437
+
438
+ if replace_flag:
439
+ batch["bg_pixel_values"] = rearrange(
440
+ torch.tensor(np.stack(bg_images[start:end]) / 127.5 - 1),
441
+ "t h w c -> 1 c t h w",
442
+ )
443
+
444
+ batch["mask_pixel_values"] = rearrange(
445
+ torch.tensor(np.stack(mask_images[start:end])[:, :, :, None]),
446
+ "t h w c -> 1 t c h w",
447
+ )
448
+
449
+
450
+ for key, value in batch.items():
451
+ if isinstance(value, torch.Tensor):
452
+ batch[key] = value.to(device=self.device, dtype=torch.bfloat16)
453
+
454
+ ref_pixel_values = batch["refer_pixel_values"]
455
+ refer_t_pixel_values = batch["refer_t_pixel_values"]
456
+ conditioning_pixel_values = batch["conditioning_pixel_values"]
457
+ face_pixel_values = batch["face_pixel_values"]
458
+
459
+ B, _, H, W = ref_pixel_values.shape
460
+ T = clip_len
461
+ lat_h = H // 8
462
+ lat_w = W // 8
463
+ lat_t = T // 4 + 1
464
+ target_shape = [lat_t + 1, lat_h, lat_w]
465
+ noise = [
466
+ torch.randn(
467
+ 16,
468
+ target_shape[0],
469
+ target_shape[1],
470
+ target_shape[2],
471
+ dtype=torch.float32,
472
+ device=self.device,
473
+ generator=seed_g,
474
+ )
475
+ ]
476
+
477
+ max_seq_len = int(math.ceil(np.prod(target_shape) // 4 / self.sp_size)) * self.sp_size
478
+ if max_seq_len % self.sp_size != 0:
479
+ raise ValueError(f"max_seq_len {max_seq_len} is not divisible by sp_size {self.sp_size}")
480
+
481
+ with (
482
+ torch.autocast(device_type=str(self.device), dtype=torch.bfloat16, enabled=True),
483
+ torch.no_grad()
484
+ ):
485
+ if sample_solver == 'unipc':
486
+ sample_scheduler = FlowUniPCMultistepScheduler(
487
+ num_train_timesteps=self.num_train_timesteps,
488
+ shift=1,
489
+ use_dynamic_shifting=False)
490
+ sample_scheduler.set_timesteps(
491
+ sampling_steps, device=self.device, shift=shift)
492
+ timesteps = sample_scheduler.timesteps
493
+ elif sample_solver == 'dpm++':
494
+ sample_scheduler = FlowDPMSolverMultistepScheduler(
495
+ num_train_timesteps=self.num_train_timesteps,
496
+ shift=1,
497
+ use_dynamic_shifting=False)
498
+ sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
499
+ timesteps, _ = retrieve_timesteps(
500
+ sample_scheduler,
501
+ device=self.device,
502
+ sigmas=sampling_sigmas)
503
+ else:
504
+ raise NotImplementedError("Unsupported solver.")
505
+
506
+ latents = noise
507
+
508
+ pose_latents_no_ref = self.vae.encode(conditioning_pixel_values.to(torch.bfloat16))
509
+ pose_latents_no_ref = torch.stack(pose_latents_no_ref)
510
+ pose_latents = torch.cat([pose_latents_no_ref], dim=2)
511
+
512
+ ref_pixel_values = rearrange(ref_pixel_values, "t c h w -> 1 c t h w")
513
+ ref_latents = self.vae.encode(ref_pixel_values.to(torch.bfloat16))
514
+ ref_latents = torch.stack(ref_latents)
515
+
516
+ mask_ref = self.get_i2v_mask(1, lat_h, lat_w, 1, device=self.device)
517
+ y_ref = torch.concat([mask_ref, ref_latents[0]]).to(dtype=torch.bfloat16, device=self.device)
518
+
519
+ img = ref_pixel_values[0, :, 0]
520
+ clip_context = self.clip.visual([img[:, None, :, :]]).to(dtype=torch.bfloat16, device=self.device)
521
+
522
+ if mask_reft_len > 0:
523
+ if replace_flag:
524
+ bg_pixel_values = batch["bg_pixel_values"]
525
+ y_reft = self.vae.encode(
526
+ [
527
+ torch.concat([refer_t_pixel_values[0, :, :mask_reft_len], bg_pixel_values[0, :, mask_reft_len:]], dim=1).to(self.device)
528
+ ]
529
+ )[0]
530
+ mask_pixel_values = 1 - batch["mask_pixel_values"]
531
+ mask_pixel_values = rearrange(mask_pixel_values, "b t c h w -> (b t) c h w")
532
+ mask_pixel_values = F.interpolate(mask_pixel_values, size=(H//8, W//8), mode='nearest')
533
+ mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b t c h w", b=1)[:,:,0]
534
+ msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, mask_pixel_values=mask_pixel_values, device=self.device)
535
+ else:
536
+ y_reft = self.vae.encode(
537
+ [
538
+ torch.concat(
539
+ [
540
+ torch.nn.functional.interpolate(refer_t_pixel_values[0, :, :mask_reft_len].cpu(),
541
+ size=(H, W), mode="bicubic"),
542
+ torch.zeros(3, T - mask_reft_len, H, W),
543
+ ],
544
+ dim=1,
545
+ ).to(self.device)
546
+ ]
547
+ )[0]
548
+ msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, device=self.device)
549
+ else:
550
+ if replace_flag:
551
+ bg_pixel_values = batch["bg_pixel_values"]
552
+ mask_pixel_values = 1 - batch["mask_pixel_values"]
553
+ mask_pixel_values = rearrange(mask_pixel_values, "b t c h w -> (b t) c h w")
554
+ mask_pixel_values = F.interpolate(mask_pixel_values, size=(H//8, W//8), mode='nearest')
555
+ mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b t c h w", b=1)[:,:,0]
556
+ y_reft = self.vae.encode(
557
+ [
558
+ torch.concat(
559
+ [
560
+ bg_pixel_values[0],
561
+ ],
562
+ dim=1,
563
+ ).to(self.device)
564
+ ]
565
+ )[0]
566
+ msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, mask_pixel_values=mask_pixel_values, device=self.device)
567
+ else:
568
+ y_reft = self.vae.encode(
569
+ [
570
+ torch.concat(
571
+ [
572
+ torch.zeros(3, T - mask_reft_len, H, W),
573
+ ],
574
+ dim=1,
575
+ ).to(self.device)
576
+ ]
577
+ )[0]
578
+ msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, device=self.device)
579
+
580
+ y_reft = torch.concat([msk_reft, y_reft]).to(dtype=torch.bfloat16, device=self.device)
581
+ y = torch.concat([y_ref, y_reft], dim=1)
582
+
583
+ arg_c = {
584
+ "context": context,
585
+ "seq_len": max_seq_len,
586
+ "clip_fea": clip_context.to(dtype=torch.bfloat16, device=self.device),
587
+ "y": [y],
588
+ "pose_latents": pose_latents,
589
+ "face_pixel_values": face_pixel_values,
590
+ }
591
+
592
+ if guide_scale > 1:
593
+ face_pixel_values_uncond = face_pixel_values * 0 - 1
594
+ arg_null = {
595
+ "context": context_null,
596
+ "seq_len": max_seq_len,
597
+ "clip_fea": clip_context.to(dtype=torch.bfloat16, device=self.device),
598
+ "y": [y],
599
+ "pose_latents": pose_latents,
600
+ "face_pixel_values": face_pixel_values_uncond,
601
+ }
602
+
603
+ for i, t in enumerate(tqdm(timesteps)):
604
+ latent_model_input = latents
605
+ timestep = [t]
606
+
607
+ timestep = torch.stack(timestep)
608
+
609
+ noise_pred_cond = TensorList(
610
+ self.noise_model(TensorList(latent_model_input), t=timestep, **arg_c)
611
+ )
612
+
613
+ if guide_scale > 1:
614
+ noise_pred_uncond = TensorList(
615
+ self.noise_model(
616
+ TensorList(latent_model_input), t=timestep, **arg_null
617
+ )
618
+ )
619
+ noise_pred = noise_pred_uncond + guide_scale * (
620
+ noise_pred_cond - noise_pred_uncond
621
+ )
622
+ else:
623
+ noise_pred = noise_pred_cond
624
+
625
+ temp_x0 = sample_scheduler.step(
626
+ noise_pred[0].unsqueeze(0),
627
+ t,
628
+ latents[0].unsqueeze(0),
629
+ return_dict=False,
630
+ generator=seed_g,
631
+ )[0]
632
+ latents[0] = temp_x0.squeeze(0)
633
+
634
+ x0 = latents
635
+
636
+ x0 = [x.to(dtype=torch.float32) for x in x0]
637
+ out_frames = torch.stack(self.vae.decode([x0[0][:, 1:]]))
638
+
639
+ if start != 0:
640
+ out_frames = out_frames[:, :, refert_num:]
641
+
642
+ all_out_frames.append(out_frames.cpu())
643
+
644
+ start += clip_len - refert_num
645
+ end += clip_len - refert_num
646
+
647
+ videos = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len]
648
+ return videos[0] if self.rank == 0 else None
tugra/configs/__init__.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024-2025 The Alibaba TUĞRA Team Authors. All rights reserved.
2
+ import copy
3
+ import os
4
+
5
+ os.environ['TOKENIZERS_PARALLELISM'] = 'false'
6
+
7
+ from .tugra_i2v_A14B import i2v_A14B
8
+ from .tugra_s2v_14B import s2v_14B
9
+ from .tugra_t2v_A14B import t2v_A14B
10
+ from .tugra_ti2v_5B import ti2v_5B
11
+ from .tugra_animate_14B import animate_14B
12
+ from .shared_config import tugra_shared_cfg
13
+
14
+ TUGRA_CONFIGS = {
15
+ 't2v-A14B': t2v_A14B,
16
+ 'i2v-A14B': i2v_A14B,
17
+ 'ti2v-5B': ti2v_5B,
18
+ 'animate-14B': animate_14B,
19
+ 's2v-14B': s2v_14B,
20
+ }
21
+
22
+ SIZE_CONFIGS = {
23
+ '720*1280': (720, 1280),
24
+ '1280*720': (1280, 720),
25
+ '480*832': (480, 832),
26
+ '832*480': (832, 480),
27
+ '704*1280': (704, 1280),
28
+ '1280*704': (1280, 704),
29
+ '1024*704': (1024, 704),
30
+ '704*1024': (704, 1024),
31
+ }
32
+
33
+ MAX_AREA_CONFIGS = {
34
+ '720*1280': 720 * 1280,
35
+ '1280*720': 1280 * 720,
36
+ '480*832': 480 * 832,
37
+ '832*480': 832 * 480,
38
+ '704*1280': 704 * 1280,
39
+ '1280*704': 1280 * 704,
40
+ '1024*704': 1024 * 704,
41
+ '704*1024': 704 * 1024,
42
+ }
43
+
44
+ SUPPORTED_SIZES = {
45
+ 't2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'),
46
+ 'i2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'),
47
+ 'ti2v-5B': ('704*1280', '1280*704'),
48
+ 's2v-14B': ('720*1280', '1280*720', '480*832', '832*480', '1024*704',
49
+ '704*1024', '704*1280', '1280*704'),
50
+ 'animate-14B': ('720*1280', '1280*720')
51
+ }
tugra/configs/__pycache__/__init__.cpython-314.pyc ADDED
Binary file (1.27 kB). View file
 
tugra/configs/__pycache__/shared_config.cpython-314.pyc ADDED
Binary file (1.06 kB). View file
 
tugra/configs/__pycache__/tugra_animate_14B.cpython-314.pyc ADDED
Binary file (1.58 kB). View file
 
tugra/configs/__pycache__/tugra_i2v_A14B.cpython-314.pyc ADDED
Binary file (1.34 kB). View file
 
tugra/configs/__pycache__/tugra_s2v_14B.cpython-314.pyc ADDED
Binary file (3 kB). View file
 
tugra/configs/__pycache__/tugra_t2v_A14B.cpython-314.pyc ADDED
Binary file (1.32 kB). View file
 
tugra/configs/__pycache__/tugra_ti2v_5B.cpython-314.pyc ADDED
Binary file (1.19 kB). View file
 
tugra/configs/__pycache__/wan_i2v_A14B.cpython-314.pyc ADDED
Binary file (1.32 kB). View file
 
tugra/configs/shared_config.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024-2025 The Alibaba TUĞRA Team Authors. All rights reserved.
2
+ import torch
3
+ from easydict import EasyDict
4
+
5
+ #------------------------ TUĞRA shared config ------------------------#
6
+ tugra_shared_cfg = EasyDict()
7
+
8
+ # t5
9
+ tugra_shared_cfg.t5_model = 'umt5_xxl'
10
+ tugra_shared_cfg.t5_dtype = torch.bfloat16
11
+ tugra_shared_cfg.text_len = 512
12
+
13
+ # transformer
14
+ tugra_shared_cfg.param_dtype = torch.bfloat16
15
+
16
+ # inference
17
+ tugra_shared_cfg.num_train_timesteps = 1000
18
+ tugra_shared_cfg.sample_fps = 16
19
+ tugra_shared_cfg.sample_neg_prompt = '色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走'
20
+ tugra_shared_cfg.frame_num = 81
tugra/configs/tugra_animate_14B.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 KAYRA. All rights reserved.
2
+ # pyrefly: ignore [missing-import]
3
+ from easydict import EasyDict
4
+
5
+ from .shared_config import tugra_shared_cfg
6
+
7
+ #------------------------ Wan animate 14B ------------------------#
8
+ animate_14B = EasyDict(__name__='Config: Wan animate 14B')
9
+ animate_14B.update(tugra_shared_cfg)
10
+
11
+ animate_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
12
+ animate_14B.t5_tokenizer = 'google/umt5-xxl'
13
+
14
+ animate_14B.clip_checkpoint = 'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
15
+ animate_14B.clip_tokenizer = 'xlm-roberta-large'
16
+ animate_14B.lora_checkpoint = 'relighting_lora.ckpt'
17
+ # vae
18
+ animate_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
19
+ animate_14B.vae_stride = (4, 8, 8)
20
+
21
+ # transformer
22
+ animate_14B.patch_size = (1, 2, 2)
23
+ animate_14B.dim = 5120
24
+ animate_14B.ffn_dim = 13824
25
+ animate_14B.freq_dim = 256
26
+ animate_14B.num_heads = 40
27
+ animate_14B.num_layers = 40
28
+ animate_14B.window_size = (-1, -1)
29
+ animate_14B.qk_norm = True
30
+ animate_14B.cross_attn_norm = True
31
+ animate_14B.eps = 1e-6
32
+ animate_14B.use_face_encoder = True
33
+ animate_14B.motion_encoder_dim = 512
34
+
35
+ # inference
36
+ animate_14B.sample_shift = 5.0
37
+ animate_14B.sample_steps = 20
38
+ animate_14B.sample_guide_scale = 1.0
39
+ animate_14B.frame_num = 77
40
+ animate_14B.sample_fps = 30
41
+ animate_14B.prompt = '视频中的人在做动作'
tugra/configs/tugra_i2v_A14B.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024-2025 The Alibaba TUĞRA Team Authors. All rights reserved.
2
+ import torch
3
+ from easydict import EasyDict
4
+
5
+ from .shared_config import tugra_shared_cfg
6
+
7
+ #------------------------ TUĞRA I2V A14B ------------------------#
8
+
9
+ i2v_A14B = EasyDict(__name__='Config: TUĞRA I2V A14B')
10
+ i2v_A14B.update(tugra_shared_cfg)
11
+
12
+ i2v_A14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
13
+ i2v_A14B.t5_tokenizer = 'google/umt5-xxl'
14
+
15
+ # vae
16
+ i2v_A14B.vae_checkpoint = 'Tugra2.1_VAE.pth'
17
+ i2v_A14B.vae_stride = (4, 8, 8)
18
+
19
+ # transformer
20
+ i2v_A14B.patch_size = (1, 2, 2)
21
+ i2v_A14B.dim = 5120
22
+ i2v_A14B.ffn_dim = 13824
23
+ i2v_A14B.freq_dim = 256
24
+ i2v_A14B.num_heads = 40
25
+ i2v_A14B.num_layers = 40
26
+ i2v_A14B.window_size = (-1, -1)
27
+ i2v_A14B.qk_norm = True
28
+ i2v_A14B.cross_attn_norm = True
29
+ i2v_A14B.eps = 1e-6
30
+ i2v_A14B.low_noise_checkpoint = 'low_noise_model'
31
+ i2v_A14B.high_noise_checkpoint = 'high_noise_model'
32
+
33
+ # inference
34
+ i2v_A14B.sample_shift = 5.0
35
+ i2v_A14B.sample_steps = 40
36
+ i2v_A14B.boundary = 0.900
37
+ i2v_A14B.sample_guide_scale = (3.5, 3.5) # low noise, high noise
tugra/configs/tugra_s2v_14B.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024-2025 The Alibaba TUĞRA Team Authors. All rights reserved.
2
+ from easydict import EasyDict
3
+
4
+ from .shared_config import tugra_shared_cfg
5
+
6
+ #------------------------ TUĞRA S2V 14B ------------------------#
7
+
8
+ s2v_14B = EasyDict(__name__='Config: TUĞRA S2V 14B')
9
+ s2v_14B.update(tugra_shared_cfg)
10
+
11
+ # t5
12
+ s2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
13
+ s2v_14B.t5_tokenizer = 'google/umt5-xxl'
14
+
15
+ # vae
16
+ s2v_14B.vae_checkpoint = 'Tugra2.1_VAE.pth'
17
+ s2v_14B.vae_stride = (4, 8, 8)
18
+
19
+ # wav2vec
20
+ s2v_14B.wav2vec = "wav2vec2-large-xlsr-53-english"
21
+
22
+ s2v_14B.num_heads = 40
23
+ # transformer
24
+ s2v_14B.transformer = EasyDict(
25
+ __name__="Config: Transformer config for TugraModel_S2V")
26
+ s2v_14B.transformer.patch_size = (1, 2, 2)
27
+ s2v_14B.transformer.dim = 5120
28
+ s2v_14B.transformer.ffn_dim = 13824
29
+ s2v_14B.transformer.freq_dim = 256
30
+ s2v_14B.transformer.num_heads = 40
31
+ s2v_14B.transformer.num_layers = 40
32
+ s2v_14B.transformer.window_size = (-1, -1)
33
+ s2v_14B.transformer.qk_norm = True
34
+ s2v_14B.transformer.cross_attn_norm = True
35
+ s2v_14B.transformer.eps = 1e-6
36
+ s2v_14B.transformer.enable_adain = True
37
+ s2v_14B.transformer.adain_mode = "attn_norm"
38
+ s2v_14B.transformer.audio_inject_layers = [
39
+ 0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39
40
+ ]
41
+ s2v_14B.transformer.zero_init = True
42
+ s2v_14B.transformer.zero_timestep = True
43
+ s2v_14B.transformer.enable_motioner = False
44
+ s2v_14B.transformer.add_last_motion = True
45
+ s2v_14B.transformer.trainable_token = False
46
+ s2v_14B.transformer.enable_tsm = False
47
+ s2v_14B.transformer.enable_framepack = True
48
+ s2v_14B.transformer.framepack_drop_mode = 'padd'
49
+ s2v_14B.transformer.audio_dim = 1024
50
+
51
+ s2v_14B.transformer.motion_frames = 73
52
+ s2v_14B.transformer.cond_dim = 16
53
+
54
+ # inference
55
+ s2v_14B.sample_neg_prompt = "画面模糊,最差质量,画面模糊,细节模糊不清,情绪激动剧烈,手快速抖动,字幕,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走"
56
+ s2v_14B.drop_first_motion = True
57
+ s2v_14B.sample_shift = 3
58
+ s2v_14B.sample_steps = 40
59
+ s2v_14B.sample_guide_scale = 4.5
tugra/configs/tugra_t2v_A14B.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024-2025 The Alibaba TUĞRA Team Authors. All rights reserved.
2
+ from easydict import EasyDict
3
+
4
+ from .shared_config import tugra_shared_cfg
5
+
6
+ #------------------------ TUĞRA T2V A14B ------------------------#
7
+
8
+ t2v_A14B = EasyDict(__name__='Config: TUĞRA T2V A14B')
9
+ t2v_A14B.update(tugra_shared_cfg)
10
+
11
+ # t5
12
+ t2v_A14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
13
+ t2v_A14B.t5_tokenizer = 'google/umt5-xxl'
14
+
15
+ # vae
16
+ t2v_A14B.vae_checkpoint = 'Tugra2.1_VAE.pth'
17
+ t2v_A14B.vae_stride = (4, 8, 8)
18
+
19
+ # transformer
20
+ t2v_A14B.patch_size = (1, 2, 2)
21
+ t2v_A14B.dim = 5120
22
+ t2v_A14B.ffn_dim = 13824
23
+ t2v_A14B.freq_dim = 256
24
+ t2v_A14B.num_heads = 40
25
+ t2v_A14B.num_layers = 40
26
+ t2v_A14B.window_size = (-1, -1)
27
+ t2v_A14B.qk_norm = True
28
+ t2v_A14B.cross_attn_norm = True
29
+ t2v_A14B.eps = 1e-6
30
+ t2v_A14B.low_noise_checkpoint = 'low_noise_model'
31
+ t2v_A14B.high_noise_checkpoint = 'high_noise_model'
32
+
33
+ # inference
34
+ t2v_A14B.sample_shift = 12.0
35
+ t2v_A14B.sample_steps = 40
36
+ t2v_A14B.boundary = 0.875
37
+ t2v_A14B.sample_guide_scale = (3.0, 4.0) # low noise, high noise