Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,16 @@ 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/I2V_exp.png filter=lfs diff=lfs merge=lfs -text
+assets/T2V_exp.png filter=lfs diff=lfs merge=lfs -text
+assets/efficiency.png filter=lfs diff=lfs merge=lfs -text
+assets/method.png filter=lfs diff=lfs merge=lfs -text
+assets/teaser.jpg filter=lfs diff=lfs merge=lfs -text
+assets/videos/more/109_seed_677347.jpg filter=lfs diff=lfs merge=lfs -text
+assets/videos/more/14_seed_876367.jpg filter=lfs diff=lfs merge=lfs -text
+temp_data/ref_imgs/a[[:space:]]brown[[:space:]]bear[[:space:]]in[[:space:]]the[[:space:]]water[[:space:]]with[[:space:]]a[[:space:]]fish[[:space:]]in[[:space:]]its[[:space:]]mouth.jpg filter=lfs diff=lfs merge=lfs -text
+temp_data/ref_imgs/a[[:space:]]close-up[[:space:]]of[[:space:]]a[[:space:]]hippopotamus[[:space:]]eating[[:space:]]grass[[:space:]]in[[:space:]]a[[:space:]]field.jpg filter=lfs diff=lfs merge=lfs -text
+temp_data/ref_imgs/a[[:space:]]sea[[:space:]]turtle[[:space:]]swimming[[:space:]]in[[:space:]]the[[:space:]]ocean[[:space:]]under[[:space:]]the[[:space:]]water.jpg filter=lfs diff=lfs merge=lfs -text
+temp_data/videos/0004e625d5bcb80130e1ea3d204e2488.mp4 filter=lfs diff=lfs merge=lfs -text
+temp_data/videos/00086ac488a1ef8833bb6b0c6714f617.mp4 filter=lfs diff=lfs merge=lfs -text
+temp_data/videos/0012a9d775ff5b5f9f8676a5970691fa.mp4 filter=lfs diff=lfs merge=lfs -text

LICENSE.txt

@@ -0,0 +1,62 @@
+Tencent is pleased to support the open source community by making prfl available.
+
+Copyright (C) 2026 Tencent. All rights reserved. 
+
+prfl is licensed under Apache-2.0. prfl does not impose any additional restrictions beyond those specified in license.
+
+Terms of the Apache-2.0 License:
+--------------------------------------------------------------------
+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:
+
+You must give any other recipients of the Work or Derivative Works a copy of this License; and
+You must cause any modified files to carry prominent notices stating that You changed the files; and
+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
+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
+
+
+

README.md

@@ -0,0 +1,295 @@
+[中文文档](./README_CN.md)
+
+# HY-Video-PRFL
+
+<div align="center">
+  <img src="assets/logo.svg"  height=100>
+
+# ⚡ HY-Video-PRFL: Video Generation Models Are Good Latent Reward Models
+
+</div>
+
+Video generation models can both create and evaluate — we enable 14B models to complete full 720P×81-frame post-training within 67GB VRAM, achieving 1.5× faster speed and 56% improvement in motion quality over traditional methods.
+
+
+<div align="center">
+  <a href="https://github.com/Tencent-Hunyuan/HY-Video-PRFL"><img src="https://img.shields.io/static/v1?label=HY-Video-PRFL%20Code&message=Github&color=blue"></a> &ensp;
+  <a href="https://hy-video-prfl.github.io/HY-VIDEO-PRFL/"><img src="https://img.shields.io/static/v1?label=Project%20Page&message=Web&color=green"></a> &ensp;
+  <a href="https://arxiv.org/pdf/2511.21541"><img src="https://img.shields.io/badge/ArXiv-2511.21541-red"></a> &ensp;
+</div>
+
+<br>
+
+![image](assets/teaser.jpg)
+
+> [**HY-Video-PRFL: Video Generation Models Are Good Latent Reward Models**](https://arxiv.org/pdf/2511.21541) <be>
+
+## 🔥🔥🔥 News!!
+
+* **Dec 07, 2025**: 👋 We release the training and inference code of HY-Video-PRFL.
+* **Nov 26, 2025**: 👋 We release the paper and project page. [[Paper](https://arxiv.org/pdf/2511.21541)] [[Project Page](https://hy-video-prfl.github.io/HY-VIDEO-PRFL/)]
+## 📑 Open-source Plan
+
+- HY-Video-PRFL
+  - [x] Training and inference code for PAVRM
+  - [x] Training and inference code for PRFL
+  
+## 📋 Table of Contents
+
+- [🔥🔥🔥 News!!](#-news)
+- [📑 Open-source Plan](#-open-source-plan)
+- [📖 Abstract](#-abstract)
+- [🏗️ Model Architecture](#-model-architecture)
+- [📊 Performance](#-performance)
+- [🎬 Case Show](#-case-show)
+- [📜 Requirements](#-requirements)
+- [🛠️ Installation](#-installation)
+- [🧱 Download Models](#-download-models)
+- [🎓 Training](#-training)
+- [🚀 Inference](#-inference)
+- [📝 Citation](#-citation)
+- [🙏 Acknowledgements](#-acknowledgements)
+
+
+
+## 📖 Abstract
+
+Reward feedback learning (ReFL) has proven effective for aligning image generation with human preferences. However, its extension to video generation faces significant challenges. Existing video reward models rely on vision-language models designed for pixel-space inputs, confining ReFL optimization to near-complete denoising steps after computationally expensive VAE decoding.
+
+**HY-Video-PRFL** introduces **Process Reward Feedback Learning (PRFL)**, a framework that conducts preference optimization entirely in latent space. We demonstrate that pre-trained video generation models are naturally suited for reward modeling in the noisy latent space, enabling efficient gradient backpropagation throughout the full denoising chain without VAE decoding.
+
+**Key advantages:**
+- ✅ Efficient latent-space optimization
+- ✅ Significant memory savings
+- ✅ 1.4X faster training compared to RGB ReFL
+- ✅ Better alignment with human preferences
+
+
+## 🏗️ Model Architecture
+
+![image](assets/method.png)
+
+**Traditional RGB ReFL** relies on vision-language models designed for pixel-space inputs, requiring expensive VAE decoding and confining optimization to late-stage denoising steps.
+
+**Our PRFL approach** leverages pre-trained video generation models as reward models in the noisy latent space. This enables:
+- Full-chain gradient backpropagation without VAE decoding
+- Early-stage supervision for motion dynamics and structure coherence
+- Substantial reductions in memory consumption and training time
+
+
+
+## 📊 Performance
+
+### Quantitative Results
+
+Our experiments demonstrate that PRFL achieves substantial motion quality improvements (with +56.00 in dynamic degree, +21.52 in human anatomy and superior alignment with human preferences) as well as significant efficiency gains (with at least 1.4X faster training and notable memory savings).
+
+#### Text-to-Video Results
+![image](assets/T2V_exp.png)
+
+#### Image-to-Video Results
+![image](assets/I2V_exp.png)
+
+#### Efficiency Comparison
+<img src="assets/efficiency.png" width="50%">
+
+## 🎬 Case Show
+
+### Text-to-Video Generation
+
+|480P Resolution|720P Resolution|
+|---|---|
+|<video src="https://github.com/user-attachments/assets/eed8d875-4b0d-43ec-b013-f0d10c2e107a" width="600" controls autoplay loop></video> <details><summary>📋 Show prompt</summary>```Two shirtless men with short dark hair are sparring in a dimly lit room. They are both wearing boxing gloves, one red and one black. One man is wearing white shorts while the other is wearing black shorts. There are several screens on the wall displaying images of buildings and people.```</details>|<video src="https://github.com/user-attachments/assets/4871a9f9-9b15-4065-8680-1c8059242707" width="600" controls autoplay loop></video> <details><summary>📋 Show prompt</summary>```A woman with fair skin, dark hair tied back, and wearing a light green t-shirt is visible against a gray background. She uses both hands to apply a white substance from below her eyes upward onto her face. Her mouth is slightly open as she spreads the cream.```</details>|
+|<video src="https://github.com/user-attachments/assets/27296444-973b-4815-b103-5a2ee06404db" width="600" controls autoplay loop></video> <details><summary>📋 Show prompt</summary>```The woman has dark eyes and is holding a black smartphone to her ear with her right hand. She is typing on the keyboard of an open silver laptop computer with her left hand. Her fingers have blue nail polish. She is sitting in front of a window covered by sheer white curtains.```</details>|<video src="https://github.com/user-attachments/assets/430ff1ca-63ae-4a67-b6fb-b010c7ceec29" width="600" controls autoplay loop></video> <details><summary>📋 Show prompt</summary>```A light-skinned man with short hair wearing a yellow baseball cap, plaid shirt, and blue overalls stands in a field of sunflowers. He holds a cut sunflower head in his left hand and touches it with his right index finger. Several other sunflowers are visible in the background, some facing away from the camera.```</details>|
+
+### Image-to-Video Generation
+
+|480P Resolution|720P Resolution|
+|---|---|
+|<img src="assets/videos/more/14_seed_876367.jpg" width="600">|<img src="assets/videos/more/109_seed_677347.jpg" width="600">|
+|<video src="https://github.com/user-attachments/assets/956f5f64-1680-45a0-8666-9fda8e253017" width="600" controls autoplay loop></video> <details><summary>📋 Show prompt</summary>```A monochromatic video capturing a cat's gaze into the camera```</details>|<video src="https://github.com/user-attachments/assets/8dfeb1ae-8b9c-45aa-899a-3b48903629f9" width="600" controls autoplay loop></video> <details><summary>📋 Show prompt</summary>```A young boy is jumping in the mud```</details>|
+|<img src="assets/videos/more/artlist_video_60ca5873a5c21ff4e4785b1997239d87_seed_561368.jpg" width="600">|<img src="assets/videos/more/real_1246_seed_277973.jpg" width="600">|
+<video src="https://github.com/user-attachments/assets/d6d48d0c-cca5-4bdc-95c7-6a2858679111" width="600" controls autoplay loop></video> <details><summary>📋 Show prompt</summary>```A family of four eats fast food at a table.```</details>| <video src="https://github.com/user-attachments/assets/f39d3382-5430-412d-999e-c4c108835b6c" width="600" controls autoplay loop></video> <details><summary>📋 Show prompt</summary>```Normal speed, Medium shot, Eye level angle, Third person viewpoint, Static camera movement, Frame-within-frame composition, Shallow depth of field, Natural light, Cinematic style, Desaturated palette with slate blue, dusty rose, and dark wood tones color palette, Dramatic atmosphere. The scene is set on a patio or veranda, framed by a stone archway. In the back, there is a large, weathered wooden gate set into a stone wall. Six people are gathered on a stone patio in front of a large wooden gate. On the right, two men are seated at a dark wooden table. An older man in a grey traditional jacket holds a cane and gestures with his right hand while speaking. A younger man in a light grey suit sits beside him, listening. On the left side of the frame, a man in a dark suit stands with his back to the camera. Next to him, a woman in a pink patterned cheongsam and a woman in a grey skirt suit are standing close together, whispering. The women then turn and smile towards the men at the table. The man in the dark suit turns to face the group, revealing a newborn baby cradled in his arms, wrapped in a pink blanket. He takes a few steps forward, holding the baby. The women look at him and the infant. The older man at the table continues to talk, now gesturing towards the man with the baby. The man holding the baby looks down at the infant as he continues to walk slowly. The table is set with white cups, plates, fruit, and a dark wooden box.```</details>|
+
+
+
+## 📜 Requirements
+
+### Hardware Requirements
+
+We recommend using  GPUs with at least 80GB of memory for better generation quality.
+
+### Software Requirements
+
+* **OS**: Linux
+* **CUDA**: 12.4
+
+## 🛠️ Installation
+
+### Step 1: Clone Repository
+```bash
+git clone https://github.com/Tencent-Hunyuan/HY-Video-PRFL.git
+cd HY-Video-PRFL
+```
+
+### Step 2: Setup Environment
+
+We recommend CUDA versions 12.4 for installation. Conda's installation instructions are available [here](https://www.anaconda.com/docs/main).
+```bash
+# Create conda environment
+conda create -n HY-Video-PRFL python==3.10
+
+# Activate environment
+conda activate HY-Video-PRFL
+
+# Install PyTorch and dependencies (CUDA 12.4)
+pip3 install torch==2.5.0 torchvision==0.20.0 torchaudio==2.5.0 --index-url https://download.pytorch.org/whl/cu121
+
+# Install additional dependencies
+pip3 install git+https://github.com/huggingface/transformers qwen-vl-utils[decord]
+pip3 install git+https://github.com/huggingface/diffusers
+pip3 install xfuser -i https://pypi.org/simple
+pip3 install flash-attn==2.5.0 --no-build-isolation
+pip3 install -e .
+pip3 install nvidia-cublas-cu12==12.4.5.8
+
+export PYTHONPATH=./
+```
+
+## 🧱 Download Models
+
+Download the pretrained models before training or inference:
+
+| Model | Resolution | Download Links | Notes |
+|-------|-----------|----------------|-------|
+| **Wan2.1-T2V-14B** | 480P & 720P | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B) <br> 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B) | Text-to-Video model |
+| **Wan2.1-I2V-14B-720P** | 720P | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P) <br> 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P) | Image-to-Video (High-res) |
+| **Wan2.1-I2V-14B-480P** | 480P | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P) <br> 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P) | Image-to-Video (Standard) |
+
+First, make sure you have installed the huggingface CLI or modelscope CLI.
+```
+pip install -U "huggingface_hub[cli]"
+pip install modelscope
+```
+Then, download the pretrained DiT and VAE checkpoints. For example, you can use the following command to download the WAN2.1 checkpoint of 720P I2V task to ```./weights``` by default.
+```
+hf download Wan-AI/Wan2.1-I2V-14B-720P --local-dir ./weights
+```
+
+## 🎓 Training
+
+### 1️⃣ Data Preprocess on single GPU
+```bash
+python3 scripts/preprocess/gen_wanx_latent.py --config configs/pre_480.yaml
+```
+
+We provide several videos in ```temp_data/videos``` as template training data and an input json file ```temp_data/temp_input_data.json```template for preprocess. ```configs/pre_480.yaml``` is for 480P latent extraction and ```configs/pre_720.yaml``` is for 720P. The ```json_path``` and ```save_dir``` in config file can be customized with your own training data.
+
+### 2️⃣ Data Annotation and Format Conversion
+
+The annotation for reward model (e.g. ```"physics_quality": 1, "human_quality": 1```) should be added in the data meta files (e.g. ```temp_data/480/meta_v1/0004e625d5bcb80130e1ea3d204e2488_meta_v1.json```). Thus we get meta file list ```temp_data/temp_data_480.list``` and ```temp_data/temp_data_720.list``` which can be used in PAVRM and PRFL training.
+
+### 3️⃣ Parallel PAVRM Training on Multiple GPUs
+
+For example, to train PAVRM with 8 GPUs, you can use the following command.
+
+```bash
+torchrun --nnodes=1 --nproc_per_node=8 --master_port 29500 scripts/pavrm/train_pavrm.py --config configs/train_pavrm_i2v_720.yaml
+```
+
+The ```meta_file_list``` and ```val_meta_file_list``` in config file can be customized with your own training and validation data. We provide several config files for different settings t2v or i2v, 480P or 720P. To be noted that, we train PAVRM with ce loss. To train PAVRM with bt loss, you can use the config file of ```configs/train_pavrm_bt_i2v_720.yaml```.
+
+### 4️⃣ Parallel PRFL Training on Multiple GPUs
+
+```bash
+torchrun --nnodes=1 --nproc_per_node=8 --master_port 29500 scripts/prfl/train_prfl.py --config configs/train_prfl_i2v_720.yaml
+```
+
+The ```meta_file_list``` in config file can be customized with your own training data, ```lrm_transformer_path```, ```lrm_mlp_path``` and ```lrm_query_attention_path``` in config file are for your reward model obtained from the previous step. We provide several config files for different settings t2v or i2v, 480P or 720P.
+
+
+## 🚀 Inference
+
+### 1️⃣ Parallel PAVRM Inference on Multiple GPUs
+
+```bash
+torchrun --nnodes=1 --nproc_per_node=8 --master_port 29500 scripts/pavrm/inference_pavrm.py --config configs/infer_pavrm_i2v_720.yaml
+```
+
+The ```val_meta_file_list``` in config file can be customized with your own inference data, ```resume_transformer_path```, ```resume_mlp_path``` and ```resume_query_attention_path``` in config file are for your reward model to be tested.
+
+### 2️⃣ Parallel PRFL Inference on Multiple GPUs
+
+The PRFL Inference is exactly same as its base model (e.g. Wan2.1).
+
+```bash
+export negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+torchrun --nnodes=1 --nproc_per_node=8 --master_port 29500 scripts/prfl/inference_prfl.py \
+    --dit_fsdp \
+    --t5_fsdp \
+    --ulysses_size 1 \
+    --task "i2v-14B"\
+    --ckpt_dir "weights/Wan2.1-I2V-14B-720P" \
+    --lora_path "" \
+    --lora_alpha 0 \
+    --dataset_path "temp_data/temp_prfl_infer_data.json" \
+    --negative_prompt "$negative_prompt" \
+    --size "1280*720" \
+    --frame_num 81 \
+    --sample_steps 40 \
+    --sample_guide_scale 5.0 \
+    --sample_shift 5.0 \
+    --teacache_thresh 0 \
+    --save_folder outputs/infer/prfl_i2v_720 \
+    --transformer_path <YOUR_CKPT_PATH> \
+    --offload_model False
+```
+
+**Parameters:**
+- `--dit_fsdp` `--t5_fsdp`: Enable FSDP for memory efficiency
+- `--task`: "t2v-14B" or "i2v-14B"
+- `--ckpt_dir`: Path to pretrained checkpoint file
+- `--lora_path` `--lora_alpha`: Path and load weight ratio for LoRA checkpoint file
+- `--dataset_path`: Path to inference dataset file
+- `--size`: Output resolution ("1280\*720" or "832\*480")
+- `--frame_num`: Number of frames to generate (default: 81)
+- `--sample_steps`: Number of inference steps (default: 40)
+- `--sample_guide_scale`: Classifier-free guidance scale (default: 5.0)
+- `--sample_shift`: Flow shift (default: 5.0)
+- `--save_folder`: Path to save generated videos
+- `--teacache_thresh`: Enable teacache
+- `--transformer_path`: Path to your PRFL checkpoint file
+- `--offload_model`: Offload to CPU to save GPU memory
+
+## 📝 Citation
+
+If you find **HY-Video-PRFL** useful for your research, please cite:
+```bibtex
+@article{mi2025video,
+  title={Video Generation Models are Good Latent Reward Models},
+  author={Mi, Xiaoyue and Yu, Wenqing and Lian, Jiesong and Jie, Shibo and Zhong, Ruizhe and Liu, Zijun and Zhang, Guozhen and Zhou, Zixiang and Xu, Zhiyong and Zhou, Yuan and Lu, Qinglin and Tang, Fan},
+  journal={arXiv preprint arXiv:2511.21541},
+  year={2025}
+}
+```
+
+
+## 🙏 Acknowledgements
+
+We sincerely thank the contributors to the following projects:
+- [HunyuanVideo](https://github.com/Tencent/HunyuanVideo)
+- [Wan2.1](https://github.com/Wan-Video/Wan2.1)
+- [ImageReward](https://github.com/THUDM/ImageReward)
+- [Diffusers](https://github.com/huggingface/diffusers)
+- [HuggingFace](https://huggingface.co)
+- [DeepSpeed](https://github.com/deepspeedai/DeepSpeed)
+
+
+
+---
+
+<div align="center">
+  
+**Star ⭐ this repo if you find it helpful!**
+
+</div>

README_CN.md

@@ -0,0 +1,284 @@
+# HY-Video-PRFL
+
+<div align="center">
+  <img src="assets/logo.svg"  height=100>
+
+# ⚡ HY-Video-PRFL: 视频生成模型是优秀的潜在奖励模型
+
+</div>
+
+视频生成模型既能创造也能评估——我们使14B模型能够在67GB显存内完成完整的720P×81帧后训练，相比传统方法实现1.5倍的速度提升和56%的运动质量改进。
+
+<div align="center">
+  <a href="https://github.com/Tencent-Hunyuan/HY-Video-PRFL"><img src="https://img.shields.io/static/v1?label=HY-Video-PRFL%20Code&message=Github&color=blue"></a> &ensp;
+  <a href="https://hy-video-prfl.github.io/HY-VIDEO-PRFL/"><img src="https://img.shields.io/static/v1?label=Project%20Page&message=Web&color=green"></a> &ensp;
+  <a href="https://arxiv.org/pdf/2511.21541"><img src="https://img.shields.io/badge/ArXiv-2511.21541-red"></a> &ensp;
+</div>
+
+<br>
+
+![image](assets/teaser.jpg)
+
+> [**HY-Video-PRFL: 视频生成模型是优秀的潜在奖励模型**](https://arxiv.org/pdf/2511.21541)
+
+## 🔥🔥🔥 最新动态！
+
+* **2025年12月7日**: 👋 我们发布了HY-Video-PRFL的训练和推理代码。
+* **2025年11月26日**: 👋 我们发布了论文和项目主页。[[论文](https://arxiv.org/pdf/2511.21541)] [[项目主页](https://hy-video-prfl.github.io/HY-VIDEO-PRFL/)]
+
+## 📑 开源计划
+
+- HY-Video-PRFL
+  - [x] PAVRM的训练和推理代码
+  - [x] PRFL的训练和推理代码
+  
+## 📋 目录
+
+- [🔥🔥🔥 最新动态！](#-最新动态)
+- [📑 开源计划](#-开源计划)
+- [📖 摘要](#-摘要)
+- [🏗️ 模型架构](#️-模型架构)
+- [📊 性能表现](#-性能表现)
+- [🎬 案例展示](#-案例展示)
+- [📜 环境要求](#-环境要求)
+- [🛠️ 安装](#️-安装)
+- [🧱 下载模型](#-下载模型)
+- [🎓 训练](#-训练)
+- [🚀 推理](#-推理)
+- [📝 引用](#-引用)
+- [🙏 致谢](#-致谢)
+
+## 📖 摘要
+
+奖励反馈学习（ReFL）已被证明能有效地使图像生成与人类偏好保持一致。然而，将其扩展到视频生成面临着重大挑战。现有的视频奖励模型依赖于为像素空间输入设计的视觉-语言模型，在计算成本高昂的VAE解码后，将ReFL优化限制在接近完成的去噪步骤。
+
+**HY-Video-PRFL** 引入了**过程奖励反馈学习（PRFL）**，这是一个完全在潜在空间中进行偏好优化的框架。我们证明了预训练的视频生成模型天然适合在噪声潜在空间中进行奖励建模，使得能够在整个去噪链中进行高效的梯度反向传播，而无需VAE解码。
+
+**核心优势:**
+- ✅ 高效的潜在空间优化
+- ✅ 显著的内存节省
+- ✅ 相比RGB ReFL快1.4倍的训练速度
+- ✅ 更好地与人类偏好保持一致
+
+## 🏗️ 模型架构
+
+![image](assets/method.png)
+
+**传统的RGB ReFL** 依赖于为像素空间输入设计的视觉-语言模型，需要昂贵的VAE解码，并将优化限制在后期去噪步骤。
+
+**我们的PRFL方法** 利用预训练的视频生成模型作为噪声潜在空间中的奖励模型。这实现了：
+- 无需VAE解码的全链梯度反向传播
+- 针对运动动态和结构一致性的早期监督
+- 大幅减少内存消耗和训练时间
+
+## 📊 性能表现
+
+### 定量结果
+
+我们的实验表明，PRFL在运动质量方面实现了显著改进（动态度提升+56.00，人体解剖结构提升+21.52，并且与人类偏好有更好的对齐），同时在效率方面也取得了显著提升（至少快1.4倍的训练速度和显著的内存节省）。
+
+#### 文本生成视频结果
+![image](assets/T2V_exp.png)
+
+#### 图像生成视频结果
+![image](assets/I2V_exp.png)
+
+#### 效率对比
+<img src="assets/efficiency.png" width="50%">
+
+## 🎬 案例展示
+
+### 文本生成视频
+
+|480P 分辨率|720P 分辨率|
+|---|---|
+|<video src="https://github.com/user-attachments/assets/eed8d875-4b0d-43ec-b013-f0d10c2e107a" width="600" controls autoplay loop></video> <details><summary>📋 展示提示词</summary>```Two shirtless men with short dark hair are sparring in a dimly lit room. They are both wearing boxing gloves, one red and one black. One man is wearing white shorts while the other is wearing black shorts. There are several screens on the wall displaying images of buildings and people.```</details>|<video src="https://github.com/user-attachments/assets/4871a9f9-9b15-4065-8680-1c8059242707" width="600" controls autoplay loop></video> <details><summary>📋 展示提示词</summary>```A woman with fair skin, dark hair tied back, and wearing a light green t-shirt is visible against a gray background. She uses both hands to apply a white substance from below her eyes upward onto her face. Her mouth is slightly open as she spreads the cream.```</details>|
+|<video src="https://github.com/user-attachments/assets/27296444-973b-4815-b103-5a2ee06404db" width="600" controls autoplay loop></video> <details><summary>📋 展示提示词</summary>```The woman has dark eyes and is holding a black smartphone to her ear with her right hand. She is typing on the keyboard of an open silver laptop computer with her left hand. Her fingers have blue nail polish. She is sitting in front of a window covered by sheer white curtains.```</details>|<video src="https://github.com/user-attachments/assets/430ff1ca-63ae-4a67-b6fb-b010c7ceec29" width="600" controls autoplay loop></video> <details><summary>📋 展示提示词</summary>```A light-skinned man with short hair wearing a yellow baseball cap, plaid shirt, and blue overalls stands in a field of sunflowers. He holds a cut sunflower head in his left hand and touches it with his right index finger. Several other sunflowers are visible in the background, some facing away from the camera.```</details>|
+
+### 图像生成视频
+
+|480P 分辨率|720P 分辨率|
+|---|---|
+|<img src="assets/videos/more/14_seed_876367.jpg" width="600">|<img src="assets/videos/more/109_seed_677347.jpg" width="600">|
+|<video src="https://github.com/user-attachments/assets/956f5f64-1680-45a0-8666-9fda8e253017" width="600" controls autoplay loop></video> <details><summary>📋 展示提示词</summary>```A monochromatic video capturing a cat's gaze into the camera```</details>|<video src="https://github.com/user-attachments/assets/8dfeb1ae-8b9c-45aa-899a-3b48903629f9" width="600" controls autoplay loop></video> <details><summary>📋 展示提示词</summary>```A young boy is jumping in the mud```</details>|
+|<img src="assets/videos/more/artlist_video_60ca5873a5c21ff4e4785b1997239d87_seed_561368.jpg" width="600">|<img src="assets/videos/more/real_1246_seed_277973.jpg" width="600">|
+<video src="https://github.com/user-attachments/assets/d6d48d0c-cca5-4bdc-95c7-6a2858679111" width="600" controls autoplay loop></video> <details><summary>📋 展示提示词</summary>```A family of four eats fast food at a table.```</details>| <video src="https://github.com/user-attachments/assets/f39d3382-5430-412d-999e-c4c108835b6c" width="600" controls autoplay loop></video> <details><summary>📋 展示提示词</summary>```Normal speed, Medium shot, Eye level angle, Third person viewpoint, Static camera movement, Frame-within-frame composition, Shallow depth of field, Natural light, Cinematic style, Desaturated palette with slate blue, dusty rose, and dark wood tones color palette, Dramatic atmosphere. The scene is set on a patio or veranda, framed by a stone archway. In the back, there is a large, weathered wooden gate set into a stone wall. Six people are gathered on a stone patio in front of a large wooden gate. On the right, two men are seated at a dark wooden table. An older man in a grey traditional jacket holds a cane and gestures with his right hand while speaking. A younger man in a light grey suit sits beside him, listening. On the left side of the frame, a man in a dark suit stands with his back to the camera. Next to him, a woman in a pink patterned cheongsam and a woman in a grey skirt suit are standing close together, whispering. The women then turn and smile towards the men at the table. The man in the dark suit turns to face the group, revealing a newborn baby cradled in his arms, wrapped in a pink blanket. He takes a few steps forward, holding the baby. The women look at him and the infant. The older man at the table continues to talk, now gesturing towards the man with the baby. The man holding the baby looks down at the infant as he continues to walk slowly. The table is set with white cups, plates, fruit, and a dark wooden box.```</details>|
+
+
+## 📜 环境要求
+
+### 硬件要求
+
+我们建议使用至少80GB显存的GPU以获得更好的生成质量。
+
+### 软件要求
+
+* **操作系统**: Linux
+* **CUDA**: 12.4
+
+## 🛠️ 安装
+
+### 步骤1: 克隆仓库
+```bash
+git clone https://github.com/Tencent-Hunyuan/HY-Video-PRFL.git
+cd HY-Video-PRFL
+```
+
+### 步骤2: 设置环境
+
+我们推荐使用CUDA 12.4版本进行安装。Conda的安装说明可在[这里](https://www.anaconda.com/docs/main)找到。
+
+```bash
+# 创建conda环境
+conda create -n HY-Video-PRFL python==3.10
+
+# 激活环境
+conda activate HY-Video-PRFL
+
+# 安装PyTorch和依赖项（CUDA 12.4）
+pip3 install torch==2.5.0 torchvision==0.20.0 torchaudio==2.5.0 --index-url https://download.pytorch.org/whl/cu121
+
+# 安装额外依赖项
+pip3 install git+https://github.com/huggingface/transformers qwen-vl-utils[decord]
+pip3 install git+https://github.com/huggingface/diffusers
+pip3 install xfuser -i https://pypi.org/simple
+pip3 install flash-attn==2.5.0 --no-build-isolation
+pip3 install -e .
+pip3 install nvidia-cublas-cu12==12.4.5.8
+
+export PYTHONPATH=./
+```
+
+## 🧱 下载模型
+
+在训练或推理前下载预训练模型：
+
+| 模型 | 分辨率 | 下载链接 | 说明 |
+|-------|-----------|----------------|-------|
+| **Wan2.1-T2V-14B** | 480P & 720P | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B) <br> 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B) | 文本生成视频模型 |
+| **Wan2.1-I2V-14B-720P** | 720P | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P) <br> 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P) | 图像生成视频（高分辨率） |
+| **Wan2.1-I2V-14B-480P** | 480P | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P) <br> 🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P) | 图像生成视频（标准） |
+
+首先，确保已安装huggingface CLI或modelscope CLI。
+```
+pip install -U "huggingface_hub[cli]"
+pip install modelscope
+```
+然后，下载预训练的DiT和VAE检查点。例如，可以使用以下命令将720P I2V任务的WAN2.1检查点下载到默认的```./weights```目录。
+```
+hf download Wan-AI/Wan2.1-I2V-14B-720P --local-dir ./weights
+```
+
+## 🎓 训练
+
+### 1️⃣ 单GPU数据预处理
+```bash
+python3 scripts/preprocess/gen_wanx_latent.py --config configs/pre_480.yaml
+```
+
+我们在```temp_data/videos```中提供了几个视频作为模板训练数据，以及用于预处理的输入json文件```temp_data/temp_input_data.json```模板。```configs/pre_480.yaml```用于480P潜在提取，```configs/pre_720.yaml```用于720P。配置文件中的```json_path```和```save_dir```可以根据自己的训练数据自定义。
+
+### 2️⃣ 数据标注和格式转换
+
+奖励模型的标注（例如```"physics_quality": 1, "human_quality": 1```）应添加到数据元文件中（例如```temp_data/480/meta_v1/0004e625d5bcb80130e1ea3d204e2488_meta_v1.json```）。这样我们得到元文件列表```temp_data/temp_data_480.list```和```temp_data/temp_data_720.list```，可用于PAVRM和PRFL训练。
+
+### 3️⃣ 多GPU并行PAVRM训练
+
+例如，要使用8个GPU训练PAVRM，可以使用以下命令。
+
+```bash
+torchrun --nnodes=1 --nproc_per_node=8 --master_port 29500 scripts/pavrm/train_pavrm.py --config configs/train_pavrm_i2v_720.yaml
+```
+
+配置文件中的```meta_file_list```和```val_meta_file_list```可以根据自己的训练和验证数据自定义。我们为不同设置（t2v或i2v，480P或720P）提供了几个配置文件。需要注意的是，我们使用ce损失训练PAVRM。要使用bt损失训练PAVRM，可以使用配置文件```configs/train_pavrm_bt_i2v_720.yaml```。
+
+### 4️⃣ 多GPU并行PRFL训练
+
+```bash
+torchrun --nnodes=1 --nproc_per_node=8 --master_port 29500 scripts/prfl/train_prfl.py --config configs/train_prfl_i2v_720.yaml
+```
+
+配置文件中的```meta_file_list```可以根据自己的训练数据自定义，配置文件中的```lrm_transformer_path```、```lrm_mlp_path```和```lrm_query_attention_path```用于从上一步获得的奖励模型。我们为不同设置（t2v或i2v，480P或720P）提供了几个配置文件。
+
+## 🚀 推理
+
+### 1️⃣ 多GPU并行PAVRM推理
+
+```bash
+torchrun --nnodes=1 --nproc_per_node=8 --master_port 29500 scripts/pavrm/inference_pavrm.py --config configs/infer_pavrm_i2v_720.yaml
+```
+
+配置文件中的```val_meta_file_list```可以根据自己的推理数据自定义，配置文件中的```resume_transformer_path```、```resume_mlp_path```和```resume_query_attention_path```用于待测试的奖励模型。
+
+### 2️⃣ 多GPU并行PRFL推理
+
+PRFL推理与其基础模型（例如Wan2.1）完全相同。
+
+```bash
+export negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+torchrun --nnodes=1 --nproc_per_node=8 --master_port 29500 scripts/prfl/inference_prfl.py \
+    --dit_fsdp \
+    --t5_fsdp \
+    --ulysses_size 1 \
+    --task "i2v-14B"\
+    --ckpt_dir "weights/Wan2.1-I2V-14B-720P" \
+    --lora_path "" \
+    --lora_alpha 0 \
+    --dataset_path "temp_data/temp_prfl_infer_data.json" \
+    --negative_prompt "$negative_prompt" \
+    --size "1280*720" \
+    --frame_num 81 \
+    --sample_steps 40 \
+    --sample_guide_scale 5.0 \
+    --sample_shift 5.0 \
+    --teacache_thresh 0 \
+    --save_folder outputs/infer/prfl_i2v_720 \
+    --transformer_path <YOUR_CKPT_PATH> \
+    --offload_model False
+```
+
+**参数说明:**
+- `--dit_fsdp` `--t5_fsdp`: 启用FSDP以提高内存效率
+- `--task`: "t2v-14B"或"i2v-14B"
+- `--ckpt_dir`: 预训练检查点文件路径
+- `--lora_path` `--lora_alpha`: LoRA检查点文件路径和加载权重比例
+- `--dataset_path`: 推理数据集文件路径
+- `--size`: 输出分辨率（"1280\*720"或"832\*480"）
+- `--frame_num`: 生成的帧数（默认：81）
+- `--sample_steps`: 推理步数（默认：40）
+- `--sample_guide_scale`: 无分类器引导比例（默认：5.0）
+- `--sample_shift`: 流偏移（默认：5.0）
+- `--save_folder`: 保存生成视频的路径
+- `--teacache_thresh`: 启用teacache
+- `--transformer_path`: PRFL检查点文件路径
+- `--offload_model`: 卸载到CPU以节省GPU内存
+
+## 📝 引用
+
+如果您发现**HY-Video-PRFL**对您的研究有用，请引用：
+```bibtex
+@article{mi2025video,
+  title={Video Generation Models are Good Latent Reward Models},
+  author={Mi, Xiaoyue and Yu, Wenqing and Lian, Jiesong and Jie, Shibo and Zhong, Ruizhe and Liu, Zijun and Zhang, Guozhen and Zhou, Zixiang and Xu, Zhiyong and Zhou, Yuan and Lu, Qinglin and Tang, Fan},
+  journal={arXiv preprint arXiv:2511.21541},
+  year={2025}
+}
+```
+
+## 🙏 致谢
+
+我们真诚感谢以下项目的贡献者：
+- [HunyuanVideo](https://github.com/Tencent/HunyuanVideo)
+- [Wan2.1](https://github.com/Wan-Video/Wan2.1)
+- [ImageReward](https://github.com/THUDM/ImageReward)
+- [Diffusers](https://github.com/huggingface/diffusers)
+- [HuggingFace](https://huggingface.co)
+- [DeepSpeed](https://github.com/deepspeedai/DeepSpeed)
+
+---
+
+<div align="center">
+  
+**如果您觉得有帮助，请给这个仓库加星 ⭐！**
+
+</div>

configs/infer_pavrm_i2v_720.yaml

@@ -0,0 +1,100 @@
+train_id: "pavrm_i2v_720"
+task: "i2v-14b-720p" # t2v-1.3b, i2v-1.3b, t2v-14b, i2v-14b-480p, i2v-14b-720p
+
+model:
+  base_path: "weights/Wan2.1-I2V-14B-720P"
+  init_transformer_path: null
+  resume_transformer_path: null
+  resume_mlp_path: null
+  resume_query_attention_path: null
+
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  val_meta_file_list:
+    - temp_data/temp_data_720.list
+
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1 #
+  sp_batch_size: 1 #4
+  num_workers: 0
+  group_frame: null
+  group_resolution: null
+
+optimizer:
+  learning_rate: 1e-6
+  learning_rate_mlp: 1e-5
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100
+  sanity_check_interval: 0
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 1
+
+eval:
+  seed: 0
+
+save:
+  output_dir: "outputs/infer"
+  sanity_check_dir: null
+
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  loss: "ce"
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/pre_480.yaml

@@ -0,0 +1,22 @@
+start_idx: 0 
+end_idx: null
+sample_n_frames: 81
+aspect_ratio: 1.73
+seed: 42
+precision: 'bf16'
+task: 'i2v'
+base_dir: 'weights/Wan2.1-I2V-14B-480P'
+resolution: [480]
+num_frames: 81
+batch_size: 1
+fps: 16
+json_path: temp_data/temp_input_data.json
+save_dir: temp_data/480
+extract_fps: 16
+model_type: 'wanx'
+vae_path: 'weights/Wan2.1-I2V-14B-480P/Wan2.1_VAE.pth'
+image_processor_path: 'weights/Wan2.1-I2V-14B-480P/xlm-roberta-large'
+image_encoder_path: 'weights/Wan2.1-I2V-14B-480P/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
+tokenizer_path: 'weights/Wan2.1-I2V-14B-480P/google/umt5-xxl'
+text_encoder_path: 'weights/Wan2.1-I2V-14B-480P/models_t5_umt5-xxl-enc-bf16.pth'
+max_sequence_length: 512

configs/pre_720.yaml

@@ -0,0 +1,22 @@
+start_idx: 0
+end_idx: null
+sample_n_frames: 81
+aspect_ratio: 1.81
+seed: 42
+precision: 'bf16'
+task: 'i2v'
+base_dir: 'weights/Wan2.1-I2V-14B-720P'
+resolution: [704]
+num_frames: 81
+batch_size: 1
+fps: 16
+json_path: temp_data/temp_input_data.json
+save_dir: temp_data/720
+extract_fps: 16
+model_type: 'wanx'
+vae_path: 'weights/Wan2.1-I2V-14B-720P/Wan2.1_VAE.pth'
+image_processor_path: 'weights/Wan2.1-I2V-14B-720P/xlm-roberta-large'
+image_encoder_path: 'weights/Wan2.1-I2V-14B-720P/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
+tokenizer_path: 'weights/Wan2.1-I2V-14B-720P/google/umt5-xxl'
+text_encoder_path: 'weights/Wan2.1-I2V-14B-720P/models_t5_umt5-xxl-enc-bf16.pth'
+max_sequence_length: 512

configs/train_pavrm_bt_i2v_720.yaml

@@ -0,0 +1,103 @@
+train_id: "pavrm_bt_i2v_720"
+task: "i2v-14b-720p" # t2v-1.3b, i2v-1.3b, t2v-14b, i2v-14b-480p, i2v-14b-720p
+
+model:
+  base_path: "weights/Wan2.1-I2V-14B-720P"
+  init_transformer_path: null
+  resume_transformer_path: null
+  resume_mlp_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_720.list
+  meta_file_lose_list:
+    - temp_data/temp_data_720.list
+  val_meta_file_list:
+    - temp_data/temp_data_720.list
+  
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1 #
+  sp_batch_size: 1 #4
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+
+optimizer:
+  learning_rate: 1e-6
+  learning_rate_mlp: 1e-5
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100 #100
+  sanity_check_interval: 0
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 1
+
+eval:
+  seed: 0
+  timestep: [201, 400, 600, 800, 1000]
+
+save:
+  output_dir: "outputs"
+  sanity_check_dir: null
+
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  loss: "bt"
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/train_pavrm_i2v_480.yaml

@@ -0,0 +1,102 @@
+train_id: "pavrm_i2v_480"
+task: "i2v-14b-480p" # t2v-1.3b, i2v-1.3b, t2v-14b, i2v-14b-480p, i2v-14b-720p
+
+model:
+  base_path: "weights/Wan2.1-I2V-14B-480P"
+  init_transformer_path: null
+  resume_transformer_path: null
+  resume_mlp_path: null
+  resume_query_attention_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_480.list
+  val_meta_file_list:
+    - temp_data/temp_data_480.list
+
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1 #
+  sp_batch_size: 1 #4
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+
+optimizer:
+  learning_rate: 1e-6
+  # learning_rate_mlp: 1e-5
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100
+  sanity_check_interval: 0
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 1
+
+eval:
+  seed: 0
+  timestep: [201, 400, 600, 800, 1000]
+
+save:
+  output_dir: "outputs"
+  sanity_check_dir: null
+
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  loss: "ce"
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/train_pavrm_i2v_720.yaml

@@ -0,0 +1,102 @@
+train_id: "pavrm_i2v_720"
+task: "i2v-14b-720p" # t2v-1.3b, i2v-1.3b, t2v-14b, i2v-14b-480p, i2v-14b-720p
+
+model:
+  base_path: "weights/Wan2.1-I2V-14B-720P"
+  init_transformer_path: null
+  resume_transformer_path: null
+  resume_mlp_path: null
+  resume_query_attention_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_720.list
+  val_meta_file_list:
+    - temp_data/temp_data_720.list
+
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1
+  sp_batch_size: 1
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+
+optimizer:
+  learning_rate: 1e-6
+  # learning_rate_mlp: 1e-5
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100
+  sanity_check_interval: 0
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 1
+
+eval:
+  seed: 0
+  timestep: [201, 400, 600, 800, 1000]
+
+save:
+  output_dir: "outputs"
+  sanity_check_dir: null
+
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  loss: "ce"
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/train_pavrm_t2v_480.yaml

@@ -0,0 +1,102 @@
+train_id: "pavrm_t2v_480"
+task: "t2v-14b" # t2v-1.3b, i2v-1.3b, t2v-14b, i2v-14b-480p, i2v-14b-720p
+
+model:
+  base_path: "weights/Wan2.1-T2V-14B"
+  init_transformer_path: null
+  resume_transformer_path: null
+  resume_mlp_path: null
+  resume_query_attention_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_480.list
+  val_meta_file_list:
+    - temp_data/temp_data_480.list
+
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1 #
+  sp_batch_size: 1 #4
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+
+optimizer:
+  learning_rate: 1e-6
+  # learning_rate_mlp: 1e-5
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100 #100
+  sanity_check_interval: 0
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 1
+
+eval:
+  seed: 0
+  timestep: [201, 400, 600, 800, 1000]
+
+save:
+  output_dir: "outputs"
+  sanity_check_dir: null
+
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  loss: "ce"
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/train_pavrm_t2v_720.yaml

@@ -0,0 +1,102 @@
+train_id: "pavrm_t2v_720"
+task: "t2v-14b" # t2v-1.3b, i2v-1.3b, t2v-14b, i2v-14b-480p, i2v-14b-720p
+
+model:
+  base_path: "weights/Wan2.1-T2V-14B"
+  init_transformer_path: null
+  resume_transformer_path: null
+  resume_mlp_path: null
+  resume_query_attention_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_720.list
+  val_meta_file_list:
+    - temp_data/temp_data_720.list
+
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1 #
+  sp_batch_size: 1 #4
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+
+optimizer:
+  learning_rate: 1e-6
+  # learning_rate_mlp: 1e-5
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100 #100
+  sanity_check_interval: 0
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 1
+
+eval:
+  seed: 0
+  timestep: [201, 400, 600, 800, 1000]
+
+save:
+  output_dir: "outputs"
+  sanity_check_dir: null
+
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  loss: "ce"
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/train_prfl_i2v_480.yaml

@@ -0,0 +1,95 @@
+train_id: "prfl_i2v_480"
+task: "i2v-14b-480p"
+
+model:
+  base_path: weights/Wan2.1-I2V-14B-480P
+  init_transformer_path: null
+  lrm_transformer_path: null
+  lrm_mlp_path: null
+  lrm_query_attention_path: null
+  resume_transformer_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_720.list
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1
+  sp_batch_size: 1
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+  # negative_prompt:
+
+optimizer:
+  learning_rate: 5e-6
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100 #100
+  sanity_check_interval: 100
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 5.
+
+save:
+  output_dir: "outputs"
+  log_dir: null
+  sanity_check_dir: null
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/train_prfl_i2v_720.yaml

@@ -0,0 +1,96 @@
+train_id: "prfl_i2v_720"
+# task: "i2v-14b-720p" # t2v-1.3b, i2v-1.3b, t2v-14b, i2v-14b-480p, i2v-14b-720p
+task: "i2v-14b-720p"
+
+model:
+  base_path: weights/Wan2.1-I2V-14B-720P
+  init_transformer_path: null
+  lrm_transformer_path: null
+  lrm_mlp_path: null
+  lrm_query_attention_path: null
+  resume_transformer_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_720.list
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1
+  sp_batch_size: 1
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+  # negative_prompt:
+
+optimizer:
+  learning_rate: 5e-6
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100
+  sanity_check_interval: 100
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 5.
+
+save:
+  output_dir: "outputs"
+  log_dir: null
+  sanity_check_dir: null
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/train_prfl_t2v_480.yaml

@@ -0,0 +1,95 @@
+train_id: "prfl_t2v_480"
+task: "t2v-14b"
+
+model:
+  base_path: weights/Wan2.1-T2V-14B
+  init_transformer_path: null
+  lrm_transformer_path: null
+  lrm_mlp_path: null
+  lrm_query_attention_path: null
+  resume_transformer_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_480.list
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1
+  sp_batch_size: 1
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+  # negative_prompt:
+
+optimizer:
+  learning_rate: 5e-6
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100 #100
+  sanity_check_interval: 100
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 5.
+
+save:
+  output_dir: "outputs"
+  log_dir: null
+  sanity_check_dir: null
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

configs/train_prfl_t2v_720.yaml

@@ -0,0 +1,96 @@
+train_id: "prfl_t2v_720"
+task: "t2v-14b"
+
+model:
+  base_path: weights/Wan2.1-T2V-14B
+  init_transformer_path: null
+  lrm_transformer_path: null
+  lrm_mlp_path: null
+  lrm_query_attention_path: null
+  resume_transformer_path: null
+  patch_size: [1, 2, 2]
+  lora:
+    use_lora: false
+    lora_rank: 128
+    target_modules: ["q", "k", "v", "o"]
+    resume_lora_path: null  # load lora ckpt if not empty
+  ema:
+    use_ema: false
+    ema_decay: 0.99
+  fsdp:
+    fsdp_sharding_startegy: full
+    use_cpu_offload: false
+  gradient_checkpointing: true
+  selective_checkpointing: 1.0
+
+extra_model:
+  vae:
+    name: Wan2.1_VAE.pth
+    vae_stride: [4, 8, 8]
+  text_encoder:
+    t5_text_len: 512
+    t5_checkpoint: models_t5_umt5-xxl-enc-bf16.pth
+    t5_tokenizer: google/umt5-xxl
+  image_encoder:
+    clip_checkpoint: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+    clip_tokenizer: xlm-roberta-large
+  scheduler:
+    flow_shift: 5.0
+    num_train_timesteps: 1000
+    weighting_scheme: uniform  # logit_normal, uniform
+    logit_mean: 0
+    logit_std: 1
+    mode_scale: 1.29
+  
+dataset:
+  meta_file_list:
+    - temp_data/temp_data_720.list
+  crop_ratio: [1, 1, 1]  # width, height
+  crop_type: "random" # center, random
+  uncond_prob: [0.1, 0.0] # prompt, image
+  sp_size: 4
+  batch_size: 1
+  sp_batch_size: 1
+  num_workers: 8
+  group_frame: null
+  group_resolution: null
+  # negative_prompt:
+
+optimizer:
+  learning_rate: 5e-6
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  weight_decay: 0.01
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+  lr_num_cycles: 1
+  lr_power: 1.0
+  max_train_steps: 1000000
+
+train:
+  seed: 110221
+  precision: bf16
+  extra_precision: bf16
+  allow_tf32: false
+  save_interval: 100
+  sanity_check_interval: 100
+  teacher_student_parallel: true
+  dpo_beta: 500
+  gradient_accumulation_steps: 5.
+
+save:
+  output_dir: "outputs"
+  log_dir: null
+  sanity_check_dir: null
+
+lrm:
+  query_attention:
+    num_queries: 1
+    num_heads: 8
+    dropout: 0.
+    return_type: query
+  feature_layer: [8]
+  pool: q_attn
+  mlp_dim: 5120
+  task: "motion_quality"
+  trainable_blocks: [0, 1, 2, 3, 4, 5, 6, 7]

diffusers_lite.egg-info/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: diffusers_lite
+Version: 0.0.1
+Author: jaspercheng
+Requires-Dist: torch>=2.4.0
+Requires-Dist: torchvision>=0.19.0
+Requires-Dist: opencv-python>=4.9.0.80
+Requires-Dist: diffusers>=0.31.0
+Requires-Dist: transformers>=4.49.0
+Requires-Dist: tokenizers>=0.20.3
+Requires-Dist: accelerate>=1.1.1
+Requires-Dist: gradio>=5.0.0
+Requires-Dist: tqdm
+Requires-Dist: imageio
+Requires-Dist: easydict
+Requires-Dist: ftfy
+Requires-Dist: dashscope
+Requires-Dist: peft
+Requires-Dist: imageio-ffmpeg
+Requires-Dist: flash_attn
+Requires-Dist: numpy
+Requires-Dist: omegaconf
+Requires-Dist: protobuf
+Requires-Dist: matplotlib
+Dynamic: author
+Dynamic: requires-dist

diffusers_lite.egg-info/SOURCES.txt

@@ -0,0 +1,7 @@
+README.md
+setup.py
+diffusers_lite.egg-info/PKG-INFO
+diffusers_lite.egg-info/SOURCES.txt
+diffusers_lite.egg-info/dependency_links.txt
+diffusers_lite.egg-info/requires.txt
+diffusers_lite.egg-info/top_level.txt

diffusers_lite.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

diffusers_lite.egg-info/requires.txt

@@ -0,0 +1,20 @@
+torch>=2.4.0
+torchvision>=0.19.0
+opencv-python>=4.9.0.80
+diffusers>=0.31.0
+transformers>=4.49.0
+tokenizers>=0.20.3
+accelerate>=1.1.1
+gradio>=5.0.0
+tqdm
+imageio
+easydict
+ftfy
+dashscope
+peft
+imageio-ffmpeg
+flash_attn
+numpy
+omegaconf
+protobuf
+matplotlib

diffusers_lite.egg-info/top_level.txt

@@ -0,0 +1 @@
+

diffusers_lite/arguments.py

@@ -0,0 +1,216 @@
+import argparse
+import random
+import sys
+
+from .wan.configs import WAN_CONFIGS
+from .wan.utils.utils import str2bool
+
+
+def args_init():
+    parser = argparse.ArgumentParser(description="diffusers lite script")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--precision", type=str, default="bf16")
+    parser.add_argument("--task", type=str, default="i2v", choices=["i2v", "t2v", "flf2v"])
+
+    ### Inference ###
+    # model
+    parser.add_argument("--base-dir", type=str, default="")
+    parser.add_argument("--transformer_path", nargs="?", const="", default="")
+
+    # lora
+    parser.add_argument("--lora-path", nargs="?", const="", default="")
+    parser.add_argument("--lora-alpha", type=float, default=1.0)
+
+    # dataset
+    parser.add_argument("--dataset-path", type=str, default="")
+    parser.add_argument("--resolution", nargs="+", type=int, default=[512])
+    parser.add_argument("--num-frames", type=int, default=81)
+    parser.add_argument("--batch-size", type=int, default=1)
+
+    # inference
+    parser.add_argument("--cfg", type=float, default=5.0)
+    parser.add_argument("--shift", type=float, default=3.0)
+    parser.add_argument("--step", type=int, default=50)
+
+    # save
+    parser.add_argument("--fps", type=int, default=16)
+    parser.add_argument("--save-dir", type=str, default="")
+
+    ### Preprocess dataset ###
+    parser.add_argument("--json-paths", nargs="+", type=str, default=[""])
+    parser.add_argument("--video-dir", type=str, default="")
+    parser.add_argument("--image-dir", type=str, default="")
+    parser.add_argument("--extract-fps", type=int, default=24)
+    # model
+    parser.add_argument("--model-type", type=str, default="wanx", choices=["wanx", "ltx"])
+    parser.add_argument("--vae-path", type=str, default="")
+    parser.add_argument("--image-processor-path", type=str, default="")
+    parser.add_argument("--image-encoder-path", type=str, default="")
+    parser.add_argument("--tokenizer-path", type=str, default="")
+    parser.add_argument("--text-encoder-path", type=str, default="")
+    parser.add_argument("--max-sequence-length", type=int, default=512)
+    parser.add_argument("--vlm-type", type=str, default="qwenvl2", choices=["qwenvl2", "qwenvl2.5", "nocap"])
+    parser.add_argument("--vlm-path", nargs="?", const="", default="")
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    # prompt
+    parser.add_argument("--caption-template", type=str, default="")
+    parser.add_argument("--instruct-sentence", type=str, default="")
+    parser.add_argument("--negative-prompt", nargs="?", const="", default="")
+    parser.add_argument("--null-caption-length", type=int, default=0)
+    # save
+    parser.add_argument("--save-interval", type=int, default=100)
+    parser.add_argument("--save-path", type=str, default="")
+
+    args = parser.parse_args()
+    return args
+
+
+def args_wan_init():
+    parser = argparse.ArgumentParser(
+        description="Generate a image or video from a text prompt or image using Wan"
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="i2v-14B",
+        choices=list(WAN_CONFIGS.keys()),
+        help="The task to run.")
+    parser.add_argument(
+        "--size",
+        type=str,
+        default="1280*720",
+        # choices=list(SIZE_CONFIGS.keys()),
+        help="The area (width*height) of the generated video. For the I2V task, the aspect ratio of the output video will follow that of the input image."
+    )
+    parser.add_argument(
+        "--frame_num",
+        type=int,
+        default=81,
+        help="How many frames to sample from a image or video. The number should be 4n+1"
+    )
+    parser.add_argument(
+        "--ckpt_dir",
+        type=str,
+        default='',
+        help="The path to the checkpoint directory.")
+    parser.add_argument(
+        "--offload_model",
+        type=str2bool,
+        default=None,
+        help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage."
+    )
+    parser.add_argument(
+        "--ulysses_size",
+        type=int,
+        default=1,
+        help="The size of the ulysses parallelism in DiT.")
+    parser.add_argument(
+        "--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_folder",
+        type=str,
+        default=None,
+        help="The folder to save the generated image or video to.")
+    parser.add_argument(
+        "--save_file",
+        type=str,
+        default=None,
+        help="The file to save the generated image or video to.")
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default=None,
+        help="The prompt to generate the image or video from.")
+    parser.add_argument(
+        "--base_seed",
+        type=int,
+        default=-1,
+        help="The seed to use for generating the image or video.")
+    parser.add_argument(
+        "--image",
+        type=str,
+        default=None,
+        help="The image to generate the video from.")
+    parser.add_argument(
+        "--sample_solver",
+        type=str,
+        default='unipc',
+        choices=['unipc', 'dpm++'],
+        help="The solver used to sample.")
+    parser.add_argument(
+        "--sample_steps", type=int, default=None, help="The sampling steps.")
+    parser.add_argument(
+        "--sample_shift",
+        type=float,
+        default=None,
+        help="Sampling shift factor for flow matching schedulers.")
+    parser.add_argument(
+        "--sample_guide_scale",
+        type=float,
+        default=6.0,
+        help="Classifier free guidance scale.")
+    parser.add_argument(
+        "--teacache_thresh",
+        type=float,
+        default=None,
+        help="The threshold for caching diffusion model steps.")
+
+    # NOTE: add by diffusers-lite to fill in blank args
+    parser.add_argument("--ddp_mode", type=bool, default=False)
+    # dataset
+    parser.add_argument("--dataset_path", type=str, default=None)
+    parser.add_argument("--resolution", nargs="+", type=int, default=[512])
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--negative_prompt", nargs="?", const="", default="")
+    # transformer
+    parser.add_argument("--transformer_path", nargs="?", const="", default="")
+    # lora
+    parser.add_argument("--lora_path", nargs="?", const="", default="")
+    parser.add_argument("--lora_alpha", type=float, default=1.0)
+    parser.add_argument("--distill_lora_path", nargs="?", const="", default="")
+    parser.add_argument("--distill_lora_alpha", type=float, default=1.0)
+
+    args = parser.parse_args()
+
+    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 "i2v" in args.task else 50
+
+    if args.sample_shift is None:
+        args.sample_shift = 5.0
+        if "i2v" in args.task:#  and args.size in ["832*480", "480*832"]
+            args.sample_shift = 3.0
+
+    # The default number of frames are 1 for text-to-image tasks and 81 for other tasks.
+    if args.frame_num is None:
+        args.frame_num = 1 if "t2i" in args.task else 81
+
+    # T2I frame_num check
+    if "t2i" in args.task:
+        assert args.frame_num == 1, f"Unsupport frame_num {args.frame_num} for task {args.task}"
+
+    args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
+        0, sys.maxsize)
+
+    return args

diffusers_lite/constants.py

@@ -0,0 +1,9 @@
+import torch
+
+PRECISION_TO_TYPE = {
+    'fp32': torch.float32,
+    'fp16': torch.float16,
+    'bf16': torch.bfloat16,
+}
+
+NULL_DIR="temp_data/null"

diffusers_lite/datasets/image2video_dataset.py

@@ -0,0 +1,448 @@
+import json
+import os
+import random
+import traceback
+from PIL import Image
+
+import torch
+import numpy as np
+from decord import VideoReader
+from easydict import EasyDict
+from einops import rearrange
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+from ..utils.data_utils import align_floor_to, align_ceil_to
+from ..constants import NULL_DIR
+
+
+class Image2VideoTrainDataset(Dataset):
+    def __init__(
+        self,
+        task="i2v-14b-480p",
+        dataset_type="wanx",
+        meta_file_list=[],
+        meta_file_lose_list=[],
+        uncond_prob=[0.0, 0.0],
+        sp_size=1,
+        patch_size=[1,2,2],
+    ):
+        self.task = task
+        self.dataset_type = dataset_type
+        self.uncond_prompt_prob = uncond_prob[0]
+        self.uncond_image_prob = uncond_prob[-1]
+        self.sp_size = sp_size
+        self.patch_size = patch_size
+        self.meta_paths = []
+
+        for meta_file in meta_file_list:
+            self.meta_paths.extend(
+                [line.strip() for line in open(meta_file, "r").readlines()]
+            )
+        if len(meta_file_lose_list) > 0:
+            self.meta_paths_lose = [] 
+            for meta_file in meta_file_lose_list:
+                self.meta_paths_lose.extend(
+                    [line.strip() for line in open(meta_file, "r").readlines()]
+                )
+
+    def __len__(self):
+        return len(self.meta_paths)
+
+    def __getitem__(self, idx):
+        try_times = 100
+        for _ in range(try_times):
+            try:
+                if self.dataset_type in ["refl"]:
+                    return self.get_batch_lrm_refl(idx)
+                elif self.dataset_type in ["lrm_ce"]:
+                    return self.get_batch_lrm_ce(idx)
+                elif self.dataset_type in ["lrm_bt_online"]:
+                    return self.get_batch_lrm_bt_online(idx)
+            except Exception as e:
+                print(
+                    f"Error details: {str(e)}-{idx}-{self.meta_paths[idx]}-{traceback.format_exc()}\n"
+                )
+                idx = np.random.randint(len(self.meta_paths))
+
+        raise RuntimeError("Too many bad data.")
+    
+    def get_batch_lrm_refl(self, idx):
+        data_json_path = self.meta_paths[idx]
+
+        with open(data_json_path, "r") as f:
+            data_dict = json.load(f)
+
+        # video
+        if 'video_vae_latent_path' in data_dict.keys():
+            latents_path = data_dict["video_vae_latent_path"]
+        elif 'vae_latent_path' in data_dict.keys():
+            latents_path = data_dict["vae_latent_path"]
+        else:
+            latents_path = data_dict["latents_path"]
+        latents = np.load(latents_path)[0]
+        latents = torch.from_numpy(latents)
+        frames = latents.shape[1]
+
+        # text states
+        if 'textshort_path' in data_dict and 'textlong_path' in data_dict:
+            text_states_path = data_dict["textshort_path"]
+            text_states_path_long = data_dict["textlong_path"]
+            prompt= data_dict["short_caption"]
+            if random.random() <= 0.7:
+                text_states_path = text_states_path_long
+                prompt=data_dict["long_caption"]
+        else:
+            text_states_path = data_dict["text_en_path"]
+            prompt= data_dict["prompt"]
+        text_states = np.load(text_states_path)[0]
+        text_states = torch.from_numpy(text_states)
+        
+        # image embeds
+        image_embeds_path = data_dict["imgclip_path"]
+        image_embeds = torch.from_numpy(np.load(image_embeds_path))
+        image_embeds = rearrange(image_embeds, "b s d -> (b s) d")
+
+        # latents condition
+        if "f1_black_path" in data_dict.keys():
+            latents_condition_path = data_dict["f1_black_path"]
+        else:
+            latents_condition_path = data_dict["latents_condition_path"]
+        latents_condition = np.load(latents_condition_path)[0]
+        latents_condition = torch.from_numpy(latents_condition)
+
+        # distill prompts
+        if "flf2v" in self.task:
+            uncond_text_states_path = os.path.join(NULL_DIR, f"wanx/uncond_flf2v.npy")
+        else:
+            uncond_text_states_path = os.path.join(NULL_DIR, f"wanx/uncond.npy")
+        uncond_text_states = np.load(uncond_text_states_path)[0]
+        uncond_text_states = torch.from_numpy(uncond_text_states)
+
+        # drop prompts
+        random_number = random.random()
+        if random_number < self.uncond_prompt_prob:
+            null_text_states_path = os.path.join(NULL_DIR, f"wanx/null.npy")
+            null_text_states = np.load(null_text_states_path)[0]
+            text_states = torch.from_numpy(null_text_states)
+
+        return latents, text_states, uncond_text_states, image_embeds, latents_condition,prompt#,inference_data_dict
+    
+    def get_batch_refl(self, idx):
+        data_json_path = self.meta_paths[idx]
+
+        with open(data_json_path, "r") as f:
+            data_dict = json.load(f)
+
+        # video
+        if 'video_vae_latent_path' in data_dict.keys():
+            latents_path = data_dict["video_vae_latent_path"]
+        elif 'vae_latent_path' in data_dict.keys():
+            latents_path = data_dict["vae_latent_path"]
+        else:
+            latents_path = data_dict["latents_path"]
+        latents = np.load(latents_path)[0]
+        latents = torch.from_numpy(latents)
+        frames = latents.shape[1]
+
+        # text states
+        text_states_path = data_dict["textshort_path"]
+        text_states_path_long = data_dict["textlong_path"]
+        prompt= data_dict["short_caption"]
+        if random.random() <= 0.7:
+            text_states_path = text_states_path_long
+            prompt=data_dict["long_caption"]
+        text_states = np.load(text_states_path)[0]
+        text_states = torch.from_numpy(text_states)
+        
+        image_embeds_path = data_dict["imgclip_path"]
+        image_embeds = torch.from_numpy(np.load(image_embeds_path))
+        image_embeds = rearrange(image_embeds, "b s d -> (b s) d")
+
+        if "f1_black_path" in data_dict.keys():
+            latents_condition_path = data_dict["f1_black_path"]
+        else:
+            latents_condition_path = data_dict["latents_condition_path"]
+        latents_condition = np.load(latents_condition_path)[0]
+        latents_condition = torch.from_numpy(latents_condition)
+
+        if "flf2v" in self.task:
+            uncond_text_states_path = os.path.join(NULL_DIR, f"wanx/uncond_flf2v.npy")
+        else:
+            uncond_text_states_path = os.path.join(NULL_DIR, f"wanx/uncond.npy")
+        uncond_text_states = np.load(uncond_text_states_path)[0]
+        uncond_text_states = torch.from_numpy(uncond_text_states)
+
+        random_number = random.random()
+        if random_number < self.uncond_prompt_prob:
+            null_text_states_path = os.path.join(NULL_DIR, f"wanx/null.npy")
+            null_text_states = np.load(null_text_states_path)[0]
+            text_states = torch.from_numpy(null_text_states)
+
+        return latents, text_states, uncond_text_states, image_embeds, latents_condition, prompt # inference_data_dict
+    
+    def get_batch_lrm_ce(self, idx):
+        data_json_path = self.meta_paths[idx]
+
+        with open(data_json_path, "r") as f:
+            data_dict = json.load(f)
+
+        source_id = data_dict["source_id"]
+
+        if 'video_vae_latent_path' in data_dict:
+            latents_path = data_dict["video_vae_latent_path"]
+        else:
+            latents_path = data_dict["vae_latent_path"]
+        
+        latents = np.load(latents_path)[0]
+        latents = torch.from_numpy(latents)
+        frames = latents.shape[1]
+        
+        if 'save_textshort_path' in data_dict:
+            text_states_path = data_dict["save_textshort_path"]
+        elif 'textshort_path' in data_dict:
+            text_states_path = data_dict["textshort_path"]
+        else:
+            text_states_path = data_dict["text_en_path"]
+        
+        text_states = np.load(text_states_path)[0]
+        text_states = torch.from_numpy(text_states)
+
+        if "image_embeds" in data_dict:
+            image_embeds_path = data_dict["image_embeds"]
+        else:
+            image_embeds_path = data_dict["imgclip_path"]
+        
+        image_embeds = torch.from_numpy(np.load(image_embeds_path))
+        image_embeds = rearrange(image_embeds, "b s d -> (b s) d")
+
+        if "f1_black_path" in data_dict:
+            latents_condition_path = data_dict["f1_black_path"]
+        else:
+            latents_condition_path = data_dict["latents_condition_path"]
+        
+        latents_condition = np.load(latents_condition_path)[0]
+        latents_condition = torch.from_numpy(latents_condition)
+        
+        if "flf2v" in self.task:
+            uncond_text_states_path = os.path.join(NULL_DIR, "wanx/uncond_flf2v.npy")
+        else:
+            uncond_text_states_path = os.path.join(NULL_DIR, "wanx/uncond.npy")
+        
+        uncond_text_states = np.load(uncond_text_states_path)[0]
+        uncond_text_states = torch.from_numpy(uncond_text_states)
+        
+        if "model" in data_dict:
+            data_from_model = data_dict["model"]
+        else:
+            data_from_model = ""
+        if "text_alignment" in data_dict:
+            text_alignment = data_dict["text_alignment"]
+        else:
+            text_alignment = 0
+        if "blur_quality" in data_dict:
+            blur_quality = data_dict["blur_quality"]
+        else:
+            blur_quality = 0
+        if "physics_quality" in data_dict:
+            physics_quality = data_dict["physics_quality"]
+        else:
+            physics_quality = 0
+        if "human_quality" in data_dict:
+            human_quality = data_dict["human_quality"]
+        else:
+            human_quality = 0
+
+        if text_alignment == "poor" or text_alignment is None: text_alignment = 0
+        if blur_quality == "poor" or blur_quality is None: blur_quality = 0 
+        if physics_quality == "poor" or physics_quality is None: physics_quality = 0 
+        if human_quality == "poor" or human_quality is None: human_quality = 0 
+        if text_alignment == "good": text_alignment = 1 
+        if blur_quality == "good": blur_quality = 1 
+        if physics_quality == "good": physics_quality = 1 
+        if human_quality == "good": human_quality = 1 
+
+        return (latents, text_states, uncond_text_states, image_embeds, latents_condition,
+                data_from_model, text_alignment, blur_quality, physics_quality, human_quality)
+    
+    def get_batch_lrm_bt_online(self, idx):
+        data_json_path = self.meta_paths[idx]
+        
+        if self.meta_paths_lose is None or len(self.meta_paths_lose) == 0:
+            raise ValueError("meta_paths_lose is None or empty. Please ensure bt=True and meta_file_list_lose is provided.")
+        
+        data_json_path_lose = self.meta_paths_lose[random.randint(0, len(self.meta_paths_lose)-1)]
+
+        with open(data_json_path, "r") as f:
+            data_dict = json.load(f)
+        with open(data_json_path_lose, "r") as f:
+            data_dict_lose = json.load(f)
+
+        if 'video_vae_latent_path' in data_dict:
+            latents_path = data_dict["video_vae_latent_path"]
+            latents_path_lose = data_dict_lose["video_vae_latent_path"]
+        else:
+            latents_path = data_dict["vae_latent_path"]
+            latents_path_lose = data_dict_lose["vae_latent_path"]
+        
+        latents = np.load(latents_path)[0]
+        latents = torch.from_numpy(latents)
+        frames = latents.shape[1]
+        latents_lose = np.load(latents_path_lose)[0]
+        latents_lose = torch.from_numpy(latents_lose)
+        frames_lose = latents_lose.shape[1]
+        assert latents.shape == latents_lose.shape, f'latents.shape {latents.shape} != latents_lose.shape {latents_lose.shape}'
+
+        if 'save_textshort_path' in data_dict:
+            text_states_path = data_dict["save_textshort_path"]
+            text_states_path_lose = data_dict_lose["save_textshort_path"]
+        elif 'textshort_path' in data_dict:
+            text_states_path_lose = data_dict_lose["textshort_path"]
+            text_states_path = data_dict["textshort_path"]
+        else:
+            text_states_path = data_dict["text_en_path"]
+            text_states_path_lose = data_dict_lose["text_en_path"]
+        
+        text_states = np.load(text_states_path)[0]
+        text_states = torch.from_numpy(text_states)
+        text_states_lose = np.load(text_states_path_lose)[0]
+        text_states_lose = torch.from_numpy(text_states_lose)
+
+        if "image_embeds" in data_dict:
+            image_embeds_path = data_dict["image_embeds"]
+            image_embeds_path_lose = data_dict_lose["image_embeds"]
+        else:
+            image_embeds_path = data_dict["imgclip_path"]
+            image_embeds_path_lose = data_dict_lose["imgclip_path"]
+        
+        image_embeds = torch.from_numpy(np.load(image_embeds_path))
+        image_embeds = rearrange(image_embeds, "b s d -> (b s) d")
+        image_embeds_lose = torch.from_numpy(np.load(image_embeds_path_lose))
+        image_embeds_lose = rearrange(image_embeds_lose, "b s d -> (b s) d")
+
+        if "f1_black_path" in data_dict:
+            latents_condition_path = data_dict["f1_black_path"]
+            latents_condition_path_lose = data_dict_lose["f1_black_path"]
+        else:
+            latents_condition_path = data_dict["latents_condition_path"]
+            latents_condition_path_lose = data_dict_lose["latents_condition_path"]
+        
+        latents_condition = np.load(latents_condition_path)[0]
+        latents_condition = torch.from_numpy(latents_condition)
+        latents_condition_lose = np.load(latents_condition_path_lose)[0]
+        latents_condition_lose = torch.from_numpy(latents_condition_lose)
+
+        if "flf2v" in self.task:
+            uncond_text_states_path = os.path.join(NULL_DIR, "wanx/uncond_flf2v.npy")
+            uncond_text_states_path_lose = os.path.join(NULL_DIR, "wanx/uncond_flf2v.npy")
+        else:
+            uncond_text_states_path = os.path.join(NULL_DIR, "wanx/uncond.npy")
+            uncond_text_states_path_lose = os.path.join(NULL_DIR, "wanx/uncond.npy")
+        
+        uncond_text_states = np.load(uncond_text_states_path)[0]
+        uncond_text_states = torch.from_numpy(uncond_text_states)
+        uncond_text_states_lose = np.load(uncond_text_states_path_lose)[0]
+        uncond_text_states_lose = torch.from_numpy(uncond_text_states_lose)
+        
+        return (latents, text_states, uncond_text_states, image_embeds, latents_condition,
+                latents_lose, text_states_lose, uncond_text_states_lose, image_embeds_lose, latents_condition_lose)
+    
+
+class Image2VideoEvalDataset(Dataset):
+    def __init__(self, file_path, resolution=(512,512), alignment=16, do_scale=True):
+        self.prompts = []
+        self.image_ids = []
+        self.image_paths = []
+        self.last_image_paths = []
+        self.seeds = []
+
+        if file_path.endswith(".txt"):
+            with open(file_path, "r") as file:
+                for line in file:
+                    prompt = line.strip()
+                    self.prompts.append(prompt)
+        
+        elif file_path.endswith(".json"):
+            with open(file_path, "r") as f:
+                datas = json.load(f)
+            for data in datas:
+                self.prompts.append(data["caption"].strip())
+                if "image_id" in data.keys():
+                    self.image_ids.append(data["image_id"])
+                if "image_path" in data.keys():
+                    self.image_paths.append(data["image_path"])
+                if "last_image_path" in data.keys():
+                    self.last_image_paths.append(data["last_image_path"])
+                if "seed" in data.keys():
+                    self.seeds.append(data["seed"])
+
+        self.resolution = resolution
+        self.alignment = alignment
+        self.do_scale = do_scale
+
+        print(f"[INFO] Load text and image dataset done, total len {len(self.prompts)}")
+
+    def __len__(self):
+        return len(self.prompts)
+
+    def __getitem__(self, index):
+        prompt = self.prompts[index]
+
+        if len(self.image_paths) > 0:
+            image_path = self.image_paths[index]
+            image_id = image_path.split("/")[-1].split(".")[0]
+
+            # Load image
+            image = Image.open(image_path).convert("RGB")
+
+            # Resize image
+            width, height = image.size
+            scale = min(min(self.resolution) / min(width, height), max(self.resolution) / max(width, height))
+
+            width_scale = align_ceil_to(int(width * scale), self.alignment)
+            height_scale = align_ceil_to(int(height * scale), self.alignment)
+
+            if not self.do_scale:
+                width_scale = width
+                height_scale = height
+
+            transform = transforms.Compose(
+                [
+                    transforms.Resize((height_scale, width_scale)),
+                    transforms.ToTensor(),
+                ]
+            )
+
+            image = transform(image)
+        else:
+            image_path = ""
+            image = ""
+            image_id = str(index)
+        
+        if len(self.image_ids) > 0:
+            image_id = self.image_ids[index]
+
+        # Load last image
+        if len(self.last_image_paths) > 0:
+            last_image_path = self.last_image_paths[index]
+            last_image = Image.open(last_image_path).convert("RGB")
+            last_image = transform(last_image)
+        else:
+            last_image = ""
+        
+        if len(self.seeds) > 0:
+            seed = self.seeds[index]
+            image_id += f'_seed_{seed}'
+        else:
+            seed = 42
+
+        return {
+            "prompt": prompt, 
+            "image": image, 
+            "last_image": last_image,
+            "image_id": image_id, 
+            "image_path": image_path,
+            "seed": seed,
+        }
+    
+    

diffusers_lite/schedulers/__init__.py

@@ -0,0 +1 @@
+from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler

diffusers_lite/schedulers/scheduling_flow_match_discrete.py

@@ -0,0 +1,275 @@
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, logging
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class FlowMatchDiscreteSchedulerOutput(BaseOutput):
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        sigma_max=1.0,
+        reverse: bool = True,
+        solver: str = "euler",
+    ):
+        # sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
+        sigmas = torch.linspace(sigma_max,0,num_train_timesteps+1)
+
+        if not reverse:
+            sigmas = sigmas.flip(0)
+
+        self.sigmas = sigmas
+        # the value fed to model
+        self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32)
+
+        self._step_index = None
+        self._begin_index = None
+
+        self.supported_solver = ["euler"]
+        if solver not in self.supported_solver:
+            raise ValueError(
+                f"Solver {solver} not supported. Supported solvers: {self.supported_solver}"
+            )
+
+        self.sigma_max = sigma_max
+
+    @property
+    def step_index(self):
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        self._begin_index = begin_index
+
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Union[str, torch.device] = None,
+        dtype: torch.Tensor = torch.float32,
+    ):
+        self.num_inference_steps = num_inference_steps
+        
+        sigmas = torch.linspace(self.sigma_max, 0, num_inference_steps + 1)
+        sigmas = (self.config.shift * sigmas) / (1 + (self.config.shift - 1) * sigmas)
+
+        if not self.config.reverse:
+            sigmas = 1 - sigmas
+
+        self.sigmas = sigmas
+        self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(
+            dtype=dtype, device=device
+        )
+
+        # Reset step index
+        self._step_index = None
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: Optional[int] = None
+    ) -> torch.Tensor:
+        return sample
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        return_dict: bool = True,
+    ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
+        if (
+            isinstance(timestep, int)
+            or isinstance(timestep, torch.IntTensor)
+            or isinstance(timestep, torch.LongTensor)
+        ):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+
+        sigma_ = self.sigmas[self.step_index + 1]
+        sigma = self.sigmas[self.step_index]
+        dt = sigma_ - sigma
+
+        if self.config.solver == "euler":
+            prev_sample = sample + model_output.to(torch.float32) * dt
+        else:
+            raise ValueError(
+                f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}"
+            )
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    def __len__(self):
+        return self.config.num_train_timesteps
+
+    def get_train_timestep_and_sigma(
+        self,
+        weighting_scheme: str = "logit_normal",   # logit_norm, uniform
+        batch_size: int = 1,
+        logit_mean: float = 0.0,
+        logit_std: float = 1.0,
+        device: Union[torch.device, str] = "cpu",
+        generator: Optional[torch.Generator] = None,
+        n_dim: int = 4,
+    ):
+        if weighting_scheme == "logit_normal":
+            # NOTE: sigma from 1 to 0.
+            u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), generator=generator)
+            u = torch.nn.functional.sigmoid(u)
+        else:
+            u = torch.rand(size=(batch_size,), generator=generator)
+
+        indices = (u * self.config.num_train_timesteps).long()
+        timestep = self.timesteps[indices].to(device=device)
+        sigma = self.sigmas[indices].to(device=device, dtype=torch.float32)
+
+        while len(sigma.shape) < n_dim:
+            sigma = sigma.unsqueeze(-1)
+
+        return timestep, sigma
+    
+    def get_train_timestep(
+        self,
+        weighting_scheme: str = "logit_normal",   # logit_norm, uniform
+        batch_size: int = 1,
+        logit_mean: float = 0.0,
+        logit_std: float = 1.0,
+        device: Union[torch.device, str] = "cpu",
+        generator: Optional[torch.Generator] = None,
+    ):
+        if weighting_scheme == "logit_normal":
+            # NOTE: sigma from 1 to 0.
+            u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), generator=generator)
+            u = torch.nn.functional.sigmoid(u)
+        else:
+            u = torch.rand(size=(batch_size,), generator=generator)
+
+        indices = (u * self.config.num_train_timesteps).long()
+        timestep = self.timesteps[indices].to(device=device)
+        return timestep
+
+    def get_train_sigma(
+        self,
+        timestep: Union[float, torch.FloatTensor],
+        n_dim: int = 4,
+        device: Union[str, torch.device] = "cpu",
+        dtype: torch.dtype = torch.float32,
+    ):
+        if isinstance(timestep, float):
+            timestep = torch.tensor([timestep], dtype=dtype)
+
+        sigmas = self.sigmas.to(device, dtype=dtype)
+        schedule_timesteps = self.timesteps.to(device)
+        timestep = timestep.to(device)
+
+        step_indices = [(schedule_timesteps == t).nonzero()[0].item() for t in timestep]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < n_dim:
+            sigma = sigma.unsqueeze(-1)
+        return sigma
+
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        sigma: Union[float, torch.FloatTensor],
+    ) -> torch.FloatTensor:
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample
+
+    def get_train_target(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+    ):
+        target = noise - original_samples
+        return target
+        
+    def get_train_loss_weighting(
+        self, 
+        sigma: torch.FloatTensor,
+    ):
+        weighting = torch.ones_like(sigma)
+        return weighting
+
+    def get_x0(
+        self,
+        model_output: torch.FloatTensor,
+        sample: torch.FloatTensor,
+        sigma_t: torch.FloatTensor,
+    ):
+        sigma_0 = torch.zeros_like(sigma_t)
+        dt = sigma_0 - sigma_t
+        prev_sample = sample + model_output.to(torch.float32) * dt
+        return prev_sample

diffusers_lite/utils/communication.py

@@ -0,0 +1,691 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Tuple
+
+import os
+import torch
+import functools
+import torch.distributed as dist
+from torch import Tensor
+
+from ..utils.parallel_states import nccl_info, get_teacher_student_parallel_state
+
+
+def broadcast(input_: torch.Tensor):
+    src = nccl_info.group_id * nccl_info.sp_size
+    dist.broadcast(input_, src=src, group=nccl_info.group)
+
+def broadcast_within_ts_unit(input_):
+    src = nccl_info.ts_unit_group_id * nccl_info.ts_unit_size
+    dist.broadcast(input_, src=src, group=nccl_info.ts_unit_group)
+
+def broadcast_global(input_: torch.Tensor):
+    dist.broadcast(input_, src=0, group=None)    
+    
+def broadcast_dict(input_: dict):
+    src = nccl_info.group_id * nccl_info.sp_size
+    for k, v in input_.items():
+        if isinstance(input_[k], torch.Tensor):
+            dist.broadcast(input_[k], src=src, group=nccl_info.group)
+
+def broadcast_dict_within_ts_unit(input_: dict):
+    src = nccl_info.ts_unit_group_id * nccl_info.ts_unit_size
+    for k, v in input_.items():
+        if isinstance(input_[k], torch.Tensor):
+            dist.broadcast(input_[k], src=src, group=nccl_info.ts_unit_group)
+
+def _all_to_all_4D(
+    input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None
+) -> torch.tensor:
+    """
+    all-to-all for QKV
+
+    Args:
+        input (torch.tensor): a tensor sharded along dim scatter dim
+        scatter_idx (int): default 1
+        gather_idx (int): default 2
+        group : torch process group
+
+    Returns:
+        torch.tensor: resharded tensor (bs, seqlen/P, hc, hs)
+    """
+    assert (
+        input.dim() == 4
+    ), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}"
+
+    seq_world_size = dist.get_world_size(group)
+
+    if scatter_idx == 2 and gather_idx == 1:
+        # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs)
+        bs, shard_seqlen, hc, hs = input.shape
+        seqlen = shard_seqlen * seq_world_size
+        shard_hc = hc // seq_world_size
+
+        # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them!
+        # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs)
+        input_t = (
+            input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs)
+            .transpose(0, 2)
+            .contiguous()
+        )
+
+        output = torch.empty_like(input_t)
+        # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single
+        # (P, seq_len/P, bs, hc/P, hs) scatter seqlen -all2all-> (P, seq_len/P, bs, hc/P, hs) scatter head
+        if seq_world_size > 1:
+            dist.all_to_all_single(output, input_t, group=group)
+            torch.cuda.synchronize()
+        else:
+            output = input_t
+        # if scattering the seq-dim, transpose the heads back to the original dimension
+        output = output.reshape(seqlen, bs, shard_hc, hs)
+
+        # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs)
+        output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs)
+
+        return output
+
+    elif scatter_idx == 1 and gather_idx == 2:
+        # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs)
+        bs, seqlen, shard_hc, hs = input.shape
+        hc = shard_hc * seq_world_size
+        shard_seqlen = seqlen // seq_world_size
+        seq_world_size = dist.get_world_size(group)
+
+        # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them!
+        # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs)
+        input_t = (
+            input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, hs)
+            .transpose(0, 3)
+            .transpose(0, 1)
+            .contiguous()
+            .reshape(seq_world_size, shard_hc, shard_seqlen, bs, hs)
+        )
+
+        output = torch.empty_like(input_t)
+        # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single
+        # (P, bs x hc/P, seqlen/P, hs) scatter seqlen -all2all-> (P, bs x seq_len/P, hc/P, hs) scatter head
+        if seq_world_size > 1:
+            dist.all_to_all_single(output, input_t, group=group)
+            torch.cuda.synchronize()
+        else:
+            output = input_t
+
+        # if scattering the seq-dim, transpose the heads back to the original dimension
+        output = output.reshape(hc, shard_seqlen, bs, hs)
+
+        # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs)
+        output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs)
+
+        return output
+    else:
+        raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2")
+
+
+class SeqAllToAll4D(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        input: Tensor,
+        scatter_idx: int,
+        gather_idx: int,
+    ) -> Tensor:
+        ctx.group = group
+        ctx.scatter_idx = scatter_idx
+        ctx.gather_idx = gather_idx
+
+        return _all_to_all_4D(input, scatter_idx, gather_idx, group=group)
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
+        return (
+            None,
+            SeqAllToAll4D.apply(
+                ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx
+            ),
+            None,
+            None,
+        )
+
+
+def all_to_all_4D(
+    input_: torch.Tensor,
+    scatter_dim: int = 2,
+    gather_dim: int = 1,
+):
+    return SeqAllToAll4D.apply(nccl_info.group, input_, scatter_dim, gather_dim)
+
+
+def _all_to_all(
+    input_: torch.Tensor,
+    world_size: int,
+    group: dist.ProcessGroup,
+    scatter_dim: int,
+    gather_dim: int,
+):
+    input_list = [
+        t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)
+    ]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
+    dist.all_to_all(output_list, input_list, group=group)
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+
+class _AllToAll(torch.autograd.Function):
+    """All-to-all communication.
+
+    Args:
+        input_: input matrix
+        process_group: communication group
+        scatter_dim: scatter dimension
+        gather_dim: gather dimension
+    """
+
+    @staticmethod
+    def forward(ctx, input_, process_group, scatter_dim, gather_dim):
+        ctx.process_group = process_group
+        ctx.scatter_dim = scatter_dim
+        ctx.gather_dim = gather_dim
+        ctx.world_size = dist.get_world_size(process_group)
+        output = _all_to_all(
+            input_, ctx.world_size, process_group, scatter_dim, gather_dim
+        )
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output = _all_to_all(
+            grad_output,
+            ctx.world_size,
+            ctx.process_group,
+            ctx.gather_dim,
+            ctx.scatter_dim,
+        )
+        return (
+            grad_output,
+            None,
+            None,
+            None,
+        )
+
+
+def all_to_all(
+    input_: torch.Tensor,
+    scatter_dim: int = 2,
+    gather_dim: int = 1,
+):
+    return _AllToAll.apply(input_, nccl_info.group, scatter_dim, gather_dim)
+
+
+class _AllGather(torch.autograd.Function):
+    """All-gather communication with autograd support.
+
+    Args:
+        input_: input tensor
+        dim: dimension along which to concatenate
+    """
+
+    @staticmethod
+    def forward(ctx, input_, dim):
+        ctx.dim = dim
+        world_size = nccl_info.sp_size
+        group = nccl_info.group
+        input_size = list(input_.size())
+
+        ctx.input_size = input_size[dim]
+
+        tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+        input_ = input_.contiguous()
+        dist.all_gather(tensor_list, input_, group=group)
+
+        output = torch.cat(tensor_list, dim=dim)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        world_size = nccl_info.sp_size
+        rank = nccl_info.rank_within_group
+        dim = ctx.dim
+        input_size = ctx.input_size
+
+        sizes = [input_size] * world_size
+
+        grad_input_list = torch.split(grad_output, sizes, dim=dim)
+        grad_input = grad_input_list[rank]
+
+        return grad_input, None
+
+
+def all_gather(input_: torch.Tensor, dim: int = 1):
+    """Performs an all-gather operation on the input tensor along the specified dimension.
+
+    Args:
+        input_ (torch.Tensor): Input tensor of shape [B, H, S, D].
+        dim (int, optional): Dimension along which to concatenate. Defaults to 1.
+
+    Returns:
+        torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'.
+    """
+    return _AllGather.apply(input_, dim)
+
+class _AllGather_TeacherStudent(torch.autograd.Function):
+    """All-gather communication with autograd support.
+
+    Args:
+        input_: input tensor
+        dim: dimension along which to concatenate
+    """
+
+    @staticmethod
+    def forward(ctx, input_, dim):
+        ctx.dim = dim
+        world_size = nccl_info.ts_unit_size
+        group = nccl_info.ts_unit_group
+        input_size = list(input_.size())
+
+        ctx.input_size = input_size[dim]
+
+        tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+        input_ = input_.contiguous()
+        dist.all_gather(tensor_list, input_, group=group)
+
+        output = torch.cat(tensor_list, dim=dim)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        world_size = nccl_info.ts_unit_size
+        rank = nccl_info.rank_within_ts_unit_group
+        dim = ctx.dim
+        input_size = ctx.input_size
+
+        sizes = [input_size] * world_size
+        grad_input_list = torch.split(grad_output, sizes, dim=dim)
+        grad_input = grad_input_list[rank]
+        return grad_input, None    
+        
+def all_gather_ts(input_: torch.Tensor, dim: int = 1):
+    """Performs an all-gather operation on the input tensor along the specified dimension.
+
+    Args:
+        input_ (torch.Tensor): Input tensor of shape [B, H, S, D].
+        dim (int, optional): Dimension along which to concatenate. Defaults to 1.
+
+    Returns:
+        torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'.
+    """
+    return _AllGather_TeacherStudent.apply(input_, dim)    
+
+    
+def prepare_sequence_parallel_data_wanx(
+    hidden_states, encoder_hidden_states, uncond_text_states, image_embeds, latents_condition
+):
+    if nccl_info.sp_size == 1:
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            uncond_text_states,
+            image_embeds,
+            latents_condition,
+        )
+
+    def prepare(hidden_states, encoder_hidden_states, uncond_text_states, image_embeds, latents_condition):
+        hidden_states = all_to_all(hidden_states, scatter_dim=2, gather_dim=0)
+        encoder_hidden_states = all_to_all(
+            encoder_hidden_states, scatter_dim=1, gather_dim=0
+        )
+        uncond_text_states = all_to_all(
+            uncond_text_states, scatter_dim=1, gather_dim=0
+        )
+        image_embeds = all_to_all(image_embeds, scatter_dim=1, gather_dim=0)
+        latents_condition = all_to_all(latents_condition, scatter_dim=2, gather_dim=0)
+
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            uncond_text_states,
+            image_embeds,
+            latents_condition,
+        )
+
+    sp_size = nccl_info.sp_size
+    frame = hidden_states.shape[2]
+    assert frame % sp_size == 0, "frame should be a multiple of sp_size"
+
+    (
+        hidden_states,
+        encoder_hidden_states,
+        uncond_text_states,
+        image_embeds,
+        latents_condition,
+    ) = prepare(
+        hidden_states,
+        encoder_hidden_states.repeat(1, sp_size, 1),
+        uncond_text_states.repeat(1, sp_size, 1),
+        image_embeds.repeat(1, sp_size, 1),
+        latents_condition,
+    )
+
+    return hidden_states, encoder_hidden_states, uncond_text_states, image_embeds, latents_condition
+
+
+def sp_parallel_dataloader_wrapper_wanx(
+    dataloader, device, train_batch_size, sp_size, train_sp_batch_size
+):
+    while True:
+        for data_item in dataloader:
+            latents, text_states, uncond_text_states, image_embeds, latents_condition = data_item
+            latents = latents.to(device)
+            text_states = text_states.to(device)
+            uncond_text_states = uncond_text_states.to(device)
+            image_embeds = image_embeds.to(device)
+            latents_condition = latents_condition.to(device)
+            frame = latents.shape[2]
+            if frame == 1:
+                yield latents, text_states, uncond_text_states, image_embeds, latents_condition
+            else:
+                latents, text_states, uncond_text_states, image_embeds, latents_condition = (
+                    prepare_sequence_parallel_data_wanx(
+                        latents, text_states, uncond_text_states, image_embeds, latents_condition
+                    )
+                )
+                assert (
+                    train_batch_size * sp_size >= train_sp_batch_size
+                ), "train_batch_size * sp_size should be greater than train_sp_batch_size"
+                for iter in range(train_batch_size * sp_size // train_sp_batch_size):
+                    st_idx = iter * train_sp_batch_size
+                    ed_idx = (iter + 1) * train_sp_batch_size
+                    yield (
+                        latents[st_idx:ed_idx],
+                        text_states[st_idx:ed_idx],
+                        uncond_text_states[st_idx:ed_idx],
+                        image_embeds[st_idx:ed_idx],
+                        latents_condition[st_idx:ed_idx],
+                    )
+
+def prepare_sequence_parallel_data_wanx_dpo(
+    hidden_states, encoder_hidden_states, uncond_text_states, image_embeds, latents_condition,latents_lose
+):
+    if nccl_info.sp_size == 1:
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            uncond_text_states,
+            image_embeds,
+            latents_condition,
+            latents_lose,
+        )
+
+    def prepare(hidden_states, encoder_hidden_states, uncond_text_states, image_embeds, latents_condition, latents_lose):
+        hidden_states = all_to_all(hidden_states, scatter_dim=2, gather_dim=0)
+        latents_lose = all_to_all(latents_lose, scatter_dim=2, gather_dim=0)
+        encoder_hidden_states = all_to_all(
+            encoder_hidden_states, scatter_dim=1, gather_dim=0
+        )
+        uncond_text_states = all_to_all(
+            uncond_text_states, scatter_dim=1, gather_dim=0
+        )
+        image_embeds = all_to_all(image_embeds, scatter_dim=1, gather_dim=0)
+        latents_condition = all_to_all(latents_condition, scatter_dim=2, gather_dim=0)
+
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            uncond_text_states,
+            image_embeds,
+            latents_condition,
+            latents_lose,
+        )
+
+    sp_size = nccl_info.sp_size
+    frame = hidden_states.shape[2]
+    assert frame % sp_size == 0, "frame should be a multiple of sp_size"
+
+    (
+        hidden_states,
+        encoder_hidden_states,
+        uncond_text_states,
+        image_embeds,
+        latents_condition,latents_lose,
+    ) = prepare(
+        hidden_states,
+        encoder_hidden_states.repeat(1, sp_size, 1),
+        uncond_text_states.repeat(1, sp_size, 1),
+        image_embeds.repeat(1, sp_size, 1),
+        latents_condition,
+        latents_lose
+    )
+
+    return hidden_states, encoder_hidden_states, uncond_text_states, image_embeds, latents_condition,latents_lose
+
+def sp_parallel_dataloader_wrapper_wanx_dpo(
+    dataloader, device, train_batch_size, sp_size, train_sp_batch_size
+):
+    while True:
+        for data_item in dataloader:
+            latents, text_states, uncond_text_states, image_embeds, latents_condition,latent_lose = data_item
+            latents = latents.to(device)
+            latents_lose = latents.to(device)
+            text_states = text_states.to(device)
+            uncond_text_states = uncond_text_states.to(device)
+            image_embeds = image_embeds.to(device)
+            latents_condition = latents_condition.to(device)
+            frame = latents.shape[2]
+            if frame == 1:
+                yield latents, text_states, uncond_text_states, image_embeds, latents_condition,latent_lose
+            else:
+                latents, text_states, uncond_text_states, image_embeds, latents_condition, latents_lose = (
+                    prepare_sequence_parallel_data_wanx_dpo(
+                        latents, text_states, uncond_text_states, image_embeds, latents_condition,latent_lose
+                    )
+                )
+                assert (
+                    train_batch_size * sp_size >= train_sp_batch_size
+                ), "train_batch_size * sp_size should be greater than train_sp_batch_size"
+                for iter in range(train_batch_size * sp_size // train_sp_batch_size):
+                    st_idx = iter * train_sp_batch_size
+                    ed_idx = (iter + 1) * train_sp_batch_size
+                    yield (
+                        latents[st_idx:ed_idx],
+                        text_states[st_idx:ed_idx],
+                        uncond_text_states[st_idx:ed_idx],
+                        image_embeds[st_idx:ed_idx],
+                        latents_condition[st_idx:ed_idx],
+                        latents_lose[st_idx:ed_idx],
+                    )
+
+def prepare_sequence_parallel_data_ltx(
+    hidden_states, encoder_hidden_states, text_mask, uncond_text_states, uncond_text_mask
+):
+    if nccl_info.sp_size == 1:
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            text_mask,
+            uncond_text_states,
+            uncond_text_mask,
+        )
+
+    def prepare(hidden_states, encoder_hidden_states, text_mask, uncond_text_states, uncond_text_mask):
+        hidden_states = all_to_all(hidden_states, scatter_dim=2, gather_dim=0)
+        encoder_hidden_states = all_to_all(
+            encoder_hidden_states, scatter_dim=1, gather_dim=0
+        )
+        text_mask = all_to_all(text_mask, scatter_dim=1, gather_dim=0)
+        uncond_text_states = all_to_all(
+            uncond_text_states, scatter_dim=1, gather_dim=0
+        )
+        uncond_text_mask = all_to_all(
+            uncond_text_mask, scatter_dim=1, gather_dim=0
+        )
+
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            text_mask,
+            uncond_text_states,
+            uncond_text_mask,
+        )
+
+    sp_size = nccl_info.sp_size
+    frame = hidden_states.shape[2]
+    assert frame % sp_size == 0, "frame should be a multiple of sp_size"
+
+    (
+        hidden_states,
+        encoder_hidden_states,
+        text_mask,
+        uncond_text_states,
+        uncond_text_mask,
+    ) = prepare(
+        hidden_states,
+        encoder_hidden_states.repeat(1, sp_size, 1),
+        text_mask.repeat(1, sp_size),
+        uncond_text_states.repeat(1, sp_size, 1),
+        uncond_text_mask.repeat(1, sp_size)
+    )
+
+    return hidden_states, encoder_hidden_states, text_mask, uncond_text_states, uncond_text_mask
+
+
+def sp_parallel_dataloader_wrapper_ltx(
+    dataloader, device, train_batch_size, sp_size, train_sp_batch_size
+):
+    while True:
+        for data_item in dataloader:
+            latents, text_states, text_mask, uncond_text_states, uncond_text_mask = data_item
+            latents = latents.to(device)
+            text_states = text_states.to(device)
+            text_mask = text_mask.to(device)
+            uncond_text_states = uncond_text_states.to(device)
+            uncond_text_mask = uncond_text_mask.to(device)
+            frame = latents.shape[2]
+            if frame == 1:
+                yield latents, text_states, text_mask, uncond_text_states, uncond_text_mask
+            else:
+                latents, text_states, text_mask, uncond_text_states, uncond_text_mask = (
+                    prepare_sequence_parallel_data_ltx(
+                        latents, text_states, text_mask, uncond_text_states, uncond_text_mask
+                    )
+                )
+                assert (
+                    train_batch_size * sp_size >= train_sp_batch_size
+                ), "train_batch_size * sp_size should be greater than train_sp_batch_size"
+                for iter in range(train_batch_size * sp_size // train_sp_batch_size):
+                    st_idx = iter * train_sp_batch_size
+                    ed_idx = (iter + 1) * train_sp_batch_size
+                    yield (
+                        latents[st_idx:ed_idx],
+                        text_states[st_idx:ed_idx],
+                        text_mask[st_idx:ed_idx],
+                        uncond_text_states[st_idx:ed_idx],
+                        uncond_text_mask[st_idx:ed_idx],
+                    )
+                    
+                    
+def parallelize_model(model):
+    original_forward = model.forward
+
+    @functools.wraps(model.__class__.forward)
+    def new_forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        text_states: torch.Tensor,
+        text_states_2: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        output_features=False,
+        output_features_stride=8,
+        attention_kwargs=None,
+        freqs_cos=None,
+        freqs_sin=None,
+        return_dict=False,
+        guidance=None,
+    ):
+        x = hidden_states
+        sp_size = nccl_info.sp_size
+        sp_rank = nccl_info.rank_within_group
+
+        if x.shape[-2] // 2 % sp_size == 0:
+            # try to split x by height
+            split_dim = -2
+        elif x.shape[-1] // 2 % sp_size == 0:
+            # try to split x by width
+            split_dim = -1
+        else:
+            raise ValueError(f"Cannot split video sequence into ulysses_degree  ({sp_size}) parts evenly")
+
+        _, _, ot, oh, ow = x.shape
+        tt, th, tw = (
+            ot // self.patch_size[0],
+            oh // self.patch_size[1],
+            ow // self.patch_size[2],
+        )
+        freqs_cos, freqs_sin = self.get_rotary_pos_embed((tt, th, tw))
+        # patch sizes for the temporal, height, and width dimensions are 1, 2, and 2.
+        temporal_size, h, w = x.shape[2], x.shape[3] // 2, x.shape[4] // 2
+
+        x = torch.chunk(x, sp_size,dim=split_dim)[sp_rank]
+
+        dim_thw = freqs_cos.shape[-1]
+        freqs_cos = freqs_cos.reshape(temporal_size, h, w, dim_thw)
+        freqs_cos = torch.chunk(freqs_cos, sp_size,dim=split_dim - 1)[sp_rank]
+        freqs_cos = freqs_cos.reshape(-1, dim_thw)
+        dim_thw = freqs_sin.shape[-1]
+        freqs_sin = freqs_sin.reshape(temporal_size, h, w, dim_thw)
+        freqs_sin = torch.chunk(freqs_sin, sp_size,dim=split_dim - 1)[sp_rank]
+        freqs_sin = freqs_sin.reshape(-1, dim_thw)
+
+        output = original_forward(
+            x,
+            timestep,
+            text_states,
+            text_states_2,
+            encoder_attention_mask,
+            output_features,
+            output_features_stride,
+            attention_kwargs,
+            freqs_cos,
+            freqs_sin,
+            return_dict,
+            guidance,
+        )
+
+        return_dict = not isinstance(output, tuple)
+        shape = (tt, th, tw)
+        if return_dict:
+            assert not output_features, "output_feature is not compatible with return_dict"
+            sample = output["x"]
+            sample = all_gather(sample, dim=split_dim)
+            output["x"] = sample
+        else:
+            sample = output[0]
+            sample = all_gather(sample, dim=split_dim)
+            if output_features:
+                features_list = output[1]
+                features_list = all_gather(features_list, dim=split_dim)
+            else:
+                features_list = None
+
+            output = (sample, features_list, shape)
+        return output
+
+    new_forward = new_forward.__get__(model)
+    model.forward = new_forward
+
+
+def all_reduce_tensor_item(item):
+    world_size = int(os.environ["WORLD_SIZE"])
+    item = item.detach().clone()
+    dist.all_reduce(item, op=dist.ReduceOp.SUM)
+    item = item / nccl_info.ts_group_size if get_teacher_student_parallel_state() else item / world_size
+    return item
+
+def broadcast_item(item, idx):
+    item_list = [item]
+    dist.broadcast_object_list(item_list, src=idx)
+    return item_list[0]

diffusers_lite/utils/data_utils.py

@@ -0,0 +1,542 @@
+import os
+import math
+import random
+from collections import Counter
+from typing import List, Optional
+
+import imageio
+import torch
+import torchvision
+import numpy as np
+from einops import rearrange
+from torch.utils.data import Sampler
+from typing import Union, Optional, Iterator, List, Callable
+import warnings
+import logging
+
+import torch
+import torch.distributed as dist
+from torch.utils.data.distributed import DistributedSampler as TorchDistributedSampler
+
+
+
+def split_list(input_list, rank=0, num_process=8):
+
+    n = len(input_list)
+    base = n // num_process
+    remainder = n % num_process
+
+    if rank < remainder:
+        start = rank * (base + 1)
+        end = start + (base + 1)
+    else:
+        start = remainder * (base + 1) + (rank - remainder) * base
+        end = start + base
+
+    local_input_list = input_list[start:end]
+
+    return local_input_list
+
+
+def align_floor_to(value, alignment):
+    return int(math.floor(value / alignment) * alignment)
+
+
+def align_ceil_to(value, alignment):
+    return int(math.ceil(value / alignment) * alignment)
+
+
+def crop_tensor(
+    latents,
+    image_latents=None,
+    crop_width_ratio=1.0,
+    crop_height_ratio=1.0,
+    crop_type="center",
+    crop_time_ratio=1.0,
+):
+    b, c, t, h, w = latents.shape
+    crop_h, crop_w = int(h * crop_height_ratio), int(w * crop_width_ratio)
+    crop_t = int(t * crop_time_ratio)
+
+    if crop_type == "center":
+        top = (h - crop_h) // 2
+        left = (w - crop_w) // 2
+    elif crop_type == "random":
+        top = random.randint(0, h - crop_h)
+        left = random.randint(0, w - crop_w)
+
+    crop_h = align_floor_to(crop_h, alignment=2)
+    crop_w = align_floor_to(crop_w, alignment=2)
+    crop_t = align_floor_to(crop_t, alignment=1)
+
+    if image_latents is not None:
+        return (
+            latents[:, :, :crop_t, top : top + crop_h, left : left + crop_w],
+            image_latents[:, :, :crop_t, top : top + crop_h, left : left + crop_w],
+        )
+    else:
+        return latents[:, :, :crop_t, top : top + crop_h, left : left + crop_w], image_latents
+    
+def crop_tensor_dpo(
+    latents,
+    latents_lose,
+    image_latents=None,
+    crop_width_ratio=1.0,
+    crop_height_ratio=1.0,
+    crop_type="center",
+    crop_time_ratio=1.0,
+):
+    b, c, t, h, w = latents.shape
+    crop_h, crop_w = int(h * crop_height_ratio), int(w * crop_width_ratio)
+    crop_t = int(t * crop_time_ratio)
+
+    if crop_type == "center":
+        top = (h - crop_h) // 2
+        left = (w - crop_w) // 2
+    elif crop_type == "random":
+        top = random.randint(0, h - crop_h)
+        left = random.randint(0, w - crop_w)
+
+    crop_h = align_floor_to(crop_h, alignment=2)
+    crop_w = align_floor_to(crop_w, alignment=2)
+    crop_t = align_floor_to(crop_t, alignment=1)
+
+    if image_latents is not None:
+        return (
+            latents[:, :, :crop_t, top : top + crop_h, left : left + crop_w],
+            image_latents[:, :, :crop_t, top : top + crop_h, left : left + crop_w],
+            latents_lose[:, :, :crop_t, top : top + crop_h, left : left + crop_w]
+        )
+    else:
+        return (latents[:, :, :crop_t, top : top + crop_h, left : left + crop_w],
+        image_latents,
+        latents_lose[:, :, :crop_t, top : top + crop_h, left : left + crop_w])
+
+
+def megabatch_frame_alignment(megabatches, lengths):
+    aligned_magabatches = []
+    for _, megabatch in enumerate(megabatches):
+        assert len(megabatch) != 0
+        len_each_megabatch = [lengths[i] for i in megabatch]
+        idx_length_dict = dict([*zip(megabatch, len_each_megabatch)])
+        count_dict = Counter(len_each_megabatch)
+
+        # mixed frame length, align megabatch inside
+        if len(count_dict) != 1:
+            sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1])
+            pick_length = sorted_by_value[-1][0]  # the highest frequency
+            candidate_batch = [
+                idx for idx, length in idx_length_dict.items() if length == pick_length
+            ]
+            random_select_batch = [
+                random.choice(candidate_batch)
+                for i in range(len(idx_length_dict) - len(candidate_batch))
+            ]
+            aligned_magabatch = candidate_batch + random_select_batch
+            aligned_magabatches.append(aligned_magabatch)
+        # already aligned megabatches
+        else:
+            aligned_magabatches.append(megabatch)
+
+    return aligned_magabatches
+
+
+def split_to_even_chunks(indices, lengths, num_chunks, batch_size):
+    """
+    Split a list of indices into `chunks` chunks of roughly equal lengths.
+    """
+
+    if len(indices) % num_chunks != 0:
+        chunks = [indices[i::num_chunks] for i in range(num_chunks)]
+    else:
+        num_indices_per_chunk = len(indices) // num_chunks
+
+        chunks = [[] for _ in range(num_chunks)]
+        chunks_lengths = [0 for _ in range(num_chunks)]
+        for index in indices:
+            shortest_chunk = chunks_lengths.index(min(chunks_lengths))
+            chunks[shortest_chunk].append(index)
+            chunks_lengths[shortest_chunk] += lengths[index]
+            if len(chunks[shortest_chunk]) == num_indices_per_chunk:
+                chunks_lengths[shortest_chunk] = float("inf")
+    # return chunks
+
+    pad_chunks = []
+    for idx, chunk in enumerate(chunks):
+        if batch_size != len(chunk):
+            assert batch_size > len(chunk)
+            if len(chunk) != 0:
+                chunk = chunk + [
+                    random.choice(chunk) for _ in range(batch_size - len(chunk))
+                ]
+            else:
+                chunk = random.choice(pad_chunks)
+                print(chunks[idx], "->", chunk)
+        pad_chunks.append(chunk)
+    return pad_chunks
+
+
+def group_frame_fun(indices, lengths):
+    # sort by num_frames
+    indices.sort(key=lambda i: lengths[i], reverse=True)
+    return indices
+
+
+def get_length_grouped_indices(
+    lengths,
+    batch_size,
+    world_size,
+    generator=None,
+    group_frame=False,
+    group_resolution=False,
+    seed=42,
+):
+    # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
+    if generator is None:
+        generator = torch.Generator().manual_seed(
+            seed
+        )  # every rank will generate a fixed order but random index
+
+    indices = torch.randperm(len(lengths), generator=generator).tolist()
+
+    # sort dataset according to frame
+    indices = group_frame_fun(indices, lengths)
+
+    # chunk dataset to megabatches
+    megabatch_size = world_size * batch_size
+    megabatches = [
+        indices[i : i + megabatch_size] for i in range(0, len(lengths), megabatch_size)
+    ]
+
+    # make sure the length in each magabatch is align with each other
+    megabatches = megabatch_frame_alignment(megabatches, lengths)
+
+    # aplit aligned megabatch into batches
+    megabatches = [
+        split_to_even_chunks(megabatch, lengths, world_size, batch_size)
+        for megabatch in megabatches
+    ]
+
+    # random megabatches to do video-image mix training
+    indices = torch.randperm(len(megabatches), generator=generator).tolist()
+    shuffled_megabatches = [megabatches[i] for i in indices]
+
+    # expand indices and return
+    return [
+        i for megabatch in shuffled_megabatches for batch in megabatch for i in batch
+    ]
+
+
+class LengthGroupedSampler(Sampler):
+    r"""
+    Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while
+    keeping a bit of randomness.
+    """
+
+    def __init__(
+        self,
+        batch_size: int,
+        rank: int,
+        world_size: int,
+        lengths: Optional[List[int]] = None,
+        group_frame=False,
+        group_resolution=False,
+        generator=None,
+    ):
+        if lengths is None:
+            raise ValueError("Lengths must be provided.")
+
+        self.batch_size = batch_size
+        self.rank = rank
+        self.world_size = world_size
+        self.lengths = lengths
+        self.group_frame = group_frame
+        self.group_resolution = group_resolution
+        self.generator = generator
+
+    def __len__(self):
+        return len(self.lengths)
+
+    def __iter__(self):
+        indices = get_length_grouped_indices(
+            self.lengths,
+            self.batch_size,
+            self.world_size,
+            group_frame=self.group_frame,
+            group_resolution=self.group_resolution,
+            generator=self.generator,
+        )
+
+        def distributed_sampler(lst, rank, batch_size, world_size):
+            result = []
+            index = rank * batch_size
+            while index < len(lst):
+                result.extend(lst[index : index + batch_size])
+                index += batch_size * world_size
+            return result
+
+        indices = distributed_sampler(
+            indices, self.rank, self.batch_size, self.world_size
+        )
+        return iter(indices)
+
+
+def save_videos_grid(videos, path, rescale=False, n_rows=1, fps=24):
+    videos = rearrange(videos, "b c t h w -> t b c h w")
+    outputs = []
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=n_rows)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+        if rescale:
+            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
+        x = torch.clamp(x, 0, 1)
+        x = (x * 255).numpy().astype(np.uint8)
+        outputs.append(x)
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    imageio.mimsave(path, outputs, fps=fps)
+
+
+class BlockDistributedSampler(TorchDistributedSampler):
+    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=False, seed=0, drop_last=False,
+                 batch_size=-1, start_index=0, align=1):
+        """
+        Args:
+            dataset: Dataset used for sampling.
+            num_replicas: Number of processes participating in distributed training.
+            rank: Rank of the current process within num_replicas.
+            shuffle: If True, the sampler will shuffle the indices.
+            seed: Random seed.
+            drop_last: If True, the sampler will drop the last batch if its size would be less than batch_size.
+            batch_size: Size of mini-batch. If callable, it should accept a tuple of (w, h) as input and return an integer
+                value as the batch size. It is useful for mix-scale(e.g., 256, 512, 1024) training.
+            start_index: Start index for the sampler.
+            align: Align the indices to the multiple of align for each dp.
+        """
+        super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last)
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        if rank >= num_replicas or rank < 0:
+            raise ValueError(
+                "Invalid rank {}, rank should be in the interval"
+                " [0, {}]".format(rank, num_replicas - 1))
+        if batch_size != -1:
+            align = batch_size
+            warnings.warn("batch_size is deprecated, please use `align` instead.")
+        if align <= 0:
+            raise ValueError(f"align should be a positive integer, but got {align}.")
+
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.drop_last = drop_last
+        self.shuffle = shuffle
+        self.seed = seed
+        self.batch_size = batch_size
+        self.align = align
+        self._start_index = start_index
+        self.recompute_sizes()
+
+    @property
+    def start_index(self):
+        return self._start_index
+
+    @start_index.setter
+    def start_index(self, value):
+        if self._start_index != value:
+            self._start_index = value
+            self.recompute_sizes()
+
+    def recompute_sizes(self):
+        self.num_samples = len(self.dataset) // self.align * self.align // self.num_replicas \
+                           - self._start_index
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        if self.shuffle:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
+        else:
+            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]
+        raw_num_samples = len(indices) // self.align * self.align // self.num_replicas
+        raw_total_size = raw_num_samples * self.num_replicas
+        indices = indices[:raw_total_size]
+
+        # subsample with start_index
+        indices = indices[self.rank * raw_num_samples + self.start_index:(self.rank + 1) * raw_num_samples]
+        assert len(indices) + self.start_index == raw_num_samples, \
+            f"{len(indices) + self.start_index} vs {raw_num_samples}"
+
+        # print(f"Iterator of BlockDistributedSampler created.")
+        # This is a sequential sampler. The shuffle operation is done by the dataset itself.
+        return iter(indices)
+
+
+class DistributedSampler(TorchDistributedSampler):
+    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=False, seed=0, drop_last=False,
+                 start_index=0):
+        super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last)
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        if rank >= num_replicas or rank < 0:
+            raise ValueError(
+                "Invalid rank {}, rank should be in the interval"
+                " [0, {}]".format(rank, num_replicas - 1))
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.drop_last = drop_last
+        self._start_index = start_index
+        self.recompute_sizes()
+        self.shuffle = shuffle
+        self.seed = seed
+
+    @property
+    def start_index(self):
+        return self._start_index
+
+    @start_index.setter
+    def start_index(self, value):
+        self._start_index = value
+        self.recompute_sizes()
+
+    def recompute_sizes(self):
+        # If the dataset length is evenly divisible by # of replicas, then there
+        # is no need to drop any data, since the dataset will be split equally.
+        if self.drop_last and (len(self.dataset) - self._start_index) % self.num_replicas != 0:  # type: ignore[arg-type]
+            # Split to nearest available length that is evenly divisible.
+            # This is to ensure each rank receives the same amount of data when
+            # using this Sampler.
+            self.num_samples = math.ceil(
+                ((len(self.dataset) - self._start_index) - self.num_replicas) / self.num_replicas  # type: ignore[arg-type]
+            )
+        else:
+            self.num_samples = math.ceil((len(self.dataset) - self._start_index) / self.num_replicas)  # type: ignore[arg-type]
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        if self.shuffle:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
+            indices = indices[self._start_index:]
+        else:
+            indices = list(range(self._start_index, len(self.dataset)))  # type: ignore[arg-type]
+
+        if not self.drop_last:
+            # add extra samples to make it evenly divisible
+            padding_size = self.total_size - len(indices)
+            if padding_size <= len(indices):
+                indices += indices[:padding_size]
+            else:
+                indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
+        else:
+            # remove tail of data to make it evenly divisible.
+            indices = indices[:self.total_size]
+        assert len(indices) == self.total_size
+
+        # subsample with start_index
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        print(f"Iterator of DistributedSamplerWithStartIndex created.")
+        return iter(indices)
+
+
+# For backward compatibility
+DistributedSamplerWithStartIndex = DistributedSampler
+
+
+def cumsum(sequence):
+    r, s = [], 0
+    for e in sequence:
+        l = len(e)
+        r.append(l + s)
+        s += l
+    return r
+
+def get_infinite_iterator(dataloader):
+    while True:
+        for batch in dataloader:
+            yield batch
+        dataloader.sampler.set_epoch(dataloader.sampler.epoch + 1)
+        print(f"epoch: {dataloader.sampler.epoch}, rank: {dataloader.sampler.rank}")
+        
+
+class VideoImageBatchIterator:
+    def __init__(self,
+                 video_dataloader, 
+                 image_dataloader = None, 
+                 sp_size = 1,
+                ):
+        assert video_dataloader is not None or image_dataloader is not None
+        self.sp_size = sp_size
+        self.video_dataloader = video_dataloader
+        self.image_dataloader = image_dataloader
+        self.video_iterator = iter(self.video_dataloader) if video_dataloader is not None else None
+        self.image_iterator = iter(self.image_dataloader) if image_dataloader is not None else None
+     
+    def get_image_batch(self):
+        try:
+            if self.sp_size > 1:
+                while True:
+                    batch = next(self.image_iterator)
+                    shape = batch[0].shape
+                    if shape[-1]/16 * shape[-2]/16 % self.sp_size == 0:
+                        break
+                    else:
+                        logging.warning(f"skipping one sample due to the shape {shape} and SP {self.sp_size} mismatching")
+            else:
+                batch = next(self.image_iterator)
+            return batch
+        except StopIteration:
+            logging.info(f"Image dataset start new epoch")
+            self.image_iterator = iter(self.image_dataloader)
+            raise StopIteration
+        
+    
+    def get_video_batch(self):
+        try:
+            if self.sp_size > 1:
+                while True:
+                    batch = next(self.video_iterator)
+                    shape = batch[0].shape # [B, C, T, H, W]
+                    if (shape[-1]/2 * shape[-2]/2 * shape[-3] % self.sp_size == 0):
+                        break
+                    else:
+                        logging.warning(f"skipping one sample due to the shape {shape} and SP {self.sp_size} mismatching")
+            else:
+                batch = next(self.video_iterator)
+                
+            return batch
+        except StopIteration:
+            logging.info(f"Video dataset start new epoch")
+            self.video_iterator = iter(self.video_dataloader)
+            return next(self.video_iterator)
+            
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.video_iterator is None:
+            return self.get_image_batch()
+        if self.image_iterator is None:
+            return self.get_video_batch()

diffusers_lite/utils/diffusion_utils.py

@@ -0,0 +1,395 @@
+import os
+import torch
+import torch.amp as amp
+import torch.nn.functional as F
+from einops import rearrange
+from safetensors.torch import load_file
+
+
+# tensor
+def expand_tensor_dims(tensor, ndim):
+    while len(tensor.shape) < ndim:
+        tensor = tensor.unsqueeze(-1)
+    return tensor
+
+
+# vae
+def vae_encode(vae, images, dtype=torch.bfloat16, vae_type="wanx"):
+    if vae_type in ["wanx"]:
+        images = batch2list(images)
+        latents = vae.encode(images)
+        latents = list2batch(latents)
+    elif vae_type in ["ltx"]:
+        with amp.autocast("cuda", dtype=dtype):
+            latents = vae.encode(images).latent_dist.sample()
+
+        latents_mean = vae.latents_mean
+        latents_std = vae.latents_std
+        scaling_factor = 1.0
+
+        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents = (latents - latents_mean) * scaling_factor / latents_std
+
+    return latents
+
+def vae_decode(vae, latents, dtype=torch.bfloat16, vae_type="wanx"):
+    if vae_type in ["wanx"]:
+        latents = batch2list(latents)
+        images = vae.decode(latents)
+        images = list2batch(images)
+
+    elif vae_type in ["ltx"]:
+        latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        scaling_factor = 1.0
+        latents = latents * latents_std / scaling_factor + latents_mean
+
+        with amp.autocast("cuda", dtype=dtype):
+            images = vae.decode(latents, return_dict=False)[0]
+
+    return images
+
+
+def image_encode(
+    image_encoder, 
+    image, 
+    last_image=None,
+    image_encoder_type="wanx"
+):
+    if image_encoder_type in ["wanx"]:
+        if image.ndim == 5:
+            image = image[:,:,0]
+        image = rearrange(image, "b c h w -> c b h w")
+
+        if last_image is not None:
+            if last_image.ndim == 5:
+                last_image = last_image[:,:,0]
+
+            last_image = rearrange(last_image, "b c h w -> c b h w")
+            image_embeds = image_encoder.visual([image, last_image])
+        else:
+            image_embeds = image_encoder.visual([image])
+
+    return image_embeds
+
+
+def pack_latents(latents, patch_size=1, patch_size_t=1):
+    # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
+    # The patch dimensions are then permuted and collapsed into the channel dimension of shape:
+    # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor).
+    # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features
+    batch_size, num_channels, num_frames, height, width = latents.shape
+    post_patch_num_frames = num_frames // patch_size_t
+    post_patch_height = height // patch_size
+    post_patch_width = width // patch_size
+    latents = latents.reshape(
+        batch_size,
+        -1,
+        post_patch_num_frames,
+        patch_size_t,
+        post_patch_height,
+        patch_size,
+        post_patch_width,
+        patch_size,
+    )
+    latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+    latents = latents.contiguous()
+    return latents
+
+
+def unpack_latents(latents, num_frames, height, width, patch_size=1, patch_size_t=1):
+    # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions)
+    # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of
+    # what happens in the `_pack_latents` method.
+    batch_size = latents.size(0)
+    latents = latents.reshape(
+        batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size
+    )
+    latents = (
+        latents.permute(0, 4, 1, 5, 2, 6, 3, 7)
+        .flatten(6, 7)
+        .flatten(4, 5)
+        .flatten(2, 3)
+    )
+    latents = latents.contiguous()
+    return latents
+
+
+# text encoder
+def prompt2states(
+    prompt,
+    text_encoder,
+    device="cuda:0",
+    tokenizer=None,
+    max_length=128,
+    text_encoder_type="wanx",
+):
+    if isinstance(prompt, str):
+        prompt = [prompt]
+
+    if text_encoder_type in ["wanx"]:
+        text_states = text_encoder(prompt, device)[0]
+        text_states = text_states.unsqueeze(0)
+        return text_states
+    elif text_encoder_type in ["ltx"]:
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_ids = text_inputs.input_ids.to(device)
+        text_mask = text_inputs.attention_mask
+        text_mask = text_mask.bool().to(device)
+        text_states = text_encoder(text_ids)[0]
+
+        return text_states, text_mask
+
+
+def load_lora_for_pipeline(
+    pipeline,
+    lora_path,
+    LORA_PREFIX_TRANSFORMER="",
+    LORA_PREFIX_TEXT_ENCODER="",
+    alpha=1.0,
+    rank=0,
+):
+    # load LoRA weight from .safetensors
+    state_dict = load_file(lora_path, device=rank)
+
+    visited = []
+
+    # directly update weight in diffusers model
+    for key in state_dict:
+        # it is suggested to print out the key, it usually will be something like below
+        # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"
+
+        # as we have set the alpha beforehand, so just skip
+        if "alpha" in key or key in visited:
+            continue
+
+        if "text" in key:
+            layer_infos = (
+                key.split(".")[0].split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_")
+            )
+            curr_layer = pipeline.text_encoder
+        else:
+            layer_infos = (
+                key.split(".")[0].split(LORA_PREFIX_TRANSFORMER + "_")[-1].split("_")
+            )
+            curr_layer = pipeline.transformer
+
+        # find the target layer
+        temp_name = layer_infos.pop(0)
+        while len(layer_infos) > -1:
+            try:
+                curr_layer = curr_layer.__getattr__(temp_name)
+                if len(layer_infos) > 0:
+                    temp_name = layer_infos.pop(0)
+                elif len(layer_infos) == 0:
+                    break
+            except Exception:
+                if len(temp_name) > 0:
+                    temp_name += "_" + layer_infos.pop(0)
+                else:
+                    temp_name = layer_infos.pop(0)
+
+        pair_keys = []
+        if "lora_down" in key:
+            pair_keys.append(key.replace("lora_down", "lora_up"))
+            pair_keys.append(key)
+        else:
+            pair_keys.append(key)
+            pair_keys.append(key.replace("lora_up", "lora_down"))
+
+        # update weight
+        if len(state_dict[pair_keys[0]].shape) == 4:
+            weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
+            weight_down = (
+                state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
+            )
+            curr_layer.weight.data += alpha * torch.mm(
+                weight_up, weight_down
+            ).unsqueeze(2).unsqueeze(3)
+        else:
+            weight_up = state_dict[pair_keys[0]].to(torch.float32)
+            weight_down = state_dict[pair_keys[1]].to(torch.float32)
+            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)
+
+        # update visited list
+        for item in pair_keys:
+            visited.append(item)
+    del state_dict
+
+    return pipeline
+
+
+def load_lora_for_model(
+    model,
+    lora_path,
+    LORA_PREFIX_TRANSFORMER="",
+    LORA_PREFIX_TEXT_ENCODER="",
+    alpha=1.0,
+    rank=0,
+):
+    # load LoRA weight from .safetensors
+    state_dict = load_file(lora_path, device="cpu")
+
+    visited = []
+
+    # directly update weight in diffusers model
+    for key in state_dict:
+        # it is suggested to print out the key, it usually will be something like below
+        # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"
+
+        # as we have set the alpha beforehand, so just skip
+        if "alpha" in key or key in visited:
+            continue
+
+        layer_infos = (
+            key.split(".")[0].split(LORA_PREFIX_TRANSFORMER + "_")[-1].split("_")
+        )
+        curr_layer = model
+
+        # find the target layer
+        temp_name = layer_infos.pop(0)
+        while len(layer_infos) > -1:
+            try:
+                curr_layer = curr_layer.__getattr__(temp_name)
+                if len(layer_infos) > 0:
+                    temp_name = layer_infos.pop(0)
+                elif len(layer_infos) == 0:
+                    break
+            except Exception:
+                if len(temp_name) > 0:
+                    temp_name += "_" + layer_infos.pop(0)
+                else:
+                    temp_name = layer_infos.pop(0)
+
+        pair_keys = []
+        if "lora_down" in key:
+            pair_keys.append(key.replace("lora_down", "lora_up"))
+            pair_keys.append(key)
+        else:
+            pair_keys.append(key)
+            pair_keys.append(key.replace("lora_up", "lora_down"))
+
+        # update weight
+        if len(state_dict[pair_keys[0]].shape) == 4:
+            weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
+            weight_down = (
+                state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
+            )
+            curr_layer.weight.data += alpha * torch.mm(
+                weight_up, weight_down
+            ).unsqueeze(2).unsqueeze(3)
+        else:
+            weight_up = state_dict[pair_keys[0]].to(torch.float32)
+            weight_down = state_dict[pair_keys[1]].to(torch.float32)
+            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)
+
+        # update visited list
+        for item in pair_keys:
+            visited.append(item)
+    del state_dict
+
+    return model
+
+
+def load_lora_state_dict(lora_dir):
+    lora_path = os.path.join(lora_dir, 'pytorch_lora_transformers_weights.safetensors')
+    lora_weights = load_file(lora_path)
+    load_lora_weights = {}
+    for key in lora_weights:
+        load_lora_weights[key.replace('.weight','.default.weight')] = lora_weights[key]
+
+    return load_lora_weights
+
+
+def transformer_zero_init(transformer):
+    for p in transformer.parameters():
+        if p.dim() > 1:
+            torch.nn.init.zeros_(p.data)
+        else:
+            torch.nn.init.normal_(p.data)
+
+    return transformer
+
+
+def prepare_video_condition_wanx(
+    vae, 
+    video, 
+    mask_strategy=[0.4, 0.25, 0.3, 0.05],
+):
+    # Get mask strategy
+    mask_id = torch.multinomial(torch.tensor(mask_strategy), num_samples=1).item()
+    bsz, _, num_frames, height, width = video.shape
+    latents_height, latents_width = height // 8, width // 8
+
+    # Get video mask
+    if mask_id == 0:
+        mask = torch.cat([
+            torch.ones(bsz, 1, 1, height, width),
+            torch.zeros(bsz, 1, num_frames-1, height, width)
+        ], dim=2)
+
+    elif mask_id == 1:
+        mid_frame = (num_frames - 1) // 2 + 1
+        mask = torch.cat([
+            torch.ones(bsz, 1, mid_frame, height, width),
+            torch.zeros(bsz, 1, num_frames-mid_frame, height, width)
+        ], dim=2)
+
+    elif mask_id == 2:
+        mask = torch.cat([
+            torch.ones(bsz, 1, 1, height, width),
+            torch.zeros(bsz, 1, num_frames-2, height, width),
+            torch.ones(bsz, 1, 1, height, width)
+        ], dim=2)
+
+    elif mask_id == 3:
+        num_masked = torch.randint(1, num_frames, (bsz,)).item()
+        indices = torch.randperm(num_frames)[:num_masked].sort().values
+        mask = torch.zeros(bsz, 1, num_frames, height, width)
+        mask[:,:, indices] = 1
+
+    # Encode video mask
+    mask = mask.to(video.device, dtype=video.dtype)
+    mask_lat_size = torch.cat([
+        torch.repeat_interleave(mask[:,:,:1,:,:], dim=2, repeats=4),
+        mask[:,:,1:,:,:],
+    ], dim=2)
+    mask_lat_size = mask_lat_size[:,:,:,::8,::8]
+    mask_lat_size = mask_lat_size.view(bsz, -1, 4, latents_height, latents_width).transpose(1,2)
+
+    # Encode video condition
+    video_condition = video * mask
+    latents_condition = torch.cat([
+        mask_lat_size,
+        vae_encode(vae, video_condition, "wanx")
+    ], dim=1)
+
+    return latents_condition
+
+
+def batch2list(batch):
+    return [item for item in batch]
+
+def list2batch(list):
+    return torch.stack(list)
+
+
+def stable_mse_loss(model_pred, target, weighting=None, threshold=50):
+    if weighting is None:
+        weighting = torch.ones_like(target)
+
+    diff = model_pred - target
+    mask = (diff.abs() <= threshold).float()
+    loss = F.mse_loss(model_pred, target, reduction="none")
+    masked_loss = weighting * mask * loss
+    masked_loss = masked_loss.mean()
+
+    return masked_loss

diffusers_lite/utils/distill_utils.py

@@ -0,0 +1,136 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from .diffusion_utils import list2batch
+
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1, ) * (len(x_shape) - 1)))
+
+def get_phase_endpoint(index, num_teacher_timesteps=32, multiphase=8):
+    interval = num_teacher_timesteps // multiphase
+    max_endpoint = num_teacher_timesteps - interval
+    
+    if index >= max_endpoint:
+        return max_endpoint
+
+    else:
+        quotient = index // interval
+        return quotient * interval
+
+class EulerSolver:
+    def __init__(self, sigmas, timesteps=1000, euler_timesteps=50):
+        # sigmas: 0.0 -> 1.0, length = 1001
+        self.num_timesteps = timesteps
+
+        step_ratio = timesteps / euler_timesteps
+        euler_timesteps = np.round(np.arange(timesteps, 0, -step_ratio)).astype(np.int64) - 1   # 999,...,0
+        self.euler_timesteps = euler_timesteps[::-1].copy() + 1 # 1,...,1000
+
+        self.sigmas = sigmas[self.euler_timesteps]  # 0.001,...,1.0
+        self.sigmas_prev = np.asarray(
+            [sigmas[0]] + sigmas[self.euler_timesteps[:-1]].tolist()    # 0.000,...,0.999
+        )
+        self.sigmas_all = sigmas.copy()
+
+        self.euler_timesteps = torch.from_numpy(self.euler_timesteps).long()
+        self.sigmas = torch.from_numpy(self.sigmas)
+        self.sigmas_prev = torch.from_numpy(self.sigmas_prev)
+        self.sigmas_all = torch.from_numpy(self.sigmas_all)
+        
+
+    def to(self, device):
+        self.euler_timesteps = self.euler_timesteps.to(device)
+        self.sigmas = self.sigmas.to(device)
+        self.sigmas_prev = self.sigmas_prev.to(device)
+        self.sigmas_all = self.sigmas_all.to(device)
+        return self
+
+    def euler_step(self, sample, model_pred, timestep_index):
+        sigma = extract_into_tensor(self.sigmas, timestep_index, model_pred.shape)
+        sigma_prev = extract_into_tensor(self.sigmas_prev, timestep_index, model_pred.shape)
+        x_prev = sample + (sigma_prev - sigma) * model_pred
+        return x_prev
+    
+    def euler_step_to_target(self, sample, model_pred, timestep_index, target_timestep_index):
+        sigma = extract_into_tensor(self.sigmas, timestep_index, model_pred.shape)
+        sigma_target = extract_into_tensor(self.sigmas_prev, target_timestep_index, model_pred.shape)
+
+        x_target = sample + (sigma_target - sigma) * model_pred
+        return x_target
+
+
+class DiscriminatorHead(nn.Module):
+    def __init__(self, in_channels=1280, reduced_channels=512):
+        super(DiscriminatorHead, self).__init__()
+        
+        # Reduce channels using 1x1 convolution
+        self.reduce_ch_conv = nn.Conv3d(in_channels, reduced_channels, kernel_size=(1, 1, 1))
+        
+        # Main convolutional layers
+        self.conv_layers = nn.Sequential(
+            nn.Conv3d(reduced_channels, reduced_channels * 2, kernel_size=(3, 3, 3), stride=(1, 2, 2)),
+            nn.LeakyReLU(0.2),
+            nn.Conv3d(reduced_channels * 2, reduced_channels * 4, kernel_size=(3, 3, 3), stride=(1, 2, 2)),
+            nn.LeakyReLU(0.2),
+            nn.Conv3d(reduced_channels * 4, reduced_channels * 8, kernel_size=(3, 3, 3), stride=(1, 2, 2)),
+            nn.LeakyReLU(0.2)
+        )
+        
+        # Global pooling
+        self.global_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
+        
+        # Fully connected layer
+        self.fc = nn.Linear(reduced_channels * 8, 1)
+
+    def forward(self, feature):
+        # Reduce channels
+        reduced_feature = self.reduce_ch_conv(feature)
+        
+        # Apply main convolutional layers
+        x = self.conv_layers(reduced_feature)
+        
+        # Global pooling
+        x = self.global_pool(x)
+
+        # Fully connected layer
+        x = x.view(x.size(0), -1)
+        out = self.fc(x)
+        
+        return out
+    
+
+class Discriminator(nn.Module):
+
+    def __init__(
+        self,
+        num_h_per_head=1,
+        selected_layers=[20,30,40],
+        adapter_channel_dims=[1280],
+    ):
+        super().__init__()
+        if isinstance(adapter_channel_dims, int):
+            adapter_channel_dims = [adapter_channel_dims]
+
+        adapter_channel_dims = adapter_channel_dims * len(selected_layers)
+        self.num_h_per_head = num_h_per_head
+        self.head_num = len(adapter_channel_dims)
+        self.heads = nn.ModuleList([
+            nn.ModuleList([DiscriminatorHead(adapter_channel) for _ in range(self.num_h_per_head)])
+            for adapter_channel in adapter_channel_dims
+        ])
+
+    def forward(self, features):
+        outputs = []
+        assert len(features) == len(self.heads)
+        for i in range(0, len(features)):
+            for h in self.heads[i]:
+                if isinstance(features[i], list):
+                    input_features = list2batch(features[i])
+                else:
+                    input_features = features[i]
+                out = h(input_features)
+                outputs.append(out)
+        return outputs

diffusers_lite/utils/fsdp_utils.py

@@ -0,0 +1,168 @@
+# ruff: noqa: E731
+import functools
+from functools import partial
+
+import torch
+from peft.utils.other import fsdp_auto_wrap_policy
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    CheckpointImpl,
+    apply_activation_checkpointing,
+    checkpoint_wrapper,
+)
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+
+from .load import get_no_split_modules
+from torch.distributed.fsdp import BackwardPrefetch
+non_reentrant_wrapper = partial(
+    checkpoint_wrapper,
+    checkpoint_impl=CheckpointImpl.NO_REENTRANT,
+)
+
+
+def apply_fsdp_checkpointing(model, no_split_modules, p=1):
+    # https://github.com/foundation-model-stack/fms-fsdp/blob/408c7516d69ea9b6bcd4c0f5efab26c0f64b3c2d/fms_fsdp/policies/ac_handler.py#L16
+    """apply activation checkpointing to model
+    returns None as model is updated directly
+    """
+    print("--> applying fdsp activation checkpointing...")
+    block_idx = 0
+    cut_off = 1 / 2
+    # when passing p as a fraction number (e.g. 1/3), it will be interpreted
+    # as a string in argv, thus we need eval("1/3") here for fractions.
+    p = eval(p) if isinstance(p, str) else p
+
+    def selective_checkpointing(submodule):
+        nonlocal block_idx
+        nonlocal cut_off
+
+        if isinstance(submodule, no_split_modules):
+            block_idx += 1
+            if block_idx * p >= cut_off:
+                cut_off += 1
+                return True
+        return False
+
+    apply_activation_checkpointing(
+        model,
+        checkpoint_wrapper_fn=non_reentrant_wrapper,
+        check_fn=selective_checkpointing,
+    )
+
+
+def get_mixed_precision(master_weight_type="fp32"):
+    weight_type = torch.float32 if master_weight_type == "fp32" else torch.bfloat16
+    mixed_precision = MixedPrecision(
+        param_dtype=weight_type,
+        # Gradient communication precision.
+        reduce_dtype=weight_type,
+        # Buffer precision.
+        buffer_dtype=weight_type,
+        cast_forward_inputs=False,
+    )
+    return mixed_precision
+
+
+def get_dit_fsdp_kwargs(
+    transformer,
+    sharding_strategy,
+    use_lora=False,
+    cpu_offload=False,
+    master_weight_type="fp32",
+):
+    no_split_modules = get_no_split_modules(transformer)
+    if use_lora:
+        auto_wrap_policy = fsdp_auto_wrap_policy
+    else:
+        auto_wrap_policy = functools.partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls=no_split_modules,
+        )
+
+    # we use float32 for fsdp but autocast during training
+    mixed_precision = get_mixed_precision(master_weight_type)
+
+    # NOTE: if no modules are split, we use NO_SHARD
+    if sharding_strategy == "full":
+        sharding_strategy = ShardingStrategy.FULL_SHARD
+    elif sharding_strategy == "hybrid_full":
+        sharding_strategy = ShardingStrategy.HYBRID_SHARD
+    elif sharding_strategy == "none":
+        sharding_strategy = ShardingStrategy.NO_SHARD
+        auto_wrap_policy = None
+    elif sharding_strategy == "hybrid_zero2":
+        sharding_strategy = ShardingStrategy._HYBRID_SHARD_ZERO2
+    elif sharding_strategy == 'shard_grad_op':
+        sharding_strategy = ShardingStrategy.SHARD_GRAD_OP
+
+    device_id = torch.cuda.current_device()
+    cpu_offload = (
+        torch.distributed.fsdp.CPUOffload(offload_params=True) if cpu_offload else None
+    )
+    fsdp_kwargs = {
+        "auto_wrap_policy": auto_wrap_policy,
+        "mixed_precision": mixed_precision,
+        "sharding_strategy": sharding_strategy,
+        "device_id": device_id,
+        "limit_all_gathers": True,
+        "cpu_offload": cpu_offload,
+    }
+
+    # Add LoRA-specific settings when LoRA is enabled
+    if len(no_split_modules) != 0 and use_lora:
+        fsdp_kwargs.update(
+            {
+                "use_orig_params": False,  # Required for LoRA memory savings
+                "sync_module_states": True,
+            }
+        )
+    elif len(no_split_modules) == 0 and use_lora:
+        fsdp_kwargs.update({"use_orig_params": True})
+    
+    return fsdp_kwargs, no_split_modules
+
+
+def get_discriminator_fsdp_kwargs(master_weight_type="fp32"):
+    auto_wrap_policy = None
+
+    # Use existing mixed precision settings
+    mixed_precision = get_mixed_precision(master_weight_type)
+    sharding_strategy = ShardingStrategy.NO_SHARD
+    device_id = torch.cuda.current_device()
+    fsdp_kwargs = {
+        "auto_wrap_policy": auto_wrap_policy,
+        "mixed_precision": mixed_precision,
+        "sharding_strategy": sharding_strategy,
+        "device_id": device_id,
+        "limit_all_gathers": True,
+    }
+
+    return fsdp_kwargs
+def get_vae_fsdp_kwargs(master_weight_type="fp32", cpu_offload=False):
+    auto_wrap_policy = None
+
+    # Use existing mixed precision settings
+    mixed_precision = get_mixed_precision(master_weight_type)
+    # sharding_strategy = ShardingStrategy.SHARD_GRAD_OP
+    sharding_strategy =  ShardingStrategy.FULL_SHARD  # 而不是SHARD_GRAD_OP
+
+    
+    # sharding_strategy = ShardingStrategy.NO_SHARD  # 注释掉的备用策略
+    device_id = torch.cuda.current_device()
+    cpu_offload = (
+        torch.distributed.fsdp.CPUOffload(offload_params=True) if cpu_offload else None
+    )
+    
+    fsdp_kwargs = {
+        "auto_wrap_policy": auto_wrap_policy,
+        "mixed_precision": mixed_precision,
+        "sharding_strategy": sharding_strategy,
+        "device_id": device_id,
+        "limit_all_gathers": True,
+        "cpu_offload": cpu_offload,  # 添加cpu_offload参数
+        "limit_all_gathers": True,
+        "use_orig_params": True,  # 保持原始参数结构
+        # "backward_prefetch": BackwardPrefetch.BACKWARD_PRE,
+    }
+
+    return fsdp_kwargs

diffusers_lite/utils/load.py

@@ -0,0 +1,12 @@
+import torch
+from peft import PeftModel
+from diffusers_lite.wan.modules.model import WanModel, WanAttentionBlock
+
+
+def get_no_split_modules(transformer):
+    while isinstance(transformer, PeftModel):
+        transformer = transformer.base_model.model
+    if isinstance(transformer, WanModel):
+        return (WanAttentionBlock, )
+    else:
+        raise ValueError(f"Unsupported transformer type: {type(transformer)}")

diffusers_lite/utils/model_utils.py

@@ -0,0 +1,175 @@
+import logging
+import json
+import os
+
+import torch
+from peft import get_peft_model_state_dict
+from safetensors.torch import save_file, load_file
+from tqdm import tqdm
+from torch.distributed.fsdp import StateDictType
+from torch.distributed.fsdp import ShardingStrategy
+from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+from .torch_utils import set_logging
+
+
+def get_kohya_state_dict(lora_layers, prefix="lora", dtype=torch.float32):
+    kohya_ss_state_dict = {}
+    for peft_key, weight in lora_layers.items():
+        kohya_key = peft_key.replace("base_model.model", prefix)
+        kohya_key = kohya_key.replace("lora_A", "lora_down")
+        kohya_key = kohya_key.replace("lora_B", "lora_up")
+        kohya_key = kohya_key.replace(".", "_", kohya_key.count(".") - 2)
+        kohya_ss_state_dict[kohya_key] = weight.to(dtype)
+
+    return kohya_ss_state_dict
+
+
+def get_diffusers_state_dict(lora_layers, dtype=torch.float32):
+    diffusers_ss_state_dict = {}
+    for peft_key, weight in lora_layers.items():
+        diffusers_key = peft_key.replace("base_model.model", "diffusion_model")
+        diffusers_ss_state_dict[diffusers_key] = weight.to(dtype)
+
+    return diffusers_ss_state_dict
+
+
+def save_lora_checkpoint(transformer, rank, output_dir, step, ema=False):
+    with FSDP.state_dict_type(
+        transformer,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+    ):
+        full_state_dict = transformer.state_dict()
+
+    if rank <= 0:
+        if ema:
+            save_dir = os.path.join(output_dir, f"checkpoint-{step}-ema")
+        else:
+            save_dir = os.path.join(output_dir, f"checkpoint-{step}")
+        os.makedirs(save_dir, exist_ok=True)
+
+        # save lora weight
+        transformer_lora_layers = get_peft_model_state_dict(
+            model=transformer, state_dict=full_state_dict
+        )
+        kohya_ss_state_dict = get_kohya_state_dict(lora_layers=transformer_lora_layers)
+        diffusers_ss_state_dict = get_diffusers_state_dict(
+            lora_layers=transformer_lora_layers
+        )
+        save_transformer_name = "pytorch_lora_transformers_weights.safetensors"
+        save_kohya_name = "pytorch_lora_kohya_weights.safetensors"
+        save_diffusers_name = "pytorch_lora_diffusers_weights.safetensors"
+
+        save_file(transformer_lora_layers, os.path.join(save_dir, save_transformer_name))
+        save_file(kohya_ss_state_dict, os.path.join(save_dir, save_kohya_name))
+        save_file(diffusers_ss_state_dict, os.path.join(save_dir, save_diffusers_name))
+
+
+def save_checkpoint(transformer, rank, output_dir, step, ema=False):
+    with FSDP.state_dict_type(
+        transformer,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+    ):
+        cpu_state = transformer.state_dict()
+
+    if rank <= 0:
+        if ema:
+            save_dir = os.path.join(output_dir, f"checkpoint-{step}-ema")
+        else:
+            save_dir = os.path.join(output_dir, f"checkpoint-{step}")
+        os.makedirs(save_dir, exist_ok=True)
+
+        max_bytes = 5 * 1024 ** 3  # 5GB
+        total_bytes = sum(v.numel() * v.element_size() for v in cpu_state.values())
+
+        if total_bytes <= max_bytes:
+            save_name = "diffusion_pytorch_model.safetensors"
+            save_file(cpu_state, os.path.join(save_dir, save_name))
+        else:
+            shard, shards, current_size = {}, [], 0
+            for k, v in sorted(cpu_state.items()):
+                tensor_size = v.numel() * v.element_size()
+                if current_size + tensor_size > max_bytes and shard:
+                    shards.append(shard)
+                    shard, current_size = {}, 0
+            
+                shard[k], current_size = v, current_size + tensor_size
+            if shard:
+                shards.append(shard)
+            
+            index_data = {
+                "metadata": {
+                    "total_size": total_bytes,
+                },
+                "weight_map": {}
+            }
+
+            for i, shard in enumerate(shards, start=1):
+                save_name = f"diffusion_pytorch_model-{i:05}-of-{len(shards):05}.safetensors"
+                save_file(shard, os.path.join(save_dir, save_name))
+                for key in shard.keys():
+                    index_data["weight_map"][key] = save_name
+        
+            with open(os.path.join(save_dir, "diffusion_pytorch_model.safetensors.index.json"), "w") as f:
+                json.dump(index_data, f, indent=2)
+            
+        config_dict = dict(transformer.config)
+        if "dtype" in config_dict:
+            del config_dict["dtype"]  # TODO
+        config_path = os.path.join(save_dir, "config.json")
+        # save dict as json
+        with open(config_path, "w") as f:
+            json.dump(config_dict, f, indent=4)
+
+def load_state_dict(model_dir, postfix=".safetensors"):
+    chunk_path_list = [os.path.join(model_dir, name) for name in os.listdir(model_dir) if name.endswith(postfix)]
+    chunk_length = len(chunk_path_list)
+
+    state_dict = {}
+    for chunk_path in tqdm(chunk_path_list, total=chunk_length):
+        if postfix == ".safetensors":
+            chunk_state_dict = load_file(chunk_path, device="cpu")
+        else:
+            chunk_state_dict = torch.load(chunk_path, map_location="cpu")
+        if "module" in chunk_state_dict.keys():
+            chunk_state_dict = chunk_state_dict["module"]
+        state_dict.update(chunk_state_dict)
+
+    return state_dict
+
+
+def print_parameters_information(model, name="Model name", rank=0):
+
+    def format_params(params):
+        if params < 1e6:
+            return f"{params} (less than 1M)"
+        elif params < 1e9:
+            return f"{params / 1e6:.2f}M"
+        else:
+            return f"{params / 1e9:.2f}B"
+
+    if model is None:
+        logging.info(f"name {name} is none objects.")
+        return
+
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+
+    param = next(model.parameters())
+    logging.info(
+        f"name [{name}] trainable params: {format_params(trainable_params)} || all params: {format_params(all_param)} || trainable%: {100 * trainable_params / all_param:.2f} || device: {param.device}, dtype: {param.dtype}."
+    )
+    logging.info(f"name [{name}] device: {param.device} || dtype: {param.dtype}.")
+
+@torch.no_grad
+def update_ema_model(transformer, ema_transformer, ema_decay):
+    for p_averaged, p_model in zip(ema_transformer.parameters(), transformer.parameters()):
+        if p_model.requires_grad:
+            p_averaged.data.mul_(ema_decay).add_(p_model.data, alpha=1 - ema_decay)

diffusers_lite/utils/network.py

@@ -0,0 +1,217 @@
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from sklearn.model_selection import train_test_split
+from diffusers.models.normalization import FP32LayerNorm
+
+class QueryAttention(nn.Module):
+    """
+    Query-based attention pooling module using PyTorch's built-in MultiheadAttention.
+    Uses learnable query vectors to attend to sequence features.
+    """
+    def __init__(self, feature_dim, num_queries=1, num_heads=8, dropout=0.1, layer_norm=False, return_type=None, product_text=False, text_dim=768):
+        super(QueryAttention, self).__init__()
+        self.feature_dim = feature_dim
+        self.num_queries = num_queries
+        self.num_heads = num_heads
+        self.layer_norm = layer_norm
+        self.return_type = return_type
+        self.product_text = product_text
+
+        # Use PyTorch's built-in MultiheadAttention
+        self.multihead_attn = nn.MultiheadAttention(
+            embed_dim=feature_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True  # Use batch_first=True for easier handling
+        )
+        # Learnable query vectors
+        self.queries = nn.Parameter(torch.randn(num_queries, feature_dim))
+        # Initialize query parameters
+        nn.init.xavier_uniform_(self.queries)
+
+        if self.layer_norm:
+            self.norm = FP32LayerNorm(feature_dim, eps=1e-6, elementwise_affine=False)
+        
+        if self.product_text:
+            self.text_proj = nn.Linear(text_dim, feature_dim)
+            nn.init.xavier_uniform_(self.text_proj.weight)
+            if self.text_proj.bias is not None:
+                nn.init.zeros_(self.text_proj.bias)
+
+        
+    def forward(self, x, e = None, text = None):
+        """
+        Args:
+            x: Input tensor of shape [batch_size, seq_len, feature_dim] or [batch_size, feature_dim]
+        Returns:
+            Pooled features of shape [batch_size, feature_dim]
+        """
+
+        if self.layer_norm:
+            x = self.norm(x)
+
+        batch_size = x.shape[0]
+        original_shape = x.shape
+        
+        # Handle different input shapes
+        if len(x.shape) == 2:  # [batch_size, feature_dim]
+            # Add sequence dimension
+            x = x.unsqueeze(1)  # [batch_size, 1, feature_dim]
+            seq_len = 1
+        elif len(x.shape) == 3:  # [batch_size, seq_len, feature_dim]
+            seq_len = x.shape[1]
+        elif len(x.shape) == 4:  # [sp_size, batch_size, seq_len, feature_dim]
+            # Handle sequence parallel case
+            sp_size, batch_size, seq_len, feature_dim = x.shape
+            x = x.view(sp_size * batch_size, seq_len, feature_dim)
+            batch_size = sp_size * batch_size
+        else:
+            raise ValueError(f"Unsupported input shape: {x.shape}")
+
+        # Expand queries to batch size
+        queries = self.queries.unsqueeze(0).expand(batch_size, -1, -1)  # [batch_size, num_queries, feature_dim]
+        if e is not None:
+            queries = queries + e.unsqueeze(0).expand(batch_size, -1, -1)
+        # Use PyTorch's MultiheadAttention
+        # query: [batch_size, num_queries, feature_dim]
+        # key, value: [batch_size, seq_len, feature_dim]
+        attended, attention_weights = self.multihead_attn(
+            query=queries,
+            key=x,
+            value=x,
+            need_weights=False  # We don't need attention weights for pooling
+        )
+        
+        # attended: [batch_size, num_queries, feature_dim]
+        
+        # If multiple queries, average them
+        if self.num_queries > 1:
+            output = attended.mean(dim=1)  # [batch_size, feature_dim]
+        else:
+            output = attended.squeeze(1)  # [batch_size, feature_dim]
+        
+        # Handle sequence parallel case
+        if len(original_shape) == 4:
+            output = output.view(sp_size, batch_size // sp_size, -1)
+            output = output.mean(dim=0)  # Average across SP devices
+        
+        if self.layer_norm:
+            output = self.norm(output)
+
+        if self.return_type == 'query':
+            output = output + queries
+
+        if self.product_text and text is not None:
+            output_product_text = torch.mul(self.text_proj(text), output)
+            return output_product_text
+        else:
+            return output
+
+class MLP(nn.Module):
+    def __init__(self, input_dim):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(input_dim, 1024)  # First hidden layer
+        self.fc2 = nn.Linear(1024, 512)        # Second hidden layer
+        self.fc3 = nn.Linear(512, 1)           # Output layer (binary classification)
+        
+        # 初始化权重，避免梯度消失
+        self._init_weights()
+    
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                # 使用Xavier初始化
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+    
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)  # 注意：这里不应用sigmoid，因为forward_siamese会处理
+        return x
+
+class MultiHead(nn.Module):
+    def __init__(self, input_dim, num_heads = 3):
+        super().__init__()
+        self.num_heads = num_heads
+        self.mlps = torch.nn.ModuleList(
+            [MLP(input_dim) for _ in range(num_heads)]
+        )
+
+    def forward_mlp(self, head_idx, x):
+        return torch.sigmoid(self.mlps[head_idx](x))
+
+    def forward(self, x):
+        out = [self.forward_mlp(h, x) for h in range(self.num_heads)]
+        return torch.stack(out)
+
+def forward_mlp(model, input):
+    return torch.sigmoid(model(input))
+
+def forward_siamese(model, input1, input2):
+    # Pass both inputs through the same model (weight sharing)
+    reward1 = model(input1)
+    reward2 = model(input2)
+    # Compute the difference between the two embeddings
+    diff = reward1 - reward2
+
+    # Use this difference for binary prediction (preference/ranking)
+    return torch.sigmoid(diff)
+
+def train_model(model, device, model_mode, X_train, y_train, X_test, y_test, epochs=3, lr=0.001, batch_size = 512, verbose=False, ealry_stopping_patience=3):
+    model = model.to(device)  # Move model to GPU
+    criterion = nn.BCELoss()  # Binary cross-entropy loss with logits for numerical stability
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+    batch_size = min(batch_size, X_train.shape[0])
+    val_losses = []
+    for epoch in range(epochs):
+        for n_batch in range(0, X_train.shape[0], batch_size):
+
+            # randomly selecte batch
+            batch_idx = torch.randperm(X_train.shape[0])[:batch_size]
+            X_batch = X_train[batch_idx]
+            y_batch = y_train[batch_idx]
+
+            model.train()
+            optimizer.zero_grad()
+
+            # Forward pass
+            if model_mode == 'clf':
+                outputs = forward_mlp(model, X_batch)
+            elif model_mode == 'siamese':
+                outputs = forward_siamese(model, X_batch[:, 0], X_batch[:, 1])
+            loss = criterion(outputs, y_batch)
+
+            # Backward pass and optimization
+            loss.backward()
+            optimizer.step()
+
+        # Evaluate on validation set
+        model.eval()
+        with torch.no_grad():
+            if model_mode == 'clf':
+                val_outputs = forward_mlp(model, X_test)
+            elif model_mode == 'siamese':
+                val_outputs = forward_siamese(model, X_test[:, 0], X_test[:, 1])
+        # early stopping?
+        val_loss = criterion(val_outputs, y_test)
+        val_losses.append(val_loss.cpu().detach().item())
+        if len(val_losses) > ealry_stopping_patience:
+            if all(val_losses[-1] > x for x in val_losses[-(ealry_stopping_patience+1):-1]):
+                if verbose:
+                    print(f"Early stopping at epoch {epoch+1}")
+                break
+        if verbose:
+            print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.cpu().detach().item()}, Val Loss: {val_loss.cpu().detach().item()}")
+        # accuracy
+        val_outputs = val_outputs.cpu().detach().numpy()
+        val_pred = (val_outputs > 0.5).astype(int)
+        accuracy = (val_pred == y_test.cpu().detach().numpy()).mean()
+        if verbose:
+            print(f"Accuracy: {accuracy}")
+
+def save_model(model, path):
+    torch.save(model.state_dict(), path)

diffusers_lite/utils/parallel_states.py

@@ -0,0 +1,141 @@
+
+import torch
+import torch.distributed as dist
+import os
+import time
+import random
+import functools
+from typing import List, Optional, Tuple, Union
+
+class COMM_INFO:
+
+    def __init__(self):
+        self.group = None
+        self.sp_size = 1
+        self.global_rank = 0
+        self.rank_within_group = 0
+        self.group_id = 0
+        
+        # the group info for teacher-student parallel
+        self.ts_group_size = 1
+        self.ts_group = None  # for fsdp data parallel communication
+        self.ts_group_id = 0
+        
+        # the group for teacher-student unit union
+        self.ts_unit_size = 1
+        self.ts_unit_group = None
+        self.rank_within_ts_unit_group = 0
+        self.ts_unit_group_id = 0
+
+
+nccl_info = COMM_INFO()
+_SEQUENCE_PARALLEL_STATE = False
+_TEACHER_STUDENT_PARALLEL_STATE = False
+
+def initialize_sequence_parallel_state(sequence_parallel_size):
+    global _SEQUENCE_PARALLEL_STATE
+    if sequence_parallel_size > 1:
+        _SEQUENCE_PARALLEL_STATE = True
+        initialize_sequence_parallel_group(sequence_parallel_size)
+    else:
+        nccl_info.sp_size = 1
+        nccl_info.global_rank = int(os.getenv("RANK", "0"))
+        nccl_info.rank_within_group = 0
+        nccl_info.group_id = int(os.getenv("RANK", "0"))
+
+
+def set_sequence_parallel_state(state):
+    global _SEQUENCE_PARALLEL_STATE
+    _SEQUENCE_PARALLEL_STATE = state
+
+
+def get_sequence_parallel_state():
+    return _SEQUENCE_PARALLEL_STATE
+
+
+def initialize_sequence_parallel_group(sequence_parallel_size):
+    """Initialize the sequence parallel group."""
+    rank = int(os.getenv("RANK", "0"))
+    world_size = int(os.getenv("WORLD_SIZE", "1"))
+    assert (
+        world_size % sequence_parallel_size == 0
+    ), "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format(
+        world_size, sequence_parallel_size
+    )
+    nccl_info.sp_size = sequence_parallel_size #序列并行size
+    nccl_info.global_rank = rank #全局rank
+    num_sequence_parallel_groups: int = world_size // sequence_parallel_size
+    for i in range(num_sequence_parallel_groups):
+        ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size)
+        group = dist.new_group(ranks)
+        if rank in ranks:
+            nccl_info.group = group
+            nccl_info.rank_within_group = rank - i * sequence_parallel_size #rank在序列并行group中的rank
+            nccl_info.group_id = i #sequence parallel group id
+
+def get_sequence_parallel_state():
+    return _SEQUENCE_PARALLEL_STATE
+
+
+def set_teacher_student_parallel_state(state):
+    global _TEACHER_STUDENT_PARALLEL_STATE
+    _TEACHER_STUDENT_PARALLEL_STATE = state
+
+
+def get_teacher_student_parallel_state():
+    return _TEACHER_STUDENT_PARALLEL_STATE
+
+
+            
+def initialize_teacher_student_parallel_state(sequence_parallel_size):
+    global _TEACHER_STUDENT_PARALLEL_STATE
+    """Initialize the teacher-student parallel group."""
+    rank = int(os.getenv("RANK", "0"))
+    world_size = int(os.getenv("WORLD_SIZE", "1"))
+    assert (
+        world_size % (2 * sequence_parallel_size) == 0
+    ), "world_size must be divisible by 2 * sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format(
+        world_size, sequence_parallel_size
+    )
+    _TEACHER_STUDENT_PARALLEL_STATE = True
+    nccl_info.global_rank = rank
+    # init ts_unit_group and assign info
+    # teacher and student must have the same sp size temporally! 
+    # In the unit, front is student, back is teacher
+    nccl_info.ts_unit_size = sequence_parallel_size * 2
+    num_teacher_student_union_groups = world_size // sequence_parallel_size // 2    
+    for j in range(num_teacher_student_union_groups):
+        ts_unit_ranks = range(j * sequence_parallel_size * 2, (j+1) * sequence_parallel_size * 2)
+        ts_unit_group = dist.new_group(ts_unit_ranks)
+        if rank in ts_unit_ranks:
+            nccl_info.ts_unit_group = ts_unit_group
+            nccl_info.ts_unit_group_id = j
+            nccl_info.rank_within_ts_unit_group = rank - j * sequence_parallel_size * 2
+    
+    # init ts_goup and assign info    
+    nccl_info.ts_group_size = world_size // 2
+    for i in range(2):
+        ranks = []
+        for j in range(num_teacher_student_union_groups):
+            ranks += range((j*2+i) * sequence_parallel_size, (j*2+i+1) * sequence_parallel_size)
+                        
+        ts_group = dist.new_group(ranks)
+        if rank in ranks:
+            nccl_info.ts_group = ts_group
+            nccl_info.ts_group_id = i
+    
+def destroy_sequence_parallel_group():
+    """Destroy the sequence parallel group."""
+    dist.destroy_process_group()
+
+def is_teacher_group():
+    if _TEACHER_STUDENT_PARALLEL_STATE:
+        return nccl_info.group_id % 2 == 1
+    else:
+        return True
+
+def is_student_group():
+    if _TEACHER_STUDENT_PARALLEL_STATE:
+        return nccl_info.group_id % 2 == 0
+    else:
+        return True

diffusers_lite/utils/torch_utils.py

@@ -0,0 +1,59 @@
+import gc
+import random
+import logging
+import os
+import sys
+import numpy as np
+import torch
+import torch.distributed as dist
+# from loguru import logger
+
+def init_dist():
+    """Initializes distributed environment."""
+    rank = int(os.environ["RANK"])
+    num_gpus = torch.cuda.device_count()
+    local_rank = rank % num_gpus
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend="nccl")
+    return local_rank
+
+
+def set_manual_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def make_contiguous(x):
+    if isinstance(x, torch.Tensor):
+        return x.contiguous()
+    elif isinstance(x, dict):
+        return {k: make_contiguous(v) for k, v in x.items()}
+    else:
+        return x
+
+
+class set_worker_seed_builder():
+    def __init__(self, global_rank):
+        self.global_rank = global_rank
+
+    def __call__(self, worker_id):
+        set_manual_seed(torch.initial_seed() % (2 ** 32 - 1))
+
+def free_memory():
+    if torch.cuda.is_available():
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
+def set_logging(local_rank):
+    if local_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)

diffusers_lite/wan/__init__.py

@@ -0,0 +1,4 @@
+from . import configs, distributed, modules
+from .image2video import WanI2V
+from .text2video import WanT2V
+from .first_last_frame2video import WanFLF2V

diffusers_lite/wan/configs/__init__.py

@@ -0,0 +1,49 @@
+# 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
+
+# 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
+}
+
+SIZE_CONFIGS = {
+    '720*1280': (720, 1280),
+    '1280*720': (1280, 720),
+    '480*832': (480, 832),
+    '832*480': (832, 480),
+    '1024*1024': (1024, 1024),
+}
+
+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()),
+}

diffusers_lite/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压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'