Upload 73 files

.gitattributes

@@ -33,3 +33,17 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/logo.png filter=lfs diff=lfs merge=lfs -text
+assets/logo2.jpeg filter=lfs diff=lfs merge=lfs -text
+assets/none_quant_efficiency.png filter=lfs diff=lfs merge=lfs -text
+assets/pipe.png filter=lfs diff=lfs merge=lfs -text
+examples/multi/1/1.WAV filter=lfs diff=lfs merge=lfs -text
+examples/multi/1/2.WAV filter=lfs diff=lfs merge=lfs -text
+examples/multi/1/multi1.png filter=lfs diff=lfs merge=lfs -text
+examples/multi/2/1.wav filter=lfs diff=lfs merge=lfs -text
+examples/multi/2/multi2.png filter=lfs diff=lfs merge=lfs -text
+examples/multi/3/1-man.WAV filter=lfs diff=lfs merge=lfs -text
+examples/multi/3/1-woman.WAV filter=lfs diff=lfs merge=lfs -text
+examples/multi/3/multi3.png filter=lfs diff=lfs merge=lfs -text
+examples/single/1.wav filter=lfs diff=lfs merge=lfs -text
+examples/single/single1.png filter=lfs diff=lfs merge=lfs -text

LICENSE.txt

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

README.md

@@ -0,0 +1,467 @@
+<div align="center">
+
+<p align="center">
+  <img src="assets/logo2.jpeg" alt="MultiTalk" width="240"/>
+</p>
+
+<h1>Let Them Talk: Audio-Driven Multi-Person Conversational Video Generation</h1>
+
+
+[Zhe Kong*](https://scholar.google.com/citations?user=4X3yLwsAAAAJ&hl=zh-CN) · [Feng Gao*](https://scholar.google.com/citations?user=lFkCeoYAAAAJ) ·[Yong Zhang](https://yzhang2016.github.io/)<sup>&#9993;</sup> · [Zhuoliang Kang](https://scholar.google.com/citations?user=W1ZXjMkAAAAJ&hl=en) · [Xiaoming Wei](https://scholar.google.com/citations?user=JXV5yrZxj5MC&hl=zh-CN) · [Xunliang Cai](https://openreview.net/profile?id=~Xunliang_Cai1)  
+
+[Guanying Chen](https://guanyingc.github.io/) · [Wenhan Luo](https://whluo.github.io/)<sup>&#9993;</sup>
+
+<sup>*</sup>Equal Contribution
+<sup>&#9993;</sup>Corresponding Authors
+
+
+<a href='https://meigen-ai.github.io/multi-talk/'><img src='https://img.shields.io/badge/Project-Page-green'></a>
+<a href='https://arxiv.org/abs/2505.22647'><img src='https://img.shields.io/badge/Technique-Report-red'></a>
+<a href='https://huggingface.co/MeiGen-AI/MeiGen-MultiTalk'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-blue'></a>
+</div>
+
+> **TL; DR:**  MultiTalk is an audio-driven multi-person conversational video generation​​. It enables the video creation of multi-person conversation 💬, singing  🎤,  interaction control 👬, and cartoon 🙊.
+
+<p align="center">
+  <img src="assets/pipe.png">
+</p>
+
+## Video Demos
+
+<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
+  <tr>
+      <td>
+          <video src="https://github.com/user-attachments/assets/e55952e6-e1b2-44a5-9887-a89307a378da" width="320" controls loop></video>
+      </td>
+      <td>
+          <video src="https://github.com/user-attachments/assets/f0396c19-d459-42aa-9d78-34fdea10de18" width="320" controls loop></video>
+      </td>
+       <td>
+          <video src="https://github.com/user-attachments/assets/3576fd04-3e5f-4933-ac7b-1c4e6a601379" width="320" controls loop></video>
+     </td>
+  </tr>
+  <tr>
+      <td>
+          <video src="https://github.com/user-attachments/assets/5589056e-3202-442d-a62a-2cad7a7ecb19" width="320" controls loop></video>
+      </td>
+      <td>
+          <video src="https://github.com/user-attachments/assets/554bfbe7-0090-492c-94be-329f5e39e175" width="320" controls loop></video>
+      </td>
+       <td>
+          <video src="https://github.com/user-attachments/assets/9e961f35-9413-4846-a806-8186d54061da" width="320" controls loop></video>
+     </td>
+  </tr>
+  <tr>
+      <td>
+          <video src="https://github.com/user-attachments/assets/342595ab-cf75-4872-8182-f20fe8c95611" width="320" controls loop></video>
+      </td>
+      <td>
+          <video src="https://github.com/user-attachments/assets/6476f9f0-35e0-4484-91a4-8aa646aa994a" width="320" controls loop></video>
+      </td>
+       <td>
+          <video src="https://github.com/user-attachments/assets/d8fc8e94-0cba-4c25-9f3a-a8d7e0a785e1" width="320" controls loop></video>
+     </td>
+  </tr>
+</table>
+
+
+
+
+
+## ✨ Key Features
+
+We propose **MultiTalk** , a novel framework for audio-driven multi-person conversational video generation. Given a multi-stream audio input, a reference image and a prompt, MultiTalk generates a video containing interactions following the prompt, with consistent lip motions aligned with the audio.
+
+> - 💬 **​​Realistic Conversations**​​ - Support single & multi-person generation
+> - 👥 ​​**​​Interactive Character Control**​​​​ - Direct virtual humans via prompts
+> - 🎤 ​​**​​Generalization Performances**​​​​ - Support the generation of cartoon character and singing 
+> - 📺 **​​​​Resolution Flexibility​​**​​: 480p & 720p output at arbitrary aspect ratios
+> - ⏱️ **Long Video Generation**: Support video generation up to 15 seconds
+
+## 🔥 Latest News
+
+* July 11, 2025: 🔥🔥 `MultiTalk` supports INT8 [quantization](https://github.com/huggingface/optimum-quanto) and [SageAttention2.2](https://github.com/thu-ml/SageAttention), and updates the CFG strategy (2 NFE per step) for FusionX LoRA, 
+* July 01, 2025: 🔥🔥 `MultiTalk` supports input audios with TTS, [FusioniX](https://huggingface.co/vrgamedevgirl84/Wan14BT2VFusioniX/blob/main/FusionX_LoRa/Wan2.1_I2V_14B_FusionX_LoRA.safetensors) and [lightx2v](https://huggingface.co/Kijai/WanVideo_comfy/blob/main/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors) LoRA acceleration (requires only 4~8 steps), and Gradio.
+* June 14, 2025: 🔥🔥 We release `MultiTalk` with support for `multi-GPU inference`, `teacache acceleration`, `APG` and `low-VRAM inference` (enabling 480P video generation on a single RTX 4090). [APG](https://arxiv.org/abs/2410.02416) is used to alleviate the color error accumulation in long video generation. TeaCache is capable of increasing speed by approximately 2~3x.
+* June 9, 2025: 🔥🔥 We release the [weights](https://huggingface.co/MeiGen-AI/MeiGen-MultiTalk) and inference code of **MultiTalk** 
+* May 29, 2025: We release the [Technique-Report](https://arxiv.org/abs/2505.22647) of **MultiTalk** 
+* May 29, 2025: We release the [project page](https://meigen-ai.github.io/multi-talk/) of **MultiTalk** 
+
+## 🌐 Community  Works
+- [Wan2GP](https://github.com/deepbeepmeep/Wan2GP): thank [deepbeepmeep](https://github.com/deepbeepmeep) for providing the project [Wan2GP](https://github.com/deepbeepmeep/Wan2GP) that enables Multitalk on very low VRAM hardware (8 GB of VRAM) and combines it with the capabilities of Vace. 
+- [Replicate](https://replicate.com/zsxkib/multitalk): thank [zsxkib](https://github.com/zsxkib) for pushing MultiTalk to Replicate platform, try it! Please refer to [cog-MultiTalk](https://github.com/zsxkib/cog-MultiTalk) for details.
+- [Gradio Demo](https://github.com/MeiGen-AI/MultiTalk): thank [fffiloni](https://github.com/fffiloni) for developing this gradio demo on Hugging Face. Please refer to the [issue](https://github.com/MeiGen-AI/MultiTalk/issues/39) for details. 
+- [ComfyUI](https://github.com/kijai/ComfyUI-WanVideoWrapper/tree/multitalk): thank [kijai](https://github.com/kijai) for integrating MultiTalk into ComfyUI-WanVideoWrapper. [Rudra](https://github.com/Rudra-ai-coder) found something interesting that MultiTalk can be combined with Wanx T2V and VACE in the [issue](https://github.com/kijai/ComfyUI-WanVideoWrapper/issues/635). 
+- [Google Colab example](https://colab.research.google.com/drive/185OyRIpJDlpnRjhBRb7FnaRlq11BLZTa?usp=sharing), an exmaple for inference on A100 provided by [Braffolk](https://github.com/Braffolk).
+
+## 📑 Todo List
+
+- [x] Release the technical report
+- [x] Inference
+- [x] Checkpoints
+- [x] Multi-GPU Inference
+- [ ] Inference acceleration
+  - [x] TeaCache
+  - [x] int8 quantization
+  - [ ] LCM distillation
+  - [ ] Sparse Attention
+- [x] Run with very low VRAM
+- [x] TTS integration
+- [x] Gradio demo
+- [ ] ComfyUI
+- [ ] 1.3B model
+
+## Quick Start
+
+### 🛠️Installation
+
+#### 1. Create a conda environment and install pytorch, xformers
+```
+conda create -n multitalk python=3.10
+conda activate multitalk
+pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu121
+pip install -U xformers==0.0.28 --index-url https://download.pytorch.org/whl/cu121
+```
+#### 2. Flash-attn installation:
+```
+pip install misaki[en]
+pip install ninja 
+pip install psutil 
+pip install packaging 
+pip install flash_attn==2.7.4.post1
+```
+
+#### 3. Other dependencies
+```
+pip install -r requirements.txt
+conda install -c conda-forge librosa
+```
+
+#### 4. FFmeg installation
+```
+conda install -c conda-forge ffmpeg
+```
+or
+```
+sudo yum install ffmpeg ffmpeg-devel
+```
+
+### 🧱Model Preparation
+
+#### 1. Model Download
+
+| Models        |                       Download Link                                           |    Notes                      |
+| --------------|-------------------------------------------------------------------------------|-------------------------------|
+| Wan2.1-I2V-14B-480P  |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P)       | Base model
+| chinese-wav2vec2-base |      🤗 [Huggingface](https://huggingface.co/TencentGameMate/chinese-wav2vec2-base)          | Audio encoder
+| Kokoro-82M      |      🤗 [Huggingface](https://huggingface.co/hexgrad/Kokoro-82M)              | TTS weights
+| MeiGen-MultiTalk      |      🤗 [Huggingface](https://huggingface.co/MeiGen-AI/MeiGen-MultiTalk)              | Our audio condition weights
+
+Download models using huggingface-cli:
+``` sh
+huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./weights/Wan2.1-I2V-14B-480P
+huggingface-cli download TencentGameMate/chinese-wav2vec2-base --local-dir ./weights/chinese-wav2vec2-base
+huggingface-cli download TencentGameMate/chinese-wav2vec2-base model.safetensors --revision refs/pr/1 --local-dir ./weights/chinese-wav2vec2-base
+huggingface-cli download hexgrad/Kokoro-82M --local-dir ./weights/Kokoro-82M
+huggingface-cli download MeiGen-AI/MeiGen-MultiTalk --local-dir ./weights/MeiGen-MultiTalk
+```
+
+#### 2. Link or Copy MultiTalk Model to Wan2.1-I2V-14B-480P Directory
+
+Link through:
+```
+mv weights/Wan2.1-I2V-14B-480P/diffusion_pytorch_model.safetensors.index.json weights/Wan2.1-I2V-14B-480P/diffusion_pytorch_model.safetensors.index.json_old
+sudo ln -s {Absolute path}/weights/MeiGen-MultiTalk/diffusion_pytorch_model.safetensors.index.json weights/Wan2.1-I2V-14B-480P/
+sudo ln -s {Absolute path}/weights/MeiGen-MultiTalk/multitalk.safetensors weights/Wan2.1-I2V-14B-480P/
+```
+
+Or, copy through:
+```
+mv weights/Wan2.1-I2V-14B-480P/diffusion_pytorch_model.safetensors.index.json weights/Wan2.1-I2V-14B-480P/diffusion_pytorch_model.safetensors.index.json_old
+cp weights/MeiGen-MultiTalk/diffusion_pytorch_model.safetensors.index.json weights/Wan2.1-I2V-14B-480P/
+cp weights/MeiGen-MultiTalk/multitalk.safetensors weights/Wan2.1-I2V-14B-480P/
+```
+### 🔑 Quick Inference
+
+Our model is compatible with both 480P and 720P resolutions. The current code only supports 480P inference. 720P inference requires multiple GPUs, and we will provide an update soon.
+> Some tips
+> - Lip synchronization accuracy:​​ Audio CFG works optimally between 3–5. Increase the audio CFG value for better synchronization.
+> - ​​Video clip length:​​ The model was trained on 81-frame videos at 25 FPS. For optimal prompt following performance, generate clips at 81 frames. Generating up to 201 frames is possible, though longer clips might reduce prompt-following performance.
+> - ​​Long video generation:​​ Audio CFG influences color tone consistency across segments. Set this value to 3 to alleviate tonal variations.
+> - Sampling steps: If you want to generate a video fast, you can decrease the sampling steps to even 10 that will not hurt the lip synchronization accuracy, but affects the motion and visual quality. More sampling steps, better video quality.
+> - TeaCache accelerate:​​ The optimal range for `--teacache_thresh` is between 0.2 and 0.5. Increasing this value can further improve acceleration, but may also lead to a decline in the quality of the generated video.
+
+#### Usage of MultiTalk
+```
+--mode streaming: long video generation.
+--mode clip: generate short video with one chunk. 
+--use_teacache: run with TeaCache.
+--size multitalk-480: generate 480P video.
+--size multitalk-720: generate 720P video.
+--use_apg: run with APG.
+--teacache_thresh: A coefficient used for TeaCache acceleration
+—-sample_text_guide_scale： When not using LoRA, the optimal value is 5. After applying LoRA, the recommended value is 1.
+—-sample_audio_guide_scale： When not using LoRA, the optimal value is 4. After applying LoRA, the recommended value is 2.
+```
+
+#### 1. Single-Person
+
+##### 1) Run with single GPU
+
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/single_example_1.json \
+    --sample_steps 40 \
+    --mode streaming \
+    --use_teacache \
+    --save_file single_long_exp
+```
+
+##### 2) Run with very low VRAM
+
+If you want run with very low VRAM, set `--num_persistent_param_in_dit 0`:
+
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/single_example_1.json \
+    --sample_steps 40 \
+    --mode streaming \
+    --num_persistent_param_in_dit 0 \
+    --use_teacache \
+    --save_file single_long_lowvram_exp
+```
+
+##### 3) Multi-GPU inference
+
+```
+GPU_NUM=8
+torchrun --nproc_per_node=$GPU_NUM --standalone generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --dit_fsdp --t5_fsdp \
+    --ulysses_size=$GPU_NUM \
+    --input_json examples/single_example_1.json \
+    --sample_steps 40 \
+    --mode streaming \
+    --use_teacache \
+    --save_file single_long_multigpu_exp
+```
+
+##### 4) Run with TTS
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/single_example_tts_1.json \
+    --sample_steps 40 \
+    --mode streaming \
+    --num_persistent_param_in_dit 0 \
+    --use_teacache \
+    --save_file single_long_lowvram_tts_exp \
+    --audio_mode tts
+```
+
+
+#### 2. Multi-Person
+
+##### 1) Run with single GPU
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/multitalk_example_2.json \
+    --sample_steps 40 \
+    --mode streaming \
+    --use_teacache \
+    --save_file multi_long_exp
+```
+##### 2) Run with very low VRAM
+
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/multitalk_example_2.json \
+    --sample_steps 40 \
+    --mode streaming \
+    --num_persistent_param_in_dit 0 \
+    --use_teacache \
+    --save_file multi_long_lowvram_exp
+```
+
+##### 3) Multi-GPU inference
+
+```
+GPU_NUM=8
+torchrun --nproc_per_node=$GPU_NUM --standalone generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --dit_fsdp --t5_fsdp --ulysses_size=$GPU_NUM \
+    --input_json examples/multitalk_example_2.json \
+    --sample_steps 40 \
+    --mode streaming --use_teacache \
+    --save_file multi_long_multigpu_exp
+```
+
+##### 4) Run with TTS
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/multitalk_example_tts_1.json \
+    --sample_steps 40 \
+    --mode streaming \
+    --num_persistent_param_in_dit 0 \
+    --use_teacache \
+    --save_file multi_long_lowvram_tts_exp \
+    --audio_mode tts
+```
+
+
+#### 3. Run with FusioniX and CausVid(Require only 4~8 steps)
+
+[FusioniX](https://huggingface.co/vrgamedevgirl84/Wan14BT2VFusioniX/blob/main/FusionX_LoRa/Wan2.1_I2V_14B_FusionX_LoRA.safetensors) require 8 steps and [lightx2v](https://huggingface.co/Kijai/WanVideo_comfy/blob/main/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors) requires only 4 steps.
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/single_example_1.json \
+    --lora_dir weights/Wan2.1_I2V_14B_FusionX_LoRA.safetensors \
+    --lora_scale 1.0 \
+    --sample_text_guide_scale 1.0 \
+    --sample_audio_guide_scale 2.0 \
+    --sample_steps 8 \
+    --mode streaming \
+    --num_persistent_param_in_dit 0 \
+    --save_file single_long_lowvram_fusionx_exp \
+    --sample_shift 2
+```
+
+or 
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/multitalk_example_2.json \
+    --lora_dir weights/Wan2.1_I2V_14B_FusionX_LoRA.safetensors \
+    --lora_scale 1.0 \
+    --sample_text_guide_scale 1.0 \
+    --sample_audio_guide_scale 2.0 \
+    --sample_steps 8 \
+    --mode streaming \
+    --num_persistent_param_in_dit 0 \
+    --save_file multi_long_lowvram_fusionx_exp \
+
+```
+
+#### 4. Run with the quantization model (Only support run with single gpu)
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/multitalk_example_2.json \
+    --sample_steps 40 \
+    --mode streaming \
+    --use_teacache \
+    --quant int8 \
+    --quant_dir weights/MeiGen-MultiTalk \
+    --num_persistent_param_in_dit 0 \
+    --save_file multi_long_lowvram_exp_quant
+```
+
+Run with lora:
+
+```
+python generate_multitalk.py \
+    --ckpt_dir weights/Wan2.1-I2V-14B-480P \
+    --wav2vec_dir 'weights/chinese-wav2vec2-base' \
+    --input_json examples/multitalk_example_1.json \
+    --quant int8 \
+    --quant_dir weights/MeiGen-MultiTalk \
+    --lora_dir weights/MeiGen-MultiTalk/quant_models/quant_model_int8_FusionX.safetensors \
+    --sample_text_guide_scale 1.0 \
+    --sample_audio_guide_scale 2.0 \
+    --sample_steps 8 \
+    --mode streaming \
+    --num_persistent_param_in_dit 0 \
+    --save_file multi_long_lowvram_fusionx_exp_quant \
+    --sample_shift 2
+```
+
+#### 5. Run with Gradio
+
+
+
+```
+python app.py \
+    --lora_dir weights/Wan2.1_I2V_14B_FusionX_LoRA.safetensors \
+    --lora_scale 1.0 \
+    --num_persistent_param_in_dit 0 \
+    --sample_shift 2
+```
+
+or
+
+```
+python app.py --num_persistent_param_in_dit 0 
+```
+
+or 
+
+```
+python app.py \
+    --quant int8 \
+    --quant_dir weights/MeiGen-MultiTalk \
+    --lora_dir weights/MeiGen-MultiTalk/quant_models/quant_model_int8_FusionX.safetensors \
+    --sample_shift 2 \
+    --num_persistent_param_in_dit 0
+```
+
+## 🚀Computational Efficiency
+
+#### 1) Non quantitative results
+
+The results are evaluated on A100 GPUs for multi-person generation. Single-person generation uses less memory and provides faster inference.
+<p align="center">
+  <img src="assets/efficiency.png">
+</p>
+TeaCache is capable of increasing speed by approximately 2~3x.
+
+#### 2) Quantitative results
+
+<p align="center">
+  <img src="assets/none_quant_efficiency.png">
+</p>
+
+
+## 📚 Citation
+
+If you find our work useful in your research, please consider citing:
+
+```
+@article{kong2025let,
+  title={Let Them Talk: Audio-Driven Multi-Person Conversational Video Generation},
+  author={Kong, Zhe and Gao, Feng and Zhang, Yong and Kang, Zhuoliang and Wei, Xiaoming and Cai, Xunliang and Chen, Guanying and Luo, Wenhan},
+  journal={arXiv preprint arXiv:2505.22647},
+  year={2025}
+}
+```
+
+## 📜 License
+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. 
+

app.py

@@ -0,0 +1,787 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import argparse
+import logging
+import os
+os.environ["no_proxy"] = "localhost,127.0.0.1,::1"
+import sys
+import json
+import warnings
+from datetime import datetime
+
+import gradio as gr
+warnings.filterwarnings('ignore')
+
+import random
+
+import torch
+import torch.distributed as dist
+from PIL import Image
+import subprocess
+
+import wan
+from wan.configs import SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
+from wan.utils.utils import cache_image, cache_video, str2bool
+from wan.utils.multitalk_utils import save_video_ffmpeg
+from kokoro import KPipeline
+from transformers import Wav2Vec2FeatureExtractor
+from src.audio_analysis.wav2vec2 import Wav2Vec2Model
+
+import librosa
+import pyloudnorm as pyln
+import numpy as np
+from einops import rearrange
+import soundfile as sf
+import re
+
+def _validate_args(args):
+    # Basic check
+    assert args.ckpt_dir is not None, "Please specify the checkpoint directory."
+    assert args.task in WAN_CONFIGS, f"Unsupport task: {args.task}"
+
+    # The default sampling steps are 40 for image-to-video tasks and 50 for text-to-video tasks.
+    if args.sample_steps is None:
+        args.sample_steps = 40
+
+    if args.sample_shift is None:
+        if args.size == 'multitalk-480':
+            args.sample_shift = 7
+        elif args.size == 'multitalk-720':
+            args.sample_shift = 11
+        else:
+            raise NotImplementedError(f'Not supported size')
+
+    args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
+        0, 99999999)
+    # Size check
+    assert args.size in SUPPORTED_SIZES[
+        args.
+        task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description="Generate a image or video from a text prompt or image using Wan"
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="multitalk-14B",
+        choices=list(WAN_CONFIGS.keys()),
+        help="The task to run.")
+    parser.add_argument(
+        "--size",
+        type=str,
+        default="multitalk-480",
+        choices=list(SIZE_CONFIGS.keys()),
+        help="The buckget size of the generated video. The aspect ratio of the output video will follow that of the input image."
+    )
+    parser.add_argument(
+        "--frame_num",
+        type=int,
+        default=81,
+        help="How many frames to be generated in one clip. The number should be 4n+1"
+    )
+    parser.add_argument(
+        "--ckpt_dir",
+        type=str,
+        default='./weights/Wan2.1-I2V-14B-480P',
+        help="The path to the Wan checkpoint directory.")
+    parser.add_argument(
+        "--quant_dir",
+        type=str,
+        default=None,
+        help="The path to the Wan quant checkpoint directory.")
+    parser.add_argument(
+        "--wav2vec_dir",
+        type=str,
+        default='./weights/chinese-wav2vec2-base',
+        help="The path to the wav2vec checkpoint directory.")
+    parser.add_argument(
+        "--lora_dir",
+        type=str,
+        nargs='+',
+        default=None,
+        help="The path to the LoRA checkpoint directory.")
+    parser.add_argument(
+        "--lora_scale",
+        type=float,
+        nargs='+',
+        default=[1.2],
+        help="Controls how much to influence the outputs with the LoRA parameters. Accepts multiple float values."
+    )
+    parser.add_argument(
+        "--offload_model",
+        type=str2bool,
+        default=None,
+        help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage."
+    )
+    parser.add_argument(
+        "--ulysses_size",
+        type=int,
+        default=1,
+        help="The size of the ulysses parallelism in DiT.")
+    parser.add_argument(
+        "--ring_size",
+        type=int,
+        default=1,
+        help="The size of the ring attention parallelism in DiT.")
+    parser.add_argument(
+        "--t5_fsdp",
+        action="store_true",
+        default=False,
+        help="Whether to use FSDP for T5.")
+    parser.add_argument(
+        "--t5_cpu",
+        action="store_true",
+        default=False,
+        help="Whether to place T5 model on CPU.")
+    parser.add_argument(
+        "--dit_fsdp",
+        action="store_true",
+        default=False,
+        help="Whether to use FSDP for DiT.")
+    parser.add_argument(
+        "--save_file",
+        type=str,
+        default=None,
+        help="The file to save the generated image or video to.")
+    parser.add_argument(
+        "--audio_save_dir",
+        type=str,
+        default='save_audio/gradio',
+        help="The path to save the audio embedding.")
+    parser.add_argument(
+        "--base_seed",
+        type=int,
+        default=42,
+        help="The seed to use for generating the image or video.")
+    parser.add_argument(
+        "--input_json",
+        type=str,
+        default='examples.json',
+        help="[meta file] The condition path to generate the video.")
+    parser.add_argument(
+        "--motion_frame",
+        type=int,
+        default=25,
+        help="Driven frame length used in the mode of long video genration.")
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="streaming",
+        choices=['clip', 'streaming'],
+        help="clip: generate one video chunk, streaming: long video generation")
+    parser.add_argument(
+        "--sample_steps", type=int, default=None, help="The sampling steps.")
+    parser.add_argument(
+        "--sample_shift",
+        type=float,
+        default=None,
+        help="Sampling shift factor for flow matching schedulers.")
+    parser.add_argument(
+        "--sample_text_guide_scale",
+        type=float,
+        default=5.0,
+        help="Classifier free guidance scale for text control.")
+    parser.add_argument(
+        "--sample_audio_guide_scale",
+        type=float,
+        default=4.0,
+        help="Classifier free guidance scale for audio control.")
+    parser.add_argument(
+        "--num_persistent_param_in_dit",
+        type=int,
+        default=None,
+        required=False,
+        help="Maximum parameter quantity retained in video memory, small number to reduce VRAM required",
+    )
+    parser.add_argument(
+        "--use_teacache",
+        action="store_true",
+        default=False,
+        help="Enable teacache for video generation."
+    )
+    parser.add_argument(
+        "--teacache_thresh",
+        type=float,
+        default=0.2,
+        help="Threshold for teacache."
+    )
+    parser.add_argument(
+        "--use_apg",
+        action="store_true",
+        default=False,
+        help="Enable adaptive projected guidance for video generation (APG)."
+    )
+    parser.add_argument(
+        "--apg_momentum",
+        type=float,
+        default=-0.75,
+        help="Momentum used in adaptive projected guidance (APG)."
+    )
+    parser.add_argument(
+        "--apg_norm_threshold",
+        type=float,
+        default=55,
+        help="Norm threshold used in adaptive projected guidance (APG)."
+    )
+    parser.add_argument(
+        "--color_correction_strength",
+        type=float,
+        default=1.0,
+        help="strength for color correction [0.0 -- 1.0]."
+    )
+
+    parser.add_argument(
+        "--quant",
+        type=str,
+        default=None,
+        help="Quantization type, must be 'int8' or 'fp8'."
+    )
+    args = parser.parse_args()
+    _validate_args(args)
+    return args
+
+
+def custom_init(device, wav2vec):    
+    audio_encoder = Wav2Vec2Model.from_pretrained(wav2vec, local_files_only=True).to(device)
+    audio_encoder.feature_extractor._freeze_parameters()
+    wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec, local_files_only=True)
+    return wav2vec_feature_extractor, audio_encoder
+
+def loudness_norm(audio_array, sr=16000, lufs=-23):
+    meter = pyln.Meter(sr)
+    loudness = meter.integrated_loudness(audio_array)
+    if abs(loudness) > 100:
+        return audio_array
+    normalized_audio = pyln.normalize.loudness(audio_array, loudness, lufs)
+    return normalized_audio
+
+def audio_prepare_multi(left_path, right_path, audio_type, sample_rate=16000):
+    if not (left_path=='None' or right_path=='None'):
+        human_speech_array1 = audio_prepare_single(left_path)
+        human_speech_array2 = audio_prepare_single(right_path)
+    elif left_path=='None':
+        human_speech_array2 = audio_prepare_single(right_path)
+        human_speech_array1 = np.zeros(human_speech_array2.shape[0])
+    elif right_path=='None':
+        human_speech_array1 = audio_prepare_single(left_path)
+        human_speech_array2 = np.zeros(human_speech_array1.shape[0])
+
+    if audio_type=='para':
+        new_human_speech1 = human_speech_array1
+        new_human_speech2 = human_speech_array2
+    elif audio_type=='add':
+        new_human_speech1 = np.concatenate([human_speech_array1[: human_speech_array1.shape[0]], np.zeros(human_speech_array2.shape[0])]) 
+        new_human_speech2 = np.concatenate([np.zeros(human_speech_array1.shape[0]), human_speech_array2[:human_speech_array2.shape[0]]])
+    sum_human_speechs = new_human_speech1 + new_human_speech2
+    return new_human_speech1, new_human_speech2, sum_human_speechs
+
+def _init_logging(rank):
+    # logging
+    if rank == 0:
+        # set format
+        logging.basicConfig(
+            level=logging.INFO,
+            format="[%(asctime)s] %(levelname)s: %(message)s",
+            handlers=[logging.StreamHandler(stream=sys.stdout)])
+    else:
+        logging.basicConfig(level=logging.ERROR)
+
+def get_embedding(speech_array, wav2vec_feature_extractor, audio_encoder, sr=16000, device='cpu'):
+    audio_duration = len(speech_array) / sr
+    video_length = audio_duration * 25 # Assume the video fps is 25
+
+    # wav2vec_feature_extractor
+    audio_feature = np.squeeze(
+        wav2vec_feature_extractor(speech_array, sampling_rate=sr).input_values
+    )
+    audio_feature = torch.from_numpy(audio_feature).float().to(device=device)
+    audio_feature = audio_feature.unsqueeze(0)
+
+    # audio encoder
+    with torch.no_grad():
+        embeddings = audio_encoder(audio_feature, seq_len=int(video_length), output_hidden_states=True)
+
+    if len(embeddings) == 0:
+        print("Fail to extract audio embedding")
+        return None
+
+    audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
+    audio_emb = rearrange(audio_emb, "b s d -> s b d")
+
+    audio_emb = audio_emb.cpu().detach()
+    return audio_emb
+
+def extract_audio_from_video(filename, sample_rate):
+    raw_audio_path = filename.split('/')[-1].split('.')[0]+'.wav'
+    ffmpeg_command = [
+        "ffmpeg",
+        "-y",
+        "-i",
+        str(filename),
+        "-vn",
+        "-acodec",
+        "pcm_s16le",
+        "-ar",
+        "16000",
+        "-ac",
+        "2",
+        str(raw_audio_path),
+    ]
+    subprocess.run(ffmpeg_command, check=True)
+    human_speech_array, sr = librosa.load(raw_audio_path, sr=sample_rate)
+    human_speech_array = loudness_norm(human_speech_array, sr)
+    os.remove(raw_audio_path)
+
+    return human_speech_array
+
+def audio_prepare_single(audio_path, sample_rate=16000):
+    ext = os.path.splitext(audio_path)[1].lower()
+    if ext in ['.mp4', '.mov', '.avi', '.mkv']:
+        human_speech_array = extract_audio_from_video(audio_path, sample_rate)
+        return human_speech_array
+    else:
+        human_speech_array, sr = librosa.load(audio_path, sr=sample_rate)
+        human_speech_array = loudness_norm(human_speech_array, sr)
+        return human_speech_array
+
+def process_tts_single(text, save_dir, voice1):    
+    s1_sentences = []
+
+    pipeline = KPipeline(lang_code='a', repo_id='weights/Kokoro-82M')
+
+    voice_tensor = torch.load(voice1, weights_only=True)
+    generator = pipeline(
+        text, voice=voice_tensor, # <= change voice here
+        speed=1, split_pattern=r'\n+'
+    )
+    audios = []
+    for i, (gs, ps, audio) in enumerate(generator):
+        audios.append(audio)
+    audios = torch.concat(audios, dim=0)
+    s1_sentences.append(audios)
+    s1_sentences = torch.concat(s1_sentences, dim=0)
+    save_path1 =f'{save_dir}/s1.wav'
+    sf.write(save_path1, s1_sentences, 24000) # save each audio file
+    s1, _ = librosa.load(save_path1, sr=16000)
+    return s1, save_path1
+    
+   
+
+def process_tts_multi(text, save_dir, voice1, voice2):
+    pattern = r'\(s(\d+)\)\s*(.*?)(?=\s*\(s\d+\)|$)'
+    matches = re.findall(pattern, text, re.DOTALL)
+    
+    s1_sentences = []
+    s2_sentences = []
+
+    pipeline = KPipeline(lang_code='a', repo_id='weights/Kokoro-82M')
+    for idx, (speaker, content) in enumerate(matches):
+        if speaker == '1':
+            voice_tensor = torch.load(voice1, weights_only=True)
+            generator = pipeline(
+                content, voice=voice_tensor, # <= change voice here
+                speed=1, split_pattern=r'\n+'
+            )
+            audios = []
+            for i, (gs, ps, audio) in enumerate(generator):
+                audios.append(audio)
+            audios = torch.concat(audios, dim=0)
+            s1_sentences.append(audios)
+            s2_sentences.append(torch.zeros_like(audios))
+        elif speaker == '2':
+            voice_tensor = torch.load(voice2, weights_only=True)
+            generator = pipeline(
+                content, voice=voice_tensor, # <= change voice here
+                speed=1, split_pattern=r'\n+'
+            )
+            audios = []
+            for i, (gs, ps, audio) in enumerate(generator):
+                audios.append(audio)
+            audios = torch.concat(audios, dim=0)
+            s2_sentences.append(audios)
+            s1_sentences.append(torch.zeros_like(audios))
+    
+    s1_sentences = torch.concat(s1_sentences, dim=0)
+    s2_sentences = torch.concat(s2_sentences, dim=0)
+    sum_sentences = s1_sentences + s2_sentences
+    save_path1 =f'{save_dir}/s1.wav'
+    save_path2 =f'{save_dir}/s2.wav'
+    save_path_sum = f'{save_dir}/sum.wav'
+    sf.write(save_path1, s1_sentences, 24000) # save each audio file
+    sf.write(save_path2, s2_sentences, 24000)
+    sf.write(save_path_sum, sum_sentences, 24000)
+
+    s1, _ = librosa.load(save_path1, sr=16000)
+    s2, _ = librosa.load(save_path2, sr=16000)
+    # sum, _ = librosa.load(save_path_sum, sr=16000)
+    return s1, s2, save_path_sum
+
+def run_graio_demo(args):
+    rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    local_rank = int(os.getenv("LOCAL_RANK", 0))
+    device = local_rank
+    _init_logging(rank)
+
+    if args.offload_model is None:
+        args.offload_model = False if world_size > 1 else True
+        logging.info(
+            f"offload_model is not specified, set to {args.offload_model}.")
+    if world_size > 1:
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(
+            backend="nccl",
+            init_method="env://",
+            rank=rank,
+            world_size=world_size)
+    else:
+        assert not (
+            args.t5_fsdp or args.dit_fsdp
+        ), f"t5_fsdp and dit_fsdp are not supported in non-distributed environments."
+        assert not (
+            args.ulysses_size > 1 or args.ring_size > 1
+        ), f"context parallel are not supported in non-distributed environments."
+
+    if args.ulysses_size > 1 or args.ring_size > 1:
+        assert args.ulysses_size * args.ring_size == world_size, f"The number of ulysses_size and ring_size should be equal to the world size."
+        from xfuser.core.distributed import (
+            init_distributed_environment,
+            initialize_model_parallel,
+        )
+        init_distributed_environment(
+            rank=dist.get_rank(), world_size=dist.get_world_size())
+
+        initialize_model_parallel(
+            sequence_parallel_degree=dist.get_world_size(),
+            ring_degree=args.ring_size,
+            ulysses_degree=args.ulysses_size,
+        )
+
+   
+    cfg = WAN_CONFIGS[args.task]
+    if args.ulysses_size > 1:
+        assert cfg.num_heads % args.ulysses_size == 0, f"`{cfg.num_heads=}` cannot be divided evenly by `{args.ulysses_size=}`."
+
+    logging.info(f"Generation job args: {args}")
+    logging.info(f"Generation model config: {cfg}")
+
+    if dist.is_initialized():
+        base_seed = [args.base_seed] if rank == 0 else [None]
+        dist.broadcast_object_list(base_seed, src=0)
+        args.base_seed = base_seed[0]
+
+    assert args.task == "multitalk-14B", 'You should choose multitalk in args.task.'
+    
+
+    
+    wav2vec_feature_extractor, audio_encoder= custom_init('cpu', args.wav2vec_dir)
+    os.makedirs(args.audio_save_dir,exist_ok=True)
+
+
+    logging.info("Creating MultiTalk pipeline.")
+    # wan_i2v = None
+    wan_i2v = wan.MultiTalkPipeline(
+        config=cfg,
+        checkpoint_dir=args.ckpt_dir,
+        quant_dir=args.quant_dir,
+        device_id=device,
+        rank=rank,
+        t5_fsdp=args.t5_fsdp,
+        dit_fsdp=args.dit_fsdp, 
+        use_usp=(args.ulysses_size > 1 or args.ring_size > 1),  
+        t5_cpu=args.t5_cpu,
+        lora_dir=args.lora_dir,
+        lora_scales=args.lora_scale,
+        quant=args.quant
+    )
+
+    if args.num_persistent_param_in_dit is not None:
+        wan_i2v.vram_management = True
+        wan_i2v.enable_vram_management(
+            num_persistent_param_in_dit=args.num_persistent_param_in_dit
+        )
+
+
+    
+    def generate_video(img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2,
+                    sd_steps, seed, text_guide_scale, audio_guide_scale, mode_selector, tts_text, resolution_select, human1_voice, human2_voice):
+        input_data = {}
+        input_data["prompt"] = img2vid_prompt
+        input_data["cond_image"] = img2vid_image
+        person = {}
+        if mode_selector == "Single Person(Local File)":
+            person['person1'] = img2vid_audio_1
+        elif mode_selector == "Single Person(TTS)":
+            tts_audio = {}
+            tts_audio['text'] = tts_text
+            tts_audio['human1_voice'] = human1_voice
+            input_data["tts_audio"] = tts_audio
+        elif mode_selector == "Multi Person(Local File, audio add)":
+            person['person1'] = img2vid_audio_1
+            person['person2'] = img2vid_audio_2
+            input_data["audio_type"] = 'add'
+        elif mode_selector == "Multi Person(Local File, audio parallel)":
+            person['person1'] = img2vid_audio_1
+            person['person2'] = img2vid_audio_2
+            input_data["audio_type"] = 'para'
+        else:
+            tts_audio = {}
+            tts_audio['text'] = tts_text
+            tts_audio['human1_voice'] = human1_voice
+            tts_audio['human2_voice'] = human2_voice
+            input_data["tts_audio"] = tts_audio
+            
+        input_data["cond_audio"] = person
+
+        if 'Local File' in mode_selector:
+            if len(input_data['cond_audio'])==2:
+                new_human_speech1, new_human_speech2, sum_human_speechs = audio_prepare_multi(input_data['cond_audio']['person1'], input_data['cond_audio']['person2'], input_data['audio_type'])
+                audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
+                audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder)
+                emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+                emb2_path = os.path.join(args.audio_save_dir, '2.pt')
+                sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+                sf.write(sum_audio, sum_human_speechs, 16000)
+                torch.save(audio_embedding_1, emb1_path)
+                torch.save(audio_embedding_2, emb2_path)
+                input_data['cond_audio']['person1'] = emb1_path
+                input_data['cond_audio']['person2'] = emb2_path
+                input_data['video_audio'] = sum_audio
+            elif len(input_data['cond_audio'])==1:
+                human_speech = audio_prepare_single(input_data['cond_audio']['person1'])
+                audio_embedding = get_embedding(human_speech, wav2vec_feature_extractor, audio_encoder)
+                emb_path = os.path.join(args.audio_save_dir, '1.pt')
+                sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+                sf.write(sum_audio, human_speech, 16000)
+                torch.save(audio_embedding, emb_path)
+                input_data['cond_audio']['person1'] = emb_path
+                input_data['video_audio'] = sum_audio
+        elif 'TTS' in mode_selector:
+            if 'human2_voice' not in input_data['tts_audio'].keys():
+                new_human_speech1, sum_audio = process_tts_single(input_data['tts_audio']['text'], args.audio_save_dir, input_data['tts_audio']['human1_voice'])
+                audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
+                emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+                torch.save(audio_embedding_1, emb1_path)
+                input_data['cond_audio']['person1'] = emb1_path
+                input_data['video_audio'] = sum_audio
+            else:
+                new_human_speech1, new_human_speech2, sum_audio = process_tts_multi(input_data['tts_audio']['text'], args.audio_save_dir, input_data['tts_audio']['human1_voice'], input_data['tts_audio']['human2_voice'])
+                audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
+                audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder)
+                emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+                emb2_path = os.path.join(args.audio_save_dir, '2.pt')
+                torch.save(audio_embedding_1, emb1_path)
+                torch.save(audio_embedding_2, emb2_path)
+                input_data['cond_audio']['person1'] = emb1_path
+                input_data['cond_audio']['person2'] = emb2_path
+                input_data['video_audio'] = sum_audio
+
+
+        # if len(input_data['cond_audio'])==2:
+        #     new_human_speech1, new_human_speech2, sum_human_speechs = audio_prepare_multi(input_data['cond_audio']['person1'], input_data['cond_audio']['person2'], input_data['audio_type'])
+        #     audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
+        #     audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder)
+        #     emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+        #     emb2_path = os.path.join(args.audio_save_dir, '2.pt')
+        #     sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+        #     sf.write(sum_audio, sum_human_speechs, 16000)
+        #     torch.save(audio_embedding_1, emb1_path)
+        #     torch.save(audio_embedding_2, emb2_path)
+        #     input_data['cond_audio']['person1'] = emb1_path
+        #     input_data['cond_audio']['person2'] = emb2_path
+        #     input_data['video_audio'] = sum_audio
+        # elif len(input_data['cond_audio'])==1:
+        #     human_speech = audio_prepare_single(input_data['cond_audio']['person1'])
+        #     audio_embedding = get_embedding(human_speech, wav2vec_feature_extractor, audio_encoder)
+        #     emb_path = os.path.join(args.audio_save_dir, '1.pt')
+        #     sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+        #     sf.write(sum_audio, human_speech, 16000)
+        #     torch.save(audio_embedding, emb_path)
+        #     input_data['cond_audio']['person1'] = emb_path
+        #     input_data['video_audio'] = sum_audio
+
+        logging.info("Generating video ...")
+        video = wan_i2v.generate(
+            input_data,
+            size_buckget=resolution_select,
+            motion_frame=args.motion_frame,
+            frame_num=args.frame_num,
+            shift=args.sample_shift,
+            sampling_steps=sd_steps,
+            text_guide_scale=text_guide_scale,
+            audio_guide_scale=audio_guide_scale,
+            seed=seed,
+            n_prompt=n_prompt,
+            offload_model=args.offload_model,
+            max_frames_num=args.frame_num if args.mode == 'clip' else 1000,
+            color_correction_strength = args.color_correction_strength,
+            extra_args=args,
+            )
+        
+
+        if args.save_file is None:
+            formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+            formatted_prompt = input_data['prompt'].replace(" ", "_").replace("/",
+                                                                        "_")[:50]
+            args.save_file = f"{args.task}_{args.size.replace('*','x') if sys.platform=='win32' else args.size}_{args.ulysses_size}_{args.ring_size}_{formatted_prompt}_{formatted_time}"
+        
+        logging.info(f"Saving generated video to {args.save_file}.mp4")
+        save_video_ffmpeg(video, args.save_file, [input_data['video_audio']], high_quality_save=False)
+        logging.info("Finished.")
+
+        return args.save_file + '.mp4'
+
+    def toggle_audio_mode(mode):
+        if 'TTS' in mode:
+            return [
+                gr.Audio(visible=False, interactive=False),
+                gr.Audio(visible=False, interactive=False),
+                gr.Textbox(visible=True, interactive=True)
+            ]
+        elif 'Single' in mode:
+            return [
+                gr.Audio(visible=True, interactive=True),
+                gr.Audio(visible=False, interactive=False),
+                gr.Textbox(visible=False, interactive=False)
+            ]
+        else:
+            return [
+                gr.Audio(visible=True, interactive=True),
+                gr.Audio(visible=True, interactive=True),
+                gr.Textbox(visible=False, interactive=False)
+            ]
+            
+        
+    with gr.Blocks() as demo:
+
+        gr.Markdown("""
+                    <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
+                        MeiGen-MultiTalk
+                    </div>
+                    <div style="text-align: center; font-size: 16px; font-weight: normal; margin-bottom: 20px;">
+                        Let Them Talk: Audio-Driven Multi-Person Conversational Video Generation.
+                    </div>
+                    <div style="display: flex; justify-content: center; gap: 10px; flex-wrap: wrap;">
+                        <a href='https://meigen-ai.github.io/multi-talk/'><img src='https://img.shields.io/badge/Project-Page-blue'></a>
+                        <a href='https://huggingface.co/MeiGen-AI/MeiGen-MultiTalk'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow'></a>
+                        <a href='https://arxiv.org/abs/2505.22647'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
+                    </div>
+
+
+                    """)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                img2vid_image = gr.Image(
+                    type="filepath",
+                    label="Upload Input Image",
+                    elem_id="image_upload",
+                )
+                img2vid_prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Describe the video you want to generate",
+                )
+                
+                
+                with gr.Accordion("Audio Options", open=True):
+                    mode_selector = gr.Radio(
+                        choices=["Single Person(Local File)", "Single Person(TTS)", "Multi Person(Local File, audio add)", "Multi Person(Local File, audio parallel)", "Multi Person(TTS)"],
+                        label="Select person and audio mode.",
+                        value="Single Person(Local File)"
+                    )
+                    resolution_select = gr.Radio(
+                        choices=["multitalk-480", "multitalk-720"],
+                        label="Select resolution.",
+                        value="multitalk-480"
+                    )
+                    img2vid_audio_1 = gr.Audio(label="Conditioning Audio for speaker 1", type="filepath", visible=True)
+                    img2vid_audio_2 = gr.Audio(label="Conditioning Audio for speaker 2", type="filepath", visible=False)
+                    tts_text = gr.Textbox(
+                        label="Text for TTS",
+                        placeholder="Refer to the format in the examples",
+                        visible=False,
+                        interactive=False
+                    )
+                    mode_selector.change(
+                        fn=toggle_audio_mode,
+                        inputs=mode_selector,
+                        outputs=[img2vid_audio_1, img2vid_audio_2, tts_text]
+                    )
+
+                with gr.Accordion("Advanced Options", open=False):
+                    with gr.Row():
+                        sd_steps = gr.Slider(
+                            label="Diffusion steps",
+                            minimum=1,
+                            maximum=1000,
+                            value=8,
+                            step=1)
+                        seed = gr.Slider(
+                            label="Seed",
+                            minimum=-1,
+                            maximum=2147483647,
+                            step=1,
+                            value=42)
+                    with gr.Row():
+                        text_guide_scale = gr.Slider(
+                            label="Text Guide scale",
+                            minimum=0,
+                            maximum=20,
+                            value=1.0,
+                            step=1)
+                        audio_guide_scale = gr.Slider(
+                            label="Audio Guide scale",
+                            minimum=0,
+                            maximum=20,
+                            value=2.0,
+                            step=1)
+                    with gr.Row():
+                        human1_voice = gr.Textbox(
+                            label="Voice for the left person",
+                            value="weights/Kokoro-82M/voices/am_adam.pt",
+                        )
+                        human2_voice = gr.Textbox(
+                            label="Voice for right person",
+                            value="weights/Kokoro-82M/voices/af_heart.pt"
+                        )
+                    # with gr.Row():
+                    n_prompt = gr.Textbox(
+                        label="Negative Prompt",
+                        placeholder="Describe the negative prompt you want to add",
+                        value="bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+                    )
+
+                run_i2v_button = gr.Button("Generate Video")
+
+            with gr.Column(scale=2):
+                result_gallery = gr.Video(
+                    label='Generated Video', interactive=False, height=600, )
+                
+                gr.Examples(
+                    examples = [
+                        ["examples/single/single1.png", "A woman is passionately singing into a professional microphone in a recording studio. She wears large black headphones and a dark cardigan over a gray top. Her long, wavy brown hair frames her face as she looks slightly upwards, her mouth open mid-song. The studio is equipped with various audio equipment, including a mixing console and a keyboard, with soundproofing panels on the walls. The lighting is warm and focused on her, creating a professional and intimate atmosphere. A close-up shot captures her expressive performance.", "Single Person(Local File)", "examples/single/1.wav", None, None],
+                        ["examples/single/single1.png", "A woman is passionately singing into a professional microphone in a recording studio. She wears large black headphones and a dark cardigan over a gray top. Her long, wavy brown hair frames her face as she looks slightly upwards, her mouth open mid-song. The studio is equipped with various audio equipment, including a mixing console and a keyboard, with soundproofing panels on the walls. The lighting is warm and focused on her, creating a professional and intimate atmosphere. A close-up shot captures her expressive performance.", "Single Person(TTS)", None, None, "Welcome to multi-talk, this is an audio-driven video generation model for multiple person."],
+                        ["examples/multi/1/multi1.png", "In a casual, intimate setting, a man and a woman are engaged in a heartfelt conversation inside a car. The man, sporting a denim jacket over a blue shirt, sits attentively with a seatbelt fastened, his gaze fixed on the woman beside him. The woman, wearing a black tank top and a denim jacket draped over her shoulders, smiles warmly, her eyes reflecting genuine interest and connection. The car's interior, with its beige seats and simple design, provides a backdrop that emphasizes their interaction. The scene captures a moment of shared understanding and connection, set against the soft, diffused light of an overcast day. A medium shot from a slightly angled perspective, focusing on their expressions and body language.", "Multi Person(Local File, audio add)", "examples/multi/1/1.WAV", "examples/multi/1/2.WAV", None],
+                        ["examples/multi/3/multi3.png", "In a cozy recording studio, a man and a woman are singing together. The man, with tousled brown hair, stands to the left, wearing a light green button-down shirt. His gaze is directed towards the woman, who is smiling warmly. She, with wavy dark hair, is dressed in a black floral dress and stands to the right, her eyes closed in enjoyment. Between them is a professional microphone, capturing their harmonious voices. The background features wooden panels and various audio equipment, creating an intimate and focused atmosphere. The lighting is soft and warm, highlighting their expressions and the intimate setting. A medium shot captures their interaction closely.", "Multi Person(Local File, audio parallel)", "examples/multi/3/1-man.WAV", "examples/multi/3/1-woman.WAV", None],
+                        ["examples/multi/1/multi1.png", "In a casual, intimate setting, a man and a woman are engaged in a heartfelt conversation inside a car. The man, sporting a denim jacket over a blue shirt, sits attentively with a seatbelt fastened, his gaze fixed on the woman beside him. The woman, wearing a black tank top and a denim jacket draped over her shoulders, smiles warmly, her eyes reflecting genuine interest and connection. The car's interior, with its beige seats and simple design, provides a backdrop that emphasizes their interaction. The scene captures a moment of shared understanding and connection, set against the soft, diffused light of an overcast day. A medium shot from a slightly angled perspective, focusing on their expressions and body language.", "Multi Person(TTS)", None, None, "(s1) do you know multi-talk? (s2) yes, I know it, that's amazing! (s1) Me too."],
+                    ],
+                    inputs = [img2vid_image, img2vid_prompt, mode_selector, img2vid_audio_1, img2vid_audio_2, tts_text],
+                )
+
+
+        run_i2v_button.click(
+            fn=generate_video,
+            inputs=[img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2,sd_steps, seed, text_guide_scale, audio_guide_scale, mode_selector, tts_text, resolution_select, human1_voice, human2_voice],
+            outputs=[result_gallery],
+        )
+    demo.launch(server_name="0.0.0.0", debug=True, server_port=8418)
+
+        
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    run_graio_demo(args)
+    

examples/multi/1/1.WAV

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8397a9b3c0add26384afe7e544e36cbc4806d8f2d7c705e11bb2897dc1bc993b
+size 315436

examples/multi/1/2.WAV

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:753120ceadbdab3ce206423a1419f73018695682787414ca2f4613306be50bfc
+size 544812

examples/multi/2/1.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51eb6a408a8b5b33a732378e2a38e7412ba273186b85c324ec6a099d23fe38af
+size 1273592

examples/multi/3/1-man.WAV

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d304fd88850d6673649d1844db2894e03bf5a775123048eebcb01ab3b79bff5e
+size 1503276

examples/multi/3/1-woman.WAV

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e1ebd7ae1587ebc7f0986f8b61e7fcc99c6fb57fbb15ab9373968e701afc8bf
+size 1503276

examples/multitalk_example_1.json

@@ -0,0 +1,13 @@
+{
+    "prompt": "In a casual, intimate setting, a man and a woman are engaged in a heartfelt conversation inside a car. The man, sporting a denim jacket over a blue shirt, sits attentively with a seatbelt fastened, his gaze fixed on the woman beside him. The woman, wearing a black tank top and a denim jacket draped over her shoulders, smiles warmly, her eyes reflecting genuine interest and connection. The car's interior, with its beige seats and simple design, provides a backdrop that emphasizes their interaction. The scene captures a moment of shared understanding and connection, set against the soft, diffused light of an overcast day. A medium shot from a slightly angled perspective, focusing on their expressions and body language.",
+    "cond_image": "examples/multi/1/multi1.png",
+    "audio_type": "add",
+    "cond_audio": {
+        "person1": "examples/multi/1/1.WAV",
+        "person2": "examples/multi/1/2.WAV"
+    },
+    "bbox": {
+        "person1": [160, 120, 1280, 1080], 
+        "person2": [160, 1320, 1280, 2280]
+    }
+}

examples/multitalk_example_2.json

@@ -0,0 +1,9 @@
+{
+    "prompt": "In a cozy recording studio, a man and a woman are singing together with passion and emotion. The man, with short brown hair, wears a light gray button-up shirt, his expression filled with concentration and warmth. The woman, with long wavy brown hair, dons a sleeveless dress adorned with small polka dots, her eyes closed as she belts out a heartfelt melody. The studio is equipped with professional microphones, and the background features soundproofing panels, creating an intimate and focused atmosphere. A close-up shot captures their expressions and the intensity of their performance.",
+    "cond_image": "examples/multi/2/multi2.png",
+    "audio_type": "para",
+    "cond_audio": {
+        "person1": "examples/multi/2/1.wav",
+        "person2": "examples/multi/2/1.wav"
+    }
+}

examples/multitalk_example_3.json

@@ -0,0 +1,9 @@
+{
+    "prompt": "In a cozy recording studio, a man and a woman are singing together. The man, with tousled brown hair, stands to the left, wearing a light green button-down shirt. His gaze is directed towards the woman, who is smiling warmly. She, with wavy dark hair, is dressed in a black floral dress and stands to the right, her eyes closed in enjoyment. Between them is a professional microphone, capturing their harmonious voices. The background features wooden panels and various audio equipment, creating an intimate and focused atmosphere. The lighting is soft and warm, highlighting their expressions and the intimate setting. A medium shot captures their interaction closely.",
+    "cond_image": "examples/multi/3/multi3.png",
+    "audio_type": "para",
+    "cond_audio": {
+        "person1": "examples/multi/3/1-man.WAV",
+        "person2": "examples/multi/3/1-woman.WAV"
+    }
+}

examples/multitalk_example_tts_1.json

@@ -0,0 +1,11 @@
+{
+    "prompt": "In a cozy recording studio, a man and a woman are singing together with passion and emotion. The man, with short brown hair, wears a light gray button-up shirt, his expression filled with concentration and warmth. The woman, with long wavy brown hair, dons a sleeveless dress adorned with small polka dots, her eyes closed as she belts out a heartfelt melody. The studio is equipped with professional microphones, and the background features soundproofing panels, creating an intimate and focused atmosphere. A close-up shot captures their expressions and the intensity of their performance.",
+    "cond_image": "examples/multi/1/multi1.png",
+    "audio_type": "para",
+    "tts_audio": {
+        "text": "(s1) do you know multi-talk? (s2) yes, I know it, that's amazing! (s1) Me too.",
+        "human1_voice": "weights/Kokoro-82M/voices/af_heart.pt",
+        "human2_voice": "weights/Kokoro-82M/voices/am_adam.pt"
+    },
+    "cond_audio":{}
+}

examples/single/1.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba2733897f561f747e6508734bff4eeee29d0a73638e5c39c0c0b806701d4e8b
+size 1888320

examples/single_example_1.json

@@ -0,0 +1,7 @@
+{
+    "prompt": "A woman is passionately singing into a professional microphone in a recording studio. She wears large black headphones and a dark cardigan over a gray top. Her long, wavy brown hair frames her face as she looks slightly upwards, her mouth open mid-song. The studio is equipped with various audio equipment, including a mixing console and a keyboard, with soundproofing panels on the walls. The lighting is warm and focused on her, creating a professional and intimate atmosphere. A close-up shot captures her expressive performance.",
+    "cond_image": "examples/single/single1.png",
+    "cond_audio": {
+        "person1": "examples/single/1.wav"
+    }
+}

examples/single_example_tts_1.json

@@ -0,0 +1,9 @@
+{
+    "prompt": "A woman is passionately singing into a professional microphone in a recording studio. She wears large black headphones and a dark cardigan over a gray top. Her long, wavy brown hair frames her face as she looks slightly upwards, her mouth open mid-song. The studio is equipped with various audio equipment, including a mixing console and a keyboard, with soundproofing panels on the walls. The lighting is warm and focused on her, creating a professional and intimate atmosphere. A close-up shot captures her expressive performance.",
+    "cond_image": "examples/single/single1.png",
+    "tts_audio": {
+        "text": "Welcome to multi-talk, this is an audio-driven video generation model for multiple person.",
+        "human1_voice": "weights/Kokoro-82M/voices/af_heart.pt"
+    },
+    "cond_audio":{}
+}

generate_multitalk.py

@@ -0,0 +1,638 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import argparse
+import logging
+import os
+import sys
+import json
+import warnings
+from datetime import datetime
+
+warnings.filterwarnings('ignore')
+
+import random
+
+import torch
+import torch.distributed as dist
+from PIL import Image
+import subprocess
+
+import wan
+from wan.configs import SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
+from wan.utils.utils import cache_image, cache_video, str2bool
+from wan.utils.multitalk_utils import save_video_ffmpeg
+from kokoro import KPipeline
+from transformers import Wav2Vec2FeatureExtractor
+from src.audio_analysis.wav2vec2 import Wav2Vec2Model
+
+import librosa
+import pyloudnorm as pyln
+import numpy as np
+from einops import rearrange
+import soundfile as sf
+import re
+
+
+def _validate_args(args):
+    # Basic check
+    assert args.ckpt_dir is not None, "Please specify the checkpoint directory."
+    assert args.task in WAN_CONFIGS, f"Unsupport task: {args.task}"
+
+    # The default sampling steps are 40 for image-to-video tasks and 50 for text-to-video tasks.
+    if args.sample_steps is None:
+        args.sample_steps = 40
+
+    if args.sample_shift is None:
+        if args.size == 'multitalk-480':
+            args.sample_shift = 7
+        elif args.size == 'multitalk-720':
+            args.sample_shift = 11
+        else:
+            raise NotImplementedError(f'Not supported size')
+
+    args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
+        0, 99999999)
+    # Size check
+    assert args.size in SUPPORTED_SIZES[
+        args.
+        task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description="Generate a image or video from a text prompt or image using Wan"
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="multitalk-14B",
+        choices=list(WAN_CONFIGS.keys()),
+        help="The task to run.")
+    parser.add_argument(
+        "--size",
+        type=str,
+        default="multitalk-480",
+        choices=list(SIZE_CONFIGS.keys()),
+        help="The buckget size of the generated video. The aspect ratio of the output video will follow that of the input image."
+    )
+    parser.add_argument(
+        "--frame_num",
+        type=int,
+        default=81,
+        help="How many frames to be generated in one clip. The number should be 4n+1"
+    )
+    parser.add_argument(
+        "--ckpt_dir",
+        type=str,
+        default=None,
+        help="The path to the Wan checkpoint directory.")
+    parser.add_argument(
+        "--quant_dir",
+        type=str,
+        default=None,
+        help="The path to the Wan quant checkpoint directory.")
+    parser.add_argument(
+        "--wav2vec_dir",
+        type=str,
+        default=None,
+        help="The path to the wav2vec checkpoint directory.")
+    parser.add_argument(
+        "--lora_dir",
+        type=str,
+        nargs='+',
+        default=None,
+        help="The paths to the LoRA checkpoint files."
+    )
+    parser.add_argument(
+        "--lora_scale",
+        type=float,
+        nargs='+',
+        default=[1.2],
+        help="Controls how much to influence the outputs with the LoRA parameters. Accepts multiple float values."
+    )
+    parser.add_argument(
+        "--offload_model",
+        type=str2bool,
+        default=None,
+        help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage."
+    )
+    parser.add_argument(
+        "--ulysses_size",
+        type=int,
+        default=1,
+        help="The size of the ulysses parallelism in DiT.")
+    parser.add_argument(
+        "--ring_size",
+        type=int,
+        default=1,
+        help="The size of the ring attention parallelism in DiT.")
+    parser.add_argument(
+        "--t5_fsdp",
+        action="store_true",
+        default=False,
+        help="Whether to use FSDP for T5.")
+    parser.add_argument(
+        "--t5_cpu",
+        action="store_true",
+        default=False,
+        help="Whether to place T5 model on CPU.")
+    parser.add_argument(
+        "--dit_fsdp",
+        action="store_true",
+        default=False,
+        help="Whether to use FSDP for DiT.")
+    parser.add_argument(
+        "--save_file",
+        type=str,
+        default=None,
+        help="The file to save the generated image or video to.")
+    parser.add_argument(
+        "--audio_save_dir",
+        type=str,
+        default='save_audio',
+        help="The path to save the audio embedding.")
+    parser.add_argument(
+        "--base_seed",
+        type=int,
+        default=42,
+        help="The seed to use for generating the image or video.")
+    parser.add_argument(
+        "--input_json",
+        type=str,
+        default='examples.json',
+        help="[meta file] The condition path to generate the video.")
+    parser.add_argument(
+        "--motion_frame",
+        type=int,
+        default=25,
+        help="Driven frame length used in the mode of long video genration.")
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="clip",
+        choices=['clip', 'streaming'],
+        help="clip: generate one video chunk, streaming: long video generation")
+    parser.add_argument(
+        "--sample_steps", type=int, default=None, help="The sampling steps.")
+    parser.add_argument(
+        "--sample_shift",
+        type=float,
+        default=None,
+        help="Sampling shift factor for flow matching schedulers.")
+    parser.add_argument(
+        "--sample_text_guide_scale",
+        type=float,
+        default=5.0,
+        help="Classifier free guidance scale for text control.")
+    parser.add_argument(
+        "--sample_audio_guide_scale",
+        type=float,
+        default=4.0,
+        help="Classifier free guidance scale for audio control.")
+    parser.add_argument(
+        "--num_persistent_param_in_dit",
+        type=int,
+        default=None,
+        required=False,
+        help="Maximum parameter quantity retained in video memory, small number to reduce VRAM required",
+    )
+    parser.add_argument(
+        "--audio_mode",
+        type=str,
+        default="localfile",
+        choices=['localfile', 'tts'],
+        help="localfile: audio from local wav file, tts: audio from TTS")
+    parser.add_argument(
+        "--use_teacache",
+        action="store_true",
+        default=False,
+        help="Enable teacache for video generation."
+    )
+    parser.add_argument(
+        "--teacache_thresh",
+        type=float,
+        default=0.2,
+        help="Threshold for teacache."
+    )
+    parser.add_argument(
+        "--use_apg",
+        action="store_true",
+        default=False,
+        help="Enable adaptive projected guidance for video generation (APG)."
+    )
+    parser.add_argument(
+        "--apg_momentum",
+        type=float,
+        default=-0.75,
+        help="Momentum used in adaptive projected guidance (APG)."
+    )
+    parser.add_argument(
+        "--apg_norm_threshold",
+        type=float,
+        default=55,
+        help="Norm threshold used in adaptive projected guidance (APG)."
+    )
+    parser.add_argument(
+        "--color_correction_strength",
+        type=float,
+        default=1.0,
+        help="strength for color correction [0.0 -- 1.0]."
+    )
+
+    parser.add_argument(
+        "--quant",
+        type=str,
+        default=None,
+        help="Quantization type, must be 'int8' or 'fp8'."
+    )
+    
+    args = parser.parse_args()
+
+    _validate_args(args)
+
+    return args
+
+def custom_init(device, wav2vec):    
+    audio_encoder = Wav2Vec2Model.from_pretrained(wav2vec, local_files_only=True).to(device)
+    audio_encoder.feature_extractor._freeze_parameters()
+    wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec, local_files_only=True)
+    return wav2vec_feature_extractor, audio_encoder
+
+def loudness_norm(audio_array, sr=16000, lufs=-23):
+    meter = pyln.Meter(sr)
+    loudness = meter.integrated_loudness(audio_array)
+    if abs(loudness) > 100:
+        return audio_array
+    normalized_audio = pyln.normalize.loudness(audio_array, loudness, lufs)
+    return normalized_audio
+
+def audio_prepare_multi(left_path, right_path, audio_type, sample_rate=16000):
+
+    if not (left_path=='None' or right_path=='None'):
+        human_speech_array1 = audio_prepare_single(left_path)
+        human_speech_array2 = audio_prepare_single(right_path)
+    elif left_path=='None':
+        human_speech_array2 = audio_prepare_single(right_path)
+        human_speech_array1 = np.zeros(human_speech_array2.shape[0])
+    elif right_path=='None':
+        human_speech_array1 = audio_prepare_single(left_path)
+        human_speech_array2 = np.zeros(human_speech_array1.shape[0])
+
+    if audio_type=='para':
+        new_human_speech1 = human_speech_array1
+        new_human_speech2 = human_speech_array2
+    elif audio_type=='add':
+        new_human_speech1 = np.concatenate([human_speech_array1[: human_speech_array1.shape[0]], np.zeros(human_speech_array2.shape[0])]) 
+        new_human_speech2 = np.concatenate([np.zeros(human_speech_array1.shape[0]), human_speech_array2[:human_speech_array2.shape[0]]])
+    sum_human_speechs = new_human_speech1 + new_human_speech2
+    return new_human_speech1, new_human_speech2, sum_human_speechs
+
+def _init_logging(rank):
+    # logging
+    if rank == 0:
+        # set format
+        logging.basicConfig(
+            level=logging.INFO,
+            format="[%(asctime)s] %(levelname)s: %(message)s",
+            handlers=[logging.StreamHandler(stream=sys.stdout)])
+    else:
+        logging.basicConfig(level=logging.ERROR)
+
+def get_embedding(speech_array, wav2vec_feature_extractor, audio_encoder, sr=16000, device='cpu'):
+    audio_duration = len(speech_array) / sr
+    video_length = audio_duration * 25 # Assume the video fps is 25
+
+    # wav2vec_feature_extractor
+    audio_feature = np.squeeze(
+        wav2vec_feature_extractor(speech_array, sampling_rate=sr).input_values
+    )
+    audio_feature = torch.from_numpy(audio_feature).float().to(device=device)
+    audio_feature = audio_feature.unsqueeze(0)
+
+    # audio encoder
+    with torch.no_grad():
+        embeddings = audio_encoder(audio_feature, seq_len=int(video_length), output_hidden_states=True)
+
+    if len(embeddings) == 0:
+        print("Fail to extract audio embedding")
+        return None
+
+    audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
+    audio_emb = rearrange(audio_emb, "b s d -> s b d")
+
+    audio_emb = audio_emb.cpu().detach()
+    return audio_emb
+
+def extract_audio_from_video(filename, sample_rate):
+    raw_audio_path = filename.split('/')[-1].split('.')[0]+'.wav'
+    ffmpeg_command = [
+        "ffmpeg",
+        "-y",
+        "-i",
+        str(filename),
+        "-vn",
+        "-acodec",
+        "pcm_s16le",
+        "-ar",
+        "16000",
+        "-ac",
+        "2",
+        str(raw_audio_path),
+    ]
+    subprocess.run(ffmpeg_command, check=True)
+    human_speech_array, sr = librosa.load(raw_audio_path, sr=sample_rate)
+    human_speech_array = loudness_norm(human_speech_array, sr)
+    os.remove(raw_audio_path)
+
+    return human_speech_array
+
+def audio_prepare_single(audio_path, sample_rate=16000):
+    ext = os.path.splitext(audio_path)[1].lower()
+    if ext in ['.mp4', '.mov', '.avi', '.mkv']:
+        human_speech_array = extract_audio_from_video(audio_path, sample_rate)
+        return human_speech_array
+    else:
+        human_speech_array, sr = librosa.load(audio_path, sr=sample_rate)
+        human_speech_array = loudness_norm(human_speech_array, sr)
+        return human_speech_array
+
+def process_tts_single(text, save_dir, voice1):    
+    s1_sentences = []
+
+    pipeline = KPipeline(lang_code='a', repo_id='weights/Kokoro-82M')
+
+    voice_tensor = torch.load(voice1, weights_only=True)
+    generator = pipeline(
+        text, voice=voice_tensor, # <= change voice here
+        speed=1, split_pattern=r'\n+'
+    )
+    audios = []
+    for i, (gs, ps, audio) in enumerate(generator):
+        audios.append(audio)
+    audios = torch.concat(audios, dim=0)
+    s1_sentences.append(audios)
+    s1_sentences = torch.concat(s1_sentences, dim=0)
+    save_path1 =f'{save_dir}/s1.wav'
+    sf.write(save_path1, s1_sentences, 24000) # save each audio file
+    s1, _ = librosa.load(save_path1, sr=16000)
+    return s1, save_path1
+    
+   
+
+def process_tts_multi(text, save_dir, voice1, voice2):
+    pattern = r'\(s(\d+)\)\s*(.*?)(?=\s*\(s\d+\)|$)'
+    matches = re.findall(pattern, text, re.DOTALL)
+    
+    s1_sentences = []
+    s2_sentences = []
+
+    pipeline = KPipeline(lang_code='a', repo_id='weights/Kokoro-82M')
+    for idx, (speaker, content) in enumerate(matches):
+        if speaker == '1':
+            voice_tensor = torch.load(voice1, weights_only=True)
+            generator = pipeline(
+                content, voice=voice_tensor, # <= change voice here
+                speed=1, split_pattern=r'\n+'
+            )
+            audios = []
+            for i, (gs, ps, audio) in enumerate(generator):
+                audios.append(audio)
+            audios = torch.concat(audios, dim=0)
+            s1_sentences.append(audios)
+            s2_sentences.append(torch.zeros_like(audios))
+        elif speaker == '2':
+            voice_tensor = torch.load(voice2, weights_only=True)
+            generator = pipeline(
+                content, voice=voice_tensor, # <= change voice here
+                speed=1, split_pattern=r'\n+'
+            )
+            audios = []
+            for i, (gs, ps, audio) in enumerate(generator):
+                audios.append(audio)
+            audios = torch.concat(audios, dim=0)
+            s2_sentences.append(audios)
+            s1_sentences.append(torch.zeros_like(audios))
+    
+    s1_sentences = torch.concat(s1_sentences, dim=0)
+    s2_sentences = torch.concat(s2_sentences, dim=0)
+    sum_sentences = s1_sentences + s2_sentences
+    save_path1 =f'{save_dir}/s1.wav'
+    save_path2 =f'{save_dir}/s2.wav'
+    save_path_sum = f'{save_dir}/sum.wav'
+    sf.write(save_path1, s1_sentences, 24000) # save each audio file
+    sf.write(save_path2, s2_sentences, 24000)
+    sf.write(save_path_sum, sum_sentences, 24000)
+
+    s1, _ = librosa.load(save_path1, sr=16000)
+    s2, _ = librosa.load(save_path2, sr=16000)
+    # sum, _ = librosa.load(save_path_sum, sr=16000)
+    return s1, s2, save_path_sum
+
+def generate(args):
+    rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    local_rank = int(os.getenv("LOCAL_RANK", 0))
+    device = local_rank
+    _init_logging(rank)
+
+    if args.offload_model is None:
+        args.offload_model = False if world_size > 1 else True
+        logging.info(
+            f"offload_model is not specified, set to {args.offload_model}.")
+    if world_size > 1:
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(
+            backend="nccl",
+            init_method="env://",
+            rank=rank,
+            world_size=world_size)
+    else:
+        assert not (
+            args.t5_fsdp or args.dit_fsdp
+        ), f"t5_fsdp and dit_fsdp are not supported in non-distributed environments."
+        assert not (
+            args.ulysses_size > 1 or args.ring_size > 1
+        ), f"context parallel are not supported in non-distributed environments."
+
+    if args.ulysses_size > 1 or args.ring_size > 1:
+        assert args.ulysses_size * args.ring_size == world_size, f"The number of ulysses_size and ring_size should be equal to the world size."
+        from xfuser.core.distributed import (
+            init_distributed_environment,
+            initialize_model_parallel,
+        )
+        init_distributed_environment(
+            rank=dist.get_rank(), world_size=dist.get_world_size())
+
+        initialize_model_parallel(
+            sequence_parallel_degree=dist.get_world_size(),
+            ring_degree=args.ring_size,
+            ulysses_degree=args.ulysses_size,
+        )
+
+    # TODO: use prompt refine
+    # if args.use_prompt_extend:
+    #     if args.prompt_extend_method == "dashscope":
+    #         prompt_expander = DashScopePromptExpander(
+    #             model_name=args.prompt_extend_model,
+    #             is_vl="i2v" in args.task or "flf2v" in args.task)
+    #     elif args.prompt_extend_method == "local_qwen":
+    #         prompt_expander = QwenPromptExpander(
+    #             model_name=args.prompt_extend_model,
+    #             is_vl="i2v" in args.task,
+    #             device=rank)
+    #     else:
+    #         raise NotImplementedError(
+    #             f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
+
+    cfg = WAN_CONFIGS[args.task]
+    if args.ulysses_size > 1:
+        assert cfg.num_heads % args.ulysses_size == 0, f"`{cfg.num_heads=}` cannot be divided evenly by `{args.ulysses_size=}`."
+
+    logging.info(f"Generation job args: {args}")
+    logging.info(f"Generation model config: {cfg}")
+
+    if dist.is_initialized():
+        base_seed = [args.base_seed] if rank == 0 else [None]
+        dist.broadcast_object_list(base_seed, src=0)
+        args.base_seed = base_seed[0]
+
+    assert args.task == "multitalk-14B", 'You should choose multitalk in args.task.'
+    
+
+    # TODO: add prompt refine
+    # img = Image.open(args.image).convert("RGB")
+    # if args.use_prompt_extend:
+    #     logging.info("Extending prompt ...")
+    #     if rank == 0:
+    #         prompt_output = prompt_expander(
+    #             args.prompt,
+    #             tar_lang=args.prompt_extend_target_lang,
+    #             image=img,
+    #             seed=args.base_seed)
+    #         if prompt_output.status == False:
+    #             logging.info(
+    #                 f"Extending prompt failed: {prompt_output.message}")
+    #             logging.info("Falling back to original prompt.")
+    #             input_prompt = args.prompt
+    #         else:
+    #             input_prompt = prompt_output.prompt
+    #         input_prompt = [input_prompt]
+    #     else:
+    #         input_prompt = [None]
+    #     if dist.is_initialized():
+    #         dist.broadcast_object_list(input_prompt, src=0)
+    #     args.prompt = input_prompt[0]
+    #     logging.info(f"Extended prompt: {args.prompt}")
+
+    # read input files
+
+    
+
+    with open(args.input_json, 'r', encoding='utf-8') as f:
+        input_data = json.load(f)
+        
+        wav2vec_feature_extractor, audio_encoder= custom_init('cpu', args.wav2vec_dir)
+        args.audio_save_dir = os.path.join(args.audio_save_dir, input_data['cond_image'].split('/')[-1].split('.')[0])
+        os.makedirs(args.audio_save_dir,exist_ok=True)
+        
+        if args.audio_mode=='localfile':
+            if len(input_data['cond_audio'])==2:
+                new_human_speech1, new_human_speech2, sum_human_speechs = audio_prepare_multi(input_data['cond_audio']['person1'], input_data['cond_audio']['person2'], input_data['audio_type'])
+                audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
+                audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder)
+                emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+                emb2_path = os.path.join(args.audio_save_dir, '2.pt')
+                sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+                sf.write(sum_audio, sum_human_speechs, 16000)
+                torch.save(audio_embedding_1, emb1_path)
+                torch.save(audio_embedding_2, emb2_path)
+                input_data['cond_audio']['person1'] = emb1_path
+                input_data['cond_audio']['person2'] = emb2_path
+                input_data['video_audio'] = sum_audio
+            elif len(input_data['cond_audio'])==1:
+                human_speech = audio_prepare_single(input_data['cond_audio']['person1'])
+                audio_embedding = get_embedding(human_speech, wav2vec_feature_extractor, audio_encoder)
+                emb_path = os.path.join(args.audio_save_dir, '1.pt')
+                sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+                sf.write(sum_audio, human_speech, 16000)
+                torch.save(audio_embedding, emb_path)
+                input_data['cond_audio']['person1'] = emb_path
+                input_data['video_audio'] = sum_audio
+        elif args.audio_mode=='tts':
+            if 'human2_voice' not in input_data['tts_audio'].keys():
+                new_human_speech1, sum_audio = process_tts_single(input_data['tts_audio']['text'], args.audio_save_dir, input_data['tts_audio']['human1_voice'])
+                audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
+                emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+                torch.save(audio_embedding_1, emb1_path)
+                input_data['cond_audio']['person1'] = emb1_path
+                input_data['video_audio'] = sum_audio
+            else:
+                new_human_speech1, new_human_speech2, sum_audio = process_tts_multi(input_data['tts_audio']['text'], args.audio_save_dir, input_data['tts_audio']['human1_voice'], input_data['tts_audio']['human2_voice'])
+                audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
+                audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder)
+                emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+                emb2_path = os.path.join(args.audio_save_dir, '2.pt')
+                torch.save(audio_embedding_1, emb1_path)
+                torch.save(audio_embedding_2, emb2_path)
+                input_data['cond_audio']['person1'] = emb1_path
+                input_data['cond_audio']['person2'] = emb2_path
+                input_data['video_audio'] = sum_audio
+
+
+    logging.info("Creating MultiTalk pipeline.")
+    wan_i2v = wan.MultiTalkPipeline(
+        config=cfg,
+        checkpoint_dir=args.ckpt_dir,
+        quant_dir=args.quant_dir,
+        device_id=device,
+        rank=rank,
+        t5_fsdp=args.t5_fsdp,
+        dit_fsdp=args.dit_fsdp, 
+        use_usp=(args.ulysses_size > 1 or args.ring_size > 1),  
+        t5_cpu=args.t5_cpu,
+        lora_dir=args.lora_dir,
+        lora_scales=args.lora_scale,
+        quant=args.quant
+    )
+
+
+    if args.num_persistent_param_in_dit is not None:
+        wan_i2v.vram_management = True
+        wan_i2v.enable_vram_management(
+            num_persistent_param_in_dit=args.num_persistent_param_in_dit
+        )
+    
+    logging.info("Generating video ...")
+    video = wan_i2v.generate(
+        input_data,
+        size_buckget=args.size,
+        motion_frame=args.motion_frame,
+        frame_num=args.frame_num,
+        shift=args.sample_shift,
+        sampling_steps=args.sample_steps,
+        text_guide_scale=args.sample_text_guide_scale,
+        audio_guide_scale=args.sample_audio_guide_scale,
+        seed=args.base_seed,
+        offload_model=args.offload_model,
+        max_frames_num=args.frame_num if args.mode == 'clip' else 1000,
+        color_correction_strength = args.color_correction_strength,
+        extra_args=args,
+        )
+    
+
+    if rank == 0:
+        
+        if args.save_file is None:
+            formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+            formatted_prompt = input_data['prompt'].replace(" ", "_").replace("/",
+                                                                        "_")[:50]
+            args.save_file = f"{args.task}_{args.size.replace('*','x') if sys.platform=='win32' else args.size}_{args.ulysses_size}_{args.ring_size}_{formatted_prompt}_{formatted_time}"
+        
+        logging.info(f"Saving generated video to {args.save_file}.mp4")
+        save_video_ffmpeg(video, args.save_file, [input_data['video_audio']], high_quality_save=False)
+        
+    logging.info("Finished.")
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    generate(args)

kokoro/__init__.py

@@ -0,0 +1,23 @@
+__version__ = '0.9.4'
+
+from loguru import logger
+import sys
+
+# Remove default handler
+logger.remove()
+
+# Add custom handler with clean format including module and line number
+logger.add(
+    sys.stderr,
+    format="<green>{time:HH:mm:ss}</green> | <cyan>{module:>16}:{line}</cyan> | <level>{level: >8}</level> | <level>{message}</level>",
+    colorize=True,
+    level="INFO" # "DEBUG" to enable logger.debug("message") and up prints 
+                 # "ERROR" to enable only logger.error("message") prints
+                 # etc
+)
+
+# Disable before release or as needed
+logger.disable("kokoro")
+
+from .model import KModel
+from .pipeline import KPipeline

kokoro/__main__.py

@@ -0,0 +1,148 @@
+"""Kokoro TTS CLI
+Example usage:
+python3 -m kokoro --text "The sky above the port was the color of television, tuned to a dead channel." -o file.wav --debug
+
+echo "Bom dia mundo, como vão vocês" > text.txt
+python3 -m kokoro -i text.txt -l p --voice pm_alex > audio.wav
+
+Common issues:
+pip not installed: `uv pip install pip`
+(Temporary workaround while https://github.com/explosion/spaCy/issues/13747 is not fixed)
+
+espeak not installed: `apt-get install espeak-ng`
+"""
+
+import argparse
+import wave
+from pathlib import Path
+from typing import Generator, TYPE_CHECKING
+
+import numpy as np
+from loguru import logger
+
+languages = [
+    "a",  # American English
+    "b",  # British English
+    "h",  # Hindi
+    "e",  # Spanish
+    "f",  # French
+    "i",  # Italian
+    "p",  # Brazilian Portuguese
+    "j",  # Japanese
+    "z",  # Mandarin Chinese
+]
+
+if TYPE_CHECKING:
+    from kokoro import KPipeline
+
+
+def generate_audio(
+    text: str, kokoro_language: str, voice: str, speed=1
+) -> Generator["KPipeline.Result", None, None]:
+    from kokoro import KPipeline
+
+    if not voice.startswith(kokoro_language):
+        logger.warning(f"Voice {voice} is not made for language {kokoro_language}")
+    pipeline = KPipeline(lang_code=kokoro_language)
+    yield from pipeline(text, voice=voice, speed=speed, split_pattern=r"\n+")
+
+
+def generate_and_save_audio(
+    output_file: Path, text: str, kokoro_language: str, voice: str, speed=1
+) -> None:
+    with wave.open(str(output_file.resolve()), "wb") as wav_file:
+        wav_file.setnchannels(1)  # Mono audio
+        wav_file.setsampwidth(2)  # 2 bytes per sample (16-bit audio)
+        wav_file.setframerate(24000)  # Sample rate
+
+        for result in generate_audio(
+            text, kokoro_language=kokoro_language, voice=voice, speed=speed
+        ):
+            logger.debug(result.phonemes)
+            if result.audio is None:
+                continue
+            audio_bytes = (result.audio.numpy() * 32767).astype(np.int16).tobytes()
+            wav_file.writeframes(audio_bytes)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-m",
+        "--voice",
+        default="af_heart",
+        help="Voice to use",
+    )
+    parser.add_argument(
+        "-l",
+        "--language",
+        help="Language to use (defaults to the one corresponding to the voice)",
+        choices=languages,
+    )
+    parser.add_argument(
+        "-o",
+        "--output-file",
+        "--output_file",
+        type=Path,
+        help="Path to output WAV file",
+        required=True,
+    )
+    parser.add_argument(
+        "-i",
+        "--input-file",
+        "--input_file",
+        type=Path,
+        help="Path to input text file (default: stdin)",
+    )
+    parser.add_argument(
+        "-t",
+        "--text",
+        help="Text to use instead of reading from stdin",
+    )
+    parser.add_argument(
+        "-s",
+        "--speed",
+        type=float,
+        default=1.0,
+        help="Speech speed",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Print DEBUG messages to console",
+    )
+    args = parser.parse_args()
+    if args.debug:
+        logger.level("DEBUG")
+    logger.debug(args)
+
+    lang = args.language or args.voice[0]
+
+    if args.text is not None and args.input_file is not None:
+        raise Exception("You cannot specify both 'text' and 'input_file'")
+    elif args.text:
+        text = args.text
+    elif args.input_file:
+        file: Path = args.input_file
+        text = file.read_text()
+    else:
+        import sys
+        print("Press Ctrl+D to stop reading input and start generating", flush=True)
+        text = '\n'.join(sys.stdin)
+
+    logger.debug(f"Input text: {text!r}")
+
+    out_file: Path = args.output_file
+    if not out_file.suffix == ".wav":
+        logger.warning("The output file name should end with .wav")
+    generate_and_save_audio(
+        output_file=out_file,
+        text=text,
+        kokoro_language=lang,
+        voice=args.voice,
+        speed=args.speed,
+    )
+
+
+if __name__ == "__main__":
+    main()

kokoro/custom_stft.py

@@ -0,0 +1,197 @@
+from attr import attr
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CustomSTFT(nn.Module):
+    """
+    STFT/iSTFT without unfold/complex ops, using conv1d and conv_transpose1d.
+
+    - forward STFT => Real-part conv1d + Imag-part conv1d
+    - inverse STFT => Real-part conv_transpose1d + Imag-part conv_transpose1d + sum
+    - avoids F.unfold, so easier to export to ONNX
+    - uses replicate or constant padding for 'center=True' to approximate 'reflect' 
+      (reflect is not supported for dynamic shapes in ONNX)
+    """
+
+    def __init__(
+        self,
+        filter_length=800,
+        hop_length=200,
+        win_length=800,
+        window="hann",
+        center=True,
+        pad_mode="replicate",  # or 'constant'
+    ):
+        super().__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_fft = filter_length
+        self.center = center
+        self.pad_mode = pad_mode
+
+        # Number of frequency bins for real-valued STFT with onesided=True
+        self.freq_bins = self.n_fft // 2 + 1
+
+        # Build window
+        assert window == 'hann', window
+        window_tensor = torch.hann_window(win_length, periodic=True, dtype=torch.float32)
+        if self.win_length < self.n_fft:
+            # Zero-pad up to n_fft
+            extra = self.n_fft - self.win_length
+            window_tensor = F.pad(window_tensor, (0, extra))
+        elif self.win_length > self.n_fft:
+            window_tensor = window_tensor[: self.n_fft]
+        self.register_buffer("window", window_tensor)
+
+        # Precompute forward DFT (real, imag)
+        # PyTorch stft uses e^{-j 2 pi k n / N} => real=cos(...), imag=-sin(...)
+        n = np.arange(self.n_fft)
+        k = np.arange(self.freq_bins)
+        angle = 2 * np.pi * np.outer(k, n) / self.n_fft  # shape (freq_bins, n_fft)
+        dft_real = np.cos(angle)
+        dft_imag = -np.sin(angle)  # note negative sign
+
+        # Combine window and dft => shape (freq_bins, filter_length)
+        # We'll make 2 conv weight tensors of shape (freq_bins, 1, filter_length).
+        forward_window = window_tensor.numpy()  # shape (n_fft,)
+        forward_real = dft_real * forward_window  # (freq_bins, n_fft)
+        forward_imag = dft_imag * forward_window
+
+        # Convert to PyTorch
+        forward_real_torch = torch.from_numpy(forward_real).float()
+        forward_imag_torch = torch.from_numpy(forward_imag).float()
+
+        # Register as Conv1d weight => (out_channels, in_channels, kernel_size)
+        # out_channels = freq_bins, in_channels=1, kernel_size=n_fft
+        self.register_buffer(
+            "weight_forward_real", forward_real_torch.unsqueeze(1)
+        )
+        self.register_buffer(
+            "weight_forward_imag", forward_imag_torch.unsqueeze(1)
+        )
+
+        # Precompute inverse DFT
+        # Real iFFT formula => scale = 1/n_fft, doubling for bins 1..freq_bins-2 if n_fft even, etc.
+        # For simplicity, we won't do the "DC/nyquist not doubled" approach here. 
+        # If you want perfect real iSTFT, you can add that logic. 
+        # This version just yields good approximate reconstruction with Hann + typical overlap.
+        inv_scale = 1.0 / self.n_fft
+        n = np.arange(self.n_fft)
+        angle_t = 2 * np.pi * np.outer(n, k) / self.n_fft  # shape (n_fft, freq_bins)
+        idft_cos = np.cos(angle_t).T  # => (freq_bins, n_fft)
+        idft_sin = np.sin(angle_t).T  # => (freq_bins, n_fft)
+
+        # Multiply by window again for typical overlap-add
+        # We also incorporate the scale factor 1/n_fft
+        inv_window = window_tensor.numpy() * inv_scale
+        backward_real = idft_cos * inv_window  # (freq_bins, n_fft)
+        backward_imag = idft_sin * inv_window
+
+        # We'll implement iSTFT as real+imag conv_transpose with stride=hop.
+        self.register_buffer(
+            "weight_backward_real", torch.from_numpy(backward_real).float().unsqueeze(1)
+        )
+        self.register_buffer(
+            "weight_backward_imag", torch.from_numpy(backward_imag).float().unsqueeze(1)
+        )
+        
+
+
+    def transform(self, waveform: torch.Tensor):
+        """
+        Forward STFT => returns magnitude, phase
+        Output shape => (batch, freq_bins, frames)
+        """
+        # waveform shape => (B, T).  conv1d expects (B, 1, T).
+        # Optional center pad
+        if self.center:
+            pad_len = self.n_fft // 2
+            waveform = F.pad(waveform, (pad_len, pad_len), mode=self.pad_mode)
+
+        x = waveform.unsqueeze(1)  # => (B, 1, T)
+        # Convolution to get real part => shape (B, freq_bins, frames)
+        real_out = F.conv1d(
+            x,
+            self.weight_forward_real,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        # Imag part
+        imag_out = F.conv1d(
+            x,
+            self.weight_forward_imag,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        # magnitude, phase
+        magnitude = torch.sqrt(real_out**2 + imag_out**2 + 1e-14)
+        phase = torch.atan2(imag_out, real_out)
+        # Handle the case where imag_out is 0 and real_out is negative to correct ONNX atan2 to match PyTorch
+        # In this case, PyTorch returns pi, ONNX returns -pi
+        correction_mask = (imag_out == 0) & (real_out < 0)
+        phase[correction_mask] = torch.pi
+        return magnitude, phase
+
+
+    def inverse(self, magnitude: torch.Tensor, phase: torch.Tensor, length=None):
+        """
+        Inverse STFT => returns waveform shape (B, T).
+        """
+        # magnitude, phase => (B, freq_bins, frames)
+        # Re-create real/imag => shape (B, freq_bins, frames)
+        real_part = magnitude * torch.cos(phase)
+        imag_part = magnitude * torch.sin(phase)
+
+        # conv_transpose wants shape (B, freq_bins, frames). We'll treat "frames" as time dimension
+        # so we do (B, freq_bins, frames) => (B, freq_bins, frames)
+        # But PyTorch conv_transpose1d expects (B, in_channels, input_length)
+        real_part = real_part  # (B, freq_bins, frames)
+        imag_part = imag_part
+
+        # real iSTFT => convolve with "backward_real", "backward_imag", and sum
+        # We'll do 2 conv_transpose calls, each giving (B, 1, time),
+        # then add them => (B, 1, time).
+        real_rec = F.conv_transpose1d(
+            real_part,
+            self.weight_backward_real,  # shape (freq_bins, 1, filter_length)
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        imag_rec = F.conv_transpose1d(
+            imag_part,
+            self.weight_backward_imag,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        # sum => (B, 1, time)
+        waveform = real_rec - imag_rec  # typical real iFFT has minus for imaginary part
+
+        # If we used "center=True" in forward, we should remove pad
+        if self.center:
+            pad_len = self.n_fft // 2
+            # Because of transposed convolution, total length might have extra samples
+            # We remove `pad_len` from start & end if possible
+            waveform = waveform[..., pad_len:-pad_len]
+
+        # If a specific length is desired, clamp
+        if length is not None:
+            waveform = waveform[..., :length]
+
+        # shape => (B, T)
+        return waveform
+
+    def forward(self, x: torch.Tensor):
+        """
+        Full STFT -> iSTFT pass: returns time-domain reconstruction.
+        Same interface as your original code.
+        """
+        mag, phase = self.transform(x)
+        return self.inverse(mag, phase, length=x.shape[-1])

kokoro/istftnet.py

@@ -0,0 +1,421 @@
+# ADAPTED from https://github.com/yl4579/StyleTTS2/blob/main/Modules/istftnet.py
+from kokoro.custom_stft import CustomSTFT
+from torch.nn.utils import weight_norm
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# https://github.com/yl4579/StyleTTS2/blob/main/Modules/utils.py
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+class AdaIN1d(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        # affine should be False, however there's a bug in the old torch.onnx.export (not newer dynamo) that causes the channel dimension to be lost if affine=False. When affine is true, there's additional learnably parameters. This shouldn't really matter setting it to True, since we're in inference mode
+        self.norm = nn.InstanceNorm1d(num_features, affine=True)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+
+class AdaINResBlock1(nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
+        super(AdaINResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList([
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                                  padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                                  padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                                  padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+        self.convs2 = nn.ModuleList([
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                                  padding=get_padding(kernel_size, 1))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                                  padding=get_padding(kernel_size, 1))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                                  padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+        self.adain1 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+        self.adain2 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+        self.alpha1 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
+        self.alpha2 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])
+
+    def forward(self, x, s):
+        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
+            xt = n1(x, s)
+            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
+            xt = c1(xt)
+            xt = n2(xt, s)
+            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+
+class TorchSTFT(nn.Module):
+    def __init__(self, filter_length=800, hop_length=200, win_length=800, window='hann'):
+        super().__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        assert window == 'hann', window
+        self.window = torch.hann_window(win_length, periodic=True, dtype=torch.float32)
+
+    def transform(self, input_data):
+        forward_transform = torch.stft(
+            input_data,
+            self.filter_length, self.hop_length, self.win_length, window=self.window.to(input_data.device),
+            return_complex=True)
+        return torch.abs(forward_transform), torch.angle(forward_transform)
+
+    def inverse(self, magnitude, phase):
+        inverse_transform = torch.istft(
+            magnitude * torch.exp(phase * 1j),
+            self.filter_length, self.hop_length, self.win_length, window=self.window.to(magnitude.device))
+        return inverse_transform.unsqueeze(-2)  # unsqueeze to stay consistent with conv_transpose1d implementation
+
+    def forward(self, input_data):
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction
+
+
+class SineGen(nn.Module):
+    """ Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(torch.pi) or cos(0)
+    """
+    def __init__(self, samp_rate, upsample_scale, harmonic_num=0,
+                 sine_amp=0.1, noise_std=0.003,
+                 voiced_threshold=0,
+                 flag_for_pulse=False):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = self.harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+        self.flag_for_pulse = flag_for_pulse
+        self.upsample_scale = upsample_scale
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = (f0 > self.voiced_threshold).type(torch.float32)
+        return uv
+
+    def _f02sine(self, f0_values):
+        """ f0_values: (batchsize, length, dim)
+            where dim indicates fundamental tone and overtones
+        """
+        # convert to F0 in rad. The interger part n can be ignored
+        # because 2 * torch.pi * n doesn't affect phase
+        rad_values = (f0_values / self.sampling_rate) % 1
+        # initial phase noise (no noise for fundamental component)
+        rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device)
+        rand_ini[:, 0] = 0
+        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
+        if not self.flag_for_pulse:
+            rad_values = F.interpolate(rad_values.transpose(1, 2), scale_factor=1/self.upsample_scale, mode="linear").transpose(1, 2)
+            phase = torch.cumsum(rad_values, dim=1) * 2 * torch.pi
+            phase = F.interpolate(phase.transpose(1, 2) * self.upsample_scale, scale_factor=self.upsample_scale, mode="linear").transpose(1, 2)
+            sines = torch.sin(phase)
+        else:
+            # If necessary, make sure that the first time step of every
+            # voiced segments is sin(pi) or cos(0)
+            # This is used for pulse-train generation
+            # identify the last time step in unvoiced segments
+            uv = self._f02uv(f0_values)
+            uv_1 = torch.roll(uv, shifts=-1, dims=1)
+            uv_1[:, -1, :] = 1
+            u_loc = (uv < 1) * (uv_1 > 0)
+            # get the instantanouse phase
+            tmp_cumsum = torch.cumsum(rad_values, dim=1)
+            # different batch needs to be processed differently
+            for idx in range(f0_values.shape[0]):
+                temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :]
+                temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :]
+                # stores the accumulation of i.phase within
+                # each voiced segments
+                tmp_cumsum[idx, :, :] = 0
+                tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum
+            # rad_values - tmp_cumsum: remove the accumulation of i.phase
+            # within the previous voiced segment.
+            i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1)
+            # get the sines
+            sines = torch.cos(i_phase * 2 * torch.pi)
+        return sines
+
+    def forward(self, f0):
+        """ sine_tensor, uv = forward(f0)
+        input F0: tensor(batchsize=1, length, dim=1)
+                  f0 for unvoiced steps should be 0
+        output sine_tensor: tensor(batchsize=1, length, dim)
+        output uv: tensor(batchsize=1, length, 1)
+        """
+        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim, device=f0.device)
+        # fundamental component
+        fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
+        # generate sine waveforms
+        sine_waves = self._f02sine(fn) * self.sine_amp
+        # generate uv signal
+        # uv = torch.ones(f0.shape)
+        # uv = uv * (f0 > self.voiced_threshold)
+        uv = self._f02uv(f0)
+        # noise: for unvoiced should be similar to sine_amp
+        #        std = self.sine_amp/3 -> max value ~ self.sine_amp
+        #        for voiced regions is self.noise_std
+        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+        noise = noise_amp * torch.randn_like(sine_waves)
+        # first: set the unvoiced part to 0 by uv
+        # then: additive noise
+        sine_waves = sine_waves * uv + noise
+        return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(nn.Module):
+    """ SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+    def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0):
+        super(SourceModuleHnNSF, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(sampling_rate, upsample_scale, harmonic_num,
+                                 sine_amp, add_noise_std, voiced_threshod)
+        # to merge source harmonics into a single excitation
+        self.l_linear = nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = nn.Tanh()
+
+    def forward(self, x):
+        """
+        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+        F0_sampled (batchsize, length, 1)
+        Sine_source (batchsize, length, 1)
+        noise_source (batchsize, length 1)
+        """
+        # source for harmonic branch
+        with torch.no_grad():
+            sine_wavs, uv, _ = self.l_sin_gen(x)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+        # source for noise branch, in the same shape as uv
+        noise = torch.randn_like(uv) * self.sine_amp / 3
+        return sine_merge, noise, uv
+
+
+class Generator(nn.Module):
+    def __init__(self, style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, disable_complex=False):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.m_source = SourceModuleHnNSF(
+                    sampling_rate=24000,
+                    upsample_scale=math.prod(upsample_rates) * gen_istft_hop_size,
+                    harmonic_num=8, voiced_threshod=10)
+        self.f0_upsamp = nn.Upsample(scale_factor=math.prod(upsample_rates) * gen_istft_hop_size)
+        self.noise_convs = nn.ModuleList()
+        self.noise_res = nn.ModuleList()
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(weight_norm(
+                nn.ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
+                                   k, u, padding=(k-u)//2)))
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel//(2**(i+1))
+            for j, (k, d) in enumerate(zip(resblock_kernel_sizes,resblock_dilation_sizes)):
+                self.resblocks.append(AdaINResBlock1(ch, k, d, style_dim))
+            c_cur = upsample_initial_channel // (2 ** (i + 1))
+            if i + 1 < len(upsample_rates):
+                stride_f0 = math.prod(upsample_rates[i + 1:])
+                self.noise_convs.append(nn.Conv1d(
+                    gen_istft_n_fft + 2, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2))
+                self.noise_res.append(AdaINResBlock1(c_cur, 7, [1,3,5], style_dim))
+            else:
+                self.noise_convs.append(nn.Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1))
+                self.noise_res.append(AdaINResBlock1(c_cur, 11, [1,3,5], style_dim))
+        self.post_n_fft = gen_istft_n_fft
+        self.conv_post = weight_norm(nn.Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.reflection_pad = nn.ReflectionPad1d((1, 0))
+        self.stft = (
+            CustomSTFT(filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft)
+            if disable_complex
+            else TorchSTFT(filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft)
+        )
+
+    def forward(self, x, s, f0):
+        with torch.no_grad():
+            f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
+            har_source, noi_source, uv = self.m_source(f0)
+            har_source = har_source.transpose(1, 2).squeeze(1)
+            har_spec, har_phase = self.stft.transform(har_source)
+            har = torch.cat([har_spec, har_phase], dim=1)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, negative_slope=0.1) 
+            x_source = self.noise_convs[i](har)
+            x_source = self.noise_res[i](x_source, s)
+            x = self.ups[i](x)
+            if i == self.num_upsamples - 1:
+                x = self.reflection_pad(x)
+            x = x + x_source
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x, s)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x, s)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :])
+        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
+        return self.stft.inverse(spec, phase)
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == 'none':
+            return x
+        else:
+            return F.interpolate(x, scale_factor=2, mode='nearest')
+
+
+class AdainResBlk1d(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2), upsample='none', dropout_p=0.0):
+        super().__init__()
+        self.actv = actv
+        self.upsample_type = upsample
+        self.upsample = UpSample1d(upsample)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out, style_dim)
+        self.dropout = nn.Dropout(dropout_p)
+        if upsample == 'none':
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
+
+    def _build_weights(self, dim_in, dim_out, style_dim):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
+        self.norm1 = AdaIN1d(style_dim, dim_in)
+        self.norm2 = AdaIN1d(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        x = self.upsample(x)
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        x = self.pool(x)
+        x = self.conv1(self.dropout(x))
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(self.dropout(x))
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) * torch.rsqrt(torch.tensor(2))
+        return out
+
+
+class Decoder(nn.Module):
+    def __init__(self, dim_in, style_dim, dim_out, 
+                 resblock_kernel_sizes,
+                 upsample_rates,
+                 upsample_initial_channel,
+                 resblock_dilation_sizes,
+                 upsample_kernel_sizes,
+                 gen_istft_n_fft, gen_istft_hop_size,
+                 disable_complex=False):
+        super().__init__()
+        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
+        self.decode = nn.ModuleList()
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 512, style_dim, upsample=True))
+        self.F0_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        self.N_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        self.asr_res = nn.Sequential(weight_norm(nn.Conv1d(512, 64, kernel_size=1)))
+        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates, 
+                                   upsample_initial_channel, resblock_dilation_sizes, 
+                                   upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, disable_complex=disable_complex)
+
+    def forward(self, asr, F0_curve, N, s):
+        F0 = self.F0_conv(F0_curve.unsqueeze(1))
+        N = self.N_conv(N.unsqueeze(1))
+        x = torch.cat([asr, F0, N], axis=1)
+        x = self.encode(x, s)
+        asr_res = self.asr_res(asr)
+        res = True
+        for block in self.decode:
+            if res:
+                x = torch.cat([x, asr_res, F0, N], axis=1)
+            x = block(x, s)
+            if block.upsample_type != "none":
+                res = False
+        x = self.generator(x, s, F0_curve)
+        return x

kokoro/model.py

@@ -0,0 +1,155 @@
+from .istftnet import Decoder
+from .modules import CustomAlbert, ProsodyPredictor, TextEncoder
+from dataclasses import dataclass
+from huggingface_hub import hf_hub_download
+from loguru import logger
+from transformers import AlbertConfig
+from typing import Dict, Optional, Union
+import json
+import torch
+import os
+
+class KModel(torch.nn.Module):
+    '''
+    KModel is a torch.nn.Module with 2 main responsibilities:
+    1. Init weights, downloading config.json + model.pth from HF if needed
+    2. forward(phonemes: str, ref_s: FloatTensor) -> (audio: FloatTensor)
+
+    You likely only need one KModel instance, and it can be reused across
+    multiple KPipelines to avoid redundant memory allocation.
+
+    Unlike KPipeline, KModel is language-blind.
+
+    KModel stores self.vocab and thus knows how to map phonemes -> input_ids,
+    so there is no need to repeatedly download config.json outside of KModel.
+    '''
+
+    MODEL_NAMES = {
+        'hexgrad/Kokoro-82M': 'kokoro-v1_0.pth',
+        'hexgrad/Kokoro-82M-v1.1-zh': 'kokoro-v1_1-zh.pth',
+    }
+
+    def __init__(
+        self,
+        repo_id: Optional[str] = None,
+        config: Union[Dict, str, None] = None,
+        model: Optional[str] = None,
+        disable_complex: bool = False
+    ):
+        super().__init__()
+        if repo_id is None:
+            repo_id = 'hexgrad/Kokoro-82M'
+            print(f"WARNING: Defaulting repo_id to {repo_id}. Pass repo_id='{repo_id}' to suppress this warning.")
+        self.repo_id = repo_id
+        if not isinstance(config, dict):
+            if not config:
+                logger.debug("No config provided, downloading from HF")
+                config = hf_hub_download(repo_id=repo_id, filename='config.json')
+            with open(config, 'r', encoding='utf-8') as r:
+                config = json.load(r)
+                logger.debug(f"Loaded config: {config}")
+        self.vocab = config['vocab']
+        self.bert = CustomAlbert(AlbertConfig(vocab_size=config['n_token'], **config['plbert']))
+        self.bert_encoder = torch.nn.Linear(self.bert.config.hidden_size, config['hidden_dim'])
+        self.context_length = self.bert.config.max_position_embeddings
+        self.predictor = ProsodyPredictor(
+            style_dim=config['style_dim'], d_hid=config['hidden_dim'],
+            nlayers=config['n_layer'], max_dur=config['max_dur'], dropout=config['dropout']
+        )
+        self.text_encoder = TextEncoder(
+            channels=config['hidden_dim'], kernel_size=config['text_encoder_kernel_size'],
+            depth=config['n_layer'], n_symbols=config['n_token']
+        )
+        self.decoder = Decoder(
+            dim_in=config['hidden_dim'], style_dim=config['style_dim'],
+            dim_out=config['n_mels'], disable_complex=disable_complex, **config['istftnet']
+        )
+        if not model:
+            try:
+                model = hf_hub_download(repo_id=repo_id, filename=KModel.MODEL_NAMES[repo_id])
+            except:
+                model = os.path.join(repo_id, 'kokoro-v1_0.pth')
+        for key, state_dict in torch.load(model, map_location='cpu', weights_only=True).items():
+            assert hasattr(self, key), key
+            try:
+                getattr(self, key).load_state_dict(state_dict)
+            except:
+                logger.debug(f"Did not load {key} from state_dict")
+                state_dict = {k[7:]: v for k, v in state_dict.items()}
+                getattr(self, key).load_state_dict(state_dict, strict=False)
+
+    @property
+    def device(self):
+        return self.bert.device
+
+    @dataclass
+    class Output:
+        audio: torch.FloatTensor
+        pred_dur: Optional[torch.LongTensor] = None
+
+    @torch.no_grad()
+    def forward_with_tokens(
+        self,
+        input_ids: torch.LongTensor,
+        ref_s: torch.FloatTensor,
+        speed: float = 1
+    ) -> tuple[torch.FloatTensor, torch.LongTensor]:
+        input_lengths = torch.full(
+            (input_ids.shape[0],), 
+            input_ids.shape[-1], 
+            device=input_ids.device,
+            dtype=torch.long
+        )
+
+        text_mask = torch.arange(input_lengths.max()).unsqueeze(0).expand(input_lengths.shape[0], -1).type_as(input_lengths)
+        text_mask = torch.gt(text_mask+1, input_lengths.unsqueeze(1)).to(self.device)
+        bert_dur = self.bert(input_ids, attention_mask=(~text_mask).int())
+        d_en = self.bert_encoder(bert_dur).transpose(-1, -2)
+        s = ref_s[:, 128:]
+        d = self.predictor.text_encoder(d_en, s, input_lengths, text_mask)
+        x, _ = self.predictor.lstm(d)
+        duration = self.predictor.duration_proj(x)
+        duration = torch.sigmoid(duration).sum(axis=-1) / speed
+        pred_dur = torch.round(duration).clamp(min=1).long().squeeze()
+        indices = torch.repeat_interleave(torch.arange(input_ids.shape[1], device=self.device), pred_dur)
+        pred_aln_trg = torch.zeros((input_ids.shape[1], indices.shape[0]), device=self.device)
+        pred_aln_trg[indices, torch.arange(indices.shape[0])] = 1
+        pred_aln_trg = pred_aln_trg.unsqueeze(0).to(self.device)
+        en = d.transpose(-1, -2) @ pred_aln_trg
+        F0_pred, N_pred = self.predictor.F0Ntrain(en, s)
+        t_en = self.text_encoder(input_ids, input_lengths, text_mask)
+        asr = t_en @ pred_aln_trg
+        audio = self.decoder(asr, F0_pred, N_pred, ref_s[:, :128]).squeeze()
+        return audio, pred_dur
+
+    def forward(
+        self,
+        phonemes: str,
+        ref_s: torch.FloatTensor,
+        speed: float = 1,
+        return_output: bool = False
+    ) -> Union['KModel.Output', torch.FloatTensor]:
+        input_ids = list(filter(lambda i: i is not None, map(lambda p: self.vocab.get(p), phonemes)))
+        logger.debug(f"phonemes: {phonemes} -> input_ids: {input_ids}")
+        assert len(input_ids)+2 <= self.context_length, (len(input_ids)+2, self.context_length)
+        input_ids = torch.LongTensor([[0, *input_ids, 0]]).to(self.device)
+        ref_s = ref_s.to(self.device)
+        audio, pred_dur = self.forward_with_tokens(input_ids, ref_s, speed)
+        audio = audio.squeeze().cpu()
+        pred_dur = pred_dur.cpu() if pred_dur is not None else None
+        logger.debug(f"pred_dur: {pred_dur}")
+        return self.Output(audio=audio, pred_dur=pred_dur) if return_output else audio
+
+class KModelForONNX(torch.nn.Module):
+    def __init__(self, kmodel: KModel):
+        super().__init__()
+        self.kmodel = kmodel
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        ref_s: torch.FloatTensor,
+        speed: float = 1
+    ) -> tuple[torch.FloatTensor, torch.LongTensor]:
+        waveform, duration = self.kmodel.forward_with_tokens(input_ids, ref_s, speed)
+        return waveform, duration

kokoro/modules.py

@@ -0,0 +1,183 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/models.py
+from .istftnet import AdainResBlk1d
+from torch.nn.utils import weight_norm
+from transformers import AlbertModel
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LinearNorm(nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
+        super(LinearNorm, self).__init__()
+        self.linear_layer = nn.Linear(in_dim, out_dim, bias=bias)
+        nn.init.xavier_uniform_(self.linear_layer.weight, gain=nn.init.calculate_gain(w_init_gain))
+
+    def forward(self, x):
+        return self.linear_layer(x)
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)):
+        super().__init__()
+        self.embedding = nn.Embedding(n_symbols, channels)
+        padding = (kernel_size - 1) // 2
+        self.cnn = nn.ModuleList()
+        for _ in range(depth):
+            self.cnn.append(nn.Sequential(
+                weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)),
+                LayerNorm(channels),
+                actv,
+                nn.Dropout(0.2),
+            ))
+        self.lstm = nn.LSTM(channels, channels//2, 1, batch_first=True, bidirectional=True)
+
+    def forward(self, x, input_lengths, m):
+        x = self.embedding(x)  # [B, T, emb]
+        x = x.transpose(1, 2)  # [B, emb, T]
+        m = m.unsqueeze(1)
+        x.masked_fill_(m, 0.0)
+        for c in self.cnn:
+            x = c(x)
+            x.masked_fill_(m, 0.0)
+        x = x.transpose(1, 2)  # [B, T, chn]
+        lengths = input_lengths if input_lengths.device == torch.device('cpu') else input_lengths.to('cpu')
+        x = nn.utils.rnn.pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
+        x = x.transpose(-1, -2)
+        x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device)
+        x_pad[:, :, :x.shape[-1]] = x
+        x = x_pad
+        x.masked_fill_(m, 0.0)
+        return x
+
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, style_dim, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+        self.fc = nn.Linear(style_dim, channels*2)
+
+    def forward(self, x, s):
+        x = x.transpose(-1, -2)
+        x = x.transpose(1, -1)
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        gamma, beta = gamma.transpose(1, -1), beta.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), eps=self.eps)
+        x = (1 + gamma) * x + beta
+        return x.transpose(1, -1).transpose(-1, -2)
+
+
+class ProsodyPredictor(nn.Module):
+    def __init__(self, style_dim, d_hid, nlayers, max_dur=50, dropout=0.1):
+        super().__init__()
+        self.text_encoder = DurationEncoder(sty_dim=style_dim, d_model=d_hid,nlayers=nlayers, dropout=dropout)
+        self.lstm = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True)
+        self.duration_proj = LinearNorm(d_hid, max_dur)
+        self.shared = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True)
+        self.F0 = nn.ModuleList()
+        self.F0.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
+        self.F0.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout))
+        self.F0.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout))
+        self.N = nn.ModuleList()
+        self.N.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
+        self.N.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout))
+        self.N.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout))
+        self.F0_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+        self.N_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+
+    def forward(self, texts, style, text_lengths, alignment, m):
+        d = self.text_encoder(texts, style, text_lengths, m)
+        m = m.unsqueeze(1)
+        lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu')
+        x = nn.utils.rnn.pack_padded_sequence(d, lengths, batch_first=True, enforce_sorted=False)
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
+        x_pad = torch.zeros([x.shape[0], m.shape[-1], x.shape[-1]], device=x.device)
+        x_pad[:, :x.shape[1], :] = x
+        x = x_pad
+        duration = self.duration_proj(nn.functional.dropout(x, 0.5, training=False))
+        en = (d.transpose(-1, -2) @ alignment)
+        return duration.squeeze(-1), en
+
+    def F0Ntrain(self, x, s):
+        x, _ = self.shared(x.transpose(-1, -2))
+        F0 = x.transpose(-1, -2)
+        for block in self.F0:
+            F0 = block(F0, s)
+        F0 = self.F0_proj(F0)
+        N = x.transpose(-1, -2)
+        for block in self.N:
+            N = block(N, s)
+        N = self.N_proj(N)
+        return F0.squeeze(1), N.squeeze(1)
+
+
+class DurationEncoder(nn.Module):
+    def __init__(self, sty_dim, d_model, nlayers, dropout=0.1):
+        super().__init__()
+        self.lstms = nn.ModuleList()
+        for _ in range(nlayers):
+            self.lstms.append(nn.LSTM(d_model + sty_dim, d_model // 2, num_layers=1, batch_first=True, bidirectional=True, dropout=dropout))
+            self.lstms.append(AdaLayerNorm(sty_dim, d_model))
+        self.dropout = dropout
+        self.d_model = d_model
+        self.sty_dim = sty_dim
+
+    def forward(self, x, style, text_lengths, m):
+        masks = m
+        x = x.permute(2, 0, 1)
+        s = style.expand(x.shape[0], x.shape[1], -1)
+        x = torch.cat([x, s], axis=-1)
+        x.masked_fill_(masks.unsqueeze(-1).transpose(0, 1), 0.0)
+        x = x.transpose(0, 1)
+        x = x.transpose(-1, -2)
+        for block in self.lstms:
+            if isinstance(block, AdaLayerNorm):
+                x = block(x.transpose(-1, -2), style).transpose(-1, -2)
+                x = torch.cat([x, s.permute(1, 2, 0)], axis=1)
+                x.masked_fill_(masks.unsqueeze(-1).transpose(-1, -2), 0.0)
+            else:
+                lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu')
+                x = x.transpose(-1, -2)
+                x = nn.utils.rnn.pack_padded_sequence(
+                    x, lengths, batch_first=True, enforce_sorted=False)
+                block.flatten_parameters()
+                x, _ = block(x)
+                x, _ = nn.utils.rnn.pad_packed_sequence(
+                    x, batch_first=True)
+                x = F.dropout(x, p=self.dropout, training=False)
+                x = x.transpose(-1, -2)
+                x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device)
+                x_pad[:, :, :x.shape[-1]] = x
+                x = x_pad
+
+        return x.transpose(-1, -2)
+
+
+# https://github.com/yl4579/StyleTTS2/blob/main/Utils/PLBERT/util.py
+class CustomAlbert(AlbertModel):
+    def forward(self, *args, **kwargs):
+        outputs = super().forward(*args, **kwargs)
+        return outputs.last_hidden_state

kokoro/pipeline.py

@@ -0,0 +1,445 @@
+from .model import KModel
+from dataclasses import dataclass
+from huggingface_hub import hf_hub_download
+from loguru import logger
+from misaki import en, espeak
+from typing import Callable, Generator, List, Optional, Tuple, Union
+import re
+import torch
+import os
+
+ALIASES = {
+    'en-us': 'a',
+    'en-gb': 'b',
+    'es': 'e',
+    'fr-fr': 'f',
+    'hi': 'h',
+    'it': 'i',
+    'pt-br': 'p',
+    'ja': 'j',
+    'zh': 'z',
+}
+
+LANG_CODES = dict(
+    # pip install misaki[en]
+    a='American English',
+    b='British English',
+
+    # espeak-ng
+    e='es',
+    f='fr-fr',
+    h='hi',
+    i='it',
+    p='pt-br',
+
+    # pip install misaki[ja]
+    j='Japanese',
+
+    # pip install misaki[zh]
+    z='Mandarin Chinese',
+)
+
+class KPipeline:
+    '''
+    KPipeline is a language-aware support class with 2 main responsibilities:
+    1. Perform language-specific G2P, mapping (and chunking) text -> phonemes
+    2. Manage and store voices, lazily downloaded from HF if needed
+
+    You are expected to have one KPipeline per language. If you have multiple
+    KPipelines, you should reuse one KModel instance across all of them.
+
+    KPipeline is designed to work with a KModel, but this is not required.
+    There are 2 ways to pass an existing model into a pipeline:
+    1. On init: us_pipeline = KPipeline(lang_code='a', model=model)
+    2. On call: us_pipeline(text, voice, model=model)
+
+    By default, KPipeline will automatically initialize its own KModel. To
+    suppress this, construct a "quiet" KPipeline with model=False.
+
+    A "quiet" KPipeline yields (graphemes, phonemes, None) without generating
+    any audio. You can use this to phonemize and chunk your text in advance.
+
+    A "loud" KPipeline _with_ a model yields (graphemes, phonemes, audio).
+    '''
+    def __init__(
+        self,
+        lang_code: str,
+        repo_id: Optional[str] = None,
+        model: Union[KModel, bool] = True,
+        trf: bool = False,
+        en_callable: Optional[Callable[[str], str]] = None,
+        device: Optional[str] = None
+    ):
+        """Initialize a KPipeline.
+        
+        Args:
+            lang_code: Language code for G2P processing
+            model: KModel instance, True to create new model, False for no model
+            trf: Whether to use transformer-based G2P
+            device: Override default device selection ('cuda' or 'cpu', or None for auto)
+                   If None, will auto-select cuda if available
+                   If 'cuda' and not available, will explicitly raise an error
+        """
+        if repo_id is None:
+            repo_id = 'hexgrad/Kokoro-82M'
+            print(f"WARNING: Defaulting repo_id to {repo_id}. Pass repo_id='{repo_id}' to suppress this warning.")
+            config=None
+        else:
+            config = os.path.join(repo_id, 'config.json')
+        self.repo_id = repo_id
+        lang_code = lang_code.lower()
+        lang_code = ALIASES.get(lang_code, lang_code)
+        assert lang_code in LANG_CODES, (lang_code, LANG_CODES)
+        self.lang_code = lang_code
+        self.model = None
+        if isinstance(model, KModel):
+            self.model = model
+        elif model:
+            if device == 'cuda' and not torch.cuda.is_available():
+                raise RuntimeError("CUDA requested but not available")
+            if device == 'mps' and not torch.backends.mps.is_available():
+                raise RuntimeError("MPS requested but not available")
+            if device == 'mps' and os.environ.get('PYTORCH_ENABLE_MPS_FALLBACK') != '1':
+                raise RuntimeError("MPS requested but fallback not enabled")
+            if device is None:
+                if torch.cuda.is_available():
+                    device = 'cuda'
+                elif os.environ.get('PYTORCH_ENABLE_MPS_FALLBACK') == '1' and torch.backends.mps.is_available():
+                    device = 'mps'
+                else:
+                    device = 'cpu'
+            try:
+                self.model = KModel(repo_id=repo_id, config=config).to(device).eval()
+            except RuntimeError as e:
+                if device == 'cuda':
+                    raise RuntimeError(f"""Failed to initialize model on CUDA: {e}. 
+                                       Try setting device='cpu' or check CUDA installation.""")
+                raise
+        self.voices = {}
+        if lang_code in 'ab':
+            try:
+                fallback = espeak.EspeakFallback(british=lang_code=='b')
+            except Exception as e:
+                logger.warning("EspeakFallback not Enabled: OOD words will be skipped")
+                logger.warning({str(e)})
+                fallback = None
+            self.g2p = en.G2P(trf=trf, british=lang_code=='b', fallback=fallback, unk='')
+        elif lang_code == 'j':
+            try:
+                from misaki import ja
+                self.g2p = ja.JAG2P()
+            except ImportError:
+                logger.error("You need to `pip install misaki[ja]` to use lang_code='j'")
+                raise
+        elif lang_code == 'z':
+            try:
+                from misaki import zh
+                self.g2p = zh.ZHG2P(
+                    version=None if repo_id.endswith('/Kokoro-82M') else '1.1',
+                    en_callable=en_callable
+                )
+            except ImportError:
+                logger.error("You need to `pip install misaki[zh]` to use lang_code='z'")
+                raise
+        else:
+            language = LANG_CODES[lang_code]
+            logger.warning(f"Using EspeakG2P(language='{language}'). Chunking logic not yet implemented, so long texts may be truncated unless you split them with '\\n'.")
+            self.g2p = espeak.EspeakG2P(language=language)
+
+    def load_single_voice(self, voice: str):
+        if voice in self.voices:
+            return self.voices[voice]
+        if voice.endswith('.pt'):
+            f = voice
+        else:
+            f = hf_hub_download(repo_id=self.repo_id, filename=f'voices/{voice}.pt')
+            if not voice.startswith(self.lang_code):
+                v = LANG_CODES.get(voice, voice)
+                p = LANG_CODES.get(self.lang_code, self.lang_code)
+                logger.warning(f'Language mismatch, loading {v} voice into {p} pipeline.')
+        pack = torch.load(f, weights_only=True)
+        self.voices[voice] = pack
+        return pack
+
+    """
+    load_voice is a helper function that lazily downloads and loads a voice:
+    Single voice can be requested (e.g. 'af_bella') or multiple voices (e.g. 'af_bella,af_jessica').
+    If multiple voices are requested, they are averaged.
+    Delimiter is optional and defaults to ','.
+    """
+    def load_voice(self, voice: Union[str, torch.FloatTensor], delimiter: str = ",") -> torch.FloatTensor:
+        if isinstance(voice, torch.FloatTensor):
+            return voice
+        if voice in self.voices:
+            return self.voices[voice]
+        logger.debug(f"Loading voice: {voice}")
+        packs = [self.load_single_voice(v) for v in voice.split(delimiter)]
+        if len(packs) == 1:
+            return packs[0]
+        self.voices[voice] = torch.mean(torch.stack(packs), dim=0)
+        return self.voices[voice]
+
+    @staticmethod
+    def tokens_to_ps(tokens: List[en.MToken]) -> str:
+        return ''.join(t.phonemes + (' ' if t.whitespace else '') for t in tokens).strip()
+
+    @staticmethod
+    def waterfall_last(
+        tokens: List[en.MToken],
+        next_count: int,
+        waterfall: List[str] = ['!.?…', ':;', ',—'],
+        bumps: List[str] = [')', '”']
+    ) -> int:
+        for w in waterfall:
+            z = next((i for i, t in reversed(list(enumerate(tokens))) if t.phonemes in set(w)), None)
+            if z is None:
+                continue
+            z += 1
+            if z < len(tokens) and tokens[z].phonemes in bumps:
+                z += 1
+            if next_count - len(KPipeline.tokens_to_ps(tokens[:z])) <= 510:
+                return z
+        return len(tokens)
+
+    @staticmethod
+    def tokens_to_text(tokens: List[en.MToken]) -> str:
+        return ''.join(t.text + t.whitespace for t in tokens).strip()
+
+    def en_tokenize(
+        self,
+        tokens: List[en.MToken]
+    ) -> Generator[Tuple[str, str, List[en.MToken]], None, None]:
+        tks = []
+        pcount = 0
+        for t in tokens:
+            # American English: ɾ => T
+            t.phonemes = '' if t.phonemes is None else t.phonemes#.replace('ɾ', 'T')
+            next_ps = t.phonemes + (' ' if t.whitespace else '')
+            next_pcount = pcount + len(next_ps.rstrip())
+            if next_pcount > 510:
+                z = KPipeline.waterfall_last(tks, next_pcount)
+                text = KPipeline.tokens_to_text(tks[:z])
+                logger.debug(f"Chunking text at {z}: '{text[:30]}{'...' if len(text) > 30 else ''}'")
+                ps = KPipeline.tokens_to_ps(tks[:z])
+                yield text, ps, tks[:z]
+                tks = tks[z:]
+                pcount = len(KPipeline.tokens_to_ps(tks))
+                if not tks:
+                    next_ps = next_ps.lstrip()
+            tks.append(t)
+            pcount += len(next_ps)
+        if tks:
+            text = KPipeline.tokens_to_text(tks)
+            ps = KPipeline.tokens_to_ps(tks)
+            yield ''.join(text).strip(), ''.join(ps).strip(), tks
+
+    @staticmethod
+    def infer(
+        model: KModel,
+        ps: str,
+        pack: torch.FloatTensor,
+        speed: Union[float, Callable[[int], float]] = 1
+    ) -> KModel.Output:
+        if callable(speed):
+            speed = speed(len(ps))
+        return model(ps, pack[len(ps)-1], speed, return_output=True)
+
+    def generate_from_tokens(
+        self,
+        tokens: Union[str, List[en.MToken]],
+        voice: str,
+        speed: float = 1,
+        model: Optional[KModel] = None
+    ) -> Generator['KPipeline.Result', None, None]:
+        """Generate audio from either raw phonemes or pre-processed tokens.
+        
+        Args:
+            tokens: Either a phoneme string or list of pre-processed MTokens
+            voice: The voice to use for synthesis
+            speed: Speech speed modifier (default: 1)
+            model: Optional KModel instance (uses pipeline's model if not provided)
+        
+        Yields:
+            KPipeline.Result containing the input tokens and generated audio
+            
+        Raises:
+            ValueError: If no voice is provided or token sequence exceeds model limits
+        """
+        model = model or self.model
+        if model and voice is None:
+            raise ValueError('Specify a voice: pipeline.generate_from_tokens(..., voice="af_heart")')
+        
+        pack = self.load_voice(voice).to(model.device) if model else None
+
+        # Handle raw phoneme string
+        if isinstance(tokens, str):
+            logger.debug("Processing phonemes from raw string")
+            if len(tokens) > 510:
+                raise ValueError(f'Phoneme string too long: {len(tokens)} > 510')
+            output = KPipeline.infer(model, tokens, pack, speed) if model else None
+            yield self.Result(graphemes='', phonemes=tokens, output=output)
+            return
+        
+        logger.debug("Processing MTokens")
+        # Handle pre-processed tokens
+        for gs, ps, tks in self.en_tokenize(tokens):
+            if not ps:
+                continue
+            elif len(ps) > 510:
+                logger.warning(f"Unexpected len(ps) == {len(ps)} > 510 and ps == '{ps}'")
+                logger.warning("Truncating to 510 characters")
+                ps = ps[:510]
+            output = KPipeline.infer(model, ps, pack, speed) if model else None
+            if output is not None and output.pred_dur is not None:
+                KPipeline.join_timestamps(tks, output.pred_dur)
+            yield self.Result(graphemes=gs, phonemes=ps, tokens=tks, output=output)
+
+    @staticmethod
+    def join_timestamps(tokens: List[en.MToken], pred_dur: torch.LongTensor):
+        # Multiply by 600 to go from pred_dur frames to sample_rate 24000
+        # Equivalent to dividing pred_dur frames by 40 to get timestamp in seconds
+        # We will count nice round half-frames, so the divisor is 80
+        MAGIC_DIVISOR = 80
+        if not tokens or len(pred_dur) < 3:
+            # We expect at least 3: <bos>, token, <eos>
+            return
+        # We track 2 counts, measured in half-frames: (left, right)
+        # This way we can cut space characters in half
+        # TODO: Is -3 an appropriate offset?
+        left = right = 2 * max(0, pred_dur[0].item() - 3)
+        # Updates:
+        # left = right + (2 * token_dur) + space_dur
+        # right = left + space_dur
+        i = 1
+        for t in tokens:
+            if i >= len(pred_dur)-1:
+                break
+            if not t.phonemes:
+                if t.whitespace:
+                    i += 1
+                    left = right + pred_dur[i].item()
+                    right = left + pred_dur[i].item()
+                    i += 1
+                continue
+            j = i + len(t.phonemes)
+            if j >= len(pred_dur):
+                break
+            t.start_ts = left / MAGIC_DIVISOR
+            token_dur = pred_dur[i: j].sum().item()
+            space_dur = pred_dur[j].item() if t.whitespace else 0
+            left = right + (2 * token_dur) + space_dur
+            t.end_ts = left / MAGIC_DIVISOR
+            right = left + space_dur
+            i = j + (1 if t.whitespace else 0)
+
+    @dataclass
+    class Result:
+        graphemes: str
+        phonemes: str
+        tokens: Optional[List[en.MToken]] = None
+        output: Optional[KModel.Output] = None
+        text_index: Optional[int] = None
+
+        @property
+        def audio(self) -> Optional[torch.FloatTensor]:
+            return None if self.output is None else self.output.audio
+
+        @property
+        def pred_dur(self) -> Optional[torch.LongTensor]:
+            return None if self.output is None else self.output.pred_dur
+
+        ### MARK: BEGIN BACKWARD COMPAT ###
+        def __iter__(self):
+            yield self.graphemes
+            yield self.phonemes
+            yield self.audio
+
+        def __getitem__(self, index):
+            return [self.graphemes, self.phonemes, self.audio][index]
+
+        def __len__(self):
+            return 3
+        #### MARK: END BACKWARD COMPAT ####
+
+    def __call__(
+        self,
+        text: Union[str, List[str]],
+        voice: Optional[str] = None,
+        speed: Union[float, Callable[[int], float]] = 1,
+        split_pattern: Optional[str] = r'\n+',
+        model: Optional[KModel] = None
+    ) -> Generator['KPipeline.Result', None, None]:
+        model = model or self.model
+        if model and voice is None:
+            raise ValueError('Specify a voice: en_us_pipeline(text="Hello world!", voice="af_heart")')
+        pack = self.load_voice(voice).to(model.device) if model else None
+        
+        # Convert input to list of segments
+        if isinstance(text, str):
+            text = re.split(split_pattern, text.strip()) if split_pattern else [text]
+            
+        # Process each segment
+        for graphemes_index, graphemes in enumerate(text):
+            if not graphemes.strip():  # Skip empty segments
+                continue
+                
+            # English processing (unchanged)
+            if self.lang_code in 'ab':
+                logger.debug(f"Processing English text: {graphemes[:50]}{'...' if len(graphemes) > 50 else ''}")
+                _, tokens = self.g2p(graphemes)
+                for gs, ps, tks in self.en_tokenize(tokens):
+                    if not ps:
+                        continue
+                    elif len(ps) > 510:
+                        logger.warning(f"Unexpected len(ps) == {len(ps)} > 510 and ps == '{ps}'")
+                        ps = ps[:510]
+                    output = KPipeline.infer(model, ps, pack, speed) if model else None
+                    if output is not None and output.pred_dur is not None:
+                        KPipeline.join_timestamps(tks, output.pred_dur)
+                    yield self.Result(graphemes=gs, phonemes=ps, tokens=tks, output=output, text_index=graphemes_index)
+            
+            # Non-English processing with chunking
+            else:
+                # Split long text into smaller chunks (roughly 400 characters each)
+                # Using sentence boundaries when possible
+                chunk_size = 400
+                chunks = []
+                
+                # Try to split on sentence boundaries first
+                sentences = re.split(r'([.!?]+)', graphemes)
+                current_chunk = ""
+                
+                for i in range(0, len(sentences), 2):
+                    sentence = sentences[i]
+                    # Add the punctuation back if it exists
+                    if i + 1 < len(sentences):
+                        sentence += sentences[i + 1]
+                        
+                    if len(current_chunk) + len(sentence) <= chunk_size:
+                        current_chunk += sentence
+                    else:
+                        if current_chunk:
+                            chunks.append(current_chunk.strip())
+                        current_chunk = sentence
+                
+                if current_chunk:
+                    chunks.append(current_chunk.strip())
+                
+                # If no chunks were created (no sentence boundaries), fall back to character-based chunking
+                if not chunks:
+                    chunks = [graphemes[i:i+chunk_size] for i in range(0, len(graphemes), chunk_size)]
+                
+                # Process each chunk
+                for chunk in chunks:
+                    if not chunk.strip():
+                        continue
+                        
+                    ps, _ = self.g2p(chunk)
+                    if not ps:
+                        continue
+                    elif len(ps) > 510:
+                        logger.warning(f'Truncating len(ps) == {len(ps)} > 510')
+                        ps = ps[:510]
+                        
+                    output = KPipeline.infer(model, ps, pack, speed) if model else None
+                    yield self.Result(graphemes=chunk, phonemes=ps, output=output, text_index=graphemes_index)

requirements.txt

@@ -0,0 +1,18 @@
+opencv-python>=4.9.0.80
+diffusers>=0.31.0
+transformers>=4.49.0
+tokenizers>=0.20.3
+accelerate>=1.1.1
+tqdm
+imageio
+easydict
+ftfy
+dashscope
+imageio-ffmpeg
+scikit-image   
+loguru
+gradio>=5.0.0
+numpy>=1.23.5,<2
+xfuser>=0.4.1
+pyloudnorm
+optimum-quanto==0.2.6

src/audio_analysis/torch_utils.py

@@ -0,0 +1,20 @@
+import torch
+import torch.nn.functional as F
+
+
+def get_mask_from_lengths(lengths, max_len=None):
+    lengths = lengths.to(torch.long)
+    if max_len is None:
+        max_len = torch.max(lengths).item()
+
+    ids = torch.arange(0, max_len).unsqueeze(0).expand(lengths.shape[0], -1).to(lengths.device)
+    mask = ids < lengths.unsqueeze(1).expand(-1, max_len)
+
+    return mask
+
+
+def linear_interpolation(features, seq_len):
+    features = features.transpose(1, 2)
+    output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear')
+    return output_features.transpose(1, 2)
+

src/audio_analysis/wav2vec2.py

@@ -0,0 +1,125 @@
+from transformers import Wav2Vec2Config, Wav2Vec2Model
+from transformers.modeling_outputs import BaseModelOutput
+
+from src.audio_analysis.torch_utils import linear_interpolation
+
+# the implementation of Wav2Vec2Model is borrowed from
+# https://github.com/huggingface/transformers/blob/HEAD/src/transformers/models/wav2vec2/modeling_wav2vec2.py
+# initialize our encoder with the pre-trained wav2vec 2.0 weights.
+class Wav2Vec2Model(Wav2Vec2Model):
+    def __init__(self, config: Wav2Vec2Config):
+        super().__init__(config)
+
+    def forward(
+        self,
+        input_values,
+        seq_len,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        self.config.output_attentions = True
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+    def feature_extract(
+        self,
+        input_values,
+        seq_len,
+    ):
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        return extract_features
+
+    def encode(
+        self,
+        extract_features,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        self.config.output_attentions = True
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+            
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

src/utils.py

@@ -0,0 +1,60 @@
+from contextlib import contextmanager
+
+import torch
+
+@contextmanager
+def init_weights_on_device(device=torch.device("meta"), include_buffers: bool = False):
+    old_register_parameter = torch.nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = torch.nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            kwargs["requires_grad"] = param.requires_grad
+            module._parameters[name] = param_cls(
+                module._parameters[name].to(device), **kwargs
+            )
+
+    def register_empty_buffer(module, name, buffer, persistent=True):
+        old_register_buffer(module, name, buffer, persistent=persistent)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(device)
+
+    def patch_tensor_constructor(fn):
+        def wrapper(*args, **kwargs):
+            kwargs["device"] = device
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    if include_buffers:
+        tensor_constructors_to_patch = {
+            torch_function_name: getattr(torch, torch_function_name)
+            for torch_function_name in ["empty", "zeros", "ones", "full"]
+        }
+    else:
+        tensor_constructors_to_patch = {}
+
+    try:
+        torch.nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            torch.nn.Module.register_buffer = register_empty_buffer
+        for torch_function_name in tensor_constructors_to_patch.keys():
+            setattr(
+                torch,
+                torch_function_name,
+                patch_tensor_constructor(getattr(torch, torch_function_name)),
+            )
+        yield
+    finally:
+        torch.nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            torch.nn.Module.register_buffer = old_register_buffer
+        for (
+            torch_function_name,
+            old_torch_function,
+        ) in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)

src/vram_management/__init__.py

@@ -0,0 +1 @@
+from .layers import *

src/vram_management/layers.py

@@ -0,0 +1,243 @@
+import copy
+
+import torch
+
+from src.utils import init_weights_on_device
+import optimum.quanto.nn.qlinear as qlinear
+
+def cast_to(weight, dtype, device):
+    r = torch.empty_like(weight, dtype=dtype, device=device)
+    r.copy_(weight)
+    return r
+
+def cast_to_device(weight, device):
+    if hasattr(weight, '__class__') and 'optimum.quanto' in str(weight.__class__):  
+        return weight.to(device)
+    else:
+        r = torch.empty_like(weight, device=device)
+        r.copy_(weight)
+        return r
+
+class AutoWrappedModule(torch.nn.Module):
+    def __init__(
+        self,
+        module: torch.nn.Module,
+        offload_dtype,
+        offload_device,
+        onload_dtype,
+        onload_device,
+        computation_dtype,
+        computation_device,
+    ):
+        super().__init__()
+        self.module = module.to(dtype=offload_dtype, device=offload_device)
+        self.offload_dtype = offload_dtype
+        self.offload_device = offload_device
+        self.onload_dtype = onload_dtype
+        self.onload_device = onload_device
+        self.computation_dtype = computation_dtype
+        self.computation_device = computation_device
+        self.state = 0
+
+    def offload(self):
+        if self.state == 1 and (
+            self.offload_dtype != self.onload_dtype
+            or self.offload_device != self.onload_device
+        ):
+            self.module.to(dtype=self.offload_dtype, device=self.offload_device)
+            self.state = 0
+
+    def onload(self):
+        if self.state == 0 and (
+            self.offload_dtype != self.onload_dtype
+            or self.offload_device != self.onload_device
+        ):
+            self.module.to(dtype=self.onload_dtype, device=self.onload_device)
+            self.state = 1
+
+    def forward(self, *args, **kwargs):
+        if (
+            self.onload_dtype == self.computation_dtype
+            and self.onload_device == self.computation_device
+        ):
+            module = self.module
+        else:
+            module = copy.deepcopy(self.module).to(
+                dtype=self.computation_dtype, device=self.computation_device
+            )
+        return module(*args, **kwargs)
+
+
+
+class AutoWrappedQLinear(qlinear.QLinear):
+    def __init__(
+        self,
+        module: qlinear.QLinear,
+        offload_dtype,
+        offload_device,
+        onload_dtype,
+        onload_device,
+        computation_dtype,
+        computation_device,
+    ):
+        with init_weights_on_device(device=torch.device("meta")):
+            super().__init__(
+                in_features=module.in_features,
+                out_features=module.out_features,
+                bias=module.bias is not None,
+                device=offload_device,
+            )
+        self.weight = module.weight
+        self.bias = module.bias
+        self.offload_device = offload_device
+
+        self.onload_device = onload_device
+        self.computation_device = computation_device
+        self.state = 0
+
+    def offload(self):
+        if self.state == 1 and (
+             self.offload_device != self.onload_device
+        ):
+            self.to(device=self.offload_device)
+            self.state = 0
+
+    def onload(self):
+        if self.state == 0 and (
+            self.offload_device != self.onload_device
+        ):
+            self.to(device=self.onload_device)
+            self.state = 1
+
+    def forward(self, x, *args, **kwargs):
+        if (
+            self.onload_device == self.computation_device
+        ):
+            
+            return torch.nn.functional.linear(x, self.weight, bias=self.bias)
+        else:
+            
+            qweight = cast_to_device(self.weight, self.computation_device)
+            bias = (
+                None
+                if self.bias is None
+                else cast_to_device(self.bias, self.computation_device)
+            )
+            return torch.nn.functional.linear(x, qweight, bias)
+
+class AutoWrappedLinear(torch.nn.Linear):
+    def __init__(
+        self,
+        module: torch.nn.Linear,
+        offload_dtype,
+        offload_device,
+        onload_dtype,
+        onload_device,
+        computation_dtype,
+        computation_device,
+    ):
+        with init_weights_on_device(device=torch.device("meta")):
+            super().__init__(
+                in_features=module.in_features,
+                out_features=module.out_features,
+                bias=module.bias is not None,
+                dtype=offload_dtype,
+                device=offload_device,
+            )
+        self.weight = module.weight
+        self.bias = module.bias
+        self.offload_dtype = offload_dtype
+        self.offload_device = offload_device
+        self.onload_dtype = onload_dtype
+        self.onload_device = onload_device
+        self.computation_dtype = computation_dtype
+        self.computation_device = computation_device
+        self.state = 0
+
+    def offload(self):
+        if self.state == 1 and (
+            self.offload_dtype != self.onload_dtype
+            or self.offload_device != self.onload_device
+        ):
+            self.to(dtype=self.offload_dtype, device=self.offload_device)
+            self.state = 0
+
+    def onload(self):
+        if self.state == 0 and (
+            self.offload_dtype != self.onload_dtype
+            or self.offload_device != self.onload_device
+        ):
+            self.to(dtype=self.onload_dtype, device=self.onload_device)
+            self.state = 1
+
+    def forward(self, x, *args, **kwargs):
+        if (
+            self.onload_dtype == self.computation_dtype
+            and self.onload_device == self.computation_device
+        ):
+            weight, bias = self.weight, self.bias
+        else:
+            weight = cast_to(
+                self.weight, self.computation_dtype, self.computation_device
+            )
+            bias = (
+                None
+                if self.bias is None
+                else cast_to(self.bias, self.computation_dtype, self.computation_device)
+            )
+        return torch.nn.functional.linear(x, weight, bias)
+
+
+def enable_vram_management_recursively(
+    model: torch.nn.Module,
+    module_map: dict,
+    module_config: dict,
+    max_num_param=None,
+    overflow_module_config: dict = None,
+    total_num_param=0,
+):
+    for name, module in model.named_children():
+        for source_module, target_module in module_map.items():
+            if isinstance(module, source_module):
+                num_param = sum(p.numel() for p in module.parameters())
+                # print(str(module) + ':' + str(num_param))
+                if (
+                    max_num_param is not None
+                    and total_num_param + num_param > max_num_param
+                ):
+                    # print(str(module) + '-->\t\t num:' + str(num_param) + "\t total:" + str(total_num_param))
+                    module_config_ = overflow_module_config
+                else:
+                    module_config_ = module_config
+                module_ = target_module(module, **module_config_)
+                setattr(model, name, module_)
+                total_num_param += num_param
+                break
+        else:
+            total_num_param = enable_vram_management_recursively(
+                module,
+                module_map,
+                module_config,
+                max_num_param,
+                overflow_module_config,
+                total_num_param,
+            )
+    return total_num_param
+
+
+def enable_vram_management(
+    model: torch.nn.Module,
+    module_map: dict,
+    module_config: dict,
+    max_num_param=None,
+    overflow_module_config: dict = None,
+):
+    enable_vram_management_recursively(
+        model,
+        module_map,
+        module_config,
+        max_num_param,
+        overflow_module_config,
+        total_num_param=0,
+    )
+    model.vram_management_enabled = True

wan/__init__.py

@@ -0,0 +1,6 @@
+from . import configs, distributed, modules
+from .first_last_frame2video import WanFLF2V
+from .image2video import WanI2V
+from .text2video import WanT2V
+from .vace import WanVace, WanVaceMP
+from .multitalk import MultiTalkPipeline

wan/configs/__init__.py

@@ -0,0 +1,58 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import copy
+import os
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+from .wan_i2v_14B import i2v_14B
+from .wan_t2v_1_3B import t2v_1_3B
+from .wan_t2v_14B import t2v_14B
+from .wan_multitalk_14B import multitalk_14B
+
+# the config of t2i_14B is the same as t2v_14B
+t2i_14B = copy.deepcopy(t2v_14B)
+t2i_14B.__name__ = 'Config: Wan T2I 14B'
+
+# the config of flf2v_14B is the same as i2v_14B
+flf2v_14B = copy.deepcopy(i2v_14B)
+flf2v_14B.__name__ = 'Config: Wan FLF2V 14B'
+flf2v_14B.sample_neg_prompt = "镜头切换，" + flf2v_14B.sample_neg_prompt
+
+WAN_CONFIGS = {
+    't2v-14B': t2v_14B,
+    't2v-1.3B': t2v_1_3B,
+    'i2v-14B': i2v_14B,
+    't2i-14B': t2i_14B,
+    'flf2v-14B': flf2v_14B,
+    'vace-1.3B': t2v_1_3B,
+    'vace-14B': t2v_14B,
+    'multitalk-14B': multitalk_14B,
+}
+
+SIZE_CONFIGS = {
+    '720*1280': (720, 1280),
+    '1280*720': (1280, 720),
+    '480*832': (480, 832),
+    '832*480': (832, 480),
+    '1024*1024': (1024, 1024),
+    'multitalk-480': (640, 640),
+    'multitalk-720': (960, 960),
+}
+
+MAX_AREA_CONFIGS = {
+    '720*1280': 720 * 1280,
+    '1280*720': 1280 * 720,
+    '480*832': 480 * 832,
+    '832*480': 832 * 480,
+}
+
+SUPPORTED_SIZES = {
+    't2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    't2v-1.3B': ('480*832', '832*480'),
+    'i2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    'flf2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    't2i-14B': tuple(SIZE_CONFIGS.keys()),
+    'vace-1.3B': ('480*832', '832*480'),
+    'vace-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    'multitalk-14B': ('multitalk-480', 'multitalk-720'),
+}

wan/configs/shared_config.py

@@ -0,0 +1,19 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+from easydict import EasyDict
+
+#------------------------ Wan shared config ------------------------#
+wan_shared_cfg = EasyDict()
+
+# t5
+wan_shared_cfg.t5_model = 'umt5_xxl'
+wan_shared_cfg.t5_dtype = torch.bfloat16
+wan_shared_cfg.text_len = 512
+
+# transformer
+wan_shared_cfg.param_dtype = torch.bfloat16
+
+# inference
+wan_shared_cfg.num_train_timesteps = 1000
+wan_shared_cfg.sample_fps = 16
+wan_shared_cfg.sample_neg_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'

wan/configs/wan_i2v_14B.py

@@ -0,0 +1,24 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan I2V 14B ------------------------#
+
+i2v_14B = EasyDict(__name__='Config: Wan I2V 14B')
+i2v_14B.update(wan_shared_cfg)
+i2v_14B.sample_neg_prompt = "镜头晃动，" + i2v_14B.sample_neg_prompt
+
+i2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+i2v_14B.t5_tokenizer = 'google/umt5-xxl'
+
+# clip
+i2v_14B.clip_model = 'clip_xlm_roberta_vit_h_14'
+i2v_14B.clip_dtype = torch.float16
+i2v_14B.clip_checkpoint = 'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
+i2v_14B.clip_tokenizer = 'xlm-roberta-large'
+
+# vae
+i2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+i2v_14B.vae_stride = (4, 8, 8)

wan/configs/wan_multitalk_14B.py

@@ -0,0 +1,36 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan I2V 14B ------------------------#
+
+multitalk_14B = EasyDict(__name__='Config: Wan MultiTalk AI2V 14B')
+multitalk_14B.update(wan_shared_cfg)
+multitalk_14B.sample_neg_prompt = 'bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards'
+
+multitalk_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+multitalk_14B.t5_tokenizer = 'google/umt5-xxl'
+
+# clip
+multitalk_14B.clip_model = 'clip_xlm_roberta_vit_h_14'
+multitalk_14B.clip_dtype = torch.float16
+multitalk_14B.clip_checkpoint = 'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
+multitalk_14B.clip_tokenizer = 'xlm-roberta-large'
+
+# vae
+multitalk_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+multitalk_14B.vae_stride = (4, 8, 8)
+
+# transformer
+multitalk_14B.patch_size = (1, 2, 2)
+multitalk_14B.dim = 5120
+multitalk_14B.ffn_dim = 13824
+multitalk_14B.freq_dim = 256
+multitalk_14B.num_heads = 40
+multitalk_14B.num_layers = 40
+multitalk_14B.window_size = (-1, -1)
+multitalk_14B.qk_norm = True
+multitalk_14B.cross_attn_norm = True
+multitalk_14B.eps = 1e-6

wan/configs/wan_t2v_14B.py

@@ -0,0 +1,29 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan T2V 14B ------------------------#
+
+t2v_14B = EasyDict(__name__='Config: Wan T2V 14B')
+t2v_14B.update(wan_shared_cfg)
+
+# t5
+t2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+t2v_14B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+t2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+t2v_14B.vae_stride = (4, 8, 8)
+
+# transformer
+t2v_14B.patch_size = (1, 2, 2)
+t2v_14B.dim = 5120
+t2v_14B.ffn_dim = 13824
+t2v_14B.freq_dim = 256
+t2v_14B.num_heads = 40
+t2v_14B.num_layers = 40
+t2v_14B.window_size = (-1, -1)
+t2v_14B.qk_norm = True
+t2v_14B.cross_attn_norm = True
+t2v_14B.eps = 1e-6

wan/configs/wan_t2v_1_3B.py

@@ -0,0 +1,29 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan T2V 1.3B ------------------------#
+
+t2v_1_3B = EasyDict(__name__='Config: Wan T2V 1.3B')
+t2v_1_3B.update(wan_shared_cfg)
+
+# t5
+t2v_1_3B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+t2v_1_3B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+t2v_1_3B.vae_checkpoint = 'Wan2.1_VAE.pth'
+t2v_1_3B.vae_stride = (4, 8, 8)
+
+# transformer
+t2v_1_3B.patch_size = (1, 2, 2)
+t2v_1_3B.dim = 1536
+t2v_1_3B.ffn_dim = 8960
+t2v_1_3B.freq_dim = 256
+t2v_1_3B.num_heads = 12
+t2v_1_3B.num_layers = 30
+t2v_1_3B.window_size = (-1, -1)
+t2v_1_3B.qk_norm = True
+t2v_1_3B.cross_attn_norm = True
+t2v_1_3B.eps = 1e-6

wan/distributed/fsdp.py

@@ -0,0 +1,43 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import gc
+from functools import partial
+
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
+from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
+from torch.distributed.utils import _free_storage
+
+
+def shard_model(
+    model,
+    device_id,
+    param_dtype=torch.bfloat16,
+    reduce_dtype=torch.float32,
+    buffer_dtype=torch.float32,
+    process_group=None,
+    sharding_strategy=ShardingStrategy.FULL_SHARD,
+    sync_module_states=True,
+):
+    model = FSDP(
+        module=model,
+        process_group=process_group,
+        sharding_strategy=sharding_strategy,
+        auto_wrap_policy=partial(
+            lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks),
+        # mixed_precision=MixedPrecision(
+        #     param_dtype=param_dtype,
+        #     reduce_dtype=reduce_dtype,
+        #     buffer_dtype=buffer_dtype),
+        device_id=device_id,
+        sync_module_states=sync_module_states)
+    return model
+
+
+def free_model(model):
+    for m in model.modules():
+        if isinstance(m, FSDP):
+            _free_storage(m._handle.flat_param.data)
+    del model
+    gc.collect()
+    torch.cuda.empty_cache()

wan/distributed/xdit_context_parallel.py

@@ -0,0 +1,550 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.cuda.amp as amp
+from xfuser.core.distributed import (
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+    get_sp_group,
+)
+from einops import rearrange
+from xfuser.core.long_ctx_attention import xFuserLongContextAttention
+import xformers.ops
+
+from ..modules.model import sinusoidal_embedding_1d
+from ..utils.multitalk_utils import get_attn_map_with_target, split_token_counts_and_frame_ids, normalize_and_scale
+from ..modules.attention import SingleStreamAttention, SingleStreamMutiAttention
+
+
+def pad_freqs(original_tensor, target_len):
+    seq_len, s1, s2 = original_tensor.shape
+    pad_size = target_len - seq_len
+    padding_tensor = torch.ones(
+        pad_size,
+        s1,
+        s2,
+        dtype=original_tensor.dtype,
+        device=original_tensor.device)
+    padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
+    return padded_tensor
+
+
+@amp.autocast(enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    """
+    x:          [B, L, N, C].
+    grid_sizes: [B, 3].
+    freqs:      [M, C // 2].
+    """
+    s, n, c = x.size(1), x.size(2), x.size(3) // 2
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) # [[N, head_dim/2], [N, head_dim/2], [N, head_dim/2]] # T H W 极坐标
+
+    # loop over samples
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(
+            s, n, -1, 2)) # [L, N, C/2] # 极坐标
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+                            dim=-1).reshape(seq_len, 1, -1) # seq_lens, 1,  3 * dim / 2 (T H W)
+
+        # apply rotary embedding
+        sp_size = get_sequence_parallel_world_size()
+        sp_rank = get_sequence_parallel_rank()
+        freqs_i = pad_freqs(freqs_i, s * sp_size)
+        s_per_rank = s
+        freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *
+                                                       s_per_rank), :, :]
+        x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
+        x_i = torch.cat([x_i, x[i, s:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).float()
+
+
+def usp_dit_forward_vace(self, x, vace_context, seq_len, kwargs):
+    # embeddings
+    c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context]
+    c = [u.flatten(2).transpose(1, 2) for u in c]
+    c = torch.cat([
+        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
+        for u in c
+    ])
+
+    # arguments
+    new_kwargs = dict(x=x)
+    new_kwargs.update(kwargs)
+
+    # Context Parallel
+    c = torch.chunk(
+        c, get_sequence_parallel_world_size(),
+        dim=1)[get_sequence_parallel_rank()]
+
+    hints = []
+    for block in self.vace_blocks:
+        c, c_skip = block(c, **new_kwargs)
+        hints.append(c_skip)
+    return hints
+
+
+def usp_dit_forward(
+    self,
+    x,
+    t,
+    context,
+    seq_len,
+    vace_context=None,
+    vace_context_scale=1.0,
+    clip_fea=None,
+    y=None,
+):
+    """
+    x:              A list of videos each with shape [C, T, H, W].
+    t:              [B].
+    context:        A list of text embeddings each with shape [L, C].
+    """
+    if self.model_type == 'i2v':
+        assert clip_fea is not None and y is not None
+    # params
+    device = self.patch_embedding.weight.device
+    if self.freqs.device != device:
+        self.freqs = self.freqs.to(device)
+
+    if self.model_type != 'vace' and y is not None:
+        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+    # embeddings
+    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+    grid_sizes = torch.stack(
+        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+    x = [u.flatten(2).transpose(1, 2) for u in x]
+    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+    assert seq_lens.max() <= seq_len
+    x = torch.cat([
+        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
+        for u in x
+    ])
+
+    # time embeddings
+    with amp.autocast(dtype=torch.float32):
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t).float())
+        e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+        assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+    # context
+    context_lens = None
+    context = self.text_embedding(
+        torch.stack([
+            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+            for u in context
+        ]))
+
+    if self.model_type != 'vace' and clip_fea is not None:
+        context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
+        context = torch.concat([context_clip, context], dim=1)
+
+    # arguments
+    kwargs = dict(
+        e=e0,
+        seq_lens=seq_lens,
+        grid_sizes=grid_sizes,
+        freqs=self.freqs,
+        context=context,
+        context_lens=context_lens)
+    
+    # Context Parallel
+    x = torch.chunk(
+        x, get_sequence_parallel_world_size(),
+        dim=1)[get_sequence_parallel_rank()]
+
+    for block in self.blocks:
+        x = block(x, **kwargs)
+
+    # head
+    x = self.head(x, e)
+
+    # Context Parallel
+    x = get_sp_group().all_gather(x, dim=1)
+
+    # unpatchify
+    x = self.unpatchify(x, grid_sizes)
+    return [u.float() for u in x]
+
+
+def usp_attn_forward(self,
+                     x,
+                     seq_lens,
+                     grid_sizes,
+                     freqs,
+                     dtype=torch.bfloat16):
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x).view(b, s, n, d)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+    q = rope_apply(q, grid_sizes, freqs)
+    k = rope_apply(k, grid_sizes, freqs)
+
+    # TODO: We should use unpaded q,k,v for attention.
+    # k_lens = seq_lens // get_sequence_parallel_world_size()
+    # if k_lens is not None:
+    #     q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0)
+    #     k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0)
+    #     v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0)
+
+    x = xFuserLongContextAttention()(
+        None,
+        query=half(q),
+        key=half(k),
+        value=half(v),
+        window_size=self.window_size)
+
+    # TODO: padding after attention.
+    # x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1)
+
+    # output
+    x = x.flatten(2)
+    x = self.o(x)
+    return x
+
+
+
+
+def usp_dit_forward_multitalk(
+    self,
+    x,
+    t,
+    context,
+    seq_len,
+    clip_fea=None,
+    y=None,
+    audio=None,
+    ref_target_masks=None,
+):
+    """
+    x:              A list of videos each with shape [C, T, H, W].
+    t:              [B].
+    context:        A list of text embeddings each with shape [L, C].
+    """
+    
+    assert clip_fea is not None and y is not None
+    # params
+    device = self.patch_embedding.weight.device
+    if self.freqs.device != device:
+        self.freqs = self.freqs.to(device)
+
+    _, T, H, W = x[0].shape
+    N_t = T // self.patch_size[0]
+    N_h = H // self.patch_size[1]
+    N_w = W // self.patch_size[2]
+
+    if y is not None:
+        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+    x[0] = x[0].to(context[0].dtype)
+
+    # embeddings
+    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+    grid_sizes = torch.stack(
+        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+    x = [u.flatten(2).transpose(1, 2) for u in x]
+    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+    assert seq_lens.max() <= seq_len
+    x = torch.cat([
+        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
+        for u in x
+    ])
+
+    # time embeddings
+    with amp.autocast(dtype=torch.float32):
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t).float())
+        e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+        assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+    # context
+    context_lens = None
+    context = self.text_embedding(
+        torch.stack([
+            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+            for u in context
+        ]))
+
+    if clip_fea is not None:
+        context_clip = self.img_emb(clip_fea)  
+        context = torch.concat([context_clip, context], dim=1)
+
+    # get audio token
+    audio_cond = audio.to(device=x.device, dtype=x.dtype)
+    first_frame_audio_emb_s = audio_cond[:, :1, ...] 
+    latter_frame_audio_emb = audio_cond[:, 1:, ...] 
+    latter_frame_audio_emb = rearrange(latter_frame_audio_emb, "b (n_t n) w s c -> b n_t n w s c", n=self.vae_scale) 
+    middle_index = self.audio_window // 2
+    latter_first_frame_audio_emb = latter_frame_audio_emb[:, :, :1, :middle_index+1, ...] 
+    latter_first_frame_audio_emb = rearrange(latter_first_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 
+    latter_last_frame_audio_emb = latter_frame_audio_emb[:, :, -1:, middle_index:, ...] 
+    latter_last_frame_audio_emb = rearrange(latter_last_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 
+    latter_middle_frame_audio_emb = latter_frame_audio_emb[:, :, 1:-1, middle_index:middle_index+1, ...] 
+    latter_middle_frame_audio_emb = rearrange(latter_middle_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 
+    latter_frame_audio_emb_s = torch.concat([latter_first_frame_audio_emb, latter_middle_frame_audio_emb, latter_last_frame_audio_emb], dim=2) 
+    audio_embedding = self.audio_proj(first_frame_audio_emb_s, latter_frame_audio_emb_s) 
+    human_num = len(audio_embedding)
+    audio_embedding = torch.concat(audio_embedding.split(1), dim=2).to(x.dtype)
+
+
+    # convert ref_target_masks to token_ref_target_masks
+    if ref_target_masks is not None:
+        ref_target_masks = ref_target_masks.unsqueeze(0).to(torch.float32) 
+        token_ref_target_masks = nn.functional.interpolate(ref_target_masks, size=(N_h, N_w), mode='nearest') 
+        token_ref_target_masks = token_ref_target_masks.squeeze(0) 
+        token_ref_target_masks = (token_ref_target_masks > 0)
+        token_ref_target_masks = token_ref_target_masks.view(token_ref_target_masks.shape[0], -1) 
+        token_ref_target_masks = token_ref_target_masks.to(x.dtype)
+    
+    if self.enable_teacache:
+        modulated_inp = e0 if self.use_ret_steps else e
+        if self.cnt%3==0: # cond
+            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
+                should_calc_cond = True
+                self.accumulated_rel_l1_distance_cond = 0
+            else:
+                rescale_func = np.poly1d(self.coefficients)
+                self.accumulated_rel_l1_distance_cond += rescale_func(((modulated_inp-self.previous_e0_cond).abs().mean() / self.previous_e0_cond.abs().mean()).cpu().item())
+                # print("accumulated_rel_l1_distance_even", self.accumulated_rel_l1_distance_even)
+                if self.accumulated_rel_l1_distance_cond < self.teacache_thresh:
+                    should_calc_cond = False
+                else:
+                    should_calc_cond = True
+                    self.accumulated_rel_l1_distance_cond = 0
+            self.previous_e0_cond = modulated_inp.clone()
+        elif self.cnt%3==1: # drop_text
+            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
+                should_calc_drop_text = True
+                self.accumulated_rel_l1_distance_drop_text = 0
+            else:
+                rescale_func = np.poly1d(self.coefficients)
+                self.accumulated_rel_l1_distance_drop_text += rescale_func(((modulated_inp-self.previous_e0_drop_text).abs().mean() / self.previous_e0_drop_text.abs().mean()).cpu().item())
+                if self.accumulated_rel_l1_distance_drop_text < self.teacache_thresh:
+                    should_calc_drop_text = False
+                else:
+                    should_calc_drop_text = True
+                    self.accumulated_rel_l1_distance_drop_text = 0
+            self.previous_e0_drop_text = modulated_inp.clone()
+        else: # uncond
+            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
+                should_calc_uncond = True
+                self.accumulated_rel_l1_distance_uncond = 0
+            else:
+                rescale_func = np.poly1d(self.coefficients)
+                self.accumulated_rel_l1_distance_uncond += rescale_func(((modulated_inp-self.previous_e0_uncond).abs().mean() / self.previous_e0_uncond.abs().mean()).cpu().item())
+                if self.accumulated_rel_l1_distance_uncond < self.teacache_thresh:
+                    should_calc_uncond = False
+                else:
+                    should_calc_uncond = True
+                    self.accumulated_rel_l1_distance_uncond = 0
+            self.previous_e0_uncond = modulated_inp.clone()
+
+    # Context Parallel
+    x = torch.chunk(
+        x, get_sequence_parallel_world_size(),
+        dim=1)[get_sequence_parallel_rank()]
+
+    # arguments
+    kwargs = dict(
+        e=e0,
+        seq_lens=seq_lens,
+        grid_sizes=grid_sizes,
+        freqs=self.freqs,
+        context=context,
+        context_lens=context_lens,
+        audio_embedding=audio_embedding,
+        ref_target_masks=token_ref_target_masks,
+        human_num=human_num,
+        )
+
+    if self.enable_teacache:
+        if self.cnt%3==0:
+            if not should_calc_cond:
+                x +=  self.previous_residual_cond
+            else:
+                ori_x = x.clone()
+                for block in self.blocks:
+                    x = block(x, **kwargs)
+                self.previous_residual_cond = x - ori_x
+        elif self.cnt%3==1:
+            if not should_calc_drop_text:
+                x +=  self.previous_residual_drop_text
+            else:
+                ori_x = x.clone()
+                for block in self.blocks:
+                    x = block(x, **kwargs)
+                self.previous_residual_drop_text = x - ori_x
+        else:
+            if not should_calc_uncond:
+                x +=  self.previous_residual_uncond
+            else:
+                ori_x = x.clone()
+                for block in self.blocks:
+                    x = block(x, **kwargs)
+                self.previous_residual_uncond = x - ori_x
+    else:
+        for block in self.blocks:
+            x = block(x, **kwargs)
+
+    # head
+    x = self.head(x, e)
+
+    # Context Parallel
+    x = get_sp_group().all_gather(x, dim=1)
+
+    # unpatchify
+    x = self.unpatchify(x, grid_sizes)
+    if self.enable_teacache:
+        self.cnt += 1
+        if self.cnt >= self.num_steps:
+            self.cnt = 0
+        
+    return torch.stack(x).float()
+
+
+def usp_attn_forward_multitalk(self,
+                     x,
+                     seq_lens,
+                     grid_sizes,
+                     freqs,
+                     dtype=torch.bfloat16,
+                     ref_target_masks=None):
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x).view(b, s, n, d)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+    q = rope_apply(q, grid_sizes, freqs)
+    k = rope_apply(k, grid_sizes, freqs)
+
+
+    x = xFuserLongContextAttention()(
+        None,
+        query=half(q),
+        key=half(k),
+        value=half(v),
+        window_size=self.window_size)
+
+
+    # output
+    x = x.flatten(2)
+    x = self.o(x)
+
+    with torch.no_grad():
+        x_ref_attn_map = get_attn_map_with_target(q.type_as(x), k.type_as(x), grid_sizes[0], 
+                                            ref_target_masks=ref_target_masks, enable_sp=True) 
+
+    return x, x_ref_attn_map
+
+
+
+
+def usp_crossattn_multi_forward_multitalk(self, 
+                                        x: torch.Tensor, 
+                                        encoder_hidden_states: torch.Tensor,  # 1, 21, 64, C
+                                        shape=None, 
+                                        x_ref_attn_map=None,
+                                        human_num=None) -> torch.Tensor:
+        
+        N_t, N_h, N_w = shape 
+        sp_size = get_sequence_parallel_world_size()
+        sp_rank = get_sequence_parallel_rank()
+        audio_tokens_per_frame = 32
+        visual_seqlen, frame_ids = split_token_counts_and_frame_ids(N_t, N_h * N_w, sp_size, sp_rank)
+        encoder_hidden_states = encoder_hidden_states[:, min(frame_ids):max(frame_ids)+1, ...]
+        encoder_hidden_states = rearrange(encoder_hidden_states, "B T N C -> B (T N) C")
+        N_a = len(frame_ids)
+        kv_seq = [audio_tokens_per_frame * human_num] * N_a
+
+        if human_num == 1:
+            return super(SingleStreamMutiAttention, self).forward(x, encoder_hidden_states, shape, enable_sp=True, kv_seq=kv_seq)
+
+
+        # get q for hidden_state
+        B, N, C = x.shape
+        q = self.q_linear(x) 
+        q_shape = (B, N, self.num_heads, self.head_dim) 
+        q = q.view(q_shape).permute((0, 2, 1, 3))
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+
+        max_values = x_ref_attn_map.max(1).values[:, None, None] 
+        min_values = x_ref_attn_map.min(1).values[:, None, None] 
+        max_min_values = torch.cat([max_values, min_values], dim=2)
+        max_min_values = get_sp_group().all_gather(max_min_values, dim=1)
+
+        human1_max_value, human1_min_value = max_min_values[0, :, 0].max(), max_min_values[0, :, 1].min()
+        human2_max_value, human2_min_value = max_min_values[1, :, 0].max(), max_min_values[1, :, 1].min()
+
+        human1 = normalize_and_scale(x_ref_attn_map[0], (human1_min_value, human1_max_value), (self.rope_h1[0], self.rope_h1[1]))
+        human2 = normalize_and_scale(x_ref_attn_map[1], (human2_min_value, human2_max_value), (self.rope_h2[0], self.rope_h2[1]))
+        back   = torch.full((x_ref_attn_map.size(1),), self.rope_bak, dtype=human1.dtype).to(human1.device)
+        max_indices = x_ref_attn_map.argmax(dim=0)
+        normalized_map = torch.stack([human1, human2, back], dim=1)
+        normalized_pos = normalized_map[range(x_ref_attn_map.size(1)), max_indices] # N 
+        q = self.rope_1d(q, normalized_pos)
+ 
+        encoder_kv = self.kv_linear(encoder_hidden_states) 
+        encoder_kv_shape = (B, encoder_hidden_states.size(1), 2, self.num_heads, self.head_dim)
+        encoder_kv = encoder_kv.view(encoder_kv_shape).permute((2, 0, 3, 1, 4)) 
+        encoder_k, encoder_v = encoder_kv.unbind(0) # B H N C
+
+        if self.qk_norm:
+            encoder_k = self.add_k_norm(encoder_k)
+
+        # position embedding for condition audio embeddings
+        per_frame = torch.zeros(audio_tokens_per_frame * human_num, dtype=encoder_k.dtype).to(encoder_k.device)
+        per_frame[:audio_tokens_per_frame] = (self.rope_h1[0] + self.rope_h1[1]) / 2
+        per_frame[audio_tokens_per_frame:] = (self.rope_h2[0] + self.rope_h2[1]) / 2
+        encoder_pos = torch.concat([per_frame]*N_a, dim=0)
+        encoder_k = self.rope_1d(encoder_k, encoder_pos)
+
+        # get attn
+        q = rearrange(q, "B H M K -> B M H K")
+        encoder_k = rearrange(encoder_k, "B H M K -> B M H K")
+        encoder_v = rearrange(encoder_v, "B H M K -> B M H K")
+        attn_bias = xformers.ops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(visual_seqlen, kv_seq)
+        x = xformers.ops.memory_efficient_attention(q, encoder_k, encoder_v, attn_bias=attn_bias, op=None,)
+        x = rearrange(x, "B M H K -> B H M K")
+
+        # linear transform
+        x_output_shape = (B, N, C)
+        x = x.transpose(1, 2) 
+        x = x.reshape(x_output_shape) 
+        x = self.proj(x) 
+        x = self.proj_drop(x)
+
+        return x

wan/first_last_frame2video.py

@@ -0,0 +1,377 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import gc
+import logging
+import math
+import os
+import random
+import sys
+import types
+from contextlib import contextmanager
+from functools import partial
+
+import numpy as np
+import torch
+import torch.cuda.amp as amp
+import torch.distributed as dist
+import torchvision.transforms.functional as TF
+from tqdm import tqdm
+
+from .distributed.fsdp import shard_model
+from .modules.clip import CLIPModel
+from .modules.model import WanModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae import WanVAE
+from .utils.fm_solvers import (
+    FlowDPMSolverMultistepScheduler,
+    get_sampling_sigmas,
+    retrieve_timesteps,
+)
+from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+
+class WanFLF2V:
+
+    def __init__(
+        self,
+        config,
+        checkpoint_dir,
+        device_id=0,
+        rank=0,
+        t5_fsdp=False,
+        dit_fsdp=False,
+        use_usp=False,
+        t5_cpu=False,
+        init_on_cpu=True,
+    ):
+        r"""
+        Initializes the image-to-video generation model components.
+
+        Args:
+            config (EasyDict):
+                Object containing model parameters initialized from config.py
+            checkpoint_dir (`str`):
+                Path to directory containing model checkpoints
+            device_id (`int`,  *optional*, defaults to 0):
+                Id of target GPU device
+            rank (`int`,  *optional*, defaults to 0):
+                Process rank for distributed training
+            t5_fsdp (`bool`, *optional*, defaults to False):
+                Enable FSDP sharding for T5 model
+            dit_fsdp (`bool`, *optional*, defaults to False):
+                Enable FSDP sharding for DiT model
+            use_usp (`bool`, *optional*, defaults to False):
+                Enable distribution strategy of USP.
+            t5_cpu (`bool`, *optional*, defaults to False):
+                Whether to place T5 model on CPU. Only works without t5_fsdp.
+            init_on_cpu (`bool`, *optional*, defaults to True):
+                Enable initializing Transformer Model on CPU. Only works without FSDP or USP.
+        """
+        self.device = torch.device(f"cuda:{device_id}")
+        self.config = config
+        self.rank = rank
+        self.use_usp = use_usp
+        self.t5_cpu = t5_cpu
+
+        self.num_train_timesteps = config.num_train_timesteps
+        self.param_dtype = config.param_dtype
+
+        shard_fn = partial(shard_model, device_id=device_id)
+        self.text_encoder = T5EncoderModel(
+            text_len=config.text_len,
+            dtype=config.t5_dtype,
+            device=torch.device('cpu'),
+            checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
+            tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
+            shard_fn=shard_fn if t5_fsdp else None,
+        )
+
+        self.vae_stride = config.vae_stride
+        self.patch_size = config.patch_size
+        self.vae = WanVAE(
+            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+            device=self.device)
+
+        self.clip = CLIPModel(
+            dtype=config.clip_dtype,
+            device=self.device,
+            checkpoint_path=os.path.join(checkpoint_dir,
+                                         config.clip_checkpoint),
+            tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
+
+        logging.info(f"Creating WanModel from {checkpoint_dir}")
+        self.model = WanModel.from_pretrained(checkpoint_dir)
+        self.model.eval().requires_grad_(False)
+
+        if t5_fsdp or dit_fsdp or use_usp:
+            init_on_cpu = False
+
+        if use_usp:
+            from xfuser.core.distributed import get_sequence_parallel_world_size
+
+            from .distributed.xdit_context_parallel import (
+                usp_attn_forward,
+                usp_dit_forward,
+            )
+            for block in self.model.blocks:
+                block.self_attn.forward = types.MethodType(
+                    usp_attn_forward, block.self_attn)
+            self.model.forward = types.MethodType(usp_dit_forward, self.model)
+            self.sp_size = get_sequence_parallel_world_size()
+        else:
+            self.sp_size = 1
+
+        if dist.is_initialized():
+            dist.barrier()
+        if dit_fsdp:
+            self.model = shard_fn(self.model)
+        else:
+            if not init_on_cpu:
+                self.model.to(self.device)
+
+        self.sample_neg_prompt = config.sample_neg_prompt
+
+    def generate(self,
+                 input_prompt,
+                 first_frame,
+                 last_frame,
+                 max_area=720 * 1280,
+                 frame_num=81,
+                 shift=16,
+                 sample_solver='unipc',
+                 sampling_steps=50,
+                 guide_scale=5.5,
+                 n_prompt="",
+                 seed=-1,
+                 offload_model=True):
+        r"""
+        Generates video frames from input first-last frame and text prompt using diffusion process.
+
+        Args:
+            input_prompt (`str`):
+                Text prompt for content generation.
+            first_frame (PIL.Image.Image):
+                Input image tensor. Shape: [3, H, W]
+            last_frame (PIL.Image.Image):
+                Input image tensor. Shape: [3, H, W]
+                [NOTE] If the sizes of first_frame and last_frame are mismatched, last_frame will be cropped & resized
+                to match first_frame.
+            max_area (`int`, *optional*, defaults to 720*1280):
+                Maximum pixel area for latent space calculation. Controls video resolution scaling
+            frame_num (`int`, *optional*, defaults to 81):
+                How many frames to sample from a video. The number should be 4n+1
+            shift (`float`, *optional*, defaults to 5.0):
+                Noise schedule shift parameter. Affects temporal dynamics
+                [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0.
+            sample_solver (`str`, *optional*, defaults to 'unipc'):
+                Solver used to sample the video.
+            sampling_steps (`int`, *optional*, defaults to 40):
+                Number of diffusion sampling steps. Higher values improve quality but slow generation
+            guide_scale (`float`, *optional*, defaults 5.0):
+                Classifier-free guidance scale. Controls prompt adherence vs. creativity
+            n_prompt (`str`, *optional*, defaults to ""):
+                Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
+            seed (`int`, *optional*, defaults to -1):
+                Random seed for noise generation. If -1, use random seed
+            offload_model (`bool`, *optional*, defaults to True):
+                If True, offloads models to CPU during generation to save VRAM
+
+        Returns:
+            torch.Tensor:
+                Generated video frames tensor. Dimensions: (C, N H, W) where:
+                - C: Color channels (3 for RGB)
+                - N: Number of frames (81)
+                - H: Frame height (from max_area)
+                - W: Frame width from max_area)
+        """
+        first_frame_size = first_frame.size
+        last_frame_size = last_frame.size
+        first_frame = TF.to_tensor(first_frame).sub_(0.5).div_(0.5).to(
+            self.device)
+        last_frame = TF.to_tensor(last_frame).sub_(0.5).div_(0.5).to(
+            self.device)
+
+        F = frame_num
+        first_frame_h, first_frame_w = first_frame.shape[1:]
+        aspect_ratio = first_frame_h / first_frame_w
+        lat_h = round(
+            np.sqrt(max_area * aspect_ratio) // self.vae_stride[1] //
+            self.patch_size[1] * self.patch_size[1])
+        lat_w = round(
+            np.sqrt(max_area / aspect_ratio) // self.vae_stride[2] //
+            self.patch_size[2] * self.patch_size[2])
+        first_frame_h = lat_h * self.vae_stride[1]
+        first_frame_w = lat_w * self.vae_stride[2]
+        if first_frame_size != last_frame_size:
+            # 1. resize
+            last_frame_resize_ratio = max(
+                first_frame_size[0] / last_frame_size[0],
+                first_frame_size[1] / last_frame_size[1])
+            last_frame_size = [
+                round(last_frame_size[0] * last_frame_resize_ratio),
+                round(last_frame_size[1] * last_frame_resize_ratio),
+            ]
+            # 2. center crop
+            last_frame = TF.center_crop(last_frame, last_frame_size)
+
+        max_seq_len = ((F - 1) // self.vae_stride[0] + 1) * lat_h * lat_w // (
+            self.patch_size[1] * self.patch_size[2])
+        max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size
+
+        seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
+        seed_g = torch.Generator(device=self.device)
+        seed_g.manual_seed(seed)
+        noise = torch.randn(
+            16, (F - 1) // 4 + 1,
+            lat_h,
+            lat_w,
+            dtype=torch.float32,
+            generator=seed_g,
+            device=self.device)
+
+        msk = torch.ones(1, 81, lat_h, lat_w, device=self.device)
+        msk[:, 1:-1] = 0
+        msk = torch.concat([
+            torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]
+        ],
+                           dim=1)
+        msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
+        msk = msk.transpose(1, 2)[0]
+
+        if n_prompt == "":
+            n_prompt = self.sample_neg_prompt
+
+        # preprocess
+        if not self.t5_cpu:
+            self.text_encoder.model.to(self.device)
+            context = self.text_encoder([input_prompt], self.device)
+            context_null = self.text_encoder([n_prompt], self.device)
+            if offload_model:
+                self.text_encoder.model.cpu()
+        else:
+            context = self.text_encoder([input_prompt], torch.device('cpu'))
+            context_null = self.text_encoder([n_prompt], torch.device('cpu'))
+            context = [t.to(self.device) for t in context]
+            context_null = [t.to(self.device) for t in context_null]
+
+        self.clip.model.to(self.device)
+        clip_context = self.clip.visual(
+            [first_frame[:, None, :, :], last_frame[:, None, :, :]])
+        if offload_model:
+            self.clip.model.cpu()
+
+        y = self.vae.encode([
+            torch.concat([
+                torch.nn.functional.interpolate(
+                    first_frame[None].cpu(),
+                    size=(first_frame_h, first_frame_w),
+                    mode='bicubic').transpose(0, 1),
+                torch.zeros(3, F - 2, first_frame_h, first_frame_w),
+                torch.nn.functional.interpolate(
+                    last_frame[None].cpu(),
+                    size=(first_frame_h, first_frame_w),
+                    mode='bicubic').transpose(0, 1),
+            ],
+                         dim=1).to(self.device)
+        ])[0]
+        y = torch.concat([msk, y])
+
+        @contextmanager
+        def noop_no_sync():
+            yield
+
+        no_sync = getattr(self.model, 'no_sync', noop_no_sync)
+
+        # evaluation mode
+        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
+
+            if sample_solver == 'unipc':
+                sample_scheduler = FlowUniPCMultistepScheduler(
+                    num_train_timesteps=self.num_train_timesteps,
+                    shift=1,
+                    use_dynamic_shifting=False)
+                sample_scheduler.set_timesteps(
+                    sampling_steps, device=self.device, shift=shift)
+                timesteps = sample_scheduler.timesteps
+            elif sample_solver == 'dpm++':
+                sample_scheduler = FlowDPMSolverMultistepScheduler(
+                    num_train_timesteps=self.num_train_timesteps,
+                    shift=1,
+                    use_dynamic_shifting=False)
+                sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
+                timesteps, _ = retrieve_timesteps(
+                    sample_scheduler,
+                    device=self.device,
+                    sigmas=sampling_sigmas)
+            else:
+                raise NotImplementedError("Unsupported solver.")
+
+            # sample videos
+            latent = noise
+
+            arg_c = {
+                'context': [context[0]],
+                'clip_fea': clip_context,
+                'seq_len': max_seq_len,
+                'y': [y],
+            }
+
+            arg_null = {
+                'context': context_null,
+                'clip_fea': clip_context,
+                'seq_len': max_seq_len,
+                'y': [y],
+            }
+
+            if offload_model:
+                torch.cuda.empty_cache()
+
+            self.model.to(self.device)
+            for _, t in enumerate(tqdm(timesteps)):
+                latent_model_input = [latent.to(self.device)]
+                timestep = [t]
+
+                timestep = torch.stack(timestep).to(self.device)
+
+                noise_pred_cond = self.model(
+                    latent_model_input, t=timestep, **arg_c)[0].to(
+                        torch.device('cpu') if offload_model else self.device)
+                if offload_model:
+                    torch.cuda.empty_cache()
+                noise_pred_uncond = self.model(
+                    latent_model_input, t=timestep, **arg_null)[0].to(
+                        torch.device('cpu') if offload_model else self.device)
+                if offload_model:
+                    torch.cuda.empty_cache()
+                noise_pred = noise_pred_uncond + guide_scale * (
+                    noise_pred_cond - noise_pred_uncond)
+
+                latent = latent.to(
+                    torch.device('cpu') if offload_model else self.device)
+
+                temp_x0 = sample_scheduler.step(
+                    noise_pred.unsqueeze(0),
+                    t,
+                    latent.unsqueeze(0),
+                    return_dict=False,
+                    generator=seed_g)[0]
+                latent = temp_x0.squeeze(0)
+
+                x0 = [latent.to(self.device)]
+                del latent_model_input, timestep
+
+            if offload_model:
+                self.model.cpu()
+                torch.cuda.empty_cache()
+
+            if self.rank == 0:
+                videos = self.vae.decode(x0)
+
+        del noise, latent
+        del sample_scheduler
+        if offload_model:
+            gc.collect()
+            torch.cuda.synchronize()
+        if dist.is_initialized():
+            dist.barrier()
+
+        return videos[0] if self.rank == 0 else None