Image-to-Video
Diffusers
Safetensors
video-generation
video diffusion transformer
joint-video-audio-generation
sparse-attention
high-resolution
Instructions to use FrancisRing/Prism with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Diffusers
How to use FrancisRing/Prism with Diffusers:
pip install -U diffusers transformers accelerate
import torch from diffusers import DiffusionPipeline from diffusers.utils import load_image, export_to_video # switch to "mps" for apple devices pipe = DiffusionPipeline.from_pretrained("FrancisRing/Prism", dtype=torch.bfloat16, device_map="cuda") pipe.to("cuda") prompt = "A man with short gray hair plays a red electric guitar." image = load_image( "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/guitar-man.png" ) output = pipe(image=image, prompt=prompt).frames[0] export_to_video(output, "output.mp4") - Notebooks
- Google Colab
- Kaggle
Update README.md
Browse files
README.md
CHANGED
|
@@ -5,447 +5,369 @@ pipeline_tag: image-to-video
|
|
| 5 |
tags:
|
| 6 |
- video-generation
|
| 7 |
- video diffusion transformer
|
| 8 |
-
- audio-
|
|
|
|
|
|
|
| 9 |
task_categories:
|
| 10 |
- image-to-video
|
| 11 |
- text-to-video
|
| 12 |
---
|
| 13 |
|
| 14 |
-
#
|
| 15 |
|
| 16 |
-
<a href='https://francis-rings.github.io/
|
| 17 |
|
| 18 |
-
|
| 19 |
<br/>
|
| 20 |
-
|
| 21 |
<br/>
|
| 22 |
-
[<sup>1</sup>Fudan University; <sup>2</sup>
|
|
|
|
| 23 |
|
| 24 |
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
| 25 |
<tr>
|
| 26 |
<td>
|
| 27 |
-
<video src="https://github.com/user-attachments/assets/
|
| 28 |
-
</td>
|
| 29 |
-
<td>
|
| 30 |
-
<video src="https://github.com/user-attachments/assets/b5902ac4-8188-4da8-b9e6-6df280690ed1" width="320" controls loop></video>
|
| 31 |
-
</td>
|
| 32 |
-
<td>
|
| 33 |
-
<video src="https://github.com/user-attachments/assets/87faa5c1-a118-4a03-a071-45f18e87e6a0" width="320" controls loop></video>
|
| 34 |
-
</td>
|
| 35 |
-
</tr>
|
| 36 |
-
<tr>
|
| 37 |
-
<td>
|
| 38 |
-
<video src="https://github.com/user-attachments/assets/531eb413-8993-4f8f-9804-e3c5ec5794d4" width="320" controls loop></video>
|
| 39 |
</td>
|
| 40 |
<td>
|
| 41 |
-
<video src="https://github.com/user-attachments/assets/
|
| 42 |
</td>
|
| 43 |
<td>
|
| 44 |
-
<video src="https://github.com/user-attachments/assets/
|
| 45 |
</td>
|
| 46 |
</tr>
|
| 47 |
<tr>
|
| 48 |
<td>
|
| 49 |
-
<video src="https://github.com/user-attachments/assets/
|
| 50 |
</td>
|
| 51 |
<td>
|
| 52 |
-
<video src="https://github.com/user-attachments/assets/
|
| 53 |
</td>
|
| 54 |
<td>
|
| 55 |
-
<video src="https://github.com/user-attachments/assets/
|
| 56 |
</td>
|
| 57 |
</tr>
|
| 58 |
</table>
|
| 59 |
|
| 60 |
-
<p style="text-align: justify;">
|
| 61 |
-
<span>Audio-driven avatar videos generated by StableAvatar, showing its power to synthesize <b>infinite-length</b> and <b>ID-preserving videos</b>. All videos are <b>directly synthesized by StableAvatar without the use of any face-related post-processing tools</b>, such as the face-swapping tool FaceFusion or face restoration models like GFP-GAN and CodeFormer.</span>
|
| 62 |
-
</p>
|
| 63 |
-
|
| 64 |
-
<p align="center">
|
| 65 |
-
<video src="https://github.com/user-attachments/assets/90691318-311e-40b9-9bd9-62db83ab1492" width="768" controls loop></video>
|
| 66 |
-
<br/>
|
| 67 |
-
<span>Comparison results between StableAvatar and state-of-the-art (SOTA) audio-driven avatar video generation models highlight the superior performance of StableAvatar in delivering <b>infinite-length, high-fidelity, identity-preserving avatar animation</b>.</span>
|
| 68 |
-
</p>
|
| 69 |
-
|
| 70 |
|
| 71 |
## Overview
|
| 72 |
|
| 73 |
<p align="center">
|
| 74 |
<img src="assets/figures/framework.jpg" alt="model architecture" width="1280"/>
|
| 75 |
</br>
|
| 76 |
-
<i>The overview of the framework of
|
| 77 |
</p>
|
| 78 |
|
| 79 |
-
|
| 80 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 81 |
|
| 82 |
## News
|
| 83 |
-
* `[
|
| 84 |
-
* `[2025-8-15]`:🔥 StableAvatar can run on Gradio Interface. Thanks @[gluttony-10](https://space.bilibili.com/893892) for the contribution!
|
| 85 |
-
* `[2025-8-15]`:🔥 StableAvatar can run on [ComfyUI](https://github.com/smthemex/ComfyUI_StableAvatar). Thanks @[smthemex](https://github.com/smthemex) for the contribution.
|
| 86 |
-
* `[2025-8-13]`:🔥 Added changes to run StableAvatar on the new Blackwell series Nvidia chips, including the RTX 6000 Pro.
|
| 87 |
-
* `[2025-8-11]`:🔥 The project page, code, technical report and [a basic model checkpoint](https://huggingface.co/FrancisRing/StableAvatar/tree/main) are released. Further lora training codes, the evaluation dataset and StableAvatar-pro will be released very soon. Stay tuned!
|
| 88 |
|
| 89 |
## 🛠️ To-Do List
|
| 90 |
-
- [x]
|
| 91 |
-
- [x]
|
| 92 |
-
- [x] Data Pre-Processing Code (
|
| 93 |
-
- [x] Data Pre-Processing Code (Vocal Separation)
|
| 94 |
- [x] Training Code
|
| 95 |
- [x] Full Finetuning Code
|
| 96 |
-
- [x]
|
| 97 |
-
- [
|
| 98 |
-
- [ ] Inference Code with Audio Native Guidance
|
| 99 |
-
- [ ] StableAvatar-pro
|
| 100 |
|
| 101 |
## 🔑 Quickstart
|
| 102 |
|
| 103 |
-
|
| 104 |
|
| 105 |
-
### 🧱 Environment
|
| 106 |
|
|
|
|
| 107 |
```
|
| 108 |
-
|
|
|
|
|
|
|
| 109 |
pip install -r requirements.txt
|
| 110 |
-
#
|
| 111 |
-
pip install
|
| 112 |
-
``
|
| 113 |
-
|
| 114 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 115 |
|
| 116 |
-
```
|
| 117 |
-
pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
|
| 118 |
-
pip install -r requirements.txt
|
| 119 |
-
# Optional to install flash_attn to accelerate attention computation
|
| 120 |
-
pip install flash_attn
|
| 121 |
-
```
|
| 122 |
-
|
| 123 |
-
### 🧱 Download weights
|
| 124 |
If you encounter connection issues with Hugging Face, you can utilize the mirror endpoint by setting the environment variable: `export HF_ENDPOINT=https://hf-mirror.com`.
|
| 125 |
Please download weights manually as follows:
|
| 126 |
```
|
| 127 |
pip install "huggingface_hub[cli]"
|
| 128 |
-
cd
|
| 129 |
mkdir checkpoints
|
| 130 |
-
huggingface-cli download FrancisRing/
|
| 131 |
```
|
| 132 |
-
All the weights should be organized
|
| 133 |
The overall file structure of this project should be organized as follows:
|
| 134 |
```
|
| 135 |
-
|
| 136 |
-
├──
|
| 137 |
-
├──
|
| 138 |
-
├──
|
| 139 |
-
├──
|
| 140 |
-
├── checkpoints
|
| 141 |
-
│
|
| 142 |
-
│
|
| 143 |
-
│
|
| 144 |
-
│
|
| 145 |
-
|
| 146 |
-
|
| 147 |
-
├──
|
| 148 |
-
├──
|
| 149 |
-
├──
|
| 150 |
-
|
| 151 |
-
|
| 152 |
-
|
| 153 |
-
|
| 154 |
-
|
| 155 |
-
|
| 156 |
-
|
| 157 |
-
|
| 158 |
-
|
| 159 |
-
|
| 160 |
-
``
|
| 161 |
-
|
| 162 |
-
|
| 163 |
-
|
| 164 |
-
|
| 165 |
-
``
|
| 166 |
-
|
| 167 |
-
|
| 168 |
-
```
|
| 169 |
-
|
| 170 |
-
|
| 171 |
-
|
| 172 |
-
|
| 173 |
-
```
|
| 174 |
-
|
| 175 |
-
``
|
| 176 |
-
|
| 177 |
-
|
| 178 |
-
|
| 179 |
-
|
| 180 |
-
|
| 181 |
-
"
|
| 182 |
-
|
| 183 |
-
|
| 184 |
-
|
| 185 |
-
|
| 186 |
-
```
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 187 |
|
| 188 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 189 |
|
| 190 |
-
|
| 191 |
-
- Wan2.1-1.3B-based StableAvatar weights have two versions: `transformer3d-square.pt` and `transformer3d-rec-vec.pt`, which are trained on two video datasets in two different resolution settings. Two versions both support generating audio-driven avatar video at three different resolution settings: 512x512, 480x832, and 832x480. You can modify `--transformer_path` in `inference.sh` to switch these two versions.
|
| 192 |
|
| 193 |
-
|
| 194 |
-
Setting `--GPU_memory_mode` to `model_cpu_offload` can significantly cut GPU memory usage, reducing it by roughly half compared to `model_full_load` mode.
|
| 195 |
-
|
| 196 |
-
- If you have multiple Gpus, you can run Multi-GPU inference to speed up by modifying "--ulysses_degree" and "--ring_degree" in `inference.sh`. For example, if you have 8 GPUs, you can set `--ulysses_degree=4` and `--ring_degree=2`. Notably, you have to ensure ulysses_degree*ring_degree=total GPU number/world-size. Moreover, you can also add `--fsdp_dit` in `inference.sh` to activate FSDP in DiT to further reduce GPU memory consumption.
|
| 197 |
-
You can fun the following command:
|
| 198 |
-
```
|
| 199 |
-
bash multiple_gpu_inference.sh
|
| 200 |
-
```
|
| 201 |
-
In my setting, 4 GPUs are utilized for inference.
|
| 202 |
-
|
| 203 |
-
The video synthesized by StableAvatar is without audio. If you want to obtain the high quality MP4 file with audio, we recommend you to leverage ffmpeg on the <b>output_path</b> as follows:
|
| 204 |
-
```
|
| 205 |
-
ffmpeg -i video_without_audio.mp4 -i /path/audio.wav -c:v copy -c:a aac -shortest /path/output_with_audio.mp4
|
| 206 |
-
```
|
| 207 |
-
|
| 208 |
-
### 🧱 Model Training
|
| 209 |
-
<b>🔥🔥It’s worth noting that if you’re looking to train a conditioned Video Diffusion Transformer (DiT) model, such as Wan2.1, this training tutorial will also be helpful.🔥🔥</b>
|
| 210 |
-
For the training dataset, it has to be organized as follows:
|
| 211 |
-
|
| 212 |
-
```
|
| 213 |
-
talking_face_data/
|
| 214 |
-
├── rec
|
| 215 |
-
│ │ ├──speech
|
| 216 |
-
│ │ │ ├──00001
|
| 217 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 218 |
-
│ │ │ │ ├──audio.wav
|
| 219 |
-
│ │ │ │ ├──images
|
| 220 |
-
│ │ │ │ │ ├──frame_0.png
|
| 221 |
-
│ │ │ │ │ ├──frame_1.png
|
| 222 |
-
│ │ │ │ │ ├──frame_2.png
|
| 223 |
-
│ │ │ │ │ ├──...
|
| 224 |
-
│ │ │ │ ├──face_masks
|
| 225 |
-
│ │ │ │ │ ├──frame_0.png
|
| 226 |
-
│ │ │ │ │ ├──frame_1.png
|
| 227 |
-
│ │ │ │ │ ├──frame_2.png
|
| 228 |
-
│ │ │ │ │ ├──...
|
| 229 |
-
│ │ │ │ ├──lip_masks
|
| 230 |
-
│ │ │ │ │ ├──frame_0.png
|
| 231 |
-
│ │ │ │ │ ├──frame_1.png
|
| 232 |
-
│ │ │ │ │ ├──frame_2.png
|
| 233 |
-
│ │ │ │ │ ├──...
|
| 234 |
-
│ │ │ ├──00002
|
| 235 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 236 |
-
│ │ │ │ ├──audio.wav
|
| 237 |
-
│ │ │ │ ├──images
|
| 238 |
-
│ │ │ │ ├──face_masks
|
| 239 |
-
│ │ │ │ ├──lip_masks
|
| 240 |
-
│ │ │ └──...
|
| 241 |
-
│ │ ├──singing
|
| 242 |
-
│ │ │ ├──00001
|
| 243 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 244 |
-
│ │ │ │ ├──audio.wav
|
| 245 |
-
│ │ │ │ ├──images
|
| 246 |
-
│ │ │ │ ├──face_masks
|
| 247 |
-
│ │ │ │ ├──lip_masks
|
| 248 |
-
│ │ │ └──...
|
| 249 |
-
│ │ ├──dancing
|
| 250 |
-
│ │ │ ├──00001
|
| 251 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 252 |
-
│ │ │ ��│ ├──audio.wav
|
| 253 |
-
│ │ │ │ ├──images
|
| 254 |
-
│ │ │ │ ├──face_masks
|
| 255 |
-
│ │ │ │ ├──lip_masks
|
| 256 |
-
│ │ │ └──...
|
| 257 |
-
├── vec
|
| 258 |
-
│ │ ├──speech
|
| 259 |
-
│ │ │ ├──00001
|
| 260 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 261 |
-
│ │ │ │ ├──audio.wav
|
| 262 |
-
│ │ │ │ ├──images
|
| 263 |
-
│ │ │ │ ├──face_masks
|
| 264 |
-
│ │ │ │ ├──lip_masks
|
| 265 |
-
│ │ │ └──...
|
| 266 |
-
│ │ ├──singing
|
| 267 |
-
│ │ │ ├──00001
|
| 268 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 269 |
-
│ │ │ │ ├──audio.wav
|
| 270 |
-
│ │ │ │ ├──images
|
| 271 |
-
│ │ │ │ ├──face_masks
|
| 272 |
-
│ │ │ │ ├──lip_masks
|
| 273 |
-
│ │ │ └──...
|
| 274 |
-
│ │ ├──dancing
|
| 275 |
-
│ │ │ ├──00001
|
| 276 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 277 |
-
│ │ │ │ ├──audio.wav
|
| 278 |
-
│ │ │ │ ├──images
|
| 279 |
-
│ │ │ │ ├──face_masks
|
| 280 |
-
│ │ │ │ ├──lip_masks
|
| 281 |
-
│ │ │ └──...
|
| 282 |
-
├── square
|
| 283 |
-
│ │ ├──speech
|
| 284 |
-
│ │ │ ├──00001
|
| 285 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 286 |
-
│ │ │ │ ├──audio.wav
|
| 287 |
-
│ │ │ │ ├──images
|
| 288 |
-
│ │ │ │ ├──face_masks
|
| 289 |
-
│ │ │ │ ├──lip_masks
|
| 290 |
-
│ │ │ └──...
|
| 291 |
-
│ │ ├──singing
|
| 292 |
-
│ │ │ ├──00001
|
| 293 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 294 |
-
│ │ │ │ ├──audio.wav
|
| 295 |
-
│ │ │ │ ├──images
|
| 296 |
-
│ │ │ │ ├──face_masks
|
| 297 |
-
│ │ │ │ ├──lip_masks
|
| 298 |
-
│ │ │ └──...
|
| 299 |
-
│ │ ├──dancing
|
| 300 |
-
│ │ │ ├──00001
|
| 301 |
-
│ │ │ │ ├──sub_clip.mp4
|
| 302 |
-
│ │ │ │ ├──audio.wav
|
| 303 |
-
│ │ │ │ ├──images
|
| 304 |
-
│ │ │ │ ├──face_masks
|
| 305 |
-
│ │ │ │ ├──lip_masks
|
| 306 |
-
│ │ │ └──...
|
| 307 |
-
├── video_rec_path.txt
|
| 308 |
-
├── video_square_path.txt
|
| 309 |
-
└── video_vec_path.txt
|
| 310 |
-
```
|
| 311 |
-
StableAvatar is trained on mixed-resolution videos, with 512x512 videos stored in `talking_face_data/square`, 480x832 videos stored in `talking_face_data/vec`, and 832x480 videos stored in `talking_face_data/rec`. Each folder in `talking_face_data/square` or `talking_face_data/rec` or `talking_face_data/vec` contains three subfolders which contains different types of videos (speech, singing, and dancing).
|
| 312 |
-
All `.png` image files are named in the format `frame_i.png`, such as `frame_0.png`, `frame_1.png`, and so on.
|
| 313 |
-
`00001`, `00002`, `00003` indicate individual video information.
|
| 314 |
-
In terms of three subfolders, `images`, `face_masks`, and `lip_masks` store RGB frames, corresponding human face masks, and corresponding human lip masks, respectively.
|
| 315 |
-
`sub_clip.mp4` and `audio.wav` refer to the corresponding RGB video of `images` and the corresponding audio file.
|
| 316 |
-
`video_square_path.txt`, `video_rec_path.txt`, and `video_vec_path.txt` record folder paths of `talking_face_data/square`, `talking_face_data/rec`, and `talking_face_data/vec`, respectively.
|
| 317 |
-
For example, the content of `video_rec_path.txt` is shown as follows:
|
| 318 |
-
```
|
| 319 |
-
path/StableAvatar/talking_face_data/rec/speech/00001
|
| 320 |
-
path/StableAvatar/talking_face_data/rec/speech/00002
|
| 321 |
-
...
|
| 322 |
-
path/StableAvatar/talking_face_data/rec/singing/00003
|
| 323 |
-
path/StableAvatar/talking_face_data/rec/singing/00004
|
| 324 |
-
...
|
| 325 |
-
path/StableAvatar/talking_face_data/rec/dancing/00005
|
| 326 |
-
path/StableAvatar/talking_face_data/rec/dancing/00006
|
| 327 |
-
...
|
| 328 |
-
```
|
| 329 |
-
If you only have raw videos, you can leverage `ffmpeg` to extract frames from raw videos (speech) and store them in the subfolder `images`.
|
| 330 |
-
```
|
| 331 |
-
ffmpeg -i raw_video_1.mp4 -q:v 1 -start_number 0 path/StableAvatar/talking_face_data/rec/speech/00001/images/frame_%d.png
|
| 332 |
-
```
|
| 333 |
-
The obtained frames are saved in `path/StableAvatar/talking_face_data/rec/speech/00001/images`.
|
| 334 |
-
|
| 335 |
-
For extracting the human face masks, please refer to [StableAnimator repo](https://github.com/Francis-Rings/StableAnimator). The Human Face Mask Extraction section in the tutorial provides off-the-shelf codes.
|
| 336 |
-
|
| 337 |
-
For extracting the human lip masks, you can run the following command:
|
| 338 |
-
```
|
| 339 |
-
pip install mediapipe
|
| 340 |
-
python lip_mask_extractor.py --folder_root="path/StableAvatar/talking_face_data/rec/singing" --start=1 --end=500
|
| 341 |
-
```
|
| 342 |
-
`--folder_root` refers to the root path of training datasets.
|
| 343 |
-
`--start` and `--end` specify the starting and ending indices of the selected training dataset. For example, `--start=1 --end=500` indicates that the human lip extraction will start at `path/StableAvatar/talking_face_data/rec/singing/00001` and end at `path/StableAvatar/talking_face_data/rec/singing/00500`.
|
| 344 |
-
|
| 345 |
-
For extraction details of corresponding audio, please refer to the Audio Extraction section.
|
| 346 |
-
When your dataset is organized exactly as outlined above, you can easily train your Wan2.1-1.3B-based StableAvatar by running the following command:
|
| 347 |
-
```
|
| 348 |
-
# Training StableAvatar on a single resolution setting (512x512) in a single machine
|
| 349 |
-
bash train_1B_square.sh
|
| 350 |
-
# Training StableAvatar on a single resolution setting (512x512) in multiple machines
|
| 351 |
-
bash train_1B_square_64.sh
|
| 352 |
-
# Training StableAvatar on a mixed resolution setting (480x832 and 832x480) in a single machine
|
| 353 |
-
bash train_1B_rec_vec.sh
|
| 354 |
-
# Training StableAvatar on a mixed resolution setting (480x832 and 832x480) in multiple machines
|
| 355 |
-
bash train_1B_rec_vec_64.sh
|
| 356 |
-
```
|
| 357 |
-
For the parameter details of `train_1B_square.sh` and `train_1B_rec_vec.sh`, `CUDA_VISIBLE_DEVICES` refers to gpu devices. In my setting, I use 4 NVIDIA A100 80G to train StableAvatar (`CUDA_VISIBLE_DEVICES=3,2,1,0`).
|
| 358 |
-
`--pretrained_model_name_or_path`, `--pretrained_wav2vec_path`, and `--output_dir` refer to the pretrained Wan2.1-1.3B path, pretrained Wav2Vec2.0 path, and the checkpoint saved path of the trained StableAvatar.
|
| 359 |
-
`--train_data_square_dir`, `--train_data_rec_dir`, and `--train_data_vec_dir` are the paths of `video_square_path.txt`, `video_rec_path.txt`, and `video_vec_path.txt`, respectively.
|
| 360 |
-
`--validation_reference_path` and `--validation_driven_audio_path` are paths of the validation reference image and the validation driven audio.
|
| 361 |
-
`--video_sample_n_frames` is the number of frames that StableAvatar processes in a single batch.
|
| 362 |
-
`--num_train_epochs` is the training epoch number. It is worth noting that the default number of training epochs is set to infinite. You can manually terminate the training process once you observe that your StableAvatar has reached its peak performance.
|
| 363 |
-
For the parameter details of `train_1B_square_64.sh` and `train_1B_rec_vec_64.sh`, we set the GPU configuration in `path/StableAvatar/accelerate_config/accelerate_config_machine_1B_multiple.yaml`. In my setting, the training setup consists of 8 nodes, each equipped with 8 NVIDIA A100 80GB GPUs, for training StableAvatar.
|
| 364 |
-
|
| 365 |
-
The overall file structure of StableAvatar at training is shown as follows:
|
| 366 |
-
```
|
| 367 |
-
StableAvatar/
|
| 368 |
-
├── accelerate_config
|
| 369 |
-
├── deepspeed_config
|
| 370 |
-
├── talking_face_data
|
| 371 |
-
├── examples
|
| 372 |
-
├── wan
|
| 373 |
-
├── checkpoints
|
| 374 |
-
│ ├── Kim_Vocal_2.onnx
|
| 375 |
-
│ ├── wav2vec2-base-960h
|
| 376 |
-
│ ├── Wan2.1-Fun-V1.1-1.3B-InP
|
| 377 |
-
│ └── StableAvatar-1.3B
|
| 378 |
-
├── inference.py
|
| 379 |
-
├── inference.sh
|
| 380 |
-
├── train_1B_square.py
|
| 381 |
-
├── train_1B_square.sh
|
| 382 |
-
├── train_1B_vec_rec.py
|
| 383 |
-
├── train_1B_vec_rec.sh
|
| 384 |
-
├── audio_extractor.py
|
| 385 |
-
├── vocal_seperator.py
|
| 386 |
-
├── requirement.txt
|
| 387 |
-
```
|
| 388 |
-
<b>It is worth noting that training StableAvatar requires approximately 50GB of VRAM due to the mixed-resolution (480x832 and 832x480) training pipeline.
|
| 389 |
-
However, if you train StableAvatar exclusively on 512x512 videos, the VRAM requirement is reduced to approximately 40GB.</b>
|
| 390 |
-
Additionally, The backgrounds of the selected training videos should remain static, as this helps the diffusion model calculate accurate reconstruction loss.
|
| 391 |
-
The audio should be clear and free from excessive background noise.
|
| 392 |
-
|
| 393 |
-
Regarding training Wan2.1-14B-based StableAvatar, you can run the following command:
|
| 394 |
-
```
|
| 395 |
-
# Training StableAvatar on a mixed resolution setting (480x832, 832x480, and 512x512) in multiple machines
|
| 396 |
-
huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P
|
| 397 |
-
huggingface-cli download Wan-AI/Wan2.1-I2V-14B-720P --local-dir ./checkpoints/Wan2.1-I2V-14B-720P # Optional
|
| 398 |
-
bash train_14B.sh
|
| 399 |
-
```
|
| 400 |
-
We utilize deepspeed stage-2 to train Wan2.1-14B-based StableAvatar. The GPU configuration can be modified in `path/StableAvatar/accelerate_config/accelerate_config_machine_14B_multiple.yaml`.
|
| 401 |
-
The deepspeed optimization configuration and deepspeed scheduler configuration are in `path/StableAvatar/deepspeed_config/zero_stage2_config.json`.
|
| 402 |
-
Notably, we observe that Wan2.1-1.3B-based StableAvatar is already capable of synthesizing infinite-length high quality avatar videos. The Wan2.1-14B backbone significantly increase the inference latency and GPU memory consumption during training, indicating limited efficiency in terms of performance-to-resource ratio.
|
| 403 |
-
|
| 404 |
-
You can also run the following commands to perform lora training:
|
| 405 |
-
```
|
| 406 |
-
# Training StableAvatar-1.3B on a mixed resolution setting (480x832 and 832x480) in a single machine
|
| 407 |
-
bash train_1B_rec_vec_lora.sh
|
| 408 |
-
# Training StableAvatar-1.3B on a mixed resolution setting (480x832 and 832x480) in multiple machines
|
| 409 |
-
bash train_1B_rec_vec_lora_64.sh
|
| 410 |
-
# Lora-Training StableAvatar-14B on a mixed resolution setting (480x832, 832x480, and 512x512) in multiple machines
|
| 411 |
-
bash train_14B_lora.sh
|
| 412 |
-
```
|
| 413 |
-
You can modify `--rank` and `--network_alpha` to control the quality of your lora training/finetuning.
|
| 414 |
-
|
| 415 |
-
If you want to train 720P Wan2.1-1.3B-based or Wan2.1-14B-based StableAvatar, you can directly modify the height and width of the dataloader (480p-->720p) in `train_1B_square.py`/`train_1B_vec_rec.py`/`train_14B.py`.
|
| 416 |
-
|
| 417 |
-
### 🧱 Model Finetuning
|
| 418 |
-
Regarding fully finetuning StableAvatar, you can add `--transformer_path="path/StableAvatar/checkpoints/StableAvatar-1.3B/transformer3d-square.pt"` to the `train_1B_rec_vec.sh` or `train_1B_rec_vec_64.sh`:
|
| 419 |
-
```
|
| 420 |
-
# Finetuning StableAvatar on a mixed resolution setting (480x832 and 832x480) in a single machine
|
| 421 |
-
bash train_1B_rec_vec.sh
|
| 422 |
-
# Finetuning StableAvatar on a mixed resolution setting (480x832 and 832x480) in multiple machines
|
| 423 |
-
bash train_1B_rec_vec_64.sh
|
| 424 |
-
```
|
| 425 |
-
In terms of lora finetuning StableAvatar, you can add `--transformer_path="path/StableAvatar/checkpoints/StableAvatar-1.3B/transformer3d-square.pt"` to the `train_1B_rec_vec_lora.sh`:
|
| 426 |
-
```
|
| 427 |
-
# Lora-Finetuning StableAvatar-1.3B on a mixed resolution setting (480x832 and 832x480) in a single machine
|
| 428 |
-
bash train_1B_rec_vec_lora.sh
|
| 429 |
-
```
|
| 430 |
-
You can modify `--rank` and `--network_alpha` to control the quality of your lora training/finetuning.
|
| 431 |
|
| 432 |
-
### 🧱
|
|
|
|
|
|
|
| 433 |
|
| 434 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 435 |
|
| 436 |
-
|
| 437 |
|
| 438 |
-
|
| 439 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 440 |
|
| 441 |
Email: francisshuyuan@gmail.com
|
| 442 |
|
| 443 |
-
If you find our work useful, <b>please consider giving a star ⭐ to this github repository
|
| 444 |
```bib
|
| 445 |
-
@article{
|
| 446 |
-
title={
|
| 447 |
-
author={Tu, Shuyuan and
|
| 448 |
-
journal={arXiv preprint arXiv:
|
| 449 |
-
year={
|
| 450 |
}
|
| 451 |
```
|
|
|
|
| 5 |
tags:
|
| 6 |
- video-generation
|
| 7 |
- video diffusion transformer
|
| 8 |
+
- joint-video-audio-generation
|
| 9 |
+
- sparse-attention
|
| 10 |
+
- high-resolution
|
| 11 |
task_categories:
|
| 12 |
- image-to-video
|
| 13 |
- text-to-video
|
| 14 |
---
|
| 15 |
|
| 16 |
+
# Prism
|
| 17 |
|
| 18 |
+
<a href='https://francis-rings.github.io/Prism'><img src='https://img.shields.io/badge/Project-Page-Green'></a> <a href='https://arxiv.org/abs/'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a> <a href='https://huggingface.co/FrancisRing/Prism/tree/main'><img src='https://img.shields.io/badge/HuggingFace-Model-orange'></a> <a href='https://www.youtube.com/watch?v=6lhvmbzvv3Y'><img src='https://img.shields.io/badge/YouTube-Watch-red?style=flat-square&logo=youtube'></a> <a href='https://www.bilibili.com/video/BV1hUt9z4EoQ'><img src='https://img.shields.io/badge/Bilibili-Watch-blue?style=flat-square&logo=bilibili'></a>
|
| 19 |
|
| 20 |
+
Prism:Dynamic Sparse Attention for Native 2K Joint Video-Audio Generation Model Training
|
| 21 |
<br/>
|
| 22 |
+
[Shuyuan Tu](https://github.com/Francis-Rings)<sup>1</sup>, [Qi Tian](https://scholar.google.com/citations?user=Ypu45nIAAAAJ&hl=zh-CN)<sup>2</sup>, [Yinming Huang](https://yinminghuang.github.io/)<sup>1</sup>, [Yue Wu](https://scholar.google.com/citations?user=1xTR6qoAAAAJ&hl=en)<sup>2</sup>, [Xintong Han](https://scholar.google.com/citations?user=FGiWOIAAAAAJ&hl=en)<sup>2</sup>, [Kaihang Pan](https://scholar.google.com/citations?user=lMQADDUAAAAJ&hl=zh-CN)<sup>3</sup>, [Weijie Kong](https://scholar.google.com/citations?user=gsOklKAAAAAJ&hl=zh-CN)<sup>2</sup>, [Jiangfeng Xiong](https://scholar.google.com/citations?user=lHbXg_0AAAAJ&hl=zh-TW)<sup>2</sup>, [Jian-Wei Zhang](https://scholar.google.com/citations?user=nF_klRIAAAAJ&hl=zh-CN)<sup>2</sup>, [Zuxuan Wu](https://scholar.google.com/citations?user=7t12hVkAAAAJ&hl=en)<sup>1</sup>, [Yu-Gang Jiang](https://scholar.google.com/citations?user=f3_FP8AAAAAJ&hl=en)<sup>1</sup>
|
| 23 |
<br/>
|
| 24 |
+
[<sup>1</sup>Fudan University; <sup>2</sup>Tencent Hunyuan; <sup>3</sup>Zhejiang University]
|
| 25 |
+
|
| 26 |
|
| 27 |
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
| 28 |
<tr>
|
| 29 |
<td>
|
| 30 |
+
<video src="https://github.com/user-attachments/assets/0f2c8471-bd25-4151-82a1-42ae631b325a" width="320" controls loop></video>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 31 |
</td>
|
| 32 |
<td>
|
| 33 |
+
<video src="https://github.com/user-attachments/assets/39e384ee-b421-447e-aaa5-2db2b4f5c4c2" width="320" controls loop></video>
|
| 34 |
</td>
|
| 35 |
<td>
|
| 36 |
+
<video src="https://github.com/user-attachments/assets/965b05cb-83a9-48b1-8ca6-783a38876db6" width="320" controls loop></video>
|
| 37 |
</td>
|
| 38 |
</tr>
|
| 39 |
<tr>
|
| 40 |
<td>
|
| 41 |
+
<video src="https://github.com/user-attachments/assets/8b3aaa07-2d3e-446b-8400-655cf11f715a" width="320" controls loop></video>
|
| 42 |
</td>
|
| 43 |
<td>
|
| 44 |
+
<video src="https://github.com/user-attachments/assets/39a57fdc-499a-46e5-bb00-1628d6d8780f" width="320" controls loop></video>
|
| 45 |
</td>
|
| 46 |
<td>
|
| 47 |
+
<video src="https://github.com/user-attachments/assets/16443caa-490f-4cf7-8115-e380ab57c106" width="320" controls loop></video>
|
| 48 |
</td>
|
| 49 |
</tr>
|
| 50 |
</table>
|
| 51 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 52 |
|
| 53 |
## Overview
|
| 54 |
|
| 55 |
<p align="center">
|
| 56 |
<img src="assets/figures/framework.jpg" alt="model architecture" width="1280"/>
|
| 57 |
</br>
|
| 58 |
+
<i>The overview of the framework of Prism.</i>
|
| 59 |
</p>
|
| 60 |
|
| 61 |
+
Native training joint video-audio generation models at higher resolutions empowers them to learn richer visual details and sharper motion dynamics.
|
| 62 |
+
However, full attention incurs quadratic cost and, as resolution increases, spreads attention over increasingly redundant tokens, diluting learning signals for informative content and disrupting pretrained priors.
|
| 63 |
+
Existing sparse attention methods either target training-free acceleration or overlook the unique structure of joint video-audio data, where cross-modal interactions are inherently concentrated around sound-producing regions.
|
| 64 |
+
To address this, we propose Prism, a dynamic sparse attention framework for natively training joint video-audio generation models at 2K.
|
| 65 |
+
In particular, Prism organizes the token sequence into spatiotemporal macro-zones, enabling the attention structure to adapt to local content.
|
| 66 |
+
For each zone, it estimates local information structure via video feature variance along the channel and feature norms from the audio-to-video cross-attention, jointly capturing how visual content varies directionally and how strongly audio influences each visual region.
|
| 67 |
+
Based on these signals, Prism dynamically assigns a tailored block shape to each zone, applying finer partitioning along axes of rapid visual content variation and strong audio-visual coupling.
|
| 68 |
+
This encourages tokens within each block to remain semantically coherent, allowing block-level features to capture both visual content and joint video-audio interaction patterns.
|
| 69 |
+
Prism further adopts a hybrid block selection strategy to dynamically determine per-query sparsity.
|
| 70 |
+
Experiments show that Prism achieves 2.5$\times$ training speedup compared to full attention, while surpassing it in generation quality.
|
| 71 |
|
| 72 |
## News
|
| 73 |
+
* `[2026-x-xx]`:🔥 The project page, training/inference code, technical report and [a preview model checkpoint](https://huggingface.co/FrancisRing/Prism/tree/main) are released. Stay tuned!
|
|
|
|
|
|
|
|
|
|
|
|
|
| 74 |
|
| 75 |
## 🛠️ To-Do List
|
| 76 |
+
- [x] Prism-preview-alpha (stable)
|
| 77 |
+
- [x] Prism-preview-beta (motion)
|
| 78 |
+
- [x] Data Pre-Processing Code (Latent Extraction and Latent Decode)
|
|
|
|
| 79 |
- [x] Training Code
|
| 80 |
- [x] Full Finetuning Code
|
| 81 |
+
- [x] Inference Code
|
| 82 |
+
- [ ] Prism-pro
|
|
|
|
|
|
|
| 83 |
|
| 84 |
## 🔑 Quickstart
|
| 85 |
|
| 86 |
+
The Preview Model supports native joint video-audio generation at 720p, 1080p, and 2K resolutions, as well as native joint video-audio training at these resolutions.
|
| 87 |
|
| 88 |
+
### 🧱 Environment Setup
|
| 89 |
|
| 90 |
+
**Prerequisites.** Prism requires Python ≥ 3.10, CUDA ≥ 12.4, and a system-level `ffmpeg` installation.
|
| 91 |
```
|
| 92 |
+
# 1 — Install PyTorch (CUDA 12.4)
|
| 93 |
+
pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu124
|
| 94 |
+
# 2 — Install core dependencies
|
| 95 |
pip install -r requirements.txt
|
| 96 |
+
# 3 — Install FlashAttention (requires CUDA)
|
| 97 |
+
pip install flash-attn --no-build-isolation
|
| 98 |
+
# 4 — Install `ffmpeg` (system-level)
|
| 99 |
+
sudo apt-get install ffmpeg
|
| 100 |
+
# 5 — Verify the installation
|
| 101 |
+
python -c "
|
| 102 |
+
import torch, diffusers, transformers, flash_attn, triton
|
| 103 |
+
from audiotools import AudioSignal
|
| 104 |
+
from decord import VideoReader
|
| 105 |
+
print(f'torch {torch.__version__}')
|
| 106 |
+
print(f'diffusers {diffusers.__version__}')
|
| 107 |
+
print(f'transformers {transformers.__version__}')
|
| 108 |
+
print(f'flash_attn {flash_attn.__version__}')
|
| 109 |
+
print(f'triton {triton.__version__}')
|
| 110 |
+
print('All imports OK ✅')"
|
| 111 |
+
```
|
| 112 |
+
**(Multi-node only) Quick cluster-wide install via `pdsh`**
|
| 113 |
+
|
| 114 |
+
If you are running on a multi-node cluster, the following one-liner installs the key
|
| 115 |
+
packages across all nodes in parallel:
|
| 116 |
+
```
|
| 117 |
+
export node_ip=$(echo ${NODE_IP_LIST} | sed 's/:8//g')
|
| 118 |
+
pdsh -f 256 -w $node_ip "pip install ipdb yunchang mmengine"
|
| 119 |
+
pdsh -f 256 -w $node_ip "pip uninstall diffusers -y && pip install diffusers==0.33.0"
|
| 120 |
+
pdsh -f 256 -w $node_ip "pip install descript-audiotools"
|
| 121 |
+
pdsh -f 256 -w $node_ip "pip install imageio[ffmpeg] imageio[pyav]"
|
| 122 |
+
pdsh -f 256 -w $node_ip "pip install decord"
|
| 123 |
+
pdsh -f 256 -w $node_ip "pip install flash-attn --no-build-isolation"
|
| 124 |
+
```
|
| 125 |
+
|
| 126 |
+
### 🧱 Download Weights
|
| 127 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 128 |
If you encounter connection issues with Hugging Face, you can utilize the mirror endpoint by setting the environment variable: `export HF_ENDPOINT=https://hf-mirror.com`.
|
| 129 |
Please download weights manually as follows:
|
| 130 |
```
|
| 131 |
pip install "huggingface_hub[cli]"
|
| 132 |
+
cd Prism
|
| 133 |
mkdir checkpoints
|
| 134 |
+
huggingface-cli download FrancisRing/Prism --local-dir ./checkpoints
|
| 135 |
```
|
| 136 |
+
All the weights should be organized as follows.
|
| 137 |
The overall file structure of this project should be organized as follows:
|
| 138 |
```
|
| 139 |
+
Prism/
|
| 140 |
+
├── configs/
|
| 141 |
+
├── hymm/
|
| 142 |
+
├── scripts/
|
| 143 |
+
├── tools/
|
| 144 |
+
├── checkpoints/
|
| 145 |
+
│ ├── pretrained_models/
|
| 146 |
+
│ │ └── MOVA-360p/ # base MOVA pretrained model
|
| 147 |
+
│ ├── preview_alpha/ # Prism-preview-alpha weights
|
| 148 |
+
│ │ └── diffusion_pytorch_model.safetensors
|
| 149 |
+
│ └── preview_beta/ # Prism-preview-beta weights
|
| 150 |
+
│ └── diffusion_pytorch_model.safetensors
|
| 151 |
+
├── mova_infer.sh
|
| 152 |
+
├── mova_infer_fsdp.sh
|
| 153 |
+
├── requirements.txt
|
| 154 |
+
└── README.md
|
| 155 |
+
```
|
| 156 |
+
|
| 157 |
+
### 🧱 Training Dataset Preparation
|
| 158 |
+
|
| 159 |
+
Training reads a plain CSV file (one row per clip). The row details are depicted as follows:
|
| 160 |
+
|
| 161 |
+
| column | required | meaning |
|
| 162 |
+
| --- | --- | --- |
|
| 163 |
+
| `latent_path` | yes | cached WAN VAE latent `.npy`, shape `[C, T, H, W]` |
|
| 164 |
+
| `video_path` | yes | source video, read only for the reference frame |
|
| 165 |
+
| `audio_path` | yes | source audio (`.m4a`, `.wav`, …) |
|
| 166 |
+
| `caption` | yes | structured JSON caption, or plain text |
|
| 167 |
+
| `num_frames` | yes† | total frame count of the source video (pixel-space) |
|
| 168 |
+
| `height` | yes† | video height in pixels |
|
| 169 |
+
| `width` | yes† | video width in pixels |
|
| 170 |
+
| `fps` | yes† | source frame rate (used for audio-video alignment) |
|
| 171 |
+
| `video_id` | no | identifier used in logs |
|
| 172 |
+
| `audio_caption` | no | separate audio prompt; defaults to `caption` |
|
| 173 |
+
| `latent_frames`, `latent_height`, `latent_width` | no | cached latent dims; when present, used for bucketing instead of `num_frames/height/width` |
|
| 174 |
+
| `best_frame_index` | no | frame offset where the audio crop starts (default `0`) |
|
| 175 |
+
| `ref_frame_index` | no | frame used as the i2v reference (default `1`) |
|
| 176 |
+
|
| 177 |
+
> † When `latent_frames/latent_height/latent_width` are present in the CSV, they take priority for bucketing and `num_frames/height/width` become optional. `fps` defaults to `dataloader_config.video_fps` (24) when absent. We recommend always including all four columns (`num_frames`, `height`, `width`, `fps`) for clarity.
|
| 178 |
+
|
| 179 |
+
`caption` keeps the structured format, so the existing sampling ratios in
|
| 180 |
+
`video_caption_sample_ratio` still control which fields get assembled into the
|
| 181 |
+
prompt:
|
| 182 |
+
|
| 183 |
+
```json
|
| 184 |
+
{"style_features": "...", "content_summary": "...", "background_audio": "...",
|
| 185 |
+
"shots": [{"time_range": ["0", "9.0"], "static_description": "...", "dynamic_description": "..."}]}
|
| 186 |
+
```
|
| 187 |
+
|
| 188 |
+
**Step 1 — Extract WAN VAE latents** (offline, once per video):
|
| 189 |
+
|
| 190 |
+
```
|
| 191 |
+
python hymm/wan_vae_latent_extraction.py \
|
| 192 |
+
--input_csv /path/i2va_raw.csv \
|
| 193 |
+
--output_dir /path/latents \
|
| 194 |
+
--model_path /path/checkpoints/pretrained_models/MOVA-360p/video_vae \
|
| 195 |
+
--sample_n_frames 289 --fps 24
|
| 196 |
+
```
|
| 197 |
+
|
| 198 |
+
**Step 2 — Assign multi-resolution and multi-duration buckets** (offline):
|
| 199 |
+
|
| 200 |
+
The temporal bucket range **must match the training config**.
|
| 201 |
+
The three resolution configs use:
|
| 202 |
+
|
| 203 |
+
| config | `latent-resolution` | `temporal-max-length` |
|
| 204 |
+
| --- | --- | --- |
|
| 205 |
+
| 720p | `720p` | `289` |
|
| 206 |
+
| 1080p | `1080p` | `289` |
|
| 207 |
+
| 2k | `2k` | `265` |
|
| 208 |
|
| 209 |
+
```bash
|
| 210 |
+
# 720p
|
| 211 |
+
python tools/build_dataset_csv.py \
|
| 212 |
+
--input /path/i2va_raw.csv \
|
| 213 |
+
--output /path/i2va_720p.csv \
|
| 214 |
+
--latent-resolution 720p \
|
| 215 |
+
--temporal-min-length 49 --temporal-max-length 289 --temporal-interval 12 \
|
| 216 |
+
--scan-latents --workers 32
|
| 217 |
+
|
| 218 |
+
# 1080p
|
| 219 |
+
python tools/build_dataset_csv.py \
|
| 220 |
+
--input /path/i2va_raw.csv \
|
| 221 |
+
--output /path/i2va_1080p.csv \
|
| 222 |
+
--latent-resolution 1080p \
|
| 223 |
+
--temporal-min-length 49 --temporal-max-length 289 --temporal-interval 12 \
|
| 224 |
+
--scan-latents --workers 32
|
| 225 |
+
|
| 226 |
+
# 2k (note: temporal-max-length is 265)
|
| 227 |
+
python tools/build_dataset_csv.py \
|
| 228 |
+
--input /path/i2va_raw.csv \
|
| 229 |
+
--output /path/i2va_2k.csv \
|
| 230 |
+
--latent-resolution 2k \
|
| 231 |
+
--temporal-min-length 49 --temporal-max-length 265 --temporal-interval 12 \
|
| 232 |
+
--scan-latents --workers 32
|
| 233 |
+
```
|
| 234 |
+
|
| 235 |
+
`--scan-latents` fills `latent_frames/height/width` from the `.npy` headers, and
|
| 236 |
+
the tool prints the per-bucket histogram so you can see how batches will fill
|
| 237 |
+
before launching.
|
| 238 |
+
The bucket columns are optional. The dataset derives them on the fly when they are
|
| 239 |
+
missing, using the same `bucket_*` settings from the config.
|
| 240 |
+
|
| 241 |
+
### 🧱 Preview Model Inference
|
| 242 |
+
|
| 243 |
+
Prism provides two inference scripts: `prism_infer.sh` for single-GPU or DeepSpeed-based multi-GPU inference, and `prism_infer_fsdp.sh` for FSDP-based multi-GPU inference that shards model parameters across GPUs (recommended for 1080p and 2K).
|
| 244 |
+
|
| 245 |
+
```bash
|
| 246 |
+
# Single-GPU / DeepSpeed multi-GPU inference (720p recommended)
|
| 247 |
+
bash prism_infer.sh
|
| 248 |
+
|
| 249 |
+
# FSDP multi-GPU inference (1080p / 2K recommended)
|
| 250 |
+
bash prism_infer_fsdp.sh
|
| 251 |
+
```
|
| 252 |
+
|
| 253 |
+
Before running, edit the shell script to set your own paths and generation settings. The core parameters are:
|
| 254 |
+
|
| 255 |
+
| Parameter | Default | Description |
|
| 256 |
+
| --- |---------------------------------------------| --- |
|
| 257 |
+
| `--ckpt` | — | Path to the pretrained MOVA base model (e.g. `checkpoints/pretrained_models/MOVA-360p`) |
|
| 258 |
+
| `--resume_ckpt` | — | Path to the Prism checkpoint (e.g. `checkpoints/preview_alpha/diffusion_pytorch_model.safetensors`) |
|
| 259 |
+
| `--config` | `configs/train/t2va_config/mova_infer.yaml` | Inference config that sets diffusion steps, default resolution, and seed |
|
| 260 |
+
| `--height` / `--width` | 480 / 848 | Output resolution. Recommended: 720p = `720 × 1280`, 1080p = `1072 × 1920`, 2K = `1440 × 2560` |
|
| 261 |
+
| `--num_frames` | `205` | Number of frames to generate; must satisfy `(n − 1) % 4 == 0`, auto-snapped if not |
|
| 262 |
+
| `--prompt` / `--audio_prompt` | — | Text prompt for video / audio generation; supports structured tags such as `<music>`, `<sfx>`, `<speech>` |
|
| 263 |
+
| `--ref_path` | — | Path to the reference image for image-to-video generation |
|
| 264 |
+
| `--output_path` | — | Path to save the generated `.mp4` video |
|
| 265 |
+
| `--sp_size` | `4` or `8` | Sequence-parallelism degree (number of GPUs sharing one sample's token sequence) |
|
| 266 |
+
| `--offload` | `cpu` | Offload frozen encoders to CPU to save GPU memory; omit for full GPU-resident inference |
|
| 267 |
+
| `--seed` | `42` | Random seed for reproducibility |
|
| 268 |
+
| `VISUAL_SHIFT` / `AUDIO_SHIFT` | `9.0` / `7.0` | Noise schedule shift for video / audio denoising; higher values allocate more steps to high-noise stages. Recommended: 720p `7.0`, 1080p `9.0–13.0`, 2K `13.0–17.0` |
|
| 269 |
+
| `CFG_SCALE` | `5.0` | Classifier-free guidance scale; `1.0` = no guidance, higher = stronger prompt adherence |
|
| 270 |
+
| `ENABLE_TILING` | `false` | Enable spatial tiling for VAE decode to avoid OOM at high resolutions (e.g. 2K); `TILE_SAMPLE_MIN_SIZE` and `TILE_SAMPLE_STRIDE` control tile size and overlap |
|
| 271 |
+
| `ENABLE_BSA` / `BSA_SPARSITY` | `true` / `0.93` or `0.85` or `0.75` | Enable Prism Block Sparse Attention and set the top-k sparsity ratio during inference |
|
| 272 |
+
| `ENABLE_IVPQ_DYNAMIC_BLOCK` | `true` | Enable dynamic block-shape assignment for inference |
|
| 273 |
|
| 274 |
+
> **Tip:** `mova_infer.sh` uses `--deepspeed` launch and is suited for 720p on a single 80 GB GPU (with `--offload cpu`). `mova_infer_fsdp.sh` uses `torchrun` with FSDP `FULL_SHARD` and is recommended for 1080p / 2K on 4+ GPUs, as it shards parameters across all GPUs rather than replicating the full model.
|
|
|
|
| 275 |
|
| 276 |
+
**❤️ Notably, for your convenience, we provide several ready-to-infer cases in the `assets/ti2va_cases` folder. Feel free to try them out and have fun! ❤️**
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 277 |
|
| 278 |
+
### 🧱 Native Joint Video-Audio Training at High-Resolution
|
| 279 |
+
|
| 280 |
+
Prism supports native joint video-audio training at 720p, 1080p, and 2K resolutions. The training pipeline uses FSDP and sequence parallelism via `torchrun` + `pdsh`, controlled by a two-level config system: a top-level YAML (`configs/train/wan_ti2va.yaml`) that holds per-resolution experiment entries (optimizer, block sparse attention flags, etc.) and a per-resolution inject-config YAML (`configs/train/t2va_config/wan_15B_ti2va_*_init_top_k_p.yaml`) that specifies the pretrained model path, loss weights, and all dataloader settings.
|
| 281 |
|
| 282 |
+
Regarding the bucketed CSV files from step2 in 🧱Training Dataset Preparation, they are the sole data manifest consumed by the trainer.
|
| 283 |
+
Each inject-config YAML has a `dataloader_config.video_csv_file` field. This is where you place your CSV path.
|
| 284 |
+
For example, the 720p inject config (`configs/train/t2va_config/wan_15B_ti2va_720p_init_top_k_p.yaml`) contains:
|
| 285 |
+
```yaml
|
| 286 |
+
dataloader_config:
|
| 287 |
+
video_csv_file:
|
| 288 |
+
- /path/to/your/i2va_720p.csv # ← put your bucketed CSV path here
|
| 289 |
+
```
|
| 290 |
+
|
| 291 |
+
Replace the path with the actual location of the CSV produced by `tools/build_dataset_csv.py`. The 1080p and 2K inject configs follow the same pattern — update `video_csv_file` in `configs/train/t2va_config/wan_15B_ti2va_1080p_init_top_k_p.yaml` and `configs/train/t2va_config/wan_15B_ti2va_2k_init_top_k_p.yaml` respectively. You may also specify multiple CSV files as a list under `video_csv_file` to merge several datasets. At training launch, the script reads the `inject_config` path from the top-level YAML, loads the inject-config, and passes `video_csv_file` to `CsvVideoAudioDataset`. The dataset reads every row, derives or looks up the `(bucket_frames, bucket_height, bucket_width)` assignment, and groups all samples by bucket key. A `DistributedBucketBatchSampler` then draws each mini-batch from a single bucket so every sample in the batch shares the exact same tensor shape, avoiding shape mismatches across data-parallel ranks. The multi-resolution spatial buckets are generated by enumerating all `(H, W)` crops whose patch count equals `(base_size / patch_size)²`, and each sample is matched to the bucket with the closest aspect ratio; the multi-duration temporal buckets are a series of frame counts of the form `4k + 1` from `temporal_min_length` to `temporal_max_length` at stride `temporal_interval`, and each sample is assigned to the largest temporal bucket that fits.
|
| 292 |
+
|
| 293 |
+
**Launch training** at each resolution with:
|
| 294 |
+
|
| 295 |
+
```bash
|
| 296 |
+
# 720p native joint video-audio training
|
| 297 |
+
bash scripts/pretrain/pdsh_train.sh "scripts/pretrain/train_wan_ti2va.sh configs/train/wan_ti2va.yaml wan_15B_ti2va_720p_init_top_k_p"
|
| 298 |
+
|
| 299 |
+
# 1080p native joint video-audio training
|
| 300 |
+
bash scripts/pretrain/pdsh_train.sh "scripts/pretrain/train_wan_ti2va.sh configs/train/wan_ti2va.yaml wan_15B_ti2va_1080p_init_top_k_p"
|
| 301 |
+
|
| 302 |
+
# 2K native joint video-audio training
|
| 303 |
+
bash scripts/pretrain/pdsh_train.sh "scripts/pretrain/train_wan_ti2va.sh configs/train/wan_ti2va.yaml wan_15B_ti2va_2k_init_top_k_p"
|
| 304 |
+
```
|
| 305 |
+
|
| 306 |
+
`pdsh_train.sh` is the inner script across all nodes listed in `$NODE_IP_LIST` via `pdsh`.
|
| 307 |
+
On each node, `train_wan_ti2va.sh` parses the named config from the top-level YAML, extracts every flag, then launches `torchrun` with `hymm/train_t2va_wan.py`.
|
| 308 |
+
|
| 309 |
+
**Core training parameters.** The training configuration is split across two YAML files. Parameters in the **top-level config** (`configs/train/wan_ti2va.yaml`) control the optimizer, parallelism, checkpointing, and Prism sparse-attention behaviour.
|
| 310 |
+
Parameters in the **inject config** (`configs/train/t2va_config/wan_15B_ti2va_*_init_top_k_p.yaml`) control the pretrained model, loss weighting, noise schedule, and dataloader.
|
| 311 |
+
|
| 312 |
+
**Top-level config** (`configs/train/wan_ti2va.yaml`):
|
| 313 |
+
|
| 314 |
+
| Parameter | Default | Description |
|
| 315 |
+
| --- |------------------------| --- |
|
| 316 |
+
| `sp_size` | 8 (1080p/2K), 4 (720p) | Sequence-parallelism degree — number of GPUs that jointly shard a single sample's token sequence to reduce per-GPU memory |
|
| 317 |
+
| `lr` | `1e-5` | Peak learning rate |
|
| 318 |
+
| `lr_scheduler` | `constant_with_warmup` | LR schedule shape; ramps linearly during warmup then holds constant |
|
| 319 |
+
| `lr_warmup_steps` | `200` | Number of steps for the linear LR warmup |
|
| 320 |
+
| `optimizer` | `adamw` | Optimizer type |
|
| 321 |
+
| `weight_decay` | `0.01` | L2 regularization weight |
|
| 322 |
+
| `max_grad_norm` | `1.0` | Maximum gradient norm for gradient clipping to stabilize training |
|
| 323 |
+
| `gradient_accumulation_steps` | `1` | Number of forward-backward micro-steps accumulated before one optimizer update; effectively multiplies the global batch size without extra memory |
|
| 324 |
+
| `fsdp_strategy` | `full` | FSDP sharding strategy — `full` shards parameters, gradients, and optimizer states across all ranks for maximum memory savings; `hybrid` (HSDP) shards within each node and replicates across nodes, reducing cross-node communication at the cost of higher per-node memory; `none` replicates the full model on every rank |
|
| 325 |
+
| `checkpointing_steps` | `50` | Save a checkpoint every N training steps |
|
| 326 |
+
| `train_full_model` | `true` | Whether all model parameters are trainable; `false` = attention-only |
|
| 327 |
+
| `offload_frozen_encoders` | `true` | Move frozen text encoder / audio VAE / video VAE back to CPU after each forward pass to free GPU memory; `false` keeps them on GPU |
|
| 328 |
+
| `empty_cache_interval` | `1` | Call `torch.cuda.empty_cache()` every N steps to release the CUDA allocator pool; `1` is safest but adds a device sync each step |
|
| 329 |
+
| `loss_spike_threshold` | `1.5` | Skip the optimizer step if loss exceeds this value, preventing catastrophic weight updates from noisy batches |
|
| 330 |
+
| `grad_norm_spike_threshold` | `1e6` | Skip the optimizer step if gradient norm exceeds this value |
|
| 331 |
+
| `enable_bsa` | `true` | Activate Prism Block Sparse Attention for video self-attention |
|
| 332 |
+
| `bsa_sparsity` | `0.75/0.80/0.85/0.93` | Top-k sparsity ratio — each query block attends to only 25% of key blocks |
|
| 333 |
+
| `bsa_chunk_3d_shape_q` / `bsa_chunk_3d_shape_k` | `4 4 4` | The 3-D macro-zone shape `(T, H, W)` that tiles the spatiotemporal volume for query / key blocking |
|
| 334 |
+
| `bsa_cdf_threshold` | `0.20` | Top-p CDF threshold for hybrid block selection; combined with top-k for adaptive per-query sparsity |
|
| 335 |
+
| `enable_ivpq_dynamic_block` | `true` | Enable the anisotropic dynamic block-shape assignment (Prism core), where each macro-zone gets a tailored block shape based on local content variance and audio-visual coupling |
|
| 336 |
|
| 337 |
+
**Inject config** (`configs/train/t2va_config/wan_15B_ti2va_*_init_top_k_p.yaml`):
|
| 338 |
|
| 339 |
+
| Parameter | Default | Description |
|
| 340 |
+
| --- | --- | --- |
|
| 341 |
+
| `pretrained_model_name_or_path` | — | Path to the pretrained MOVA base model weights |
|
| 342 |
+
| `video_loss_weight` / `audio_loss_weight` | `0.75` / `0.25` | Balance between the video and audio reconstruction losses |
|
| 343 |
+
| `visual_shift` / `audio_shift` | `9.0` / `5.0` | Per-tower noise schedule shift in the flow-matching formulation |
|
| 344 |
+
| `video_csv_file` | — | Path(s) to the bucketed CSV manifest produced by `tools/build_dataset_csv.py` |
|
| 345 |
+
| `video_size` | `[960,960]` / `[1440,1440]` / `[1920,1920]` | Spatial resolution for 720p / 1080p / 2K |
|
| 346 |
+
| `sample_n_frames` | `289` (720p/1080p), `265` (2K) | Maximum number of pixel-space frames per training clip |
|
| 347 |
+
| `video_micro_batch_size` | `1` | Per-GPU batch size within a data-parallel rank |
|
| 348 |
+
| `video_uncond_p` / `audio_uncond_p` | `0.1` / `0.2` | Classifier-free guidance dropout probability for video / audio captions |
|
| 349 |
+
| `latent_resolution` | `720p` / `1080p` / `2k` | Shorthand that auto-selects the spatial bucket base size |
|
| 350 |
+
| `video_bucket_temporal_min_length` | `49` | Shortest temporal bucket (frames) |
|
| 351 |
+
| `video_bucket_temporal_max_length` | `289` (720p/1080p), `265` (2K) | Longest temporal bucket (frames) |
|
| 352 |
+
| `video_bucket_temporal_interval` | `12` | Stride between successive temporal buckets |
|
| 353 |
+
|
| 354 |
+
|
| 355 |
+
### 🧱 VRAM Requirement
|
| 356 |
+
|
| 357 |
+
For native high-resolution training, 720p requires at least 32 × NVIDIA 80 GB GPUs (e.g. A100 or H100), and 1080p / 2K each require at least 64 × NVIDIA 80 GB GPUs. For inference, 720p can run on a single NVIDIA 80 GB GPU, while 1080p and 2K each require 4 or more NVIDIA 80 GB GPUs.
|
| 358 |
+
|
| 359 |
+
## ⭐ Contact
|
| 360 |
+
|
| 361 |
+
If you have any suggestions or find our work helpful, feel free to contact me.
|
| 362 |
|
| 363 |
Email: francisshuyuan@gmail.com
|
| 364 |
|
| 365 |
+
If you find our work useful, <b>please consider giving a star ⭐ to this github repository and citing it ❤️</b>:
|
| 366 |
```bib
|
| 367 |
+
@article{tu2026prism,
|
| 368 |
+
title={Prism:Dynamic Sparse Attention for Native 2K Joint Video-Audio Generation Model Training},
|
| 369 |
+
author={Tu, Shuyuan and Tian, Qi and Huang, Yinming and Wu, Yue and Han, Xintong and Pan, Kaihang and Kong, Weijie and Xiong, Jiangfeng and Zhang, Jian-Wei and Wu, Zuxuan and Jiang, Yu-Gang},
|
| 370 |
+
journal={arXiv preprint arXiv:},
|
| 371 |
+
year={2026}
|
| 372 |
}
|
| 373 |
```
|