fix name typo

upload github pytest workflow
init test file
2024-12-29 04:53:45 +00:00 · 2024-12-29 04:50:59 +00:00 · 2024-12-29 04:40:14 +00:00 · 2024-12-25 21:50:57 -08:00 · 2024-12-25 12:40:16 -08:00 · 2024-12-19 13:36:43 -08:00
99 changed files with 13524 additions and 784 deletions
@@ -0,0 +1,29 @@
+name: 🐞 Bug report
+description: Create a report to help us reproduce and fix the bug
+title: "[Bug] "
+labels: ['Bug']
+
+body:
+- type: textarea
+  attributes:
+    label: Environment
+    description: |
+      Please share your environment with us. You can run the command **python fastvideo/utils/env_utils.py** and copy-paste its output below.
+    placeholder: FastVideo version, platform, python version, cuda version...
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Describe the bug
+    description: A clear and concise description of what the bug is.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Reproduction
+    description: |
+      What command or script did you run? Which **model** are you using?
+    placeholder: |
+      A placeholder for the command.
+  validations:
+    required: true
@@ -0,0 +1,17 @@
+name: 🚀 Feature request
+description: Suggest an idea for this project
+title: "[Feature] "
+
+body:
+- type: textarea
+  attributes:
+    label: Motivation
+    description: |
+      A clear and concise description of the motivation of the feature.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Related resources
+    description: |
+      If there is an official code release or third-party implementations, please also provide the information here, which would be very helpful.
@@ -0,0 +1,27 @@
+name: Run Tests
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'  # or any version you need
+
+      - name: Install dependencies
+        run: |
+          pip install --upgrade pip
+          pip install -e .
+      - name: Run Pytest
+        run: |
+          pytest
@@ -20,7 +20,6 @@ wandb/
 *.pt
 cache_dir/
 wandb/
-test*
 sample_video*
 sample_image*
 512*
@@ -1,189 +1,155 @@
-# FastVideo
-
 <div align="center">
-  <a href=""><img src="https://img.shields.io/static/v1?label=API:H100&message=Replicate&color=pink"></a> &ensp;
-  <a href=""><img src="https://img.shields.io/static/v1?label=Discuss&message=Discord&color=purple&logo=discord"></a> &ensp;
-</div>
-<br>
-<div align="center">
-<img src=assets/logo.png width="50%"/>
+<img src=assets/logo.jpg width="30%"/>
 </div>

-FastVideo is a scalable framework for post-training video diffusion models, addressing the growing challenges of fine-tuning, distillation, and inference as model sizes and sequence lengths increase. As a first step, it provides an efficient script for distilling and fine-tuning the 10B Mochi model, with plans to expand features and support for more models.
+FastVideo is a lightweight framework for accelerating large video diffusion models.

-### Features

- FastMochi, a distilled Mochi model that can generate videos with merely 8 sampling steps.
- Finetuning with FSDP (both master weight and ema weight), sequence parallelism, and selective gradient checkpointing.
- LoRA coupled with pecomputed the latents and text embedding for minumum memory consumption.
- Finetuning with both image and videos.
+https://github.com/user-attachments/assets/5fbc4596-56d6-43aa-98e0-da472cf8e26c
+
+
+
+
+<p align="center">
+    🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank">FastMochi</a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank">FastHunyuan</a>  | 🎮 <a href="https://discord.gg/REBzDQTWWt" target="_blank"> Discord </a> | 🕹️ <a href="https://replicate.com/lucataco/fast-hunyuan-video" target="_blank"> Replicate </a> 
+</p> 
+
+
+FastVideo currently offers: (with more to come)
+
+- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
+- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+- Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
+- Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
+
+Dev in progress and highly experimental.
+
+## 🎥 More Demos
+
+Fast-Hunyuan comparison with original Hunyuan, achieving an 8X diffusion speed boost with the FastVideo framework.
+
+https://github.com/user-attachments/assets/064ac1d2-11ed-4a0c-955b-4d412a96ef30
+
+Comparison between OpenAI Sora, original Hunyuan and FastHunyuan
+
+https://github.com/user-attachments/assets/d323b712-3f68-42b2-952b-94f6a49c4836
+
+Comparison between original FastHunyuan, LLM-INT8 quantized FastHunyuan and NF4 quantized FastHunyuan
+
+https://github.com/user-attachments/assets/cf89efb5-5f68-4949-a085-f41c1ef26c94

 ## Change Log
-
-
- ```2024/12/13```: `FastVideo` v0.0.1 is released.
-
-
-## Fast and High-Quality Text-to-video Generation
-
-<table style="margin-left: auto; margin-right: auto; border: none;">
-  <tr>
-    <td>
-      <img src="assets/8steps/mochi-demo.gif" width="640" alt="Mochi Demo">
-    </td>
-  </tr>
-  <tr>
-    <td style="text-align:center;">
-      Get 8X diffusion boost for Mochi with FastVideo
-    </td>
-  </tr>
-</table>
-
-
-
-## Table of Contents
-
-Jump to a specific section:
-
- [🔧 Installation](#-installation)
- [🚀 Inference](#-inference)
- [🧱 Data Preprocess](#-data-preprocess)
- [🎯 Distill](#-distill)
- [⚡ Finetune](#-lora-finetune)
+- ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
+- ```2024/12/17```: `FastVideo` v1.0 is released.


 ## 🔧 Installation
-
+The code is tested on Python 3.10.0, CUDA 12.1 and H100.
 ```
-conda create -n fastvideo python=3.10.0 -y && conda activate fastvideo
-pip3 install torch==2.5.0 torchvision --index-url https://download.pytorch.org/whl/cu121
-pip install packaging ninja && pip install flash-attn==2.7.0.post2 --no-build-isolation 
-pip install "git+https://github.com/huggingface/diffusers.git@bf64b32652a63a1865a0528a73a13652b201698b"
-git clone https://github.com/hao-ai-lab/FastVideo.git
-cd FastVideo && pip install -e .
+./env_setup.sh fastvideo
 ```

-
-
-
 ## 🚀 Inference

-Use [scripts/download_hf.py](scripts/download_hf.py) to download the hugging-face style model to a local directory. Use it like this:
+### Inference FastHunyuan on single RTX4090
+We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
 ```bash
-python scripts/download_hf.py --repo_id=FastVideo/FastMochi --local_dir=data/FastMochi --repo_type=model
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
+# CLI inference
+bash scripts/inference/inference_diffusers_hunyuan.sh
 ```
+For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):


-Start the gradio UI with
+| Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
+|--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
+| BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
+| INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
+| NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
+           
+
+
+For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
+
+### FastHunyuan
+```bash
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
+# CLI inference
+bash scripts/inference/inference_hunyuan.sh
 ```
-python fastvideo/demo/gradio_web_demo.py --model_path data/FastMochi
+You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).
+
+### FastMochi
+
+```bash
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffusers --local_dir=data/FastMochi-diffusers --repo_type=model
+# CLI inference
+bash scripts/inference/inference_mochi_sp.sh
 ```

-We also provide CLI inference script featured with sequence parallelism.
-
-```
-export NUM_GPUS=4
-
-torchrun --nnodes=1 --nproc_per_node=$NUM_GPUS \
-    fastvideo/sample/sample_t2v_mochi.py \
-    --model_path data/FastMochi \
-    --prompt_path assets/prompt.txt \
-    --num_frames 163 \
-    --height 480 \
-    --width 848 \
-    --num_inference_steps 8 \
-    --guidance_scale 1.5 \
-    --output_path outputs_video/demo_video \
-    --seed 12345 \
-    --scheduler_type "pcm_linear_quadratic" \
-    --linear_threshold 0.1 \
-    --linear_range 0.75
-```
-
-For the mochi style, simply following the scripts list in mochi repo.
-
-```
-git clone https://github.com/genmoai/mochi.git
-cd mochi
-
-# install env
-...
-
-python3 ./demos/cli.py --model_dir weights/ --cpu_offload
-```
-## 🧱 Data Preprocess
-
-To reduce the memory cost and time consumption caused by VAE parsing during the distillation and finetune, we offload the VAE preprocess media part to Data Preprocess section.
-For data preprocess, we need to prepare a source folder for the media we wish to use and a json file for the source information of these media. 
-For image media, the json item needs to follow the following format:
-```
-{
-    "path": "0.jpg",
-    "cap": ["captions"]
-}
-```
-For video media, the json item needs to follow the following format:
-```
-{
-    "path": "1.mp4",
-    "resolution": {
-      "width": 848,
-      "height": 480
-    },
-    "fps": 30.0,
-    "duration": 6.033333333333333,
-    "cap": [
-      "caption"
-    ]
-  }
-```
-Use a txt file to contain the source folder for media and the json file for meta information
-
-```
-path_to_media_source_foder,path_to_json_file
-```
-Adjust the `DATA_MERGE_PATH` and `OUTPUT_DIR` in `./scripts/finetune_data_gen.sh` correspondingly and run
-```
-bash ./scripts/finetune_data_gen.sh
-```
-The preprocessed data will be put into the `OUTPUT_DIR` and the `videos2caption.json` can be used in finetune and distill scripts.

 ## 🎯 Distill
-
-## 💰Hardware requirement
-
-  VRAM is required for both distill 10B mochi model
-
-
-We provide a dataset example here. First download testing data. Use [scripts/download_hf.py](scripts/download_hf.py) to download the data to a local directory. Use it like this:
+Our distillation recipe is based on [Phased Consistency Model](https://github.com/G-U-N/Phased-Consistency-Model). We did not find significant improvement using multi-phase distillation, so we keep the one phase setup similar to the original latent consistency model's recipe.
+We use the [MixKit](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/all_mixkit) dataset for distillation. To avoid running the text encoder and VAE during training, we preprocess all data to generate text embeddings and VAE latents.
+Preprocessing instructions can be found [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide preprocessed data that can be downloaded directly using the following command:
 ```bash
-python scripts/download_hf.py --repo_id=FastVideo/Mochi-425-Data --local_dir=data/Mochi-425-Data --repo_type=dataset
-python scripts/download_hf.py --repo_id=FastVideo/validation_embeddings --local_dir=data/validation_embeddings --repo_type=dataset
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/HD-Mixkit-Finetune-Hunyuan --local_dir=data/HD-Mixkit-Finetune-Hunyuan --repo_type=dataset
 ```
-
-Then the distillation can be launched by:
-
+Next, download the original model weights with:
+```bash
+python scripts/huggingface/download_hf.py --repo_id=genmo/mochi-1-preview --local_dir=data/mochi --repo_type=model # original mochi
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model # original hunyuan
 ```
-bash scripts/distill_t2v.sh
+To launch the distillation process, use the following commands:
 ```
-
-
-## ⚡ Lora Finetune
-
-
-## 💰Hardware requirement
-
-  VRAM is required for both distill 10B mochi model
-
-To launch finetuning, you will first need to prepare data in the following formats.
-
-
-
-Then the finetuning can be launched by:
-
+bash scripts/distill/distill_mochi.sh # for mochi
+bash scripts/distill/distill_hunyuan.sh # for hunyuan
 ```
-bash scripts/lora_finetune.sh
+We also provide an optional script for distillation with adversarial loss, located at `fastvideo/distill_adv.py`. Although we tried adversarial loss, we did not observe significant improvements.
+## Finetune
+### ⚡ Full Finetune
+Ensure your data is prepared and preprocessed in the format specified in [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide a mochi preprocessed Black Myth Wukong data that can be downloaded directly:
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/Mochi-Black-Myth --local_dir=data/Mochi-Black-Myth --repo_type=dataset
 ```
+Download the original model weights as specificed in [Distill Section](#-distill):
+
+Then you can run the finetune with:
+```
+bash scripts/finetune/finetune_mochi.sh # for mochi
+```
+**Note that for finetuning, we did not tune the hyperparameters in the provided script**
+### ⚡ Lora Finetune
+Currently, we only provide Lora Finetune for Mochi model, the command for Lora Finetune is
+```
+bash scripts/finetune/finetune_mochi_lora.sh
+```
+### Minimum Hardware Requirement
+- 40 GB GPU memory each for 2 GPUs with lora
+- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.
+### Finetune with Both Image and Video
+Our codebase support finetuning with both image and video. 
+```bash
+bash scripts/finetune/finetune_hunyuan.sh
+bash scripts/finetune/finetune_mochi_lora_mix.sh
+```
+For Image-Video Mixture Fine-tuning, make sure to enable the --group_frame option in your script.
+
+## 📑 Development Plan
+
+- More distillation methods
+  - [ ] Add Distribution Matching Distillation
+- More models support
+  - [ ] Add CogvideoX model
+- Code update
+  - [ ] fp8 support
+  - [ ] faster load model and save model support

 ## Acknowledgement
-We learned from and reused code from the following projects: [PCM](https://github.com/G-U-N/Phased-Consistency-Model), [diffusers](https://github.com/huggingface/diffusers), and [OpenSoraPlan](https://github.com/PKU-YuanGroup/Open-Sora-Plan).
+We learned and reused code from the following projects: [PCM](https://github.com/G-U-N/Phased-Consistency-Model), [diffusers](https://github.com/huggingface/diffusers), [OpenSoraPlan](https://github.com/PKU-YuanGroup/Open-Sora-Plan), and [xDiT](https://github.com/xdit-project/xDiT).
+
+We thank MBZUAI and Anyscale for their support throughout this project.
@@ -1,9 +1,8 @@
-A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand's movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough.
-A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks. The cars glint in the late afternoon sun, their deep crimson and gold accents lending a touch of elegance. The tracks carve a precarious path along the cliffside, revealing glimpses of a roaring river far below. Inside, passengers peer out the large windows, their faces lit with awe as the landscape unfolds.
-A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background. Strings of fairy lights hang above, casting a warm, golden glow over the scene. Groups of people gather around high tables, their laughter blending with the soft rhythm of live jazz. The aroma of freshly mixed cocktails and charred appetizers wafts through the air, mingling with the cool night breeze.
-En "The Matrix", Neo, interpretado por Keanu Reeves, personifica la lucha contra un sistema opresor a través de su icónica imagen, que incluye unos anteojos oscuros. Estos lentes no son solo un accesorio de moda; representan una barrera entre la realidad y la percepción. Al usar estos anteojos, Neo se sumerge en un mundo donde la verdad se oculta detrás de ilusiones y engaños. La oscuridad de los lentes simboliza la ignorancia y el control que las máquinas tienen sobre la humanidad, mientras que su propia búsqueda de la verdad lo lleva a descubrir sus auténticos poderes. La escena en que se los pone se convierte en un momento crucial, marcando su transformación de un simple programador a "El Elegido". Esta imagen se ha convertido en un ícono cultural, encapsulando el mensaje de que, al enfrentar la oscuridad, podemos encontrar la luz que nos guía hacia la libertad. Así, los anteojos de Neo se convierten en un símbolo de resistencia y autoconocimiento en un mundo manipulado.
-Medium close up. Low-angle shot. A woman in a 1950s retro dress sits in a diner bathed in neon light, surrounded by classic decor and lively chatter. The camera starts with a medium shot of her sitting at the counter, then slowly zooms in as she blows a shiny pink bubblegum bubble. The bubble swells dramatically before popping with a soft, playful burst. The scene is vibrant and nostalgic, evoking the fun and carefree spirit of the 1950s.
-Will Smith eats noodles.
-A short clip of the blonde woman taking a sip from her whiskey glass, her eyes locking with the camera as she smirks playfully. The background shows a group of people laughing and enjoying the party, with vibrant neon signs illuminating the space. The shot is taken in a way that conveys the feeling of a tipsy, carefree night out. The camera then zooms in on her face as she winks, creating a cheeky, flirtatious vibe.
-A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robot's immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1
-A chimpanzee lead vocalist singing into a microphone on stage. The camera zooms in to show him singing. There is a spotlight on him.
+Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.
+A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature.
+A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.
+A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones.
+A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1
+fox in the forest close-up quickly turned its head to the left
+Man walking his dog in the woods on a hot sunny day
+A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.
@@ -0,0 +1,24 @@
+# Configuration for Cog ⚙️
+# Reference: https://cog.run/yaml
+
+build:
+  gpu: true
+  cuda: "12.1"
+  python_version: "3.10"
+  python_packages:
+    - "torch==2.4.0"
+    - "torchvision"
+    - "ninja==1.11.1.3"
+    - "transformers==4.46.1"
+    - "git+https://github.com/huggingface/diffusers.git@bf64b32652a63a1865a0528a73a13652b201698b"
+    - "accelerate==1.0.1"
+    - "safetensors==0.4.5"
+    - "peft==0.13.2"
+    - "packaging==24.2"
+    - "git+https://github.com/hao-ai-lab/FastVideo"
+
+  run:
+    - FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation
+    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/latest/download/pget_$(uname -s)_$(uname -m)" && chmod +x /usr/local/bin/pget
+
+predict: "predict.py:Predictor"
@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
-from fastvideo.model.pipeline_mochi import MochiPipeline
-from fastvideo.model.modeling_mochi import MochiTransformer3DModel
+from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
+from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
 from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import export_to_video
 from fastvideo.distill.solver import PCMFMScheduler
@@ -13,20 +13,20 @@ import argparse
 def init_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--prompts", nargs="+", default=[])
-    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument("--num_frames", type=int, default=25)
    parser.add_argument("--height", type=int, default=480)
    parser.add_argument("--width", type=int, default=848)
-    parser.add_argument("--num_inference_steps", type=int, default=64)
+    parser.add_argument("--num_inference_steps", type=int, default=8)
    parser.add_argument("--guidance_scale", type=float, default=4.5)
    parser.add_argument("--model_path", type=str, default="data/mochi")
-    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--seed", type=int, default=12345)
    parser.add_argument("--transformer_path", type=str, default=None)
-    parser.add_argument("--scheduler_type", type=str, default="euler")
+    parser.add_argument("--scheduler_type", type=str, default="pcm_linear_quadratic")
    parser.add_argument("--lora_checkpoint_dir", type=str, default=None)
    parser.add_argument("--shift", type=float, default=8.0)
-    parser.add_argument("--num_euler_timesteps", type=int, default=100)
-    parser.add_argument("--linear_threshold", type=float, default=0.025)
-    parser.add_argument("--linear_range", type=float, default=0.5)
+    parser.add_argument("--num_euler_timesteps", type=int, default=50)
+    parser.add_argument("--linear_threshold", type=float, default=0.1)
+    parser.add_argument("--linear_range", type=float, default=0.75)
    parser.add_argument("--cpu_offload", action="store_true")
    return parser.parse_args()

@@ -36,11 +36,12 @@ def load_model(args):
    if args.scheduler_type == "euler":
        scheduler = FlowMatchEulerDiscreteScheduler()
    else:
+        linear_quadratic = True if "linear_quadratic" in args.scheduler_type else False
        scheduler = PCMFMScheduler(
            1000,
            args.shift,
            args.num_euler_timesteps,
-            False,
+            linear_quadratic,
            args.linear_threshold,
            args.linear_range,
        )
@@ -56,9 +57,9 @@ def load_model(args):
        args.model_path, transformer=transformer, scheduler=scheduler
    )
    pipe.enable_vae_tiling()
-    pipe.to(device)
-    if args.cpu_offload:
-        pipe.enable_model_cpu_offload()
+    # pipe.to(device)
+    # if args.cpu_offload:
+    pipe.enable_sequential_cpu_offload()
    return pipe


@@ -77,8 +78,6 @@ def generate_video(
    if randomize_seed:
        seed = torch.randint(0, 1000000, (1,)).item()

-    pipe = load_model(args)
-    print("load model successfully")
    generator = torch.Generator(device="cuda").manual_seed(seed)

    if not use_negative_prompt:
@@ -108,9 +107,10 @@ examples = [
 ]

 args = init_args()
-
+pipe = load_model(args)
+print("load model successfully")
 with gr.Blocks() as demo:
-    gr.Markdown("# Mochi Video Generation Demo")
+    gr.Markdown("# Fastvideo Mochi Video Generation Demo")

    with gr.Group():
        with gr.Row():
@@ -141,19 +141,19 @@ with gr.Blocks() as demo:
            with gr.Row():
                num_frames = gr.Slider(
                    label="Number of Frames",
-                    minimum=8,
-                    maximum=256,
+                    minimum=21,
+                    maximum=163,
                    value=args.num_frames,
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
-                    maximum=20,
+                    maximum=12,
                    value=args.guidance_scale,
                )
                num_inference_steps = gr.Slider(
                    label="Inference Steps",
-                    minimum=10,
+                    minimum=4,
                    maximum=100,
                    value=args.num_inference_steps,
                )
@@ -201,4 +201,4 @@ with gr.Blocks() as demo:
    )

 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)
@@ -0,0 +1,68 @@
+
+
+
+## 🧱 Data Preprocess
+
+To save GPU memory, we precompute text embeddings and VAE latents to eliminate the need to load the text encoder and VAE during training.
+
+
+We provide a sample dataset to help you get started. Download the source media using the following command:
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/Image-Vid-Finetune-Src --local_dir=data/Image-Vid-Finetune-Src --repo_type=dataset
+```
+To preprocess the dataset for fine-tuning or distillation, run:
+```
+bash scripts/preprocess/preprocess_mochi_data.sh # for mochi
+bash scripts/preprocess/preprocess_hunyuan_data.sh # for hunyuan
+```
+
+The preprocessed dataset will be stored in `Image-Vid-Finetune-Mochi` or `Image-Vid-Finetune-HunYuan` correspondingly.
+
+### Process your own dataset
+
+If you wish to create your own dataset for finetuning or distillation, please structure you video dataset in the following format:
+
+path_to_dataset_folder/
+├── media/
+│   ├── 0.jpg
+│   ├── 1.mp4
+│   ├── 2.jpg
+├── video2caption.json
+└── merge.txt
+
+Format the JSON file as a list, where each item represents a media source:
+
+For image media,
+```
+{
+    "path": "0.jpg",
+    "cap": ["captions"]
+}
+```
+For video media, 
+```
+{
+    "path": "1.mp4",
+    "resolution": {
+      "width": 848,
+      "height": 480
+    },
+    "fps": 30.0,
+    "duration": 6.033333333333333,
+    "cap": [
+      "caption"
+    ]
+  }
+```
+
+Use a txt file (merge.txt) to contain the source folder for media and the JSON file for meta information:
+
+```
+path_to_media_source_foder,path_to_json_file
+```
+
+Adjust the `DATA_MERGE_PATH` and `OUTPUT_DIR` in `scripts/preprocess/preprocess_****_data.sh` accordingly and run:
+```
+bash scripts/preprocess/preprocess_****_data.sh
+```
+The preprocessed data will be put into the `OUTPUT_DIR` and the `videos2caption.json` can be used in finetune and distill scripts.
@@ -0,0 +1,10 @@
+#!/bin/bash
+
+# install torch
+pip install torch==2.5.0 torchvision --index-url https://download.pytorch.org/whl/cu121
+
+# install FA2 and diffusers
+pip install packaging ninja && pip install flash-attn==2.7.0.post2 --no-build-isolation 
+
+# install fastvideo
+pip install -e .
@@ -18,9 +18,7 @@ from tqdm import tqdm

 class T5dataset(Dataset):
    def __init__(
-        self,
-        json_path,
-        vae_debug,
+        self, json_path, vae_debug,
    ):
        self.json_path = json_path
        self.vae_debug = vae_debug
@@ -67,9 +65,7 @@ def main(args):
    os.makedirs(os.path.join(args.output_dir, "prompt_embed"), exist_ok=True)
    os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True)

-    latents_json_path = os.path.join(
-        args.output_dir, "videos2caption_temp_replace.json"
-    )
+    latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
    train_dataset = T5dataset(latents_json_path, args.vae_debug)
    text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
    vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
@@ -38,7 +38,7 @@ def main(args):
        dist.init_process_group(
            backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
        )
-    vae, autocast_type = load_vae(args.model_type, args.model_path)
+    vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
    vae.enable_tiling()
    os.makedirs(args.output_dir, exist_ok=True)
    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
@@ -7,10 +7,7 @@ import random

 class LatentDataset(Dataset):
    def __init__(
-        self,
-        json_path,
-        num_latent_t,
-        cfg_rate,
+        self, json_path, num_latent_t, cfg_rate,
    ):
        # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path
        self.json_path = json_path
@@ -288,10 +288,7 @@ class LongSideResizeVideo:
    """

    def __init__(
-        self,
-        size,
-        skip_low_resolution=False,
-        interpolation_mode="bilinear",
+        self, size, skip_low_resolution=False, interpolation_mode="bilinear",
    ):
        self.size = size
        self.skip_low_resolution = skip_low_resolution
@@ -330,10 +327,7 @@ class CenterCropResizeVideo:
    """

    def __init__(
-        self,
-        size,
-        top_crop=False,
-        interpolation_mode="bilinear",
+        self, size, top_crop=False, interpolation_mode="bilinear",
    ):
        if len(size) != 2:
            raise ValueError(
@@ -374,9 +368,7 @@ class UCFCenterCropVideo:
    """

    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
+        self, size, interpolation_mode="bilinear",
    ):
        if isinstance(size, tuple):
            if len(size) != 2:
@@ -413,9 +405,7 @@ class KineticsRandomCropResizeVideo:
    """

    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
+        self, size, interpolation_mode="bilinear",
    ):
        if isinstance(size, tuple):
            if len(size) != 2:
@@ -436,9 +426,7 @@ class KineticsRandomCropResizeVideo:

 class CenterCropVideo:
    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
+        self, size, interpolation_mode="bilinear",
    ):
        if isinstance(size, tuple):
            if len(size) != 2:
@@ -27,10 +27,8 @@ import wandb
 from accelerate.utils import set_seed
 from tqdm.auto import tqdm
 from fastvideo.utils.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing
-from diffusers import (
-    FlowMatchEulerDiscreteScheduler,
-)
-from fastvideo.utils.load import  load_transformer
+from diffusers import FlowMatchEulerDiscreteScheduler
+from fastvideo.utils.load import load_transformer
 from fastvideo.distill.solver import EulerSolver, extract_into_tensor
 from copy import deepcopy
 from diffusers.optimization import get_scheduler
@@ -39,9 +37,7 @@ from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_func
 import torch.distributed as dist
 from safetensors.torch import save_file
 from peft import LoraConfig
-from torch.distributed.fsdp import (
-    FullyShardedDataParallel as FSDP,
-)
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from fastvideo.utils.checkpoint import (
    save_checkpoint,
    save_lora_checkpoint,
@@ -242,7 +238,7 @@ def distill_one_step(
        # loss = loss.mean()
        loss = (
            torch.mean(
-                torch.sqrt((model_pred.float() - target.float()) ** 2 + huber_c**2)
+                torch.sqrt((model_pred.float() - target.float()) ** 2 + huber_c ** 2)
                - huber_c
            )
            / gradient_accumulation_steps
@@ -368,19 +364,10 @@ def main(args):
        transformer._no_split_modules = no_split_modules
        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)

-    transformer = FSDP(
-        transformer,
-        **fsdp_kwargs,
-    )
-    teacher_transformer = FSDP(
-        teacher_transformer,
-        **fsdp_kwargs,
-    )
+    transformer = FSDP(transformer, **fsdp_kwargs,)
+    teacher_transformer = FSDP(teacher_transformer, **fsdp_kwargs,)
    if args.use_ema:
-        ema_transformer = FSDP(
-            ema_transformer,
-            **fsdp_kwargs,
-        )
+        ema_transformer = FSDP(ema_transformer, **fsdp_kwargs,)
    main_print(f"--> model loaded")

    if args.gradient_checkpointing:
@@ -58,11 +58,7 @@ class DiscriminatorHead(nn.Module):

 class Discriminator(nn.Module):
    def __init__(
-        self,
-        stride=8,
-        num_h_per_head=1,
-        adapter_channel_dims=[3072],
-        total_layers = 48,
+        self, stride=8, num_h_per_head=1, adapter_channel_dims=[3072], total_layers=48,
    ):
        super().__init__()
        adapter_channel_dims = adapter_channel_dims * (total_layers // stride)
@@ -275,12 +275,7 @@ class EulerSolver:
        return x_prev

    def euler_style_multiphase_pred(
-        self,
-        sample,
-        model_pred,
-        timestep_index,
-        multiphase,
-        is_target=False,
+        self, sample, model_pred, timestep_index, multiphase, is_target=False,
    ):
        inference_indices = np.linspace(
            0, len(self.euler_timesteps), num=multiphase, endpoint=False
@@ -22,7 +22,7 @@ from torch.distributed.fsdp import (
    StateDictType,
    FullStateDictConfig,
 )
-from fastvideo.utils.load import  load_transformer
+from fastvideo.utils.load import load_transformer

 from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
 import json
@@ -37,9 +37,7 @@ from fastvideo.utils.fsdp_util import (
    get_discriminator_fsdp_kwargs,
 )
 import diffusers
-from diffusers import (
-    FlowMatchEulerDiscreteScheduler,
-)
+from diffusers import FlowMatchEulerDiscreteScheduler
 from fastvideo.distill.discriminator import Discriminator
 from fastvideo.distill.solver import EulerSolver, extract_into_tensor
 from copy import deepcopy
@@ -49,9 +47,7 @@ from diffusers.utils import check_min_version
 from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_function
 import torch.distributed as dist
 from peft import LoraConfig
-from torch.distributed.fsdp import (
-    FullyShardedDataParallel as FSDP,
-)
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from fastvideo.utils.checkpoint import (
    save_checkpoint,
    save_lora_checkpoint,
@@ -77,7 +73,7 @@ def gan_d_loss(
    encoder_hidden_states,
    encoder_attention_mask,
    weight,
-    discriminator_head_stride
+    discriminator_head_stride,
 ):
    loss = 0.0
    # collate sample_fake and sample_real
@@ -119,7 +115,7 @@ def gan_g_loss(
    encoder_hidden_states,
    encoder_attention_mask,
    weight,
-    discriminator_head_stride
+    discriminator_head_stride,
 ):
    loss = 0.0
    features = teacher_transformer(
@@ -131,9 +127,7 @@ def gan_g_loss(
        output_features_stride=discriminator_head_stride,
        return_dict=False,
    )[1]
-    fake_outputs = discriminator(
-        features,
-    )
+    fake_outputs = discriminator(features,)
    for fake_output in fake_outputs:
        loss += torch.mean(weight * torch.relu(1 - fake_output.float())) / (
            discriminator.head_num * discriminator.num_h_per_head
@@ -162,7 +156,7 @@ def distill_one_step_adv(
    not_apply_cfg_solver,
    distill_cfg,
    adv_weight,
-    discriminator_head_stride
+    discriminator_head_stride,
 ):
    optimizer.zero_grad()
    discriminator_optimizer.zero_grad()
@@ -278,7 +272,7 @@ def distill_one_step_adv(

    huber_c = 0.001
    g_loss = torch.mean(
-        torch.sqrt((model_pred.float() - target.float()) ** 2 + huber_c**2) - huber_c
+        torch.sqrt((model_pred.float() - target.float()) ** 2 + huber_c ** 2) - huber_c
    )
    discriminator.requires_grad_(False)
    with torch.autocast("cuda", dtype=torch.bfloat16):
@@ -290,7 +284,7 @@ def distill_one_step_adv(
            encoder_hidden_states.float(),
            encoder_attention_mask,
            1.0,
-            discriminator_head_stride
+            discriminator_head_stride,
        )
    g_loss += g_gan_loss
    g_loss.backward()
@@ -362,7 +356,10 @@ def main(args):
        torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16,
    )
    teacher_transformer = deepcopy(transformer)
-    discriminator = Discriminator(args.discriminator_head_stride, total_layers = 48 if args.model_type =="mochi" else 40)
+    discriminator = Discriminator(
+        args.discriminator_head_stride,
+        total_layers=48 if args.model_type == "mochi" else 40,
+    )

    if args.use_lora:
        transformer.requires_grad_(False)
@@ -400,18 +397,9 @@ def main(args):
        transformer._no_split_modules = no_split_modules
        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)

-    transformer = FSDP(
-        transformer,
-        **fsdp_kwargs,
-    )
-    teacher_transformer = FSDP(
-        teacher_transformer,
-        **fsdp_kwargs,
-    )
-    discriminator = FSDP(
-        discriminator,
-        **discriminator_fsdp_kwargs,
-    )
+    transformer = FSDP(transformer, **fsdp_kwargs,)
+    teacher_transformer = FSDP(teacher_transformer, **fsdp_kwargs,)
+    discriminator = FSDP(discriminator, **discriminator_fsdp_kwargs,)
    main_print(f"--> model loaded")

    if args.gradient_checkpointing:
@@ -582,6 +570,7 @@ def main(args):
            if step <= int(phase_step):
                return int(phase)
        return phase
+
    for i in range(init_steps):
        _ = next(loader)
    for step in range(init_steps + 1, args.max_train_steps + 1):
@@ -614,7 +603,7 @@ def main(args):
            args.not_apply_cfg_solver,
            args.distill_cfg,
            args.adv_weight,
-            args.discriminator_head_stride
+            args.discriminator_head_stride,
        )

        step_time = time.time() - start_time
@@ -653,7 +642,7 @@ def main(args):
                )
            else:
                # Your existing checkpoint saving code
-                # TODO 
+                # TODO
                # save_checkpoint_generator_discriminator(
                #     transformer,
                #     optimizer,
@@ -663,7 +652,9 @@ def main(args):
                #     args.output_dir,
                #     step,
                # )
-                save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
+                save_checkpoint(
+                    transformer, rank, args.output_dir, args.max_train_steps
+                )
            main_print(f"--> checkpoint saved at step {step}")
            dist.barrier()
        if args.log_validation and step % args.validation_steps == 0:
@@ -680,7 +671,6 @@ def main(args):
                linear_range=args.linear_range,
                ema=False,
            )
-            

    if args.use_lora:
        save_lora_checkpoint(
@@ -48,10 +48,7 @@ PROMPT_TEMPLATE_ENCODE_VIDEO = (
 NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion"

 PROMPT_TEMPLATE = {
-    "dit-llm-encode": {
-        "template": PROMPT_TEMPLATE_ENCODE,
-        "crop_start": 36,
-    },
+    "dit-llm-encode": {"template": PROMPT_TEMPLATE_ENCODE, "crop_start": 36,},
    "dit-llm-encode-video": {
        "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
        "crop_start": 95,
@@ -876,8 +876,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):

        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_func_kwargs(
-            self.scheduler.step,
-            {"generator": generator, "eta": eta},
+            self.scheduler.step, {"generator": generator, "eta": eta},
        )

        target_dtype = PRECISION_TO_TYPE[self.args.precision]
@@ -195,10 +195,7 @@ def add_denoise_schedule_args(parser: argparse.ArgumentParser):
        help="If reverse, learning/sampling from t=1 -> t=0.",
    )
    group.add_argument(
-        "--flow-solver",
-        type=str,
-        default="euler",
-        help="Solver for flow matching.",
+        "--flow-solver", type=str, default="euler", help="Solver for flow matching.",
    )
    group.add_argument(
        "--use-linear-quadratic-schedule",
@@ -360,16 +357,10 @@ def add_parallel_args(parser: argparse.ArgumentParser):

    # ======================== Model loads ========================
    group.add_argument(
-        "--ulysses-degree",
-        type=int,
-        default=1,
-        help="Ulysses degree.",
+        "--ulysses-degree", type=int, default=1, help="Ulysses degree.",
    )
    group.add_argument(
-        "--ring-degree",
-        type=int,
-        default=1,
-        help="Ulysses degree.",
+        "--ring-degree", type=int, default=1, help="Ulysses degree.",
    )

    return parser
@@ -20,7 +20,7 @@ from fastvideo.models.hunyuan.text_encoder import TextEncoder
 from fastvideo.models.hunyuan.utils.data_utils import align_to
 from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler
 from fastvideo.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline
-
+from safetensors.torch import load_file as safetensors_load_file
 from fastvideo.utils.parallel_states import (
    initialize_sequence_parallel_state,
    nccl_info,
@@ -56,7 +56,9 @@ class Inference(object):
        self.device = (
            device
            if device is not None
-            else "cuda" if torch.cuda.is_available() else "cpu"
+            else "cuda"
+            if torch.cuda.is_available()
+            else "cpu"
        )
        self.logger = logger
        self.parallel_args = parallel_args
@@ -238,7 +240,16 @@ class Inference(object):
        if not model_path.exists():
            raise ValueError(f"model_path not exists: {model_path}")
        logger.info(f"Loading torch model {model_path}...")
-        state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
+        if model_path.suffix == ".safetensors":
+            # Use safetensors library for .safetensors files
+            state_dict = safetensors_load_file(model_path)
+        elif model_path.suffix == ".pt":
+            # Use torch for .pt files
+            state_dict = torch.load(
+                model_path, map_location=lambda storage, loc: storage
+            )
+        else:
+            raise ValueError(f"Unsupported file format: {model_path}")

        if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict):
            bare_model = False
@@ -414,7 +425,8 @@ class HunyuanVideoSampler(Inference):
            raise ValueError(
                f"Seed must be an integer, a list of integers, or None, got {seed}."
            )
-        generator = [torch.Generator(self.device).manual_seed(seed) for seed in seeds]
+        # Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
+        generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds]
        out_dict["seeds"] = seeds

        # ========================================================================
@@ -12,12 +12,7 @@ from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad


 def attention(
-    q,
-    k,
-    v,
-    drop_rate=0,
-    attn_mask=None,
-    causal=False,
+    q, k, v, drop_rate=0, attn_mask=None, causal=False,
 ):

    qkv = torch.stack([q, k, v], dim=2)
@@ -18,9 +18,7 @@ from .modulate_layers import ModulateDiT, modulate, apply_gate
 from .token_refiner import SingleTokenRefiner
 from fastvideo.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_embed

-from fastvideo.utils.parallel_states import (
-    nccl_info,
-)
+from fastvideo.utils.parallel_states import nccl_info


 class MMDoubleStreamBlock(nn.Module):
@@ -272,7 +270,7 @@ class MMSingleStreamBlock(nn.Module):
        head_dim = hidden_size // heads_num
        mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
        self.mlp_hidden_dim = mlp_hidden_dim
-        self.scale = qk_scale or head_dim**-0.5
+        self.scale = qk_scale or head_dim ** -0.5

        # qkv and mlp_in
        self.linear1 = nn.Linear(
@@ -602,7 +600,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
        timestep: torch.LongTensor,
        encoder_attention_mask: torch.Tensor,
        output_features=False,
-        output_features_stride = 8,
+        output_features_stride=8,
        attention_kwargs: Optional[Dict[str, Any]] = None,
        return_dict: bool = False,
        guidance=None,
@@ -135,10 +135,7 @@ class IndividualTokenRefiner(nn.Module):
        )

    def forward(
-        self,
-        x: torch.Tensor,
-        c: torch.LongTensor,
-        mask: Optional[torch.Tensor] = None,
+        self, x: torch.Tensor, c: torch.LongTensor, mask: Optional[torch.Tensor] = None,
    ):
        mask = mask.clone().bool()
        # avoid attention weight become NaN
@@ -10,9 +10,7 @@ def preprocess_text_encoder_tokenizer(args):

    processor = AutoProcessor.from_pretrained(args.input_dir)
    model = LlavaForConditionalGeneration.from_pretrained(
-        args.input_dir,
-        torch_dtype=torch.float16,
-        low_cpu_mem_usage=True,
+        args.input_dir, torch_dtype=torch.float16, low_cpu_mem_usage=True,
    ).to(0)

    model.language_model.save_pretrained(f"{args.output_dir}")
@@ -229,54 +229,54 @@ def convert_diffusers_vae_to_mochi(state_dict):
    # Convert down_blocks
    down_block_layers = [3, 4, 6]
    for block in range(3):
-        encoder_state_dict[f"layers.{block+4}.layers.0.weight"] = (
-            original_state_dict.pop(f"{prefix}down_blocks.{block}.conv_in.conv.weight")
-        )
+        encoder_state_dict[
+            f"layers.{block+4}.layers.0.weight"
+        ] = original_state_dict.pop(f"{prefix}down_blocks.{block}.conv_in.conv.weight")
        encoder_state_dict[f"layers.{block+4}.layers.0.bias"] = original_state_dict.pop(
            f"{prefix}down_blocks.{block}.conv_in.conv.bias"
        )

        for i in range(down_block_layers[block]):
            # Convert resnets
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.weight"] = (
-                original_state_dict.pop(
-                    f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"
-                )
+            encoder_state_dict[
+                f"layers.{block+4}.layers.{i+1}.stack.0.weight"
+            ] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"
            )
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.bias"] = (
-                original_state_dict.pop(
-                    f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"
-                )
+            encoder_state_dict[
+                f"layers.{block+4}.layers.{i+1}.stack.0.bias"
+            ] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"
            )
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.weight"] = (
-                original_state_dict.pop(
-                    f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight"
-                )
+            encoder_state_dict[
+                f"layers.{block+4}.layers.{i+1}.stack.2.weight"
+            ] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight"
            )
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.bias"] = (
-                original_state_dict.pop(
-                    f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias"
-                )
+            encoder_state_dict[
+                f"layers.{block+4}.layers.{i+1}.stack.2.bias"
+            ] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias"
            )
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.weight"] = (
-                original_state_dict.pop(
-                    f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"
-                )
+            encoder_state_dict[
+                f"layers.{block+4}.layers.{i+1}.stack.3.weight"
+            ] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"
            )
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.bias"] = (
-                original_state_dict.pop(
-                    f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"
-                )
+            encoder_state_dict[
+                f"layers.{block+4}.layers.{i+1}.stack.3.bias"
+            ] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"
            )
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.weight"] = (
-                original_state_dict.pop(
-                    f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight"
-                )
+            encoder_state_dict[
+                f"layers.{block+4}.layers.{i+1}.stack.5.weight"
+            ] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight"
            )
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.bias"] = (
-                original_state_dict.pop(
-                    f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias"
-                )
+            encoder_state_dict[
+                f"layers.{block+4}.layers.{i+1}.stack.5.bias"
+            ] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias"
            )

            # Convert attentions
@@ -348,18 +348,18 @@ def convert_diffusers_vae_to_mochi(state_dict):
        qkv_weight = torch.cat([q, k, v], dim=0)
        encoder_state_dict[f"layers.{i+7}.attn_block.attn.qkv.weight"] = qkv_weight

-        encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.weight"] = (
-            original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_out.0.weight")
-        )
-        encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.bias"] = (
-            original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_out.0.bias")
-        )
-        encoder_state_dict[f"layers.{i+7}.attn_block.norm.weight"] = (
-            original_state_dict.pop(f"{prefix}block_out.norms.{i}.norm_layer.weight")
-        )
-        encoder_state_dict[f"layers.{i+7}.attn_block.norm.bias"] = (
-            original_state_dict.pop(f"{prefix}block_out.norms.{i}.norm_layer.bias")
-        )
+        encoder_state_dict[
+            f"layers.{i+7}.attn_block.attn.out.weight"
+        ] = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_out.0.weight")
+        encoder_state_dict[
+            f"layers.{i+7}.attn_block.attn.out.bias"
+        ] = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_out.0.bias")
+        encoder_state_dict[
+            f"layers.{i+7}.attn_block.norm.weight"
+        ] = original_state_dict.pop(f"{prefix}block_out.norms.{i}.norm_layer.weight")
+        encoder_state_dict[
+            f"layers.{i+7}.attn_block.norm.bias"
+        ] = original_state_dict.pop(f"{prefix}block_out.norms.{i}.norm_layer.bias")

    # Convert output layers
    encoder_state_dict["output_norm.weight"] = original_state_dict.pop(
@@ -413,45 +413,45 @@ def convert_diffusers_vae_to_mochi(state_dict):
    up_block_layers = [6, 4, 3]
    for block in range(3):
        for i in range(up_block_layers[block]):
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.weight"] = (
-                original_state_dict.pop(
-                    f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"
-                )
+            decoder_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.0.weight"
+            ] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"
            )
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.bias"] = (
-                original_state_dict.pop(
-                    f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"
-                )
+            decoder_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.0.bias"
+            ] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"
            )
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.weight"] = (
-                original_state_dict.pop(
-                    f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight"
-                )
+            decoder_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.2.weight"
+            ] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight"
            )
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.bias"] = (
-                original_state_dict.pop(
-                    f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias"
-                )
+            decoder_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.2.bias"
+            ] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias"
            )
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.weight"] = (
-                original_state_dict.pop(
-                    f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"
-                )
+            decoder_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.3.weight"
+            ] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"
            )
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.bias"] = (
-                original_state_dict.pop(
-                    f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"
-                )
+            decoder_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.3.bias"
+            ] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"
            )
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.weight"] = (
-                original_state_dict.pop(
-                    f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight"
-                )
+            decoder_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.5.weight"
+            ] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight"
            )
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.bias"] = (
-                original_state_dict.pop(
-                    f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias"
-                )
+            decoder_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.5.bias"
+            ] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias"
            )
        decoder_state_dict[f"blocks.{block+1}.proj.weight"] = original_state_dict.pop(
            f"{prefix}up_blocks.{block}.proj.weight"
@@ -622,7 +622,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
        timestep: torch.LongTensor,
        encoder_attention_mask: torch.Tensor,
        output_features=False,
-        output_features_stride = 8,
+        output_features_stride=8,
        attention_kwargs: Optional[Dict[str, Any]] = None,
        return_dict: bool = False,
    ) -> torch.Tensor:
@@ -67,11 +67,7 @@ class MochiRMSNorm(nn.Module):

 class MochiLayerNormContinuous(nn.Module):
    def __init__(
-        self,
-        embedding_dim: int,
-        conditioning_embedding_dim: int,
-        eps=1e-5,
-        bias=True,
+        self, embedding_dim: int, conditioning_embedding_dim: int, eps=1e-5, bias=True,
    ):
        super().__init__()

@@ -81,9 +77,7 @@ class MochiLayerNormContinuous(nn.Module):
        self.norm = MochiModulatedRMSNorm(eps=eps)

    def forward(
-        self,
-        x: torch.Tensor,
-        conditioning_embedding: torch.Tensor,
+        self, x: torch.Tensor, conditioning_embedding: torch.Tensor,
    ) -> torch.Tensor:
        input_dtype = x.dtype

@@ -36,7 +36,7 @@ from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput
 from einops import rearrange
 from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
 from fastvideo.utils.communications import all_gather
-
+from diffusers.loaders import Mochi1LoraLoaderMixin
 if is_torch_xla_available():
    import torch_xla.core.xla_model as xm

@@ -85,13 +85,13 @@ def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
    ]
    threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
    quadratic_steps = num_steps - linear_steps
-    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
+    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps ** 2)
    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (
-        quadratic_steps**2
+        quadratic_steps ** 2
    )
-    const = quadratic_coef * (linear_steps**2)
+    const = quadratic_coef * (linear_steps ** 2)
    quadratic_sigma_schedule = [
-        quadratic_coef * (i**2) + linear_coef * i + const
+        quadratic_coef * (i ** 2) + linear_coef * i + const
        for i in range(linear_steps, num_steps)
    ]
    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule
@@ -165,7 +165,7 @@ def retrieve_timesteps(
    return timesteps, num_inference_steps


-class MochiPipeline(DiffusionPipeline):
+class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
    r"""
    The mochi pipeline for text-to-video generation.

@@ -502,7 +502,8 @@ class MochiPipeline(DiffusionPipeline):
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        latents = latents.to(dtype)
        return latents

    @property
@@ -533,8 +534,8 @@ class MochiPipeline(DiffusionPipeline):
        negative_prompt: Optional[Union[str, List[str]]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
-        num_frames: int = 16,
-        num_inference_steps: int = 28,
+        num_frames: int = 19,
+        num_inference_steps: int = 64,
        timesteps: List[int] = None,
        guidance_scale: float = 4.5,
        num_videos_per_prompt: Optional[int] = 1,
@@ -711,17 +712,11 @@ class MochiPipeline(DiffusionPipeline):
        # check if of type FlowMatchEulerDiscreteScheduler
        if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
            timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler,
-                num_inference_steps,
-                device,
-                timesteps,
-                sigmas,
+                self.scheduler, num_inference_steps, device, timesteps, sigmas,
            )
        else:
            timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler,
-                num_inference_steps,
-                device,
+                self.scheduler, num_inference_steps, device,
            )
        num_warmup_steps = max(
            len(timesteps) - num_inference_steps * self.scheduler.order, 0
@@ -729,6 +724,7 @@ class MochiPipeline(DiffusionPipeline):
        self._num_timesteps = len(timesteps)

        # 6. Denoising loop
+        self._progress_bar_config = {"disable": nccl_info.rank_within_group != 0}
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                if self.interrupt:
@@ -1,74 +1,105 @@
-import os
-import imageio
-import time
-from einops import rearrange
-
 import torch
-import torchvision
+from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel, BitsAndBytesConfig
+import imageio as iio 
+import math
 import numpy as np
-from pathlib import Path
-from loguru import logger
-from datetime import datetime
+import io
+import time
 import argparse
-from diffusers.utils import export_to_video
+import os

-from fastvideo.models.hunyuan.utils.file_utils import save_videos_grid
-from fastvideo.models.hunyuan.inference import HunyuanVideoSampler
+def export_to_video_bytes(fps, frames):
+    request = iio.core.Request("<bytes>", mode="w", extension=".mp4")
+    pyavobject = iio.plugins.pyav.PyAVPlugin(request)
+    if isinstance(frames, np.ndarray):
+        frames = (np.array(frames) * 255).astype('uint8')
+    else:
+        frames = np.array(frames)
+    new_bytes = pyavobject.write(frames, codec="libx264", fps=fps)
+    out_bytes = io.BytesIO(new_bytes)
+    return out_bytes

+def export_to_video(frames, path, fps):
+    video_bytes = export_to_video_bytes(fps, frames)
+    video_bytes.seek(0)
+    with open(path, "wb") as f:
+        f.write(video_bytes.getbuffer())

 def main(args):
-    print(args)
-    models_root_path = Path(args.model_path)
-    if not models_root_path.exists():
-        raise ValueError(f"`models_root` not exists: {models_root_path}")
+    torch.manual_seed(args.seed)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    prompt_template = {
+        "template": (
+            "<|start_header_cid|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
+            "1. The main content and theme of the video."
+            "2. The color, shape, size, texture, quantity, text, and spatial relationships of the contents, including objects, people, and anything else."
+            "3. Actions, events, behaviors temporal relationships, physical movement changes of the contents."
+            "4. Background environment, light, style, atmosphere, and qualities."
+            "5. Camera angles, movements, and transitions used in the video."
+            "6. Thematic and aesthetic concepts associated with the scene, i.e. realistic, futuristic, fairy tale, etc<|eot_id|>"
+            "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"
+        ),
+        "crop_start": 95,
+    }
+    
+    
+    model_id = args.model_path

-    # Create save folder to save the samples
-    save_path = args.output_path
-    os.makedirs(os.path.dirname(save_path), exist_ok=True)
-
-    # Load models
-    hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(
-        models_root_path, args=args
-    )
-
-    # Get the updated args
-    args = hunyuan_video_sampler.args
-
-    # Start sampling
-    samples = []
-    for prompt in args.prompts:
-        outputs = hunyuan_video_sampler.predict(
-            prompt=prompt,
-            height=args.height,
-            width=args.width,
-            video_length=args.num_frames,
-            seed=args.seed,
-            negative_prompt=args.neg_prompt,
-            infer_steps=args.num_inference_steps,
-            guidance_scale=args.guidance_scale,
-            num_videos_per_prompt=args.num_videos,
-            flow_shift=args.flow_shift,
-            batch_size=args.batch_size,
-            embedded_guidance_scale=args.embedded_cfg_scale,
+    if args.quantization == "nf4":
+        quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_quant_type="nf4", llm_int8_skip_modules=["proj_out", "norm_out"])
+        transformer = HunyuanVideoTransformer3DModel.from_pretrained(
+            model_id, subfolder="transformer/" ,torch_dtype=torch.bfloat16, quantization_config=quantization_config
        )
-        samples.append(outputs["samples"][0])
-
-    for prompt, video in zip(args.prompts, samples):
-        videos = rearrange(video.unsqueeze(0), "b c t h w -> t b c h w")
-        outputs = []
-        for x in videos:
-            x = torchvision.utils.make_grid(x, nrow=6)
-            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-            outputs.append((x * 255).numpy().astype(np.uint8))
-        os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
-        imageio.mimsave(args.output_path + f"{prompt[:100]}.mp4", outputs, fps=args.fps)
-
+    if args.quantization == "int8":
+        quantization_config = BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=["proj_out", "norm_out"])
+        transformer = HunyuanVideoTransformer3DModel.from_pretrained(
+            model_id, subfolder="transformer/" ,torch_dtype=torch.bfloat16, quantization_config=quantization_config
+        )
+    elif not args.quantization:
+        transformer = HunyuanVideoTransformer3DModel.from_pretrained(
+            model_id, subfolder="transformer/" ,torch_dtype=torch.bfloat16
+        ).to(device)
+    
+    print("Max vram for read transofrmer:", round(torch.cuda.max_memory_allocated(device="cuda") / 1024 ** 3, 3), "GiB")
+    torch.cuda.reset_max_memory_allocated(device)
+    
+    if not args.cpu_offload:
+        pipe = HunyuanVideoPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(device)
+        pipe.transformer = transformer
+    else:
+        pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.bfloat16)
+    torch.cuda.reset_max_memory_allocated(device)
+    pipe.scheduler._shift = args.flow_shift
+    pipe.vae.enable_tiling()
+    if args.cpu_offload:
+        pipe.enable_model_cpu_offload()
+    print("Max vram for init pipeline:", round(torch.cuda.max_memory_allocated(device="cuda") / 1024 ** 3, 3), "GiB")
+    with open(args.prompt) as f:
+        prompts = f.readlines()
+    
+    generator = torch.Generator("cpu").manual_seed(args.seed)
+    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
+    torch.cuda.reset_max_memory_allocated(device)
+    for prompt in prompts:
+        start_time = time.perf_counter()
+        output = pipe(
+            prompt=prompt,
+            height = args.height,
+            width = args.width,
+            num_frames = args.num_frames,
+            prompt_template=prompt_template,
+            num_inference_steps = args.num_inference_steps,
+            generator=generator,
+        ).frames[0]
+        export_to_video(output, os.path.join(args.output_path, f"{prompt[:100]}.mp4"), fps=args.fps)
+        print("Time:", round(time.perf_counter() - start_time, 2), "seconds")
+        print("Max vram for denoise:", round(torch.cuda.max_memory_allocated(device="cuda") / 1024 ** 3, 3), "GiB")

 if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    # Basic parameters
-    parser.add_argument("--prompts", nargs="+", default=[])
+    parser.add_argument("--prompt", type=str, help="prompt file for inference")
    parser.add_argument("--num_frames", type=int, default=16)
    parser.add_argument("--height", type=int, default=256)
    parser.add_argument("--width", type=int, default=256)
@@ -76,7 +107,8 @@ if __name__ == "__main__":
    parser.add_argument("--model_path", type=str, default="data/hunyuan")
    parser.add_argument("--output_path", type=str, default="./outputs/video")
    parser.add_argument("--fps", type=int, default=24)
-
+    parser.add_argument("--quantization", type=str, default=None)
+    parser.add_argument("--cpu_offload", action="store_true")
    # Additional parameters
    parser.add_argument(
        "--denoise-type",
@@ -164,7 +196,7 @@ if __name__ == "__main__":
    parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill")
    parser.add_argument("--latent-channels", type=int, default=16)
    parser.add_argument(
-        "--precision", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]
+        "--precision", type=str, default="bf16", choices=["fp32", "fp16", "bf16", "fp8"]
    )
    parser.add_argument(
        "--rope-theta", type=int, default=256, help="Theta used in RoPE."
@@ -205,4 +237,4 @@ if __name__ == "__main__":
    parser.add_argument("--text-len-2", type=int, default=77)

    args = parser.parse_args()
-    main(args)
+    main(args)
@@ -58,7 +58,11 @@ def main(args):

    # Start sampling
    samples = []
-    for prompt in args.prompts:
+
+    with open(args.prompt) as f:
+        prompts = f.readlines()
+
+    for prompt in prompts:
        outputs = hunyuan_video_sampler.predict(
            prompt=prompt,
            height=args.height,
@@ -73,24 +77,23 @@ def main(args):
            batch_size=args.batch_size,
            embedded_guidance_scale=args.embedded_cfg_scale,
        )
-        samples.append(outputs["samples"][0])
-
-    for prompt, video in zip(args.prompts, samples):
-        videos = rearrange(video.unsqueeze(0), "b c t h w -> t b c h w")
+        videos = rearrange(outputs["samples"], "b c t h w -> t b c h w")
        outputs = []
        for x in videos:
            x = torchvision.utils.make_grid(x, nrow=6)
            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
            outputs.append((x * 255).numpy().astype(np.uint8))
        os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
-        imageio.mimsave(args.output_path + f"{prompt[:100]}.mp4", outputs, fps=args.fps)
+        imageio.mimsave(
+            os.path.join(args.output_path, f"{prompt[:100]}.mp4"), outputs, fps=args.fps
+        )


 if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    # Basic parameters
-    parser.add_argument("--prompts", nargs="+", default=[])
+    parser.add_argument("--prompt", type=str, help="prompt file for inference")
    parser.add_argument("--num_frames", type=int, default=16)
    parser.add_argument("--height", type=int, default=256)
    parser.add_argument("--width", type=int, default=256)
@@ -39,7 +39,7 @@ def main(args):
    initialize_distributed()
    print(nccl_info.sp_size)
    device = torch.cuda.current_device()
-    generator = torch.Generator(device).manual_seed(args.seed)
+    # Peiyuan: GPU seed will cause A100 and H100 to produce different results .....
    weight_dtype = torch.bfloat16
    if args.scheduler_type == "euler":
        scheduler = FlowMatchEulerDiscreteScheduler()
@@ -107,9 +107,9 @@ def main(args):
        encoder_attention_mask = None

    if prompts is not None:
-        videos = []
        with torch.autocast("cuda", dtype=torch.bfloat16):
            for prompt in prompts:
+                generator = torch.Generator("cpu").manual_seed(args.seed)
                video = pipe(
                    prompt=[prompt],
                    height=args.height,
@@ -119,9 +119,17 @@ def main(args):
                    guidance_scale=args.guidance_scale,
                    generator=generator,
                ).frames
-                videos.append(video[0])
+                if nccl_info.global_rank <= 0:
+                    os.makedirs(args.output_path, exist_ok=True)
+                    suffix = prompt.split(".")[0]
+                    export_to_video(
+                        video[0],
+                        os.path.join(args.output_path, f"{suffix}.mp4"),
+                        fps=30,
+                    )
    else:
        with torch.autocast("cuda", dtype=torch.bfloat16):
+            generator = torch.Generator("cpu").manual_seed(args.seed)
            videos = pipe(
                prompt_embeds=prompt_embeds,
                prompt_attention_mask=encoder_attention_mask,
@@ -133,16 +141,7 @@ def main(args):
                generator=generator,
            ).frames

-    if nccl_info.global_rank <= 0:
-        if prompts is not None:
-            # mkdir
-            os.makedirs(args.output_path, exist_ok=True)
-            for video, prompt in zip(videos, prompts):
-                suffix = prompt.split(".")[0]
-                export_to_video(
-                    video, os.path.join(args.output_path, f"{suffix}.mp4"), fps=30
-                )
-        else:
+        if nccl_info.global_rank <= 0:
            export_to_video(videos[0], args.output_path + ".mp4", fps=30)


@@ -30,9 +30,8 @@ from accelerate.utils import set_seed
 from tqdm.auto import tqdm
 from fastvideo.utils.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing
 from diffusers.utils import convert_unet_state_dict_to_peft
-from diffusers import (
-    FlowMatchEulerDiscreteScheduler,
-)
+from diffusers import FlowMatchEulerDiscreteScheduler
+from fastvideo.utils.load import load_transformer
 from diffusers.optimization import get_scheduler
 from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
 from diffusers.utils import check_min_version
@@ -40,9 +39,7 @@ from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_func
 import torch.distributed as dist
 from safetensors.torch import save_file, load_file
 from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict
-from torch.distributed.fsdp import (
-    FullyShardedDataParallel as FSDP,
-)
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from fastvideo.utils.checkpoint import (
    save_checkpoint,
    save_lora_checkpoint,
@@ -102,8 +99,9 @@ def get_sigmas(noise_scheduler, device, timesteps, n_dim=4, dtype=torch.float32)
    return sigma


-def train_one_step_mochi(
+def train_one_step(
    transformer,
+    model_type,
    optimizer,
    lr_scheduler,
    loader,
@@ -127,7 +125,7 @@ def train_one_step_mochi(
            latents_attention_mask,
            encoder_attention_mask,
        ) = next(loader)
-        latents = normalize_mochi_dit_input(latents)
+        latents = normalize_dit_input(model_type, latents)
        batch_size = latents.shape[0]
        noise = torch.randn_like(latents)
        u = compute_density_for_timestep_sampling(
@@ -218,15 +216,15 @@ def main(args):

    main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
    # keep the master weight to float32
-    transformer = MochiTransformer3DModel.from_pretrained(
+    transformer = load_transformer(
+        args.model_type,
+        args.dit_model_name_or_path,
        args.pretrained_model_name_or_path,
-        subfolder="transformer",
-        torch_dtype=(
-            torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16
-        ),
+        torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16,
    )

    if args.use_lora:
+        assert args.model_type == "mochi", "LoRA is only supported for Mochi model."
        transformer.requires_grad_(False)
        transformer_lora_config = LoraConfig(
            r=args.lora_rank,
@@ -264,7 +262,8 @@ def main(args):
    main_print(
        f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}"
    )
-    fsdp_kwargs = get_dit_fsdp_kwargs(
+    fsdp_kwargs, no_split_modules = get_dit_fsdp_kwargs(
+        transformer,
        args.fsdp_sharding_startegy,
        args.use_lora,
        args.use_cpu_offload,
@@ -275,17 +274,18 @@ def main(args):
        transformer.config.lora_rank = args.lora_rank
        transformer.config.lora_alpha = args.lora_alpha
        transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
-        transformer._no_split_modules = ["MochiTransformerBlock"]
+        transformer._no_split_modules = [
+            no_split_module.__name__ for no_split_module in no_split_modules
+        ]
        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)

-    transformer = FSDP(
-        transformer,
-        **fsdp_kwargs,
-    )
+    transformer = FSDP(transformer, **fsdp_kwargs,)
    main_print(f"--> model loaded")

    if args.gradient_checkpointing:
-        apply_fsdp_checkpointing(transformer, args.selective_checkpointing)
+        apply_fsdp_checkpointing(
+            transformer, no_split_modules, args.selective_checkpointing
+        )

    # Set model as trainable.
    transformer.train()
@@ -411,8 +411,9 @@ def main(args):
        next(loader)
    for step in range(init_steps + 1, args.max_train_steps + 1):
        start_time = time.time()
-        loss, grad_norm = train_one_step_mochi(
+        loss, grad_norm = train_one_step(
            transformer,
+            args.model_type,
            optimizer,
            lr_scheduler,
            loader,
@@ -479,7 +480,9 @@ def main(args):

 if __name__ == "__main__":
    parser = argparse.ArgumentParser()
-
+    parser.add_argument(
+        "--model_type", type=str, default="mochi", help="The type of model to train."
+    )
    # dataset & dataloader
    parser.add_argument("--data_json_path", type=str, required=True)
    parser.add_argument("--num_frames", type=int, default=163)
@@ -503,6 +506,7 @@ if __name__ == "__main__":

    # text encoder & vae & diffusion model
    parser.add_argument("--pretrained_model_name_or_path", type=str)
+    parser.add_argument("--dit_model_name_or_path", type=str, default=None)
    parser.add_argument("--cache_dir", type=str, default="./cache_dir")

    # diffusion setting
@@ -28,10 +28,7 @@ def save_checkpoint(model, optimizer, rank, output_dir, step, discriminator=Fals
        FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True),
    ):
        cpu_state = model.state_dict()
-        optim_state = FSDP.optim_state_dict(
-            model,
-            optimizer,
-        )
+        optim_state = FSDP.optim_state_dict(model, optimizer,)

    # todo move to get_state_dict
    save_dir = os.path.join(output_dir, f"checkpoint-{step}")
@@ -55,13 +52,7 @@ def save_checkpoint(model, optimizer, rank, output_dir, step, discriminator=Fals


 def save_checkpoint_generator_discriminator(
-    model,
-    optimizer,
-    discriminator,
-    discriminator_optimizer,
-    rank,
-    output_dir,
-    step,
+    model, optimizer, discriminator, discriminator_optimizer, rank, output_dir, step,
 ):
    with FSDP.state_dict_type(
        model,
@@ -230,10 +221,7 @@ def save_lora_checkpoint(transformer, optimizer, rank, output_dir, step):
        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
    ):
        full_state_dict = transformer.state_dict()
-        lora_optim_state = FSDP.optim_state_dict(
-            transformer,
-            optimizer,
-        )
+        lora_optim_state = FSDP.optim_state_dict(transformer, optimizer,)

    if rank <= 0:
        save_dir = os.path.join(output_dir, f"lora-checkpoint-{step}")
@@ -132,9 +132,7 @@ class SeqAllToAll4D(torch.autograd.Function):


 def all_to_all_4D(
-    input_: torch.Tensor,
-    scatter_dim: int = 2,
-    gather_dim: int = 1,
+    input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1,
 ):
    return SeqAllToAll4D.apply(nccl_info.group, input_, scatter_dim, gather_dim)

@@ -193,9 +191,7 @@ class _AllToAll(torch.autograd.Function):


 def all_to_all(
-    input_: torch.Tensor,
-    scatter_dim: int = 2,
-    gather_dim: int = 1,
+    input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1,
 ):
    return _AllToAll.apply(input_, nccl_info.group, scatter_dim, gather_dim)

@@ -0,0 +1,40 @@
+import platform
+
+import accelerate
+import peft
+import torch
+import transformers
+from transformers.utils import is_torch_cuda_available, is_torch_npu_available
+
+
+VERSION = "1.2.0"
+
+
+if __name__ == "__main__":
+    info = {
+        "FastVideo version": VERSION,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "PyTorch version": torch.__version__,
+        "Transformers version": transformers.__version__,
+        "Accelerate version": accelerate.__version__,
+        "PEFT version": peft.__version__,
+    }
+
+    if is_torch_cuda_available():
+        info["PyTorch version"] += " (GPU)"
+        info["GPU type"] = torch.cuda.get_device_name()
+
+    if is_torch_npu_available():
+        info["PyTorch version"] += " (NPU)"
+        info["NPU type"] = torch.npu.get_device_name()
+        info["CANN version"] = torch.version.cann
+
+    try:
+        import bitsandbytes
+
+        info["Bitsandbytes version"] = bitsandbytes.__version__
+    except Exception:
+        pass
+
+    print("\n" + "\n".join([f"- {key}: {value}" for key, value in info.items()]) + "\n")
@@ -29,8 +29,7 @@ import functools


 non_reentrant_wrapper = partial(
-    checkpoint_wrapper,
-    checkpoint_impl=CheckpointImpl.NO_REENTRANT,
+    checkpoint_wrapper, checkpoint_impl=CheckpointImpl.NO_REENTRANT,
 )

 check_fn = lambda submodule: isinstance(submodule, MochiTransformerBlock)
@@ -91,8 +90,7 @@ def get_dit_fsdp_kwargs(
        auto_wrap_policy = fsdp_auto_wrap_policy
    else:
        auto_wrap_policy = functools.partial(
-            transformer_auto_wrap_policy,
-            transformer_layer_cls=no_split_modules,
+            transformer_auto_wrap_policy, transformer_layer_cls=no_split_modules,
        )

    # we use float32 for fsdp but autocast during training
@@ -49,10 +49,7 @@ PROMPT_TEMPLATE_ENCODE_VIDEO = (
 NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion"

 PROMPT_TEMPLATE = {
-    "dit-llm-encode": {
-        "template": PROMPT_TEMPLATE_ENCODE,
-        "crop_start": 36,
-    },
+    "dit-llm-encode": {"template": PROMPT_TEMPLATE_ENCODE, "crop_start": 36,},
    "dit-llm-encode-video": {
        "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
        "crop_start": 95,
@@ -180,13 +177,13 @@ class MochiTextEncoderWrapper(nn.Module):
            os.path.join(pretrained_model_name_or_path, "text_encoder")
        ).to(device)
        self.tokenizer = AutoTokenizer.from_pretrained(
-            os.path.join(pretrained_model_name_or_path, "text_encoder")
+            os.path.join(pretrained_model_name_or_path, "tokenizer")
        )
        self.max_sequence_length = 256

    def encode_prompt(self, prompt):
        device = self.text_encoder.device
-        dtype = self.dtype
+        dtype = self.text_encoder.dtype

        prompt = [prompt] if isinstance(prompt, str) else prompt
        batch_size = len(prompt)
@@ -274,10 +271,7 @@ def load_transformer(
            )
    elif model_type == "hunyuan":
        transformer = HYVideoDiffusionTransformer(
-            in_channels=16,
-            out_channels=16,
-            **hunyuan_config,
-            dtype=master_weight_type,
+            in_channels=16, out_channels=16, **hunyuan_config, dtype=master_weight_type,
        )
        transformer = load_hunyuan_state_dict(transformer, dit_model_name_or_path)
    else:
@@ -107,17 +107,11 @@ def sample_validation_video(
    sigmas = np.array(sigmas)
    if scheduler_type == "euler":
        timesteps, num_inference_steps = retrieve_timesteps(
-            scheduler,
-            num_inference_steps,
-            device,
-            timesteps,
-            sigmas,
+            scheduler, num_inference_steps, device, timesteps, sigmas,
        )
    else:
        timesteps, num_inference_steps = retrieve_timesteps(
-            scheduler,
-            num_inference_steps,
-            device,
+            scheduler, num_inference_steps, device,
        )
    num_warmup_steps = max(len(timesteps) - num_inference_steps * scheduler.order, 0)

@@ -1,12 +0,0 @@
-python3 fastvideo/sample/sample_t2v_hunyuan_no_sp.py \
-    --height 500 \
-    --width 700 \
-    --num_frames 29 \
-    --num_inference_steps 50 \
-    --guidance_scale 1 \
-    --embedded_cfg_scale 6 \
-    --flow-reverse \
-    --prompts "A cat walks on the grass, realistic style." \
-    --prompts "A dog runs in the park, realistic style." \
-    --seed 42 \
-    --output_path outputs_video/hunyuan/
@@ -1,39 +0,0 @@
-num_gpus=4
-
-torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
-    fastvideo/sample/sample_t2v_hunyuan.py \
-    --height 512 \
-    --width 512 \
-    --num_frames 29 \
-    --num_inference_steps 4 \
-    --guidance_scale 1 \
-    --embedded_cfg_scale 6 \
-    --flow_shift 17 \
-    --flow-reverse \
-    --prompts "A man on stage claps his hands together while facing the audience. The audience, visible in the foreground, holds up mobile devices to record the event, capturing the moment from various angles. The background features a large banner with text identifying the man on stage. Throughout the sequence, the man's expression remains engaged and directed towards the audience. The camera angle remains constant, focusing on capturing the interaction between the man on stage and the audience."\
-    --seed 12345 \
-    --output_path outputs_video/hunyuan/
-
-
-tensor(-0.1065, device='cuda:0', dtype=torch.float16)
-tensor(-0.0034, device='cuda:2', dtype=torch.float16)
-
-tensor(-0.0230, device='cuda:0', dtype=torch.float16)
->>> weight[0, :768].mean()
-tensor(-0.0367, device='cuda:0', dtype=torch.float16)
-
-num_gpus=1
-
-torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
-    fastvideo/sample/sample_t2v_hunyuan.py \
-    --height 480 \
-    --width 848 \
-    --num_frames 93 \
-    --num_inference_steps 50 \
-    --guidance_scale 1 \
-    --embedded_cfg_scale 6 \
-    --flow_shift 17 \
-    --flow-reverse \
-    --prompts "A man on stage claps his hands together while facing the audience. The audience, visible in the foreground, holds up mobile devices to record the event, capturing the moment from various angles. The background features a large banner with text identifying the man on stage. Throughout the sequence, the man's expression remains engaged and directed towards the audience. The camera angle remains constant, focusing on capturing the interaction between the man on stage and the audience."\
-    --seed 12345 \
-    --output_path outputs_video/hunyuan/
@@ -1,20 +0,0 @@
-
-
-
-num_gpus=4
-
-torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
-    fastvideo/sample/sample_t2v_mochi.py \
-    --model_path data/mochi \
-    --prompt_embed_path "data/synthetic_debug2/prompt_embed/2.pt" \
-    --encoder_attention_mask_path "data/synthetic_debug2/prompt_attention_mask/1.pt" \
-    --num_frames  163 \
-    --height 480 \
-    --width 848 \
-    --num_inference_steps 32 \
-    --guidance_scale 4.5 \
-    --output_path outputs_video/debug \
-    --shift 8 \
-    --seed 12345 \
-    --scheduler_type "pcm_linear_quadratic" 
-
@@ -1,11 +0,0 @@
-python fastvideo/sample/sample_t2v_mochi_no_sp.py \
-    --model_path data/mochi \
-    --prompts  "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand's movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough."  \
-    --num_frames 163 \
-    --height 480 \
-    --width 848 \
-    --num_inference_steps 64 \
-    --guidance_scale 0.0 \
-    --seed 12346 \
-    --transformer_path data/outputs/debug/checkpoint-100/transformer \
-    --output_path outputs_video/single_no_guidance
@@ -0,0 +1,130 @@
+# Prediction interface for Cog ⚙️
+# https://cog.run/python
+
+from cog import BasePredictor, Input, Path
+import os
+import time
+import torch
+import imageio
+import argparse
+import subprocess
+import torchvision
+import numpy as np
+from einops import rearrange
+
+MODEL_CACHE = 'FastHunyuan'
+os.environ['MODEL_BASE'] = './'+MODEL_CACHE
+from fastvideo.models.hunyuan.inference import HunyuanVideoSampler
+MODEL_URL = "https://weights.replicate.delivery/default/FastVideo/FastHunyuan/model.tar"
+
+def download_weights(url, dest):
+    start = time.time()
+    print("downloading url: ", url)
+    print("downloading to: ", dest)
+    subprocess.check_call(["pget", "-xf", url, dest], close_fds=False)
+    print("downloading took: ", time.time() - start)
+
+class Predictor(BasePredictor):
+    def setup(self):
+        """Load the model into memory"""
+        print("Model Base: " + os.environ['MODEL_BASE'])
+        # Download weights
+        if not os.path.exists(MODEL_CACHE):
+            download_weights(MODEL_URL, MODEL_CACHE)
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        args = argparse.Namespace(
+            num_frames=125,
+            height=720,
+            width=1280,
+            num_inference_steps=6,
+            fps=24,
+            denoise_type='flow',
+            seed=1024,
+            neg_prompt=None,
+            guidance_scale=1.0,
+            embedded_cfg_scale=6.0,
+            flow_shift=17,
+            batch_size=1,
+            num_videos=1,
+            load_key='module',
+            use_cpu_offload=False,
+            dit_weight='FastHunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt',
+            reproduce=True,
+            disable_autocast=False,
+            flow_reverse=True,
+            flow_solver='euler',
+            use_linear_quadratic_schedule=False,
+            linear_schedule_end=25,
+            model='HYVideo-T/2-cfgdistill',
+            latent_channels=16,
+            precision='bf16',
+            rope_theta=256,
+            vae='884-16c-hy',
+            vae_precision='fp16',
+            vae_tiling=True,
+            text_encoder='llm',
+            text_encoder_precision='fp16',
+            text_states_dim=4096,
+            text_len=256,
+            tokenizer='llm',
+            prompt_template='dit-llm-encode',
+            prompt_template_video='dit-llm-encode-video',
+            hidden_state_skip_layer=2,
+            apply_final_norm=False,
+            text_encoder_2='clipL',
+            text_encoder_precision_2='fp16',
+            text_states_dim_2=768,
+            tokenizer_2='clipL',
+            text_len_2=77,
+            model_path=MODEL_CACHE,
+        )
+        self.model = HunyuanVideoSampler.from_pretrained(MODEL_CACHE, args=args)
+
+    def predict(
+        self,
+        prompt: str = Input(description="Text prompt for video generation", default="A cat walks on the grass, realistic style."),
+        negative_prompt: str = Input(description="Text prompt to specify what you don't want in the video.", default=""),
+        width: int = Input(description="Width of output video", default=1280, ge=256),
+        height: int = Input(description="Height of output video", default=720, ge=256),
+        num_frames: int = Input(description="Number of frames to generate", default=125, ge=16),
+        num_inference_steps: int = Input(description="Number of denoising steps", default=6, ge=1, le=50),
+        guidance_scale: float = Input(description="Classifier free guidance scale", default=1.0, ge=0.1, le=10.0),
+        embedded_cfg_scale: float = Input(description="Embedded classifier free guidance scale", default=6.0, ge=0.1, le=10.0),
+        flow_shift: int = Input(description="Flow shift parameter", default=17, ge=1, le=20),
+        fps: int = Input(description="Frames per second of output video", default=24, ge=1, le=60),
+        seed: int = Input(description="0 for Random seed. Set for reproducible generation", default=0),
+    ) -> Path:
+        """Run video generation"""
+        if seed <=0:
+            seed = int.from_bytes(os.urandom(2), "big")
+        print(f"Using seed: {seed}")
+            
+        outputs = self.model.predict(
+            prompt=prompt,
+            height=height,
+            width=width,
+            video_length=num_frames,
+            seed=seed,
+            negative_prompt=negative_prompt,
+            infer_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            embedded_guidance_scale=embedded_cfg_scale,
+            flow_shift=flow_shift,
+            flow_reverse=True,
+            batch_size=1,
+            num_videos_per_prompt=1,
+        )
+
+        # Process output video
+        videos = rearrange(outputs["samples"], "b c t h w -> t b c h w")
+        frames = []
+        for x in videos:
+            x = torchvision.utils.make_grid(x, nrow=6)
+            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+            frames.append((x * 255).numpy().astype(np.uint8))
+
+        # Save video
+        output_path = Path("/tmp/output.mp4")
+        imageio.mimsave(str(output_path), frames, fps=fps)
+        return Path(output_path)
@@ -21,7 +21,7 @@ dependencies = [
 "timm==1.0.11", "torchdiffeq==0.2.4", "torchmetrics==1.5.1", "tqdm==4.66.5", "urllib3==2.2.0", "uvicorn==0.32.0", 
 "scikit-video==1.1.11", "imageio-ffmpeg==0.5.1", "sentencepiece==0.2.0", "beautifulsoup4==4.12.3", "ftfy==6.3.0", 
 "moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3", 
-"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1", "einops==0.8.0", "wheel==0.44.0", "loguru"]
+"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1", "einops==0.8.0", "wheel==0.44.0", "loguru", "diffusers==0.32.0", "bitsandbytes", "pytest", "requests-mock"]


 [tool.setuptools.packages.find]
@@ -1,25 +1,25 @@
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
-export WANDB_API_KEY=4f6de3765d6464f43e0506ec7d785641af645e73

+DATA_DIR=./data

-torchrun --nnodes 1 --nproc_per_node 4\
-    fastvideo/distill_adv.py\
+torchrun --nnodes 1 --nproc_per_node 8\
+    fastvideo/distill.py\
    --seed 42\
-    --pretrained_model_name_or_path data/hunyuan\
-    --dit_model_name_or_path data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt\
+    --pretrained_model_name_or_path $DATA_DIR/hunyuan\
+    --dit_model_name_or_path $DATA_DIR/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt\
    --model_type "hunyuan" \
-    --cache_dir "data/.cache"\
-    --data_json_path "data/Hunyuan-Distill-Data/videos2caption.json"\
-    --validation_prompt_dir "data/Hunyuan-Distill-Data/validation"\
+    --cache_dir "$DATA_DIR/.cache"\
+    --data_json_path "$DATA_DIR/Hunyuan-30K-Distill-Data/videos2caption.json"\
+    --validation_prompt_dir "$DATA_DIR/Hunyuan-Distill-Data/validation"\
    --gradient_checkpointing\
    --train_batch_size=1\
-    --num_latent_t 1\
+    --num_latent_t 24\
    --sp_size 1\
    --train_sp_batch_size 1\
    --dataloader_num_workers 4\
    --gradient_accumulation_steps=1\
-    --max_train_steps=640\
+    --max_train_steps=2000\
    --learning_rate=1e-6\
    --mixed_precision="bf16"\
    --checkpointing_steps=64\
@@ -30,8 +30,8 @@ torchrun --nnodes 1 --nproc_per_node 4\
    --ema_start_step 0\
    --cfg 0.0\
    --log_validation\
-    --output_dir="outputs/debug"\
-    --tracker_project_name PCM \
+    --output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_32"\
+    --tracker_project_name Hunyuan_Distill \
    --num_frames  93 \
    --shift 17 \
    --validation_guidance_scale "1.0" \
@@ -0,0 +1,38 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+torchrun --nnodes 1 --nproc_per_node 4 \
+    fastvideo/distill.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/mochi \
+    --model_type "mochi" \
+    --cache_dir data/.cache \
+    --data_json_path data/Merge-30k-Data/video2caption.json \
+    --validation_prompt_dir data/Image-Vid-Finetune-Mochi/validation \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 28 \
+    --sp_size 4 \
+    --train_sp_batch_size 2 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=4000 \
+    --learning_rate=1e-6 \
+    --mixed_precision=bf16 \
+    --checkpointing_steps=64 \
+    --validation_steps 1 \
+    --validation_sampling_steps 8 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --log_validation \
+    --output_dir="data/outputs/lq_euler_50_thres0.1_lrg_0.75_phase1_lr1e-6_repro" \
+    --tracker_project_name PCM \
+    --num_frames 163 \
+    --scheduler_type pcm_linear_quadratic \
+    --validation_guidance_scale 0.5,1.5,2.5 \
+    --num_euler_timesteps 50 \
+    --linear_quadratic_threshold 0.1 \
+    --linear_range 0.75 \
+    --multi_phased_distill_schedule 4000-1
@@ -1,5 +0,0 @@
-## How to Distill Hunyuan 
-
-python scripts/download_hf.py --repo_id FastVideo/hunyuan --local_dir data/hunyuan --repo_type model
-python scripts/download_hf.py --repo_id FastVideo/Hunyuan-Distill-Data --local_dir data/Hunyuan-Distill-Data --repo_type=dataset
-
@@ -0,0 +1,36 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+torchrun --nnodes 1 --nproc_per_node 8 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/hunyuan \
+    --dit_model_name_or_path data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt\
+    --model_type "hunyuan" \
+    --cache_dir data/.cache \
+    --data_json_path data/Image-Vid-Finetune-HunYuan/videos2caption.json \
+    --validation_prompt_dir data/Image-Vid-Finetune-HunYuan/validation \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 24 \
+    --sp_size 4 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=2000 \
+    --learning_rate=5e-6 \
+    --mixed_precision=bf16 \
+    --checkpointing_steps=200 \
+    --validation_steps 100 \
+    --validation_sampling_steps 64 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir=data/outputs/HSH-Taylor-Finetune-Hunyuan \
+    --tracker_project_name HSH-Taylor-Finetune-Hunyuan \
+    --num_frames 93 \
+    --validation_guidance_scale "1.0" \
+    --group_frame
@@ -0,0 +1,32 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+torchrun --nnodes 1 --nproc_per_node 4 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/mochi \
+    --cache_dir data/.cache \
+    --data_json_path data/Mochi-Black-Myth/videos2caption.json \
+    --validation_prompt_dir data/Mochi-Black-Myth/validation \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 16 \
+    --sp_size 1 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=2000 \
+    --learning_rate=5e-6 \
+    --mixed_precision=bf16 \
+    --checkpointing_steps=200 \
+    --validation_steps 100 \
+    --validation_sampling_steps 64 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir=data/outputs/Black-Myth-Finetune \
+    --tracker_project_name Black-Myth-Finetune \
+    --num_frames 93 
@@ -0,0 +1,37 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+torchrun --nnodes 1 --nproc_per_node 2 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/mochi \
+    --cache_dir data/.cache \
+    --data_json_path data/Mochi-Black-Myth/videos2caption.json \
+    --validation_prompt_dir data/Mochi-Black-Myth/validation \
+    --gradient_checkpointing \
+    --train_batch_size 1 \
+    --num_latent_t 14 \
+    --sp_size 2 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 1 \
+    --gradient_accumulation_steps 2 \
+    --max_train_steps 2000 \
+    --learning_rate 5e-6 \
+    --mixed_precision bf16 \
+    --checkpointing_steps 200 \
+    --validation_steps 100 \
+    --validation_sampling_steps 64 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir=data/outputs/Black-Myth-Lora-FT \
+    --tracker_project_name Black-Myth-Lora-Finetune \
+    --num_frames 91 \
+    --lora_rank 128 \
+    --lora_alpha 256 \
+    --master_weight_type "bf16" \
+    --use_lora \
+    --use_cpu_offload
@@ -0,0 +1,37 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+CUDA_VISIBLE_DEVICES=5 torchrun --nnodes 1 --nproc_per_node 1 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/mochi \
+    --cache_dir data/.cache \
+    --data_json_path data/Image-Vid-Finetune-Mochi/videos2caption.json \
+    --validation_prompt_dir data/Image-Vid-Finetune-Mochi/validation \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 14 \
+    --sp_size 1 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 1 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=2000 \
+    --learning_rate=5e-6 \
+    --mixed_precision=bf16 \
+    --checkpointing_steps=200 \
+    --validation_steps 100 \
+    --validation_sampling_steps 64 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir=data/outputs/HSH-Taylor-Finetune-Lora \
+    --tracker_project_name HSH-Taylor-Finetune-Lora \
+    --num_frames 91 \
+    --group_frame \
+    --lora_rank 128 \
+    --lora_alpha 256 \
+    --master_weight_type "bf16" \
+    --use_lora 
@@ -1,36 +0,0 @@
-torchrun --nnodes 4 --nproc_per_node 4 \
-    --node_rank=3 \
-    --rdzv_id=456 \
-    --rdzv_backend=c10d \
-    --rdzv_endpoint=172.23.30.15:29500 \
-    fastvideo/train.py \
-    --seed 42 \
-    --pretrained_model_name_or_path data/mochi \
-    --cache_dir "data/.cache" \
-    --data_json_path "data/BLACK-MYTH-Finetune-Dataset/videos2caption.json" \
-    --validation_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/validation_prompt_embed_mask" \
-    --uncond_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/uncond_prompt_embed_mask" \
-    --gradient_checkpointing \
-    --train_batch_size=1 \
-    --num_latent_t 16 \
-    --sp_size 4 \
-    --train_sp_batch_size 4 \
-    --dataloader_num_workers 4 \
-    --gradient_accumulation_steps=1 \
-    --max_train_steps=4000 \
-    --learning_rate=2e-5 \
-    --mixed_precision="bf16" \
-    --checkpointing_steps=500 \
-    --validation_steps=100 \
-    --validation_sampling_steps 64 \
-    --checkpoints_total_limit 3 \
-    --allow_tf32 \
-    --ema_start_step 0 \
-    --cfg 0.0 \
-    --ema_decay 0.999 \
-    --log_validation \
-    --output_dir="data/outputs/black_myth_correct_video_mask" \
-    --weighting_scheme "uniform" \
-    --num_frames 91 \
-    --selective_checkpointing 1.0
-
@@ -1,13 +0,0 @@
-num_gpus=1
-
-torchrun --nproc_per_node=$num_gpus fastvideo/sample/generate_synthetic.py \
-    --model_path data/mochi \
-    --num_frames 1 \
-    --height 480 \
-    --width 848 \
-    --num_inference_steps 4 \
-    --guidance_scale 4.5 \
-    --prompt_path "data/prompt.txt" \
-    --dataset_output_dir data/synthetic_debug2
-
-    
@@ -0,0 +1,9 @@
+from huggingface_hub import HfApi
+
+api = HfApi()
+
+api.upload_folder(
+    folder_path="data/Black-Myth-Taylor-Src",
+    repo_id="FastVideo/Image-Vid-Finetune-Src",
+    repo_type="dataset",
+)
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+num_gpus=1
+export MODEL_BASE="data/FastHunyuan-diffusers"
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 12345 \
+    fastvideo/sample/sample_t2v_diffusers_hunyuan.py \
+    --height 720 \
+    --width 1280 \
+    --num_frames 45 \
+    --num_inference_steps 6 \
+    --guidance_scale 1 \
+    --embedded_cfg_scale 6 \
+    --flow_shift 17 \
+    --flow-reverse \
+    --prompt ./assets/prompt.txt \
+    --seed 1024 \
+    --output_path outputs_video/hunyuan_quant/nf4/ \
+    --model_path $MODEL_BASE \
+    --quantization "nf4" \
+    --cpu_offload
@@ -0,0 +1,19 @@
+#!/bin/bash
+
+num_gpus=4
+export MODEL_BASE=data/FastHunyuan
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+    fastvideo/sample/sample_t2v_hunyuan.py \
+    --height 720 \
+    --width 1280 \
+    --num_frames 125 \
+    --num_inference_steps 6 \
+    --guidance_scale 1 \
+    --embedded_cfg_scale 6 \
+    --flow_shift 17 \
+    --flow-reverse \
+    --prompt ./assets/prompt.txt \
+    --seed 1024 \
+    --output_path outputs_video/hunyuan/cfg6/ \
+    --model_path $MODEL_BASE \
+    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt
@@ -0,0 +1,19 @@
+#!/bin/bash
+
+num_gpus=4
+
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+    fastvideo/sample/sample_t2v_mochi.py \
+    --model_path data/FastMochi-diffusers \
+    --prompt_path "assets/prompt.txt" \
+    --num_frames 163 \
+    --height 480 \
+    --width 848 \
+    --num_inference_steps 8 \
+    --guidance_scale 1.5 \
+    --output_path outputs_video/mochi_sp/ \
+    --seed 1024 \
+    --scheduler_type "pcm_linear_quadratic" \
+    --linear_threshold 0.1 \
+    --linear_range 0.75
+
@@ -1,24 +1,23 @@
 # export WANDB_MODE="offline"
-GPU_NUM=8
+GPU_NUM=1 # 2,4,8
 MODEL_PATH="data/hunyuan"
 MODEL_TYPE="hunyuan"
-DATA_MERGE_PATH="data/Mixkit-All-Clips/merge.txt"
-OUTPUT_DIR="data/Hunyuan-Mixkit-Data"
+DATA_MERGE_PATH="data/Image-Vid-Finetune-Src/merge.txt"
+OUTPUT_DIR="data/Image-Vid-Finetune-HunYuan"
 VALIDATION_PATH="assets/prompt.txt"
-# torchrun --nproc_per_node=$GPU_NUM \
-#     fastvideo/data_preprocess/preprocess_vae_latents.py \
-#     --model_path $MODEL_PATH \
-#     --data_merge_path $DATA_MERGE_PATH \
-#     --train_batch_size=1 \
-#     --max_height=480 \
-#     --max_width=848 \
-#     --num_frames=93 \
-#     --dataloader_num_workers 1 \
-#     --output_dir=$OUTPUT_DIR \
-#     --model_type $MODEL_TYPE \
-#     --train_fps 24 
-

+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_vae_latents.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --train_batch_size=1 \
+    --max_height=480 \
+    --max_width=848 \
+    --num_frames=93 \
+    --dataloader_num_workers 1 \
+    --output_dir=$OUTPUT_DIR \
+    --model_type $MODEL_TYPE \
+    --train_fps 24 

 torchrun --nproc_per_node=$GPU_NUM \
    fastvideo/data_preprocess/preprocess_text_embeddings.py \
@@ -26,8 +25,6 @@ torchrun --nproc_per_node=$GPU_NUM \
    --model_path $MODEL_PATH \
    --output_dir=$OUTPUT_DIR 

-
-
 torchrun --nproc_per_node=1 \
    fastvideo/data_preprocess/preprocess_validation_text_embeddings.py \
    --model_type $MODEL_TYPE \
@@ -0,0 +1,33 @@
+# export WANDB_MODE="offline"
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="data/FastMochi-diffusers"
+MODEL_TYPE="mochi"
+DATA_MERGE_PATH="data/Image-Vid-Finetune-Src/merge.txt"
+OUTPUT_DIR="data/Image-Vid-Finetune-Mochi"
+VALIDATION_PATH="assets/prompt.txt"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_vae_latents.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --train_batch_size=1 \
+    --max_height=480 \
+    --max_width=848 \
+    --num_frames=93 \
+    --dataloader_num_workers 1 \
+    --output_dir=$OUTPUT_DIR \
+    --model_type $MODEL_TYPE \
+    --train_fps 24 
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_text_embeddings.py \
+    --model_type $MODEL_TYPE \
+    --model_path $MODEL_PATH \
+    --output_dir=$OUTPUT_DIR 
+
+torchrun --nproc_per_node=1 \
+    fastvideo/data_preprocess/preprocess_validation_text_embeddings.py \
+    --model_type $MODEL_TYPE \
+    --model_path $MODEL_PATH \
+    --output_dir=$OUTPUT_DIR \
+    --validation_prompt_txt $VALIDATION_PATH
@@ -0,0 +1,10 @@
+import torch
+from tests.pipelines.hunyuan_video.test_hunyuan_video import HunyuanVideoPipelineFastTests
+from tests.pipelines.mochi.test_mochi import MochiPipelineFastTests
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+HunyuanTestCls =  HunyuanVideoPipelineFastTests()
+HunyuanTestCls.test_inference()
+
+MochiTestCls = MochiPipelineFastTests()
+MochiTestCls.test_inference()
@@ -1,9 +0,0 @@
-from huggingface_hub import HfApi
-
-api = HfApi()
-
-api.upload_folder(
-    folder_path="/ephemeral/hao.zhang/codefolder/FastVideo-OSP/data/Hunyuan-Mixkit-Data",
-    repo_id="FastVideo/Hunyuan-Distill-Data",
-    repo_type="dataset",
-)
@@ -0,0 +1,184 @@
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import sys
+import unittest
+
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast
+
+from diffusers import (
+    AutoencoderKLHunyuanVideo,
+    FlowMatchEulerDiscreteScheduler,
+    HunyuanVideoPipeline,
+    HunyuanVideoTransformer3DModel,
+)
+from diffusers.utils.testing_utils import (
+    floats_tensor,
+    require_peft_backend,
+    skip_mps,
+)
+
+
+sys.path.append(".")
+
+from utils import PeftLoraLoaderMixinTests  # noqa: E402
+
+
+@require_peft_backend
+@skip_mps
+class HunyuanVideoLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
+    pipeline_class = HunyuanVideoPipeline
+    scheduler_cls = FlowMatchEulerDiscreteScheduler
+    scheduler_classes = [FlowMatchEulerDiscreteScheduler]
+    scheduler_kwargs = {}
+
+    transformer_kwargs = {
+        "in_channels": 4,
+        "out_channels": 4,
+        "num_attention_heads": 2,
+        "attention_head_dim": 10,
+        "num_layers": 1,
+        "num_single_layers": 1,
+        "num_refiner_layers": 1,
+        "patch_size": 1,
+        "patch_size_t": 1,
+        "guidance_embeds": True,
+        "text_embed_dim": 16,
+        "pooled_projection_dim": 8,
+        "rope_axes_dim": (2, 4, 4),
+    }
+    transformer_cls = HunyuanVideoTransformer3DModel
+    vae_kwargs = {
+        "in_channels": 3,
+        "out_channels": 3,
+        "latent_channels": 4,
+        "down_block_types": (
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+        ),
+        "up_block_types": (
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+        ),
+        "block_out_channels": (8, 8, 8, 8),
+        "layers_per_block": 1,
+        "act_fn": "silu",
+        "norm_num_groups": 4,
+        "scaling_factor": 0.476986,
+        "spatial_compression_ratio": 8,
+        "temporal_compression_ratio": 4,
+        "mid_block_add_attention": True,
+    }
+    vae_cls = AutoencoderKLHunyuanVideo
+    has_two_text_encoders = True
+    tokenizer_cls, tokenizer_id, tokenizer_subfolder = (
+        LlamaTokenizerFast,
+        "hf-internal-testing/tiny-random-hunyuanvideo",
+        "tokenizer",
+    )
+    tokenizer_2_cls, tokenizer_2_id, tokenizer_2_subfolder = (
+        CLIPTokenizer,
+        "hf-internal-testing/tiny-random-hunyuanvideo",
+        "tokenizer_2",
+    )
+    text_encoder_cls, text_encoder_id, text_encoder_subfolder = (
+        LlamaModel,
+        "hf-internal-testing/tiny-random-hunyuanvideo",
+        "text_encoder",
+    )
+    text_encoder_2_cls, text_encoder_2_id, text_encoder_2_subfolder = (
+        CLIPTextModel,
+        "hf-internal-testing/tiny-random-hunyuanvideo",
+        "text_encoder_2",
+    )
+
+    @property
+    def output_shape(self):
+        return (1, 9, 32, 32, 3)
+
+    def get_dummy_inputs(self, with_generator=True):
+        batch_size = 1
+        sequence_length = 16
+        num_channels = 4
+        num_frames = 9
+        num_latent_frames = 3  # (num_frames - 1) // temporal_compression_ratio + 1
+        sizes = (4, 4)
+
+        generator = torch.manual_seed(0)
+        noise = floats_tensor((batch_size, num_latent_frames, num_channels) + sizes)
+        input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator)
+
+        pipeline_inputs = {
+            "prompt": "",
+            "num_frames": num_frames,
+            "num_inference_steps": 1,
+            "guidance_scale": 6.0,
+            "height": 32,
+            "width": 32,
+            "max_sequence_length": sequence_length,
+            "prompt_template": {"template": "{}", "crop_start": 0},
+            "output_type": "np",
+        }
+        if with_generator:
+            pipeline_inputs.update({"generator": generator})
+
+        return noise, input_ids, pipeline_inputs
+
+    def test_simple_inference_with_text_lora_denoiser_fused_multi(self):
+        super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3)
+
+    def test_simple_inference_with_text_denoiser_lora_unfused(self):
+        super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3)
+
+    # TODO(aryan): Fix the following test
+    @unittest.skip("This test fails with an error I haven't been able to debug yet.")
+    def test_simple_inference_save_pretrained(self):
+        pass
+
+    @unittest.skip("Not supported in HunyuanVideo.")
+    def test_simple_inference_with_text_denoiser_block_scale(self):
+        pass
+
+    @unittest.skip("Not supported in HunyuanVideo.")
+    def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self):
+        pass
+
+    @unittest.skip("Not supported in HunyuanVideo.")
+    def test_modify_padding_mode(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.")
+    def test_simple_inference_with_partial_text_lora(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.")
+    def test_simple_inference_with_text_lora(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.")
+    def test_simple_inference_with_text_lora_and_scale(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.")
+    def test_simple_inference_with_text_lora_fused(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.")
+    def test_simple_inference_with_text_lora_save_load(self):
+        pass
@@ -0,0 +1,138 @@
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import sys
+import unittest
+
+import torch
+from transformers import AutoTokenizer, T5EncoderModel
+
+from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel
+from diffusers.utils.testing_utils import (
+    floats_tensor,
+    require_peft_backend,
+    skip_mps,
+)
+
+
+sys.path.append(".")
+
+from utils import PeftLoraLoaderMixinTests  # noqa: E402
+
+
+@require_peft_backend
+@skip_mps
+class MochiLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
+    pipeline_class = MochiPipeline
+    scheduler_cls = FlowMatchEulerDiscreteScheduler
+    scheduler_classes = [FlowMatchEulerDiscreteScheduler]
+    scheduler_kwargs = {}
+
+    transformer_kwargs = {
+        "patch_size": 2,
+        "num_attention_heads": 2,
+        "attention_head_dim": 8,
+        "num_layers": 2,
+        "pooled_projection_dim": 16,
+        "in_channels": 12,
+        "out_channels": None,
+        "qk_norm": "rms_norm",
+        "text_embed_dim": 32,
+        "time_embed_dim": 4,
+        "activation_fn": "swiglu",
+        "max_sequence_length": 16,
+    }
+    transformer_cls = MochiTransformer3DModel
+    vae_kwargs = {
+        "latent_channels": 12,
+        "out_channels": 3,
+        "encoder_block_out_channels": (32, 32, 32, 32),
+        "decoder_block_out_channels": (32, 32, 32, 32),
+        "layers_per_block": (1, 1, 1, 1, 1),
+    }
+    vae_cls = AutoencoderKLMochi
+    tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-random-t5"
+    text_encoder_cls, text_encoder_id = T5EncoderModel, "hf-internal-testing/tiny-random-t5"
+
+    text_encoder_target_modules = ["q", "k", "v", "o"]
+
+    @property
+    def output_shape(self):
+        return (1, 7, 16, 16, 3)
+
+    def get_dummy_inputs(self, with_generator=True):
+        batch_size = 1
+        sequence_length = 16
+        num_channels = 4
+        num_frames = 7
+        num_latent_frames = 3
+        sizes = (2, 2)
+
+        generator = torch.manual_seed(0)
+        noise = floats_tensor((batch_size, num_latent_frames, num_channels) + sizes)
+        input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator)
+
+        pipeline_inputs = {
+            "prompt": "dance monkey",
+            "num_frames": num_frames,
+            "num_inference_steps": 4,
+            "guidance_scale": 6.0,
+            # Cannot reduce because convolution kernel becomes bigger than sample
+            "height": 16,
+            "width": 16,
+            "max_sequence_length": sequence_length,
+            "output_type": "np",
+        }
+        if with_generator:
+            pipeline_inputs.update({"generator": generator})
+
+        return noise, input_ids, pipeline_inputs
+
+    def test_simple_inference_with_text_lora_denoiser_fused_multi(self):
+        super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3)
+
+    def test_simple_inference_with_text_denoiser_lora_unfused(self):
+        super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3)
+
+    @unittest.skip("Not supported in Mochi.")
+    def test_simple_inference_with_text_denoiser_block_scale(self):
+        pass
+
+    @unittest.skip("Not supported in Mochi.")
+    def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self):
+        pass
+
+    @unittest.skip("Not supported in Mochi.")
+    def test_modify_padding_mode(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in Mochi.")
+    def test_simple_inference_with_partial_text_lora(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in Mochi.")
+    def test_simple_inference_with_text_lora(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in Mochi.")
+    def test_simple_inference_with_text_lora_and_scale(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in Mochi.")
+    def test_simple_inference_with_text_lora_fused(self):
+        pass
+
+    @unittest.skip("Text encoder LoRA is not supported in Mochi.")
+    def test_simple_inference_with_text_lora_save_load(self):
+        pass
@@ -0,0 +1,184 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import unittest
+
+import torch
+
+from diffusers import AutoencoderKLHunyuanVideo
+from diffusers.utils.testing_utils import (
+    enable_full_determinism,
+    floats_tensor,
+    torch_device,
+)
+
+from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin
+
+
+enable_full_determinism()
+
+
+class AutoencoderKLHunyuanVideoTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
+    model_class = AutoencoderKLHunyuanVideo
+    main_input_name = "sample"
+    base_precision = 1e-2
+
+    def get_autoencoder_kl_hunyuan_video_config(self):
+        return {
+            "in_channels": 3,
+            "out_channels": 3,
+            "latent_channels": 4,
+            "down_block_types": (
+                "HunyuanVideoDownBlock3D",
+                "HunyuanVideoDownBlock3D",
+                "HunyuanVideoDownBlock3D",
+                "HunyuanVideoDownBlock3D",
+            ),
+            "up_block_types": (
+                "HunyuanVideoUpBlock3D",
+                "HunyuanVideoUpBlock3D",
+                "HunyuanVideoUpBlock3D",
+                "HunyuanVideoUpBlock3D",
+            ),
+            "block_out_channels": (8, 8, 8, 8),
+            "layers_per_block": 1,
+            "act_fn": "silu",
+            "norm_num_groups": 4,
+            "scaling_factor": 0.476986,
+            "spatial_compression_ratio": 8,
+            "temporal_compression_ratio": 4,
+            "mid_block_add_attention": True,
+        }
+
+    @property
+    def dummy_input(self):
+        batch_size = 2
+        num_frames = 9
+        num_channels = 3
+        sizes = (16, 16)
+
+        image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
+
+        return {"sample": image}
+
+    @property
+    def input_shape(self):
+        return (3, 9, 16, 16)
+
+    @property
+    def output_shape(self):
+        return (3, 9, 16, 16)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = self.get_autoencoder_kl_hunyuan_video_config()
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    def test_enable_disable_tiling(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+
+        inputs_dict.update({"return_dict": False})
+
+        torch.manual_seed(0)
+        output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        torch.manual_seed(0)
+        model.enable_tiling()
+        output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertLess(
+            (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(),
+            0.5,
+            "VAE tiling should not affect the inference results",
+        )
+
+        torch.manual_seed(0)
+        model.disable_tiling()
+        output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertEqual(
+            output_without_tiling.detach().cpu().numpy().all(),
+            output_without_tiling_2.detach().cpu().numpy().all(),
+            "Without tiling outputs should match with the outputs when tiling is manually disabled.",
+        )
+
+    def test_enable_disable_slicing(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+
+        inputs_dict.update({"return_dict": False})
+
+        torch.manual_seed(0)
+        output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        torch.manual_seed(0)
+        model.enable_slicing()
+        output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertLess(
+            (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(),
+            0.5,
+            "VAE slicing should not affect the inference results",
+        )
+
+        torch.manual_seed(0)
+        model.disable_slicing()
+        output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertEqual(
+            output_without_slicing.detach().cpu().numpy().all(),
+            output_without_slicing_2.detach().cpu().numpy().all(),
+            "Without slicing outputs should match with the outputs when slicing is manually disabled.",
+        )
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {
+            "HunyuanVideoDecoder3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoEncoder3D",
+            "HunyuanVideoMidBlock3D",
+            "HunyuanVideoUpBlock3D",
+        }
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+    # We need to overwrite this test because the base test does not account length of down_block_types
+    def test_forward_with_norm_groups(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        init_dict["norm_num_groups"] = 16
+        init_dict["block_out_channels"] = (16, 16, 16, 16)
+
+        model = self.model_class(**init_dict)
+        model.to(torch_device)
+        model.eval()
+
+        with torch.no_grad():
+            output = model(**inputs_dict)
+
+            if isinstance(output, dict):
+                output = output.to_tuple()[0]
+
+        self.assertIsNotNone(output)
+        expected_shape = inputs_dict["sample"].shape
+        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
+
+    @unittest.skip("Unsupported test.")
+    def test_outputs_equivalence(self):
+        pass
@@ -0,0 +1,468 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import gc
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from diffusers import AutoencoderKL
+from diffusers.utils.import_utils import is_xformers_available
+from diffusers.utils.testing_utils import (
+    backend_empty_cache,
+    enable_full_determinism,
+    floats_tensor,
+    load_hf_numpy,
+    require_torch_accelerator,
+    require_torch_accelerator_with_fp16,
+    require_torch_gpu,
+    skip_mps,
+    slow,
+    torch_all_close,
+    torch_device,
+)
+
+from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin
+
+
+enable_full_determinism()
+
+
+class AutoencoderKLTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
+    model_class = AutoencoderKL
+    main_input_name = "sample"
+    base_precision = 1e-2
+
+    def get_autoencoder_kl_config(self, block_out_channels=None, norm_num_groups=None):
+        block_out_channels = block_out_channels or [2, 4]
+        norm_num_groups = norm_num_groups or 2
+        init_dict = {
+            "block_out_channels": block_out_channels,
+            "in_channels": 3,
+            "out_channels": 3,
+            "down_block_types": ["DownEncoderBlock2D"] * len(block_out_channels),
+            "up_block_types": ["UpDecoderBlock2D"] * len(block_out_channels),
+            "latent_channels": 4,
+            "norm_num_groups": norm_num_groups,
+        }
+        return init_dict
+
+    @property
+    def dummy_input(self):
+        batch_size = 4
+        num_channels = 3
+        sizes = (32, 32)
+
+        image = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
+
+        return {"sample": image}
+
+    @property
+    def input_shape(self):
+        return (3, 32, 32)
+
+    @property
+    def output_shape(self):
+        return (3, 32, 32)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = self.get_autoencoder_kl_config()
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    def test_enable_disable_tiling(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+
+        inputs_dict.update({"return_dict": False})
+
+        torch.manual_seed(0)
+        output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        torch.manual_seed(0)
+        model.enable_tiling()
+        output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertLess(
+            (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(),
+            0.5,
+            "VAE tiling should not affect the inference results",
+        )
+
+        torch.manual_seed(0)
+        model.disable_tiling()
+        output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertEqual(
+            output_without_tiling.detach().cpu().numpy().all(),
+            output_without_tiling_2.detach().cpu().numpy().all(),
+            "Without tiling outputs should match with the outputs when tiling is manually disabled.",
+        )
+
+    def test_enable_disable_slicing(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+
+        inputs_dict.update({"return_dict": False})
+
+        torch.manual_seed(0)
+        output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        torch.manual_seed(0)
+        model.enable_slicing()
+        output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertLess(
+            (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(),
+            0.5,
+            "VAE slicing should not affect the inference results",
+        )
+
+        torch.manual_seed(0)
+        model.disable_slicing()
+        output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertEqual(
+            output_without_slicing.detach().cpu().numpy().all(),
+            output_without_slicing_2.detach().cpu().numpy().all(),
+            "Without slicing outputs should match with the outputs when slicing is manually disabled.",
+        )
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"Decoder", "Encoder", "UNetMidBlock2D"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+    def test_from_pretrained_hub(self):
+        model, loading_info = AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy", output_loading_info=True)
+        self.assertIsNotNone(model)
+        self.assertEqual(len(loading_info["missing_keys"]), 0)
+
+        model.to(torch_device)
+        image = model(**self.dummy_input)
+
+        assert image is not None, "Make sure output is not None"
+
+    def test_output_pretrained(self):
+        model = AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy")
+        model = model.to(torch_device)
+        model.eval()
+
+        # Keep generator on CPU for non-CUDA devices to compare outputs with CPU result tensors
+        generator_device = "cpu" if not torch_device.startswith("cuda") else "cuda"
+        if torch_device != "mps":
+            generator = torch.Generator(device=generator_device).manual_seed(0)
+        else:
+            generator = torch.manual_seed(0)
+
+        image = torch.randn(
+            1,
+            model.config.in_channels,
+            model.config.sample_size,
+            model.config.sample_size,
+            generator=torch.manual_seed(0),
+        )
+        image = image.to(torch_device)
+        with torch.no_grad():
+            output = model(image, sample_posterior=True, generator=generator).sample
+
+        output_slice = output[0, -1, -3:, -3:].flatten().cpu()
+
+        # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
+        # the expected output slices are not the same for CPU and GPU.
+        if torch_device == "mps":
+            expected_output_slice = torch.tensor(
+                [
+                    -4.0078e-01,
+                    -3.8323e-04,
+                    -1.2681e-01,
+                    -1.1462e-01,
+                    2.0095e-01,
+                    1.0893e-01,
+                    -8.8247e-02,
+                    -3.0361e-01,
+                    -9.8644e-03,
+                ]
+            )
+        elif generator_device == "cpu":
+            expected_output_slice = torch.tensor(
+                [
+                    -0.1352,
+                    0.0878,
+                    0.0419,
+                    -0.0818,
+                    -0.1069,
+                    0.0688,
+                    -0.1458,
+                    -0.4446,
+                    -0.0026,
+                ]
+            )
+        else:
+            expected_output_slice = torch.tensor(
+                [
+                    -0.2421,
+                    0.4642,
+                    0.2507,
+                    -0.0438,
+                    0.0682,
+                    0.3160,
+                    -0.2018,
+                    -0.0727,
+                    0.2485,
+                ]
+            )
+
+        self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2))
+
+
+@slow
+class AutoencoderKLIntegrationTests(unittest.TestCase):
+    def get_file_format(self, seed, shape):
+        return f"gaussian_noise_s={seed}_shape={'_'.join([str(s) for s in shape])}.npy"
+
+    def tearDown(self):
+        # clean up the VRAM after each test
+        super().tearDown()
+        gc.collect()
+        backend_empty_cache(torch_device)
+
+    def get_sd_image(self, seed=0, shape=(4, 3, 512, 512), fp16=False):
+        dtype = torch.float16 if fp16 else torch.float32
+        image = torch.from_numpy(load_hf_numpy(self.get_file_format(seed, shape))).to(torch_device).to(dtype)
+        return image
+
+    def get_sd_vae_model(self, model_id="CompVis/stable-diffusion-v1-4", fp16=False):
+        revision = "fp16" if fp16 else None
+        torch_dtype = torch.float16 if fp16 else torch.float32
+
+        model = AutoencoderKL.from_pretrained(
+            model_id,
+            subfolder="vae",
+            torch_dtype=torch_dtype,
+            revision=revision,
+        )
+        model.to(torch_device)
+
+        return model
+
+    def get_generator(self, seed=0):
+        generator_device = "cpu" if not torch_device.startswith("cuda") else "cuda"
+        if torch_device != "mps":
+            return torch.Generator(device=generator_device).manual_seed(seed)
+        return torch.manual_seed(seed)
+
+    @parameterized.expand(
+        [
+            # fmt: off
+            [
+                33,
+                [-0.1556, 0.9848, -0.0410, -0.0642, -0.2685, 0.8381, -0.2004, -0.0700],
+                [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824],
+            ],
+            [
+                47,
+                [-0.2376, 0.1200, 0.1337, -0.4830, -0.2504, -0.0759, -0.0486, -0.4077],
+                [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131],
+            ],
+            # fmt: on
+        ]
+    )
+    def test_stable_diffusion(self, seed, expected_slice, expected_slice_mps):
+        model = self.get_sd_vae_model()
+        image = self.get_sd_image(seed)
+        generator = self.get_generator(seed)
+
+        with torch.no_grad():
+            sample = model(image, generator=generator, sample_posterior=True).sample
+
+        assert sample.shape == image.shape
+
+        output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
+        expected_output_slice = torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice)
+
+        assert torch_all_close(output_slice, expected_output_slice, atol=3e-3)
+
+    @parameterized.expand(
+        [
+            # fmt: off
+            [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
+            [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
+            # fmt: on
+        ]
+    )
+    @require_torch_accelerator_with_fp16
+    def test_stable_diffusion_fp16(self, seed, expected_slice):
+        model = self.get_sd_vae_model(fp16=True)
+        image = self.get_sd_image(seed, fp16=True)
+        generator = self.get_generator(seed)
+
+        with torch.no_grad():
+            sample = model(image, generator=generator, sample_posterior=True).sample
+
+        assert sample.shape == image.shape
+
+        output_slice = sample[-1, -2:, :2, -2:].flatten().float().cpu()
+        expected_output_slice = torch.tensor(expected_slice)
+
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-2)
+
+    @parameterized.expand(
+        [
+            # fmt: off
+            [
+                33,
+                [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814],
+                [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824],
+            ],
+            [
+                47,
+                [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085],
+                [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131],
+            ],
+            # fmt: on
+        ]
+    )
+    def test_stable_diffusion_mode(self, seed, expected_slice, expected_slice_mps):
+        model = self.get_sd_vae_model()
+        image = self.get_sd_image(seed)
+
+        with torch.no_grad():
+            sample = model(image).sample
+
+        assert sample.shape == image.shape
+
+        output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
+        expected_output_slice = torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice)
+
+        assert torch_all_close(output_slice, expected_output_slice, atol=3e-3)
+
+    @parameterized.expand(
+        [
+            # fmt: off
+            [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
+            [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
+            # fmt: on
+        ]
+    )
+    @require_torch_accelerator
+    @skip_mps
+    def test_stable_diffusion_decode(self, seed, expected_slice):
+        model = self.get_sd_vae_model()
+        encoding = self.get_sd_image(seed, shape=(3, 4, 64, 64))
+
+        with torch.no_grad():
+            sample = model.decode(encoding).sample
+
+        assert list(sample.shape) == [3, 3, 512, 512]
+
+        output_slice = sample[-1, -2:, :2, -2:].flatten().cpu()
+        expected_output_slice = torch.tensor(expected_slice)
+
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-3)
+
+    @parameterized.expand(
+        [
+            # fmt: off
+            [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
+            [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
+            # fmt: on
+        ]
+    )
+    @require_torch_accelerator_with_fp16
+    def test_stable_diffusion_decode_fp16(self, seed, expected_slice):
+        model = self.get_sd_vae_model(fp16=True)
+        encoding = self.get_sd_image(seed, shape=(3, 4, 64, 64), fp16=True)
+
+        with torch.no_grad():
+            sample = model.decode(encoding).sample
+
+        assert list(sample.shape) == [3, 3, 512, 512]
+
+        output_slice = sample[-1, -2:, :2, -2:].flatten().float().cpu()
+        expected_output_slice = torch.tensor(expected_slice)
+
+        assert torch_all_close(output_slice, expected_output_slice, atol=5e-3)
+
+    @parameterized.expand([(13,), (16,), (27,)])
+    @require_torch_gpu
+    @unittest.skipIf(
+        not is_xformers_available(),
+        reason="xformers is not required when using PyTorch 2.0.",
+    )
+    def test_stable_diffusion_decode_xformers_vs_2_0_fp16(self, seed):
+        model = self.get_sd_vae_model(fp16=True)
+        encoding = self.get_sd_image(seed, shape=(3, 4, 64, 64), fp16=True)
+
+        with torch.no_grad():
+            sample = model.decode(encoding).sample
+
+        model.enable_xformers_memory_efficient_attention()
+        with torch.no_grad():
+            sample_2 = model.decode(encoding).sample
+
+        assert list(sample.shape) == [3, 3, 512, 512]
+
+        assert torch_all_close(sample, sample_2, atol=1e-1)
+
+    @parameterized.expand([(13,), (16,), (37,)])
+    @require_torch_gpu
+    @unittest.skipIf(
+        not is_xformers_available(),
+        reason="xformers is not required when using PyTorch 2.0.",
+    )
+    def test_stable_diffusion_decode_xformers_vs_2_0(self, seed):
+        model = self.get_sd_vae_model()
+        encoding = self.get_sd_image(seed, shape=(3, 4, 64, 64))
+
+        with torch.no_grad():
+            sample = model.decode(encoding).sample
+
+        model.enable_xformers_memory_efficient_attention()
+        with torch.no_grad():
+            sample_2 = model.decode(encoding).sample
+
+        assert list(sample.shape) == [3, 3, 512, 512]
+
+        assert torch_all_close(sample, sample_2, atol=1e-2)
+
+    @parameterized.expand(
+        [
+            # fmt: off
+            [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
+            [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
+            # fmt: on
+        ]
+    )
+    def test_stable_diffusion_encode_sample(self, seed, expected_slice):
+        model = self.get_sd_vae_model()
+        image = self.get_sd_image(seed)
+        generator = self.get_generator(seed)
+
+        with torch.no_grad():
+            dist = model.encode(image).latent_dist
+            sample = dist.sample(generator=generator)
+
+        assert list(sample.shape) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]
+
+        output_slice = sample[0, -1, -3:, -3:].flatten().cpu()
+        expected_output_slice = torch.tensor(expected_slice)
+
+        tolerance = 3e-3 if torch_device != "mps" else 1e-2
+        assert torch_all_close(output_slice, expected_output_slice, atol=tolerance)
@@ -0,0 +1,86 @@
+def get_autoencoder_kl_config(block_out_channels=None, norm_num_groups=None):
+    block_out_channels = block_out_channels or [2, 4]
+    norm_num_groups = norm_num_groups or 2
+    init_dict = {
+        "block_out_channels": block_out_channels,
+        "in_channels": 3,
+        "out_channels": 3,
+        "down_block_types": ["DownEncoderBlock2D"] * len(block_out_channels),
+        "up_block_types": ["UpDecoderBlock2D"] * len(block_out_channels),
+        "latent_channels": 4,
+        "norm_num_groups": norm_num_groups,
+    }
+    return init_dict
+
+
+def get_asym_autoencoder_kl_config(block_out_channels=None, norm_num_groups=None):
+    block_out_channels = block_out_channels or [2, 4]
+    norm_num_groups = norm_num_groups or 2
+    init_dict = {
+        "in_channels": 3,
+        "out_channels": 3,
+        "down_block_types": ["DownEncoderBlock2D"] * len(block_out_channels),
+        "down_block_out_channels": block_out_channels,
+        "layers_per_down_block": 1,
+        "up_block_types": ["UpDecoderBlock2D"] * len(block_out_channels),
+        "up_block_out_channels": block_out_channels,
+        "layers_per_up_block": 1,
+        "act_fn": "silu",
+        "latent_channels": 4,
+        "norm_num_groups": norm_num_groups,
+        "sample_size": 32,
+        "scaling_factor": 0.18215,
+    }
+    return init_dict
+
+
+def get_autoencoder_tiny_config(block_out_channels=None):
+    block_out_channels = (len(block_out_channels) * [32]) if block_out_channels is not None else [32, 32]
+    init_dict = {
+        "in_channels": 3,
+        "out_channels": 3,
+        "encoder_block_out_channels": block_out_channels,
+        "decoder_block_out_channels": block_out_channels,
+        "num_encoder_blocks": [b // min(block_out_channels) for b in block_out_channels],
+        "num_decoder_blocks": [b // min(block_out_channels) for b in reversed(block_out_channels)],
+    }
+    return init_dict
+
+
+def get_consistency_vae_config(block_out_channels=None, norm_num_groups=None):
+    block_out_channels = block_out_channels or [2, 4]
+    norm_num_groups = norm_num_groups or 2
+    return {
+        "encoder_block_out_channels": block_out_channels,
+        "encoder_in_channels": 3,
+        "encoder_out_channels": 4,
+        "encoder_down_block_types": ["DownEncoderBlock2D"] * len(block_out_channels),
+        "decoder_add_attention": False,
+        "decoder_block_out_channels": block_out_channels,
+        "decoder_down_block_types": ["ResnetDownsampleBlock2D"] * len(block_out_channels),
+        "decoder_downsample_padding": 1,
+        "decoder_in_channels": 7,
+        "decoder_layers_per_block": 1,
+        "decoder_norm_eps": 1e-05,
+        "decoder_norm_num_groups": norm_num_groups,
+        "encoder_norm_num_groups": norm_num_groups,
+        "decoder_num_train_timesteps": 1024,
+        "decoder_out_channels": 6,
+        "decoder_resnet_time_scale_shift": "scale_shift",
+        "decoder_time_embedding_type": "learned",
+        "decoder_up_block_types": ["ResnetUpsampleBlock2D"] * len(block_out_channels),
+        "scaling_factor": 1,
+        "latent_channels": 4,
+    }
+
+
+def get_autoencoder_oobleck_config(block_out_channels=None):
+    init_dict = {
+        "encoder_hidden_size": 12,
+        "decoder_channels": 12,
+        "decoder_input_channels": 6,
+        "audio_channels": 2,
+        "downsampling_ratios": [2, 4],
+        "channel_multiples": [1, 2],
+    }
+    return init_dict
@@ -0,0 +1,66 @@
+import inspect
+
+from diffusers.utils import is_flax_available
+from diffusers.utils.testing_utils import require_flax
+
+
+if is_flax_available():
+    import jax
+
+
+@require_flax
+class FlaxModelTesterMixin:
+    def test_output(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        model = self.model_class(**init_dict)
+        variables = model.init(inputs_dict["prng_key"], inputs_dict["sample"])
+        jax.lax.stop_gradient(variables)
+
+        output = model.apply(variables, inputs_dict["sample"])
+
+        if isinstance(output, dict):
+            output = output.sample
+
+        self.assertIsNotNone(output)
+        expected_shape = inputs_dict["sample"].shape
+        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
+
+    def test_forward_with_norm_groups(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        init_dict["norm_num_groups"] = 16
+        init_dict["block_out_channels"] = (16, 32)
+
+        model = self.model_class(**init_dict)
+        variables = model.init(inputs_dict["prng_key"], inputs_dict["sample"])
+        jax.lax.stop_gradient(variables)
+
+        output = model.apply(variables, inputs_dict["sample"])
+
+        if isinstance(output, dict):
+            output = output.sample
+
+        self.assertIsNotNone(output)
+        expected_shape = inputs_dict["sample"].shape
+        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
+
+    def test_deprecated_kwargs(self):
+        has_kwarg_in_model_class = "kwargs" in inspect.signature(self.model_class.__init__).parameters
+        has_deprecated_kwarg = len(self.model_class._deprecated_kwargs) > 0
+
+        if has_kwarg_in_model_class and not has_deprecated_kwarg:
+            raise ValueError(
+                f"{self.model_class} has `**kwargs` in its __init__ method but has not defined any deprecated kwargs"
+                " under the `_deprecated_kwargs` class attribute. Make sure to either remove `**kwargs` if there are"
+                " no deprecated arguments or add the deprecated argument with `_deprecated_kwargs ="
+                " [<deprecated_argument>]`"
+            )
+
+        if not has_kwarg_in_model_class and has_deprecated_kwarg:
+            raise ValueError(
+                f"{self.model_class} doesn't have `**kwargs` in its __init__ method but has defined deprecated kwargs"
+                " under the `_deprecated_kwargs` class attribute. Make sure to either add the `**kwargs` argument to"
+                f" {self.model_class}.__init__ if there are deprecated arguments or remove the deprecated argument"
+                " from `_deprecated_kwargs = [<deprecated_argument>]`"
+            )
@@ -0,0 +1,169 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import unittest
+
+import torch
+
+from diffusers import FluxTransformer2DModel
+from diffusers.models.attention_processor import FluxIPAdapterJointAttnProcessor2_0
+from diffusers.models.embeddings import ImageProjection
+from diffusers.utils.testing_utils import enable_full_determinism, torch_device
+
+from ..test_modeling_common import ModelTesterMixin
+
+
+enable_full_determinism()
+
+
+def create_flux_ip_adapter_state_dict(model):
+    # "ip_adapter" (cross-attention weights)
+    ip_cross_attn_state_dict = {}
+    key_id = 0
+
+    for name in model.attn_processors.keys():
+        if name.startswith("single_transformer_blocks"):
+            continue
+
+        joint_attention_dim = model.config["joint_attention_dim"]
+        hidden_size = model.config["num_attention_heads"] * model.config["attention_head_dim"]
+        sd = FluxIPAdapterJointAttnProcessor2_0(
+            hidden_size=hidden_size, cross_attention_dim=joint_attention_dim, scale=1.0
+        ).state_dict()
+        ip_cross_attn_state_dict.update(
+            {
+                f"{key_id}.to_k_ip.weight": sd["to_k_ip.0.weight"],
+                f"{key_id}.to_v_ip.weight": sd["to_v_ip.0.weight"],
+                f"{key_id}.to_k_ip.bias": sd["to_k_ip.0.bias"],
+                f"{key_id}.to_v_ip.bias": sd["to_v_ip.0.bias"],
+            }
+        )
+
+        key_id += 1
+
+    # "image_proj" (ImageProjection layer weights)
+
+    image_projection = ImageProjection(
+        cross_attention_dim=model.config["joint_attention_dim"],
+        image_embed_dim=model.config["pooled_projection_dim"],
+        num_image_text_embeds=4,
+    )
+
+    ip_image_projection_state_dict = {}
+    sd = image_projection.state_dict()
+    ip_image_projection_state_dict.update(
+        {
+            "proj.weight": sd["image_embeds.weight"],
+            "proj.bias": sd["image_embeds.bias"],
+            "norm.weight": sd["norm.weight"],
+            "norm.bias": sd["norm.bias"],
+        }
+    )
+
+    del sd
+    ip_state_dict = {}
+    ip_state_dict.update({"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict})
+    return ip_state_dict
+
+
+class FluxTransformerTests(ModelTesterMixin, unittest.TestCase):
+    model_class = FluxTransformer2DModel
+    main_input_name = "hidden_states"
+    # We override the items here because the transformer under consideration is small.
+    model_split_percents = [0.7, 0.6, 0.6]
+
+    # Skip setting testing with default: AttnProcessor
+    uses_custom_attn_processor = True
+
+    @property
+    def dummy_input(self):
+        batch_size = 1
+        num_latent_channels = 4
+        num_image_channels = 3
+        height = width = 4
+        sequence_length = 48
+        embedding_dim = 32
+
+        hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(torch_device)
+        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
+        pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(torch_device)
+        text_ids = torch.randn((sequence_length, num_image_channels)).to(torch_device)
+        image_ids = torch.randn((height * width, num_image_channels)).to(torch_device)
+        timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "img_ids": image_ids,
+            "txt_ids": text_ids,
+            "pooled_projections": pooled_prompt_embeds,
+            "timestep": timestep,
+        }
+
+    @property
+    def input_shape(self):
+        return (16, 4)
+
+    @property
+    def output_shape(self):
+        return (16, 4)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
+            "patch_size": 1,
+            "in_channels": 4,
+            "num_layers": 1,
+            "num_single_layers": 1,
+            "attention_head_dim": 16,
+            "num_attention_heads": 2,
+            "joint_attention_dim": 32,
+            "pooled_projection_dim": 32,
+            "axes_dims_rope": [4, 4, 8],
+        }
+
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    def test_deprecated_inputs_img_txt_ids_3d(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        model = self.model_class(**init_dict)
+        model.to(torch_device)
+        model.eval()
+
+        with torch.no_grad():
+            output_1 = model(**inputs_dict).to_tuple()[0]
+
+        # update inputs_dict with txt_ids and img_ids as 3d tensors (deprecated)
+        text_ids_3d = inputs_dict["txt_ids"].unsqueeze(0)
+        image_ids_3d = inputs_dict["img_ids"].unsqueeze(0)
+
+        assert text_ids_3d.ndim == 3, "text_ids_3d should be a 3d tensor"
+        assert image_ids_3d.ndim == 3, "img_ids_3d should be a 3d tensor"
+
+        inputs_dict["txt_ids"] = text_ids_3d
+        inputs_dict["img_ids"] = image_ids_3d
+
+        with torch.no_grad():
+            output_2 = model(**inputs_dict).to_tuple()[0]
+
+        self.assertEqual(output_1.shape, output_2.shape)
+        self.assertTrue(
+            torch.allclose(output_1, output_2, atol=1e-5),
+            msg="output with deprecated inputs (img_ids and txt_ids as 3d torch tensors) are not equal as them as 2d inputs",
+        )
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"FluxTransformer2DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
@@ -0,0 +1,89 @@
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import unittest
+
+import torch
+
+from diffusers import HunyuanVideoTransformer3DModel
+from diffusers.utils.testing_utils import enable_full_determinism, torch_device
+
+from ..test_modeling_common import ModelTesterMixin
+
+
+enable_full_determinism()
+
+
+class HunyuanVideoTransformer3DTests(ModelTesterMixin, unittest.TestCase):
+    model_class = HunyuanVideoTransformer3DModel
+    main_input_name = "hidden_states"
+    uses_custom_attn_processor = True
+
+    @property
+    def dummy_input(self):
+        batch_size = 1
+        num_channels = 4
+        num_frames = 1
+        height = 16
+        width = 16
+        text_encoder_embedding_dim = 16
+        pooled_projection_dim = 8
+        sequence_length = 12
+
+        hidden_states = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device)
+        timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device)
+        encoder_hidden_states = torch.randn((batch_size, sequence_length, text_encoder_embedding_dim)).to(torch_device)
+        pooled_projections = torch.randn((batch_size, pooled_projection_dim)).to(torch_device)
+        encoder_attention_mask = torch.ones((batch_size, sequence_length)).to(torch_device)
+        guidance = torch.randint(0, 1000, size=(batch_size,)).to(torch_device, dtype=torch.float32)
+
+        return {
+            "hidden_states": hidden_states,
+            "timestep": timestep,
+            "encoder_hidden_states": encoder_hidden_states,
+            "pooled_projections": pooled_projections,
+            "encoder_attention_mask": encoder_attention_mask,
+            "guidance": guidance,
+        }
+
+    @property
+    def input_shape(self):
+        return (4, 1, 16, 16)
+
+    @property
+    def output_shape(self):
+        return (4, 1, 16, 16)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
+            "in_channels": 4,
+            "out_channels": 4,
+            "num_attention_heads": 2,
+            "attention_head_dim": 10,
+            "num_layers": 1,
+            "num_single_layers": 1,
+            "num_refiner_layers": 1,
+            "patch_size": 1,
+            "patch_size_t": 1,
+            "guidance_embeds": True,
+            "text_embed_dim": 16,
+            "pooled_projection_dim": 8,
+            "rope_axes_dim": (2, 4, 4),
+        }
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"HunyuanVideoTransformer3DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
@@ -0,0 +1,86 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import unittest
+
+import torch
+
+from diffusers import MochiTransformer3DModel
+from diffusers.utils.testing_utils import enable_full_determinism, torch_device
+
+from ..test_modeling_common import ModelTesterMixin
+
+
+enable_full_determinism()
+
+
+class MochiTransformerTests(ModelTesterMixin, unittest.TestCase):
+    model_class = MochiTransformer3DModel
+    main_input_name = "hidden_states"
+    uses_custom_attn_processor = True
+    # Overriding it because of the transformer size.
+    model_split_percents = [0.7, 0.6, 0.6]
+
+    @property
+    def dummy_input(self):
+        batch_size = 2
+        num_channels = 4
+        num_frames = 2
+        height = 16
+        width = 16
+        embedding_dim = 16
+        sequence_length = 16
+
+        hidden_states = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device)
+        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
+        encoder_attention_mask = torch.ones((batch_size, sequence_length)).bool().to(torch_device)
+        timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "timestep": timestep,
+            "encoder_attention_mask": encoder_attention_mask,
+        }
+
+    @property
+    def input_shape(self):
+        return (4, 2, 16, 16)
+
+    @property
+    def output_shape(self):
+        return (4, 2, 16, 16)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
+            "patch_size": 2,
+            "num_attention_heads": 2,
+            "attention_head_dim": 8,
+            "num_layers": 2,
+            "pooled_projection_dim": 16,
+            "in_channels": 4,
+            "out_channels": None,
+            "qk_norm": "rms_norm",
+            "text_embed_dim": 16,
+            "time_embed_dim": 4,
+            "activation_fn": "swiglu",
+            "max_sequence_length": 16,
+        }
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"MochiTransformer3DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
@@ -0,0 +1,373 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import gc
+import math
+import unittest
+
+import torch
+
+from diffusers import UNet2DModel
+from diffusers.utils import logging
+from diffusers.utils.testing_utils import (
+    enable_full_determinism,
+    floats_tensor,
+    require_torch_accelerator,
+    slow,
+    torch_all_close,
+    torch_device,
+)
+
+from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin
+
+
+logger = logging.get_logger(__name__)
+
+enable_full_determinism()
+
+
+class Unet2DModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
+    model_class = UNet2DModel
+    main_input_name = "sample"
+
+    @property
+    def dummy_input(self):
+        batch_size = 4
+        num_channels = 3
+        sizes = (32, 32)
+
+        noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
+        time_step = torch.tensor([10]).to(torch_device)
+
+        return {"sample": noise, "timestep": time_step}
+
+    @property
+    def input_shape(self):
+        return (3, 32, 32)
+
+    @property
+    def output_shape(self):
+        return (3, 32, 32)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
+            "block_out_channels": (4, 8),
+            "norm_num_groups": 2,
+            "down_block_types": ("DownBlock2D", "AttnDownBlock2D"),
+            "up_block_types": ("AttnUpBlock2D", "UpBlock2D"),
+            "attention_head_dim": 3,
+            "out_channels": 3,
+            "in_channels": 3,
+            "layers_per_block": 2,
+            "sample_size": 32,
+        }
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    def test_mid_block_attn_groups(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        init_dict["add_attention"] = True
+        init_dict["attn_norm_num_groups"] = 4
+
+        model = self.model_class(**init_dict)
+        model.to(torch_device)
+        model.eval()
+
+        self.assertIsNotNone(
+            model.mid_block.attentions[0].group_norm, "Mid block Attention group norm should exist but does not."
+        )
+        self.assertEqual(
+            model.mid_block.attentions[0].group_norm.num_groups,
+            init_dict["attn_norm_num_groups"],
+            "Mid block Attention group norm does not have the expected number of groups.",
+        )
+
+        with torch.no_grad():
+            output = model(**inputs_dict)
+
+            if isinstance(output, dict):
+                output = output.to_tuple()[0]
+
+        self.assertIsNotNone(output)
+        expected_shape = inputs_dict["sample"].shape
+        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {
+            "AttnUpBlock2D",
+            "AttnDownBlock2D",
+            "UNetMidBlock2D",
+            "UpBlock2D",
+            "DownBlock2D",
+        }
+
+        # NOTE: unlike UNet2DConditionModel, UNet2DModel does not currently support tuples for `attention_head_dim`
+        attention_head_dim = 8
+        block_out_channels = (16, 32)
+
+        super().test_gradient_checkpointing_is_applied(
+            expected_set=expected_set, attention_head_dim=attention_head_dim, block_out_channels=block_out_channels
+        )
+
+
+class UNetLDMModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
+    model_class = UNet2DModel
+    main_input_name = "sample"
+
+    @property
+    def dummy_input(self):
+        batch_size = 4
+        num_channels = 4
+        sizes = (32, 32)
+
+        noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
+        time_step = torch.tensor([10]).to(torch_device)
+
+        return {"sample": noise, "timestep": time_step}
+
+    @property
+    def input_shape(self):
+        return (4, 32, 32)
+
+    @property
+    def output_shape(self):
+        return (4, 32, 32)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
+            "sample_size": 32,
+            "in_channels": 4,
+            "out_channels": 4,
+            "layers_per_block": 2,
+            "block_out_channels": (32, 64),
+            "attention_head_dim": 32,
+            "down_block_types": ("DownBlock2D", "DownBlock2D"),
+            "up_block_types": ("UpBlock2D", "UpBlock2D"),
+        }
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    def test_from_pretrained_hub(self):
+        model, loading_info = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update", output_loading_info=True)
+
+        self.assertIsNotNone(model)
+        self.assertEqual(len(loading_info["missing_keys"]), 0)
+
+        model.to(torch_device)
+        image = model(**self.dummy_input).sample
+
+        assert image is not None, "Make sure output is not None"
+
+    @require_torch_accelerator
+    def test_from_pretrained_accelerate(self):
+        model, _ = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update", output_loading_info=True)
+        model.to(torch_device)
+        image = model(**self.dummy_input).sample
+
+        assert image is not None, "Make sure output is not None"
+
+    @require_torch_accelerator
+    def test_from_pretrained_accelerate_wont_change_results(self):
+        # by default model loading will use accelerate as `low_cpu_mem_usage=True`
+        model_accelerate, _ = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update", output_loading_info=True)
+        model_accelerate.to(torch_device)
+        model_accelerate.eval()
+
+        noise = torch.randn(
+            1,
+            model_accelerate.config.in_channels,
+            model_accelerate.config.sample_size,
+            model_accelerate.config.sample_size,
+            generator=torch.manual_seed(0),
+        )
+        noise = noise.to(torch_device)
+        time_step = torch.tensor([10] * noise.shape[0]).to(torch_device)
+
+        arr_accelerate = model_accelerate(noise, time_step)["sample"]
+
+        # two models don't need to stay in the device at the same time
+        del model_accelerate
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        model_normal_load, _ = UNet2DModel.from_pretrained(
+            "fusing/unet-ldm-dummy-update", output_loading_info=True, low_cpu_mem_usage=False
+        )
+        model_normal_load.to(torch_device)
+        model_normal_load.eval()
+        arr_normal_load = model_normal_load(noise, time_step)["sample"]
+
+        assert torch_all_close(arr_accelerate, arr_normal_load, rtol=1e-3)
+
+    def test_output_pretrained(self):
+        model = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update")
+        model.eval()
+        model.to(torch_device)
+
+        noise = torch.randn(
+            1,
+            model.config.in_channels,
+            model.config.sample_size,
+            model.config.sample_size,
+            generator=torch.manual_seed(0),
+        )
+        noise = noise.to(torch_device)
+        time_step = torch.tensor([10] * noise.shape[0]).to(torch_device)
+
+        with torch.no_grad():
+            output = model(noise, time_step).sample
+
+        output_slice = output[0, -1, -3:, -3:].flatten().cpu()
+        # fmt: off
+        expected_output_slice = torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800])
+        # fmt: on
+
+        self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-3))
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"DownBlock2D", "UNetMidBlock2D", "UpBlock2D"}
+
+        # NOTE: unlike UNet2DConditionModel, UNet2DModel does not currently support tuples for `attention_head_dim`
+        attention_head_dim = 32
+        block_out_channels = (32, 64)
+
+        super().test_gradient_checkpointing_is_applied(
+            expected_set=expected_set, attention_head_dim=attention_head_dim, block_out_channels=block_out_channels
+        )
+
+
+class NCSNppModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
+    model_class = UNet2DModel
+    main_input_name = "sample"
+
+    @property
+    def dummy_input(self, sizes=(32, 32)):
+        batch_size = 4
+        num_channels = 3
+
+        noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
+        time_step = torch.tensor(batch_size * [10]).to(dtype=torch.int32, device=torch_device)
+
+        return {"sample": noise, "timestep": time_step}
+
+    @property
+    def input_shape(self):
+        return (3, 32, 32)
+
+    @property
+    def output_shape(self):
+        return (3, 32, 32)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
+            "block_out_channels": [32, 64, 64, 64],
+            "in_channels": 3,
+            "layers_per_block": 1,
+            "out_channels": 3,
+            "time_embedding_type": "fourier",
+            "norm_eps": 1e-6,
+            "mid_block_scale_factor": math.sqrt(2.0),
+            "norm_num_groups": None,
+            "down_block_types": [
+                "SkipDownBlock2D",
+                "AttnSkipDownBlock2D",
+                "SkipDownBlock2D",
+                "SkipDownBlock2D",
+            ],
+            "up_block_types": [
+                "SkipUpBlock2D",
+                "SkipUpBlock2D",
+                "AttnSkipUpBlock2D",
+                "SkipUpBlock2D",
+            ],
+        }
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    @slow
+    def test_from_pretrained_hub(self):
+        model, loading_info = UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256", output_loading_info=True)
+        self.assertIsNotNone(model)
+        self.assertEqual(len(loading_info["missing_keys"]), 0)
+
+        model.to(torch_device)
+        inputs = self.dummy_input
+        noise = floats_tensor((4, 3) + (256, 256)).to(torch_device)
+        inputs["sample"] = noise
+        image = model(**inputs)
+
+        assert image is not None, "Make sure output is not None"
+
+    @slow
+    def test_output_pretrained_ve_mid(self):
+        model = UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256")
+        model.to(torch_device)
+
+        batch_size = 4
+        num_channels = 3
+        sizes = (256, 256)
+
+        noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device)
+        time_step = torch.tensor(batch_size * [1e-4]).to(torch_device)
+
+        with torch.no_grad():
+            output = model(noise, time_step).sample
+
+        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
+        # fmt: off
+        expected_output_slice = torch.tensor([-4836.2178, -6487.1470, -3816.8196, -7964.9302, -10966.3037, -20043.5957, 8137.0513, 2340.3328, 544.6056])
+        # fmt: on
+
+        self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2))
+
+    def test_output_pretrained_ve_large(self):
+        model = UNet2DModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update")
+        model.to(torch_device)
+
+        batch_size = 4
+        num_channels = 3
+        sizes = (32, 32)
+
+        noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device)
+        time_step = torch.tensor(batch_size * [1e-4]).to(torch_device)
+
+        with torch.no_grad():
+            output = model(noise, time_step).sample
+
+        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
+        # fmt: off
+        expected_output_slice = torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256])
+        # fmt: on
+
+        self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2))
+
+    def test_forward_with_norm_groups(self):
+        # not required for this model
+        pass
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {
+            "UNetMidBlock2D",
+        }
+
+        block_out_channels = (32, 64, 64, 64)
+
+        super().test_gradient_checkpointing_is_applied(
+            expected_set=expected_set, block_out_channels=block_out_channels
+        )
+
+    def test_effective_gradient_checkpointing(self):
+        super().test_effective_gradient_checkpointing(skip={"time_proj.weight"})
@@ -0,0 +1,104 @@
+import gc
+import unittest
+
+from parameterized import parameterized
+
+from diffusers import FlaxUNet2DConditionModel
+from diffusers.utils import is_flax_available
+from diffusers.utils.testing_utils import load_hf_numpy, require_flax, slow
+
+
+if is_flax_available():
+    import jax
+    import jax.numpy as jnp
+
+
+@slow
+@require_flax
+class FlaxUNet2DConditionModelIntegrationTests(unittest.TestCase):
+    def get_file_format(self, seed, shape):
+        return f"gaussian_noise_s={seed}_shape={'_'.join([str(s) for s in shape])}.npy"
+
+    def tearDown(self):
+        # clean up the VRAM after each test
+        super().tearDown()
+        gc.collect()
+
+    def get_latents(self, seed=0, shape=(4, 4, 64, 64), fp16=False):
+        dtype = jnp.bfloat16 if fp16 else jnp.float32
+        image = jnp.array(load_hf_numpy(self.get_file_format(seed, shape)), dtype=dtype)
+        return image
+
+    def get_unet_model(self, fp16=False, model_id="CompVis/stable-diffusion-v1-4"):
+        dtype = jnp.bfloat16 if fp16 else jnp.float32
+        revision = "bf16" if fp16 else None
+
+        model, params = FlaxUNet2DConditionModel.from_pretrained(
+            model_id, subfolder="unet", dtype=dtype, revision=revision
+        )
+        return model, params
+
+    def get_encoder_hidden_states(self, seed=0, shape=(4, 77, 768), fp16=False):
+        dtype = jnp.bfloat16 if fp16 else jnp.float32
+        hidden_states = jnp.array(load_hf_numpy(self.get_file_format(seed, shape)), dtype=dtype)
+        return hidden_states
+
+    @parameterized.expand(
+        [
+            # fmt: off
+            [83, 4, [-0.2323, -0.1304, 0.0813, -0.3093, -0.0919, -0.1571, -0.1125, -0.5806]],
+            [17, 0.55, [-0.0831, -0.2443, 0.0901, -0.0919, 0.3396, 0.0103, -0.3743, 0.0701]],
+            [8, 0.89, [-0.4863, 0.0859, 0.0875, -0.1658, 0.9199, -0.0114, 0.4839, 0.4639]],
+            [3, 1000, [-0.5649, 0.2402, -0.5518, 0.1248, 1.1328, -0.2443, -0.0325, -1.0078]],
+            # fmt: on
+        ]
+    )
+    def test_compvis_sd_v1_4_flax_vs_torch_fp16(self, seed, timestep, expected_slice):
+        model, params = self.get_unet_model(model_id="CompVis/stable-diffusion-v1-4", fp16=True)
+        latents = self.get_latents(seed, fp16=True)
+        encoder_hidden_states = self.get_encoder_hidden_states(seed, fp16=True)
+
+        sample = model.apply(
+            {"params": params},
+            latents,
+            jnp.array(timestep, dtype=jnp.int32),
+            encoder_hidden_states=encoder_hidden_states,
+        ).sample
+
+        assert sample.shape == latents.shape
+
+        output_slice = jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())), dtype=jnp.float32)
+        expected_output_slice = jnp.array(expected_slice, dtype=jnp.float32)
+
+        # Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, in the same hardware
+        assert jnp.allclose(output_slice, expected_output_slice, atol=1e-2)
+
+    @parameterized.expand(
+        [
+            # fmt: off
+            [83, 4, [0.1514, 0.0807, 0.1624, 0.1016, -0.1896, 0.0263, 0.0677, 0.2310]],
+            [17, 0.55, [0.1164, -0.0216, 0.0170, 0.1589, -0.3120, 0.1005, -0.0581, -0.1458]],
+            [8, 0.89, [-0.1758, -0.0169, 0.1004, -0.1411, 0.1312, 0.1103, -0.1996, 0.2139]],
+            [3, 1000, [0.1214, 0.0352, -0.0731, -0.1562, -0.0994, -0.0906, -0.2340, -0.0539]],
+            # fmt: on
+        ]
+    )
+    def test_stabilityai_sd_v2_flax_vs_torch_fp16(self, seed, timestep, expected_slice):
+        model, params = self.get_unet_model(model_id="stabilityai/stable-diffusion-2", fp16=True)
+        latents = self.get_latents(seed, shape=(4, 4, 96, 96), fp16=True)
+        encoder_hidden_states = self.get_encoder_hidden_states(seed, shape=(4, 77, 1024), fp16=True)
+
+        sample = model.apply(
+            {"params": params},
+            latents,
+            jnp.array(timestep, dtype=jnp.int32),
+            encoder_hidden_states=encoder_hidden_states,
+        ).sample
+
+        assert sample.shape == latents.shape
+
+        output_slice = jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())), dtype=jnp.float32)
+        expected_output_slice = jnp.array(expected_slice, dtype=jnp.float32)
+
+        # Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, on the same hardware
+        assert jnp.allclose(output_slice, expected_output_slice, atol=1e-2)
@@ -0,0 +1,126 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc.
+#
+# 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 typing import Tuple
+
+import torch
+
+from diffusers.utils.testing_utils import (
+    floats_tensor,
+    require_torch,
+    require_torch_accelerator_with_training,
+    torch_all_close,
+    torch_device,
+)
+from diffusers.utils.torch_utils import randn_tensor
+
+
+@require_torch
+class UNetBlockTesterMixin:
+    @property
+    def dummy_input(self):
+        return self.get_dummy_input()
+
+    @property
+    def output_shape(self):
+        if self.block_type == "down":
+            return (4, 32, 16, 16)
+        elif self.block_type == "mid":
+            return (4, 32, 32, 32)
+        elif self.block_type == "up":
+            return (4, 32, 64, 64)
+
+        raise ValueError(f"'{self.block_type}' is not a supported block_type. Set it to 'up', 'mid', or 'down'.")
+
+    def get_dummy_input(
+        self,
+        include_temb=True,
+        include_res_hidden_states_tuple=False,
+        include_encoder_hidden_states=False,
+        include_skip_sample=False,
+    ):
+        batch_size = 4
+        num_channels = 32
+        sizes = (32, 32)
+
+        generator = torch.manual_seed(0)
+        device = torch.device(torch_device)
+        shape = (batch_size, num_channels) + sizes
+        hidden_states = randn_tensor(shape, generator=generator, device=device)
+        dummy_input = {"hidden_states": hidden_states}
+
+        if include_temb:
+            temb_channels = 128
+            dummy_input["temb"] = randn_tensor((batch_size, temb_channels), generator=generator, device=device)
+
+        if include_res_hidden_states_tuple:
+            generator_1 = torch.manual_seed(1)
+            dummy_input["res_hidden_states_tuple"] = (randn_tensor(shape, generator=generator_1, device=device),)
+
+        if include_encoder_hidden_states:
+            dummy_input["encoder_hidden_states"] = floats_tensor((batch_size, 32, 32)).to(torch_device)
+
+        if include_skip_sample:
+            dummy_input["skip_sample"] = randn_tensor(((batch_size, 3) + sizes), generator=generator, device=device)
+
+        return dummy_input
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
+            "in_channels": 32,
+            "out_channels": 32,
+            "temb_channels": 128,
+        }
+        if self.block_type == "up":
+            init_dict["prev_output_channel"] = 32
+
+        if self.block_type == "mid":
+            init_dict.pop("out_channels")
+
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
+
+    def test_output(self, expected_slice):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        unet_block = self.block_class(**init_dict)
+        unet_block.to(torch_device)
+        unet_block.eval()
+
+        with torch.no_grad():
+            output = unet_block(**inputs_dict)
+
+        if isinstance(output, Tuple):
+            output = output[0]
+
+        self.assertEqual(output.shape, self.output_shape)
+
+        output_slice = output[0, -1, -3:, -3:]
+        expected_slice = torch.tensor(expected_slice).to(torch_device)
+        assert torch_all_close(output_slice.flatten(), expected_slice, atol=5e-3)
+
+    @require_torch_accelerator_with_training
+    def test_training(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        model = self.block_class(**init_dict)
+        model.to(torch_device)
+        model.train()
+        output = model(**inputs_dict)
+
+        if isinstance(output, Tuple):
+            output = output[0]
+
+        device = torch.device(torch_device)
+        noise = randn_tensor(output.shape, device=device)
+        loss = torch.nn.functional.mse_loss(output, noise)
+        loss.backward()
@@ -0,0 +1,213 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import os
+import unittest
+
+import pytest
+
+from diffusers import __version__
+from diffusers.utils import deprecate
+from diffusers.utils.testing_utils import str_to_bool
+
+
+# Used to test the hub
+USER = "__DUMMY_TRANSFORMERS_USER__"
+ENDPOINT_STAGING = "https://hub-ci.huggingface.co"
+
+# Not critical, only usable on the sandboxed CI instance.
+TOKEN = "hf_94wBhPGp6KrrTH3KDchhKpRxZwd6dmHWLL"
+
+
+class DeprecateTester(unittest.TestCase):
+    higher_version = ".".join([str(int(__version__.split(".")[0]) + 1)] + __version__.split(".")[1:])
+    lower_version = "0.0.1"
+
+    def test_deprecate_function_arg(self):
+        kwargs = {"deprecated_arg": 4}
+
+        with self.assertWarns(FutureWarning) as warning:
+            output = deprecate("deprecated_arg", self.higher_version, "message", take_from=kwargs)
+
+        assert output == 4
+        assert (
+            str(warning.warning)
+            == f"The `deprecated_arg` argument is deprecated and will be removed in version {self.higher_version}."
+            " message"
+        )
+
+    def test_deprecate_function_arg_tuple(self):
+        kwargs = {"deprecated_arg": 4}
+
+        with self.assertWarns(FutureWarning) as warning:
+            output = deprecate(("deprecated_arg", self.higher_version, "message"), take_from=kwargs)
+
+        assert output == 4
+        assert (
+            str(warning.warning)
+            == f"The `deprecated_arg` argument is deprecated and will be removed in version {self.higher_version}."
+            " message"
+        )
+
+    def test_deprecate_function_args(self):
+        kwargs = {"deprecated_arg_1": 4, "deprecated_arg_2": 8}
+        with self.assertWarns(FutureWarning) as warning:
+            output_1, output_2 = deprecate(
+                ("deprecated_arg_1", self.higher_version, "Hey"),
+                ("deprecated_arg_2", self.higher_version, "Hey"),
+                take_from=kwargs,
+            )
+        assert output_1 == 4
+        assert output_2 == 8
+        assert (
+            str(warning.warnings[0].message)
+            == "The `deprecated_arg_1` argument is deprecated and will be removed in version"
+            f" {self.higher_version}. Hey"
+        )
+        assert (
+            str(warning.warnings[1].message)
+            == "The `deprecated_arg_2` argument is deprecated and will be removed in version"
+            f" {self.higher_version}. Hey"
+        )
+
+    def test_deprecate_function_incorrect_arg(self):
+        kwargs = {"deprecated_arg": 4}
+
+        with self.assertRaises(TypeError) as error:
+            deprecate(("wrong_arg", self.higher_version, "message"), take_from=kwargs)
+
+        assert "test_deprecate_function_incorrect_arg in" in str(error.exception)
+        assert "line" in str(error.exception)
+        assert "got an unexpected keyword argument `deprecated_arg`" in str(error.exception)
+
+    def test_deprecate_arg_no_kwarg(self):
+        with self.assertWarns(FutureWarning) as warning:
+            deprecate(("deprecated_arg", self.higher_version, "message"))
+
+        assert (
+            str(warning.warning)
+            == f"`deprecated_arg` is deprecated and will be removed in version {self.higher_version}. message"
+        )
+
+    def test_deprecate_args_no_kwarg(self):
+        with self.assertWarns(FutureWarning) as warning:
+            deprecate(
+                ("deprecated_arg_1", self.higher_version, "Hey"),
+                ("deprecated_arg_2", self.higher_version, "Hey"),
+            )
+        assert (
+            str(warning.warnings[0].message)
+            == f"`deprecated_arg_1` is deprecated and will be removed in version {self.higher_version}. Hey"
+        )
+        assert (
+            str(warning.warnings[1].message)
+            == f"`deprecated_arg_2` is deprecated and will be removed in version {self.higher_version}. Hey"
+        )
+
+    def test_deprecate_class_obj(self):
+        class Args:
+            arg = 5
+
+        with self.assertWarns(FutureWarning) as warning:
+            arg = deprecate(("arg", self.higher_version, "message"), take_from=Args())
+
+        assert arg == 5
+        assert (
+            str(warning.warning)
+            == f"The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message"
+        )
+
+    def test_deprecate_class_objs(self):
+        class Args:
+            arg = 5
+            foo = 7
+
+        with self.assertWarns(FutureWarning) as warning:
+            arg_1, arg_2 = deprecate(
+                ("arg", self.higher_version, "message"),
+                ("foo", self.higher_version, "message"),
+                ("does not exist", self.higher_version, "message"),
+                take_from=Args(),
+            )
+
+        assert arg_1 == 5
+        assert arg_2 == 7
+        assert (
+            str(warning.warning)
+            == f"The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message"
+        )
+        assert (
+            str(warning.warnings[0].message)
+            == f"The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message"
+        )
+        assert (
+            str(warning.warnings[1].message)
+            == f"The `foo` attribute is deprecated and will be removed in version {self.higher_version}. message"
+        )
+
+    def test_deprecate_incorrect_version(self):
+        kwargs = {"deprecated_arg": 4}
+
+        with self.assertRaises(ValueError) as error:
+            deprecate(("wrong_arg", self.lower_version, "message"), take_from=kwargs)
+
+        assert (
+            str(error.exception)
+            == "The deprecation tuple ('wrong_arg', '0.0.1', 'message') should be removed since diffusers' version"
+            f" {__version__} is >= {self.lower_version}"
+        )
+
+    def test_deprecate_incorrect_no_standard_warn(self):
+        with self.assertWarns(FutureWarning) as warning:
+            deprecate(("deprecated_arg", self.higher_version, "This message is better!!!"), standard_warn=False)
+
+        assert str(warning.warning) == "This message is better!!!"
+
+    def test_deprecate_stacklevel(self):
+        with self.assertWarns(FutureWarning) as warning:
+            deprecate(("deprecated_arg", self.higher_version, "This message is better!!!"), standard_warn=False)
+        assert str(warning.warning) == "This message is better!!!"
+        assert "diffusers/tests/others/test_utils.py" in warning.filename
+
+
+def parse_flag_from_env(key, default=False):
+    try:
+        value = os.environ[key]
+    except KeyError:
+        # KEY isn't set, default to `default`.
+        _value = default
+    else:
+        # KEY is set, convert it to True or False.
+        try:
+            _value = str_to_bool(value)
+        except ValueError:
+            # More values are supported, but let's keep the message simple.
+            raise ValueError(f"If set, {key} must be yes or no.")
+    return _value
+
+
+_run_staging = parse_flag_from_env("HUGGINGFACE_CO_STAGING", default=False)
+
+
+def is_staging_test(test_case):
+    """
+    Decorator marking a test as a staging test.
+
+    Those tests will run using the staging environment of huggingface.co instead of the real model hub.
+    """
+    if not _run_staging:
+        return unittest.skip("test is staging test")(test_case)
+    else:
+        return pytest.mark.is_staging_test()(test_case)
@@ -0,0 +1,331 @@
+# Copyright 2024 The HuggingFace Team.
+#
+# 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.
+
+import inspect
+import unittest
+
+import numpy as np
+import torch
+from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer, LlamaConfig, LlamaModel, LlamaTokenizer
+
+from diffusers import (
+    AutoencoderKLHunyuanVideo,
+    FlowMatchEulerDiscreteScheduler,
+    HunyuanVideoPipeline,
+    HunyuanVideoTransformer3DModel,
+)
+from diffusers.utils.testing_utils import (
+    enable_full_determinism,
+    torch_device,
+)
+
+from ..test_pipelines_common import PipelineTesterMixin, to_np
+
+
+enable_full_determinism()
+
+
+class HunyuanVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+    pipeline_class = HunyuanVideoPipeline
+    params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"])
+    batch_params = frozenset(["prompt"])
+    required_optional_params = frozenset(
+        [
+            "num_inference_steps",
+            "generator",
+            "latents",
+            "return_dict",
+            "callback_on_step_end",
+            "callback_on_step_end_tensor_inputs",
+        ]
+    )
+
+    # there is no xformers processor for Flux
+    test_xformers_attention = False
+
+    def get_dummy_components(self):
+        torch.manual_seed(0)
+        transformer = HunyuanVideoTransformer3DModel(
+            in_channels=4,
+            out_channels=4,
+            num_attention_heads=2,
+            attention_head_dim=10,
+            num_layers=1,
+            num_single_layers=1,
+            num_refiner_layers=1,
+            patch_size=1,
+            patch_size_t=1,
+            guidance_embeds=True,
+            text_embed_dim=16,
+            pooled_projection_dim=8,
+            rope_axes_dim=(2, 4, 4),
+        )
+
+        torch.manual_seed(0)
+        vae = AutoencoderKLHunyuanVideo(
+            in_channels=3,
+            out_channels=3,
+            latent_channels=4,
+            down_block_types=(
+                "HunyuanVideoDownBlock3D",
+                "HunyuanVideoDownBlock3D",
+                "HunyuanVideoDownBlock3D",
+                "HunyuanVideoDownBlock3D",
+            ),
+            up_block_types=(
+                "HunyuanVideoUpBlock3D",
+                "HunyuanVideoUpBlock3D",
+                "HunyuanVideoUpBlock3D",
+                "HunyuanVideoUpBlock3D",
+            ),
+            block_out_channels=(8, 8, 8, 8),
+            layers_per_block=1,
+            act_fn="silu",
+            norm_num_groups=4,
+            scaling_factor=0.476986,
+            spatial_compression_ratio=8,
+            temporal_compression_ratio=4,
+            mid_block_add_attention=True,
+        )
+
+        torch.manual_seed(0)
+        scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0)
+
+        llama_text_encoder_config = LlamaConfig(
+            bos_token_id=0,
+            eos_token_id=2,
+            hidden_size=16,
+            intermediate_size=37,
+            layer_norm_eps=1e-05,
+            num_attention_heads=4,
+            num_hidden_layers=2,
+            pad_token_id=1,
+            vocab_size=1000,
+            hidden_act="gelu",
+            projection_dim=32,
+        )
+        clip_text_encoder_config = CLIPTextConfig(
+            bos_token_id=0,
+            eos_token_id=2,
+            hidden_size=8,
+            intermediate_size=37,
+            layer_norm_eps=1e-05,
+            num_attention_heads=4,
+            num_hidden_layers=2,
+            pad_token_id=1,
+            vocab_size=1000,
+            hidden_act="gelu",
+            projection_dim=32,
+        )
+
+        torch.manual_seed(0)
+        text_encoder = LlamaModel(llama_text_encoder_config)
+        tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/tiny-random-LlamaForCausalLM")
+
+        torch.manual_seed(0)
+        text_encoder_2 = CLIPTextModel(clip_text_encoder_config)
+        tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
+
+        components = {
+            "transformer": transformer,
+            "vae": vae,
+            "scheduler": scheduler,
+            "text_encoder": text_encoder,
+            "text_encoder_2": text_encoder_2,
+            "tokenizer": tokenizer,
+            "tokenizer_2": tokenizer_2,
+        }
+        return components
+
+    def get_dummy_inputs(self, device, seed=0):
+        if str(device).startswith("mps"):
+            generator = torch.manual_seed(seed)
+        else:
+            generator = torch.Generator(device=device).manual_seed(seed)
+
+        # Cannot test with dummy prompt because tokenizers are not configured correctly.
+        # TODO(aryan): create dummy tokenizers and using from hub
+        inputs = {
+            "prompt": "",
+            "prompt_template": {
+                "template": "{}",
+                "crop_start": 0,
+            },
+            "generator": generator,
+            "num_inference_steps": 2,
+            "guidance_scale": 4.5,
+            "height": 16,
+            "width": 16,
+            # 4 * k + 1 is the recommendation
+            "num_frames": 9,
+            "max_sequence_length": 16,
+            "output_type": "pt",
+        }
+        return inputs
+
+    def test_inference(self):
+        device = "cpu"
+
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        pipe.to(device)
+        pipe.set_progress_bar_config(disable=None)
+
+        inputs = self.get_dummy_inputs(device)
+        video = pipe(**inputs).frames
+        generated_video = video[0]
+
+        self.assertEqual(generated_video.shape, (9, 3, 16, 16))
+        expected_video = torch.randn(9, 3, 16, 16)
+        max_diff = np.abs(generated_video - expected_video).max()
+        self.assertLessEqual(max_diff, 1e10)
+
+    def test_callback_inputs(self):
+        sig = inspect.signature(self.pipeline_class.__call__)
+        has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
+        has_callback_step_end = "callback_on_step_end" in sig.parameters
+
+        if not (has_callback_tensor_inputs and has_callback_step_end):
+            return
+
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        pipe = pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+        self.assertTrue(
+            hasattr(pipe, "_callback_tensor_inputs"),
+            f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
+        )
+
+        def callback_inputs_subset(pipe, i, t, callback_kwargs):
+            # iterate over callback args
+            for tensor_name, tensor_value in callback_kwargs.items():
+                # check that we're only passing in allowed tensor inputs
+                assert tensor_name in pipe._callback_tensor_inputs
+
+            return callback_kwargs
+
+        def callback_inputs_all(pipe, i, t, callback_kwargs):
+            for tensor_name in pipe._callback_tensor_inputs:
+                assert tensor_name in callback_kwargs
+
+            # iterate over callback args
+            for tensor_name, tensor_value in callback_kwargs.items():
+                # check that we're only passing in allowed tensor inputs
+                assert tensor_name in pipe._callback_tensor_inputs
+
+            return callback_kwargs
+
+        inputs = self.get_dummy_inputs(torch_device)
+
+        # Test passing in a subset
+        inputs["callback_on_step_end"] = callback_inputs_subset
+        inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
+        output = pipe(**inputs)[0]
+
+        # Test passing in a everything
+        inputs["callback_on_step_end"] = callback_inputs_all
+        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
+        output = pipe(**inputs)[0]
+
+        def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
+            is_last = i == (pipe.num_timesteps - 1)
+            if is_last:
+                callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
+            return callback_kwargs
+
+        inputs["callback_on_step_end"] = callback_inputs_change_tensor
+        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
+        output = pipe(**inputs)[0]
+        assert output.abs().sum() < 1e10
+
+    def test_attention_slicing_forward_pass(
+        self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
+    ):
+        if not self.test_attention_slicing:
+            return
+
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        for component in pipe.components.values():
+            if hasattr(component, "set_default_attn_processor"):
+                component.set_default_attn_processor()
+        pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        generator_device = "cpu"
+        inputs = self.get_dummy_inputs(generator_device)
+        output_without_slicing = pipe(**inputs)[0]
+
+        pipe.enable_attention_slicing(slice_size=1)
+        inputs = self.get_dummy_inputs(generator_device)
+        output_with_slicing1 = pipe(**inputs)[0]
+
+        pipe.enable_attention_slicing(slice_size=2)
+        inputs = self.get_dummy_inputs(generator_device)
+        output_with_slicing2 = pipe(**inputs)[0]
+
+        if test_max_difference:
+            max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
+            max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
+            self.assertLess(
+                max(max_diff1, max_diff2),
+                expected_max_diff,
+                "Attention slicing should not affect the inference results",
+            )
+
+    def test_vae_tiling(self, expected_diff_max: float = 0.2):
+        # Seems to require higher tolerance than the other tests
+        expected_diff_max = 0.6
+        generator_device = "cpu"
+        components = self.get_dummy_components()
+
+        pipe = self.pipeline_class(**components)
+        pipe.to("cpu")
+        pipe.set_progress_bar_config(disable=None)
+
+        # Without tiling
+        inputs = self.get_dummy_inputs(generator_device)
+        inputs["height"] = inputs["width"] = 128
+        output_without_tiling = pipe(**inputs)[0]
+
+        # With tiling
+        pipe.vae.enable_tiling(
+            tile_sample_min_height=96,
+            tile_sample_min_width=96,
+            tile_sample_stride_height=64,
+            tile_sample_stride_width=64,
+        )
+        inputs = self.get_dummy_inputs(generator_device)
+        inputs["height"] = inputs["width"] = 128
+        output_with_tiling = pipe(**inputs)[0]
+
+        self.assertLess(
+            (to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
+            expected_diff_max,
+            "VAE tiling should not affect the inference results",
+        )
+
+    # TODO(aryan): Create a dummy gemma model with smol vocab size
+    @unittest.skip(
+        "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+    )
+    def test_inference_batch_consistent(self):
+        pass
+
+    @unittest.skip(
+        "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+    )
+    def test_inference_batch_single_identical(self):
+        pass
@@ -0,0 +1,299 @@
+# Copyright 2024 The HuggingFace Team.
+#
+# 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.
+
+import gc
+import inspect
+import unittest
+
+import numpy as np
+import torch
+from transformers import AutoTokenizer, T5EncoderModel
+
+from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel
+from diffusers.utils.testing_utils import (
+    enable_full_determinism,
+    numpy_cosine_similarity_distance,
+    require_torch_gpu,
+    slow,
+    torch_device,
+)
+
+from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
+from ..test_pipelines_common import PipelineTesterMixin, to_np
+
+
+enable_full_determinism()
+
+
+class MochiPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+    pipeline_class = MochiPipeline
+    params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
+    batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
+    image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
+    image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
+    required_optional_params = frozenset(
+        [
+            "num_inference_steps",
+            "generator",
+            "latents",
+            "return_dict",
+            "callback_on_step_end",
+            "callback_on_step_end_tensor_inputs",
+        ]
+    )
+    test_xformers_attention = False
+
+    def get_dummy_components(self):
+        torch.manual_seed(0)
+        transformer = MochiTransformer3DModel(
+            patch_size=2,
+            num_attention_heads=2,
+            attention_head_dim=8,
+            num_layers=2,
+            pooled_projection_dim=16,
+            in_channels=12,
+            out_channels=None,
+            qk_norm="rms_norm",
+            text_embed_dim=32,
+            time_embed_dim=4,
+            activation_fn="swiglu",
+            max_sequence_length=16,
+        )
+        transformer.pos_frequencies.data = transformer.pos_frequencies.new_full(transformer.pos_frequencies.shape, 0)
+
+        torch.manual_seed(0)
+        vae = AutoencoderKLMochi(
+            latent_channels=12,
+            out_channels=3,
+            encoder_block_out_channels=(32, 32, 32, 32),
+            decoder_block_out_channels=(32, 32, 32, 32),
+            layers_per_block=(1, 1, 1, 1, 1),
+        )
+
+        torch.manual_seed(0)
+        scheduler = FlowMatchEulerDiscreteScheduler()
+        text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
+
+        components = {
+            "transformer": transformer,
+            "vae": vae,
+            "scheduler": scheduler,
+            "text_encoder": text_encoder,
+            "tokenizer": tokenizer,
+        }
+        return components
+
+    def get_dummy_inputs(self, device, seed=0):
+        if str(device).startswith("mps"):
+            generator = torch.manual_seed(seed)
+        else:
+            generator = torch.Generator(device=device).manual_seed(seed)
+        inputs = {
+            "prompt": "dance monkey",
+            "negative_prompt": "",
+            "generator": generator,
+            "num_inference_steps": 2,
+            "guidance_scale": 4.5,
+            "height": 16,
+            "width": 16,
+            # 6 * k + 1 is the recommendation
+            "num_frames": 7,
+            "max_sequence_length": 16,
+            "output_type": "pt",
+        }
+        return inputs
+
+    def test_inference(self):
+        device = "cpu"
+
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        pipe.to(device)
+        pipe.set_progress_bar_config(disable=None)
+
+        inputs = self.get_dummy_inputs(device)
+        video = pipe(**inputs).frames
+        generated_video = video[0]
+
+        self.assertEqual(generated_video.shape, (7, 3, 16, 16))
+        expected_video = torch.randn(7, 3, 16, 16)
+        max_diff = np.abs(generated_video - expected_video).max()
+        self.assertLessEqual(max_diff, 1e10)
+
+    def test_callback_inputs(self):
+        sig = inspect.signature(self.pipeline_class.__call__)
+        has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
+        has_callback_step_end = "callback_on_step_end" in sig.parameters
+
+        if not (has_callback_tensor_inputs and has_callback_step_end):
+            return
+
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        pipe = pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+        self.assertTrue(
+            hasattr(pipe, "_callback_tensor_inputs"),
+            f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
+        )
+
+        def callback_inputs_subset(pipe, i, t, callback_kwargs):
+            # iterate over callback args
+            for tensor_name, tensor_value in callback_kwargs.items():
+                # check that we're only passing in allowed tensor inputs
+                assert tensor_name in pipe._callback_tensor_inputs
+
+            return callback_kwargs
+
+        def callback_inputs_all(pipe, i, t, callback_kwargs):
+            for tensor_name in pipe._callback_tensor_inputs:
+                assert tensor_name in callback_kwargs
+
+            # iterate over callback args
+            for tensor_name, tensor_value in callback_kwargs.items():
+                # check that we're only passing in allowed tensor inputs
+                assert tensor_name in pipe._callback_tensor_inputs
+
+            return callback_kwargs
+
+        inputs = self.get_dummy_inputs(torch_device)
+
+        # Test passing in a subset
+        inputs["callback_on_step_end"] = callback_inputs_subset
+        inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
+        output = pipe(**inputs)[0]
+
+        # Test passing in a everything
+        inputs["callback_on_step_end"] = callback_inputs_all
+        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
+        output = pipe(**inputs)[0]
+
+        def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
+            is_last = i == (pipe.num_timesteps - 1)
+            if is_last:
+                callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
+            return callback_kwargs
+
+        inputs["callback_on_step_end"] = callback_inputs_change_tensor
+        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
+        output = pipe(**inputs)[0]
+        assert output.abs().sum() < 1e10
+
+    def test_inference_batch_single_identical(self):
+        self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3)
+
+    def test_attention_slicing_forward_pass(
+        self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
+    ):
+        if not self.test_attention_slicing:
+            return
+
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        for component in pipe.components.values():
+            if hasattr(component, "set_default_attn_processor"):
+                component.set_default_attn_processor()
+        pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        generator_device = "cpu"
+        inputs = self.get_dummy_inputs(generator_device)
+        output_without_slicing = pipe(**inputs)[0]
+
+        pipe.enable_attention_slicing(slice_size=1)
+        inputs = self.get_dummy_inputs(generator_device)
+        output_with_slicing1 = pipe(**inputs)[0]
+
+        pipe.enable_attention_slicing(slice_size=2)
+        inputs = self.get_dummy_inputs(generator_device)
+        output_with_slicing2 = pipe(**inputs)[0]
+
+        if test_max_difference:
+            max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
+            max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
+            self.assertLess(
+                max(max_diff1, max_diff2),
+                expected_max_diff,
+                "Attention slicing should not affect the inference results",
+            )
+
+    def test_vae_tiling(self, expected_diff_max: float = 0.2):
+        generator_device = "cpu"
+        components = self.get_dummy_components()
+
+        pipe = self.pipeline_class(**components)
+        pipe.to("cpu")
+        pipe.set_progress_bar_config(disable=None)
+
+        # Without tiling
+        inputs = self.get_dummy_inputs(generator_device)
+        inputs["height"] = inputs["width"] = 128
+        output_without_tiling = pipe(**inputs)[0]
+
+        # With tiling
+        pipe.vae.enable_tiling(
+            tile_sample_min_height=96,
+            tile_sample_min_width=96,
+            tile_sample_stride_height=64,
+            tile_sample_stride_width=64,
+        )
+        inputs = self.get_dummy_inputs(generator_device)
+        inputs["height"] = inputs["width"] = 128
+        output_with_tiling = pipe(**inputs)[0]
+
+        self.assertLess(
+            (to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
+            expected_diff_max,
+            "VAE tiling should not affect the inference results",
+        )
+
+
+@slow
+@require_torch_gpu
+class MochiPipelineIntegrationTests(unittest.TestCase):
+    prompt = "A painting of a squirrel eating a burger."
+
+    def setUp(self):
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def tearDown(self):
+        super().tearDown()
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def test_mochi(self):
+        generator = torch.Generator("cpu").manual_seed(0)
+
+        pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", torch_dtype=torch.float16)
+        pipe.enable_model_cpu_offload()
+        prompt = self.prompt
+
+        videos = pipe(
+            prompt=prompt,
+            height=480,
+            width=848,
+            num_frames=19,
+            generator=generator,
+            num_inference_steps=2,
+            output_type="pt",
+        ).frames
+
+        video = videos[0]
+        expected_video = torch.randn(1, 16, 480, 848, 3).numpy()
+
+        max_diff = numpy_cosine_similarity_distance(video, expected_video)
+        assert max_diff < 1e-3, f"Max diff is too high. got {video}"
@@ -0,0 +1,129 @@
+# These are canonical sets of parameters for different types of pipelines.
+# They are set on subclasses of `PipelineTesterMixin` as `params` and
+# `batch_params`.
+#
+# If a pipeline's set of arguments has minor changes from one of the common sets
+# of arguments, do not make modifications to the existing common sets of arguments.
+# I.e. a text to image pipeline with non-configurable height and width arguments
+# should set its attribute as `params = TEXT_TO_IMAGE_PARAMS - {'height', 'width'}`.
+
+TEXT_TO_IMAGE_PARAMS = frozenset(
+    [
+        "prompt",
+        "height",
+        "width",
+        "guidance_scale",
+        "negative_prompt",
+        "prompt_embeds",
+        "negative_prompt_embeds",
+        "cross_attention_kwargs",
+    ]
+)
+
+TEXT_TO_IMAGE_BATCH_PARAMS = frozenset(["prompt", "negative_prompt"])
+
+TEXT_TO_IMAGE_IMAGE_PARAMS = frozenset([])
+
+IMAGE_TO_IMAGE_IMAGE_PARAMS = frozenset(["image"])
+
+IMAGE_VARIATION_PARAMS = frozenset(
+    [
+        "image",
+        "height",
+        "width",
+        "guidance_scale",
+    ]
+)
+
+IMAGE_VARIATION_BATCH_PARAMS = frozenset(["image"])
+
+TEXT_GUIDED_IMAGE_VARIATION_PARAMS = frozenset(
+    [
+        "prompt",
+        "image",
+        "height",
+        "width",
+        "guidance_scale",
+        "negative_prompt",
+        "prompt_embeds",
+        "negative_prompt_embeds",
+    ]
+)
+
+TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS = frozenset(["prompt", "image", "negative_prompt"])
+
+TEXT_GUIDED_IMAGE_INPAINTING_PARAMS = frozenset(
+    [
+        # Text guided image variation with an image mask
+        "prompt",
+        "image",
+        "mask_image",
+        "height",
+        "width",
+        "guidance_scale",
+        "negative_prompt",
+        "prompt_embeds",
+        "negative_prompt_embeds",
+    ]
+)
+
+TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS = frozenset(["prompt", "image", "mask_image", "negative_prompt"])
+
+IMAGE_INPAINTING_PARAMS = frozenset(
+    [
+        # image variation with an image mask
+        "image",
+        "mask_image",
+        "height",
+        "width",
+        "guidance_scale",
+    ]
+)
+
+IMAGE_INPAINTING_BATCH_PARAMS = frozenset(["image", "mask_image"])
+
+IMAGE_GUIDED_IMAGE_INPAINTING_PARAMS = frozenset(
+    [
+        "example_image",
+        "image",
+        "mask_image",
+        "height",
+        "width",
+        "guidance_scale",
+    ]
+)
+
+IMAGE_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS = frozenset(["example_image", "image", "mask_image"])
+
+CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS = frozenset(["class_labels"])
+
+CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS = frozenset(["class_labels"])
+
+UNCONDITIONAL_IMAGE_GENERATION_PARAMS = frozenset(["batch_size"])
+
+UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS = frozenset([])
+
+UNCONDITIONAL_AUDIO_GENERATION_PARAMS = frozenset(["batch_size"])
+
+UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS = frozenset([])
+
+TEXT_TO_AUDIO_PARAMS = frozenset(
+    [
+        "prompt",
+        "audio_length_in_s",
+        "guidance_scale",
+        "negative_prompt",
+        "prompt_embeds",
+        "negative_prompt_embeds",
+        "cross_attention_kwargs",
+    ]
+)
+
+TEXT_TO_AUDIO_BATCH_PARAMS = frozenset(["prompt", "negative_prompt"])
+TOKENS_TO_AUDIO_GENERATION_PARAMS = frozenset(["input_tokens"])
+
+TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS = frozenset(["input_tokens"])
+
+TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS = frozenset(["prompt_embeds"])
+
+VIDEO_TO_VIDEO_BATCH_PARAMS = frozenset(["prompt", "negative_prompt", "video"])
@@ -1,38 +0,0 @@
-import torch
-from fastvideo.models.hunyuan.diffusion.pipelines.pipeline_hunyuan_video import (
-    HunyuanVideoPipeline,
-)
-from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer
-from fastvideo.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D
-
-
-transformer = HYVideoDiffusionTransformer.from_pretrained(
-    "data/hyvideo-diffusers", torch_dtype=torch.bfloat16, subfolder="transformer"
-)
-vae = AutoencoderKLCausal3D.from_pretrained(
-    "data/hyvideo-diffusers", torch_dtype=torch.float16, subfolder="vae"
-)
-
-pipe = HunyuanVideoPipeline.from_pretrained(
-    "data/hyvideo-diffusers", transformer=transformer, vae=vae
-)
-pipe = pipe.to("cuda")
-pipe.vae.enable_tiling()
-
-prompt = "Close-up, A little girl wearing a red hoodie in winter strikes a match. The sky is dark, there is a layer of snow on the ground, and it is still snowing lightly. The flame of the match flickers, illuminating the girl's face intermittently."
-
-result = pipe(
-    prompt,
-    height=512,
-    width=512,
-    video_length=29,
-)
-
-import PIL.Image
-from diffusers.utils import export_to_video
-
-output = result.videos[0].permute(1, 2, 3, 0).detach().cpu().numpy()
-output = (output * 255).clip(0, 255).astype("uint8")
-output = [PIL.Image.fromarray(x) for x in output]
-
-export_to_video(output, "output.mp4", fps=24)
Author	SHA1	Message	Date
rlsu9	f3afc8382e	fix name typo	2024-12-29 04:53:45 +00:00
rlsu9	d256054508	upload github pytest workflow	2024-12-29 04:50:59 +00:00
rlsu9	2664c54649	init test file	2024-12-29 04:40:14 +00:00
Hangliang Ding	d467c7cd35	[Minor] Adding issue template. (#114 )	2024-12-25 21:50:57 -08:00
Zhang Peiyuanandrlsu9	88b2583c2c	[feat]:Single 4090 inference for fasthunyuan (#104 ) Co-authored-by: rlsu9 <r3su@ucsd.edu>	2024-12-25 12:40:16 -08:00
rlsu9	a730e43d5f	Update README.md layout	2024-12-19 13:36:43 -08:00
Brian Chen	edf116fa46	fix lora checkpoint saving issue (#97 )	2024-12-19 08:42:59 -08:00
Luis Catacora	de3cefb5e5	Add Replicate demo and API (#93 )	2024-12-18 19:56:09 -08:00
Hangliang Ding	e087e85e09	Adding Development plan	2024-12-18 16:46:14 +08:00
Your Name	e1b998b6ef	merge	2024-12-17 12:48:16 -08:00
rlsu9	fb49c93dbc	Update README.md	2024-12-17 12:29:03 -08:00
rlsu9	172f4802b4	Update README.md	2024-12-17 12:28:08 -08:00
rlsu9	24e57fafc9	Update README.md	2024-12-17 12:26:17 -08:00
Your Name	6debd46482	merge docs	2024-12-17 12:20:42 -08:00
rlsu9	f7dc36f7ea	Update README.md	2024-12-17 12:13:33 -08:00
Brian Chen	a0fb954f56	Update README.md fix typo	2024-12-17 15:09:10 -05:00
rlsu9	053106922c	Update README.md	2024-12-17 11:43:49 -08:00
Zhang Peiyuan rlsu9 rlsu9	a57122c519	Rlsu lora readme (#86 ) Co-authored-by: rlsu9 <r3su@ucsd.edu> Co-authored-by: rlsu9 <147024991+rlsu9@users.noreply.github.com>	2024-12-17 11:37:07 -08:00
Zhang Peiyuanandrlsu9	b393570e45	Update README (#85 ) Co-authored-by: rlsu9 <r3su@ucsd.edu>	2024-12-16 17:06:14 -08:00
Zhang Peiyuanandrlsu9	285635e8c0	Clean up (#84 ) Co-authored-by: rlsu9 <r3su@ucsd.edu>	2024-12-15 20:33:11 -08:00
Zhang Peiyuanandrlsu9	58cfd71b5e	Cleanup Co-authored-by: rlsu9 <r3su@ucsd.edu>	2024-12-15 17:03:29 -08:00
Hangliang Dingandrlsu9	3bf892b6ab	update release readme (#81 ) Co-authored-by: rlsu9 <r3su@ucsd.edu>	2024-12-15 22:24:13 +08:00