syn

Add wan dit

.github/ISSUE_TEMPLATE/1-bug-report.yml

@@ -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

.github/ISSUE_TEMPLATE/2-feature-request.yml

@@ -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.

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1 @@
+blank_issues_enabled: false

.github/workflows/codespell.yml

@@ -0,0 +1,45 @@
+name: codespell
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the main branch
+  push:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - "**/*.md"
+      - "**/*.rst"
+      - pyproject.toml
+      - requirements-lint.txt
+      - .github/workflows/codespell.yml
+  pull_request:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - "**/*.md"
+      - "**/*.rst"
+      - pyproject.toml
+      - requirements-lint.txt
+      - .github/workflows/codespell.yml
+
+jobs:
+  codespell:
+    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: |
+          python -m pip install --upgrade pip
+          pip install -r requirements-lint.txt
+      - name: Spelling check with codespell
+        run: |
+          # Refer to the above environment variable here
+          codespell --toml pyproject.toml $CODESPELL_EXCLUDES

.github/workflows/fastvideo-publish.yml

@@ -0,0 +1,70 @@
+name: Publish FastVideo to PyPI on Version Change
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - 'pyproject.toml'  # Trigger when pyproject.toml changes
+
+jobs:
+  check-version-change:
+    runs-on: ubuntu-latest
+    outputs:
+      version-changed: ${{ steps.check-version.outputs.changed }}
+      new-version: ${{ steps.check-version.outputs.new-version }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 2
+
+      - name: Check if version changed
+        id: check-version
+        run: |
+          # Get current commit's version
+          NEW_VERSION=$(grep -oP 'version\s*=\s*"\K[^"]+' pyproject.toml)
+          echo "New version: $NEW_VERSION"
+          
+          # Get previous version from git history
+          OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP 'version\s*=\s*"\K[^"]+' || echo "0.0.0")
+          echo "Old version: $OLD_VERSION"
+          
+          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
+            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
+            echo "changed=true" >> $GITHUB_OUTPUT
+            echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT
+          else
+            echo "Version did not change"
+            echo "changed=false" >> $GITHUB_OUTPUT
+          fi
+
+  build-publish-main:
+    needs: check-version-change
+    if: needs.check-version-change.outputs.version-changed == 'true'
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write  # Needed for OIDC Trusted Publishing
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+
+      - name: Install build dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build twine wheel
+
+      - name: Build package
+        run: |
+          python -m build
+          
+      - name: Publish release distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: dist/

.github/workflows/ruff.yml

@@ -0,0 +1,50 @@
+name: ruff
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the main branch
+  push:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - pyproject.toml
+      - requirements-lint.txt
+      - .github/workflows/matchers/ruff.json
+      - .github/workflows/ruff.yml
+  pull_request:
+    branches:
+      - main
+    # This workflow is only relevant when one of the following files changes.
+    # However, we have github configured to expect and require this workflow
+    # to run and pass before github with auto-merge a pull request. Until github
+    # allows more flexible auto-merge policy, we can just run this on every PR.
+    # It doesn't take that long to run, anyway.
+    #paths:
+    #  - "**/*.py"
+    #  - pyproject.toml
+    #  - requirements-lint.txt
+    #  - .github/workflows/matchers/ruff.json
+    #  - .github/workflows/ruff.yml
+
+jobs:
+  ruff:
+    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: |
+          python -m pip install --upgrade pip
+          pip install -r requirements-lint.txt
+      - name: Analysing the code with ruff
+        run: |
+          ruff check .
+      - name: Run isort
+        run: |
+          isort . --check-only

.github/workflows/sta-publish.yml

@@ -0,0 +1,221 @@
+name: Publish Sliding Tile Attention Kernel to PyPI on Version Change
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - "csrc/sliding_tile_attention/setup.py"
+
+jobs:
+  check-version-change:
+    runs-on: ubuntu-latest
+    outputs:
+      version-changed: ${{ steps.check-version.outputs.changed }}
+      new-version: ${{ steps.check-version.outputs.new-version }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 2
+
+      - name: Check if version changed
+        id: check-version
+        run: |
+          cd csrc/sliding_tile_attention
+          # Get current commit's version
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
+          echo "New version: $NEW_VERSION"
+          
+          # Get previous version from git history
+          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          echo "Old version: $OLD_VERSION"
+          
+          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
+            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
+            echo "changed=true" >> $GITHUB_OUTPUT
+            echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT
+          else
+            echo "Version did not change"
+            echo "changed=false" >> $GITHUB_OUTPUT
+          fi
+
+  build_wheels:
+    name: Build Wheel
+    needs: check-version-change
+    if: needs.check-version-change.outputs.version-changed == 'true'
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+          # Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
+          # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
+          os: [ubuntu-22.04]
+          python-version: ['3.10', '3.11', '3.12', '3.13']
+          torch-version: ['2.5.1', '2.6.0']
+          cuda-version: ['12.4.1', '12.5.1', '12.6.3']
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install CUDA ${{ matrix.cuda-version }}
+        uses: Jimver/cuda-toolkit@v0.2.21
+        id: cuda-toolkit
+        with:
+          cuda: ${{ matrix.cuda-version }}
+          linux-local-args: '["--toolkit"]'
+          method: 'network'
+
+      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
+        run: |
+          sudo apt update
+          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
+          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+          # Allow Git to Access Safe Directory
+          git config --global --add safe.directory /__w/FastVideo/FastVideo
+
+          # Set CUDA environment variables
+          export CUDA_HOME=/usr/local/cuda-${{ matrix.cuda-version }}
+          export PATH=${CUDA_HOME}/bin:${PATH}
+          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+          
+          # Verify installation
+          gcc --version
+          g++ --version
+          clang-11 --version
+          nvcc --version
+
+      - name: Install PyTorch ${{ matrix.torch-version }}+cu${{ matrix.cuda-version }}
+        run: |
+          pip install --upgrade pip
+          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
+          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
+          pip install typing-extensions==4.12.2
+          # We want to figure out the CUDA version to download pytorch
+          # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
+          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
+          export TORCH_CUDA_VERSION=124
+          pip install --no-cache-dir torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
+          nvcc --version
+          python --version
+          python -c "import torch; print('PyTorch:', torch.__version__)"
+          python -c "import torch; print('CUDA:', torch.version.cuda)"
+          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
+
+      - name: Build wheel
+        run: |
+          # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
+          # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
+          # However this still fails so I'm using a newer version of setuptools
+          pip install setuptools
+          pip install ninja packaging wheel
+
+          cd csrc/sliding_tile_attention  # Move into the correct folder
+          git submodule update --init --recursive tk # Ensure ThunderKittens submodule is initialized
+          python setup.py bdist_wheel --dist-dir=dist
+
+      - name: Rename wheel file
+        run: |
+          cd csrc/sliding_tile_attention
+
+          CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
+          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2)
+          # Get the correct version format
+          tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION}
+          wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
+          # Rename with version information
+          ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
+          echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
+
+      - name: Upload wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ env.wheel_name }}
+          path: csrc/sliding_tile_attention/dist/*.whl
+          retention-days: 90
+
+  publish_package:
+    name: Publish package
+    needs: [build_wheels]
+    if: needs.check-version-change.outputs.version-changed == 'true'
+    runs-on: ubuntu-22.04
+    permissions:
+      id-token: write  # Needed for OIDC Trusted Publishing
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+      
+      - name: Install CUDA 12.4.1
+        uses: Jimver/cuda-toolkit@v0.2.21
+        id: cuda-toolkit
+        with:
+          cuda: 12.4.1
+          linux-local-args: '["--toolkit"]'
+          method: 'network'
+          sub-packages: '["nvcc"]'
+
+      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
+        run: |
+          sudo apt update
+          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
+          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+          # Allow Git to Access Safe Directory
+          git config --global --add safe.directory /__w/FastVideo/FastVideo
+
+          # Set CUDA environment variables
+          export CUDA_HOME=/usr/local/cuda-12.4.1
+          export PATH=${CUDA_HOME}/bin:${PATH}
+          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+          
+          # Verify installation
+          gcc --version
+          g++ --version
+          clang-11 --version
+          nvcc --version
+
+      - name: Install PyTorch 2.5.1+cu12.4.1
+        run: |
+          pip install --upgrade pip
+          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
+          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
+          pip install typing-extensions==4.12.2
+          # We want to figure out the CUDA version to download pytorch
+          # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
+          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
+          export TORCH_CUDA_VERSION=124
+          pip install --no-cache-dir torch==2.5.1 --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
+          nvcc --version
+          python --version
+          python -c "import torch; print('PyTorch:', torch.__version__)"
+          python -c "import torch; print('CUDA:', torch.version.cuda)"
+          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
+
+      - name: Build source distribution
+        run: |
+          # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
+          # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
+          # However this still fails so I'm using a newer version of setuptools
+          pip install setuptools
+          pip install ninja packaging wheel
+
+          cd csrc/sliding_tile_attention  # Move into the correct folder
+          git submodule update --init --recursive tk # Ensure ThunderKittens submodule is initialized
+          python setup.py sdist --dist-dir=dist
+
+      - name: Publish release distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: csrc/sliding_tile_attention/dist/

.github/workflows/test.yml

@@ -0,0 +1,33 @@
+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: |
+          python -m pip install --upgrade pip setuptools wheel
+          pip install torch
+          pip install packaging ninja
+          # remove st-attn dependency because no cuda environment
+          sed -i '/st_attn/d' pyproject.toml
+          pip install -e .
+          pip install pytest
+
+      - name: Run Pytest
+        run: |
+          pytest --ignore csrc/sliding_tile_attention/test

.github/workflows/yapf.yml

@@ -0,0 +1,38 @@
+name: yapf
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the main branch
+  push:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - .github/workflows/yapf.yml
+  pull_request:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - .github/workflows/yapf.yml
+
+jobs:
+  yapf:
+    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: |
+          python -m pip install --upgrade pip
+          pip install yapf==0.32.0
+          pip install toml==0.10.2
+      - name: Running yapf
+        run: |
+          yapf --diff --recursive .

.gitignore

@@ -1,4 +1,3 @@
-ucf101_stride4x4x4
 __pycache__
 *.mp4
 .ipynb_checkpoints
@@ -8,10 +7,8 @@ results/
 build/
 fastvideo.egg-info/
 wandb/
-.idea
 *.ipynb
 *.jpg
-*.mp3
 *.safetensors
 *.mp4
 *.png
@@ -20,29 +17,6 @@ wandb/
 *.pt
 cache_dir/
 wandb/
-test*
-sample_video*
-sample_image*
-512*
-720*
-1024*
-debug*
-private*
-caption*
-*deepspeed*
-revised*
-129f*
-all*
-read*
-YSH*
-*pick*
-*ysh*
-hw*
-257f*
-513f*
-taming*
-221hw*
-65x512x512
 runs/
 samples/
 *validation/
@@ -52,3 +26,11 @@ outputs_video
 sbatch.sh
 *.out
 env
+dist/
+*.o
+**/build/
+**.egg-info
+**.pyc
+**.egg
+**.txt
+**.json

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "csrc/sliding_tile_attention/tk"]
+	path = csrc/sliding_tile_attention/tk
+	url = https://github.com/HazyResearch/ThunderKittens.git

LICENSE

@@ -184,18 +184,4 @@
       comment syntax for the file format. We also recommend that a
       file or class name and description of purpose be included on the
       same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright [2023] Lightning AI
-
-   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.
+      identification within third-party archives.

README.md

@@ -1,162 +1,230 @@
-# 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.
+
+
+<p align="center">
+    🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank">FastHunyuan</a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank">FastMochi</a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> Slack </a> 
+</p> 
+
+
+
+
+
+https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
+
+
+
+FastVideo currently offers: (with more to come)
+
+- [NEW!] [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
+- 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.
 
-### 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.
 
 ## Change Log
-
-
-- ```2024/12/06```: `FastMochi` v0.0.1 is released.
-
-
-## Fast and High-Quality Text-to-video Generation
-
-### 8-Step Results of FastMochi
-
-<table class="center">
-  <td><img src=assets/8steps/1.gif width="320"></td></td>
-  <td><img src=assets/8steps/2.gif width="320"></td></td></td>
-  <tr>
-  <td style="text-align:center;" width="320">tmp</td>
-  <td style="text-align:center;" width="320">tmp</td>
-  <tr>
-</table >
-
-
-## Table of Contents
-
-Jump to a specific section:
-
-- [🔧 Installation](#-installation)
-- [🚀 Inference](#-inference)
-- [🎯 Distill](#-distill)
-- [⚡ Finetune](#-lora-finetune)
+- ```2025/02/20```: FastVideo now supports STA on [StepVideo](https://github.com/stepfun-ai/Step-Video-T2V) with 3.4X speedup!
+- ```2025/02/18```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
+- ```2025/01/13```: Support Lora finetuning for HunyuanVideo.
+- ```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.4 and H100.
 ```
-conda create -n fastmochi python=3.10.0 -y && conda activate fastmochi
-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
 ```
-
-
-
+To try Sliding Tile Attention (optional), please follow the instruction in [csrc/sliding_tile_attention/README.md](csrc/sliding_tile_attention/README.md) to install STA.
 
 ## 🚀 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 StepVideo with Sliding Tile Attention 
+First, download the model:
+```
+python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model 
+```
+Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204×768×768 (currently, this is the only supported shape).
 ```bash
-python scripts/download_hf.py --repo_id=FastVideo/FastMochi --local_dir=data/FastMochi --repo_type=model
+sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA
+sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
 ```
 
-
-Start the gradio UI with
+### Inference HunyuanVideo with Sliding Tile Attention
+First, download the model:
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model 
 ```
-python fastvideo/demo/gradio_web_demo.py --model_path data/FastMochi
+We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
+```bash
+sh scripts/inference/inference_hunyuan_STA.sh
 ```
+### Video Demos using STA + Teacache
+Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.
 
-We also provide CLI inference script featured with sequence parallelism.
-
+### 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
+# 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_hunyuan_hf_quantization.sh
 ```
-export NUM_GPUS=4
+For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
 
-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
+
+| 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
 ```
+You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).
 
-For the mochi style, simply following the scripts list in mochi repo.
+### FastMochi
 
-```
-git clone https://github.com/genmoai/mochi.git
-cd mochi
-
-# install env
-...
-
-python3 ./demos/cli.py --model_dir weights/ --cpu_offload
+```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
 ```
 
 
 ## 🎯 Distill
-
-## 💰Hardware requirement
-
--  VRAM is required for both distill 10B mochi model
-
-To launch distillation, you will first need to prepare data in the following formats
-
+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
-asset/example_data
-├── AAA.txt
-├── AAA.png
-├── BCC.txt
-├── BCC.png
-├── ......
-├── CCC.txt
-└── CCC.png
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/HD-Mixkit-Finetune-Hunyuan --local_dir=data/HD-Mixkit-Finetune-Hunyuan --repo_type=dataset
 ```
-
-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:
+Next, download the original model weights with:
 ```bash
-python scripts/download_hf.py --repo_id=Stealths-Video/Merge-425-Data --local_dir=data/Merge-425-Data --repo_type=dataset
-python scripts/download_hf.py --repo_id=Stealths-Video/validation_embeddings --local_dir=data/validation_embeddings --repo_type=dataset
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model # original hunyuan
+python scripts/huggingface/download_hf.py --repo_id=genmo/mochi-1-preview --local_dir=data/mochi --repo_type=model # original mochi
 ```
-
-Then the distillation can be launched by:
-
+To launch the distillation process, use the following commands:
 ```
-bash scripts/distill_t2v.sh
+bash scripts/distill/distill_hunyuan.sh # for hunyuan
+bash scripts/distill/distill_mochi.sh # for mochi
 ```
-
-
-## ⚡ 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:
-
+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
 ```
-bash scripts/lora_finetune.sh
+Download the original model weights as specified 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 
+
+Hunyuan supports Lora fine-tuning of videos up to 720p. Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through:
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight --local_dir=data/Hunyuan-Black-Myth-Wukong-lora-weight --repo_type=model
+```
+#### 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.  
+
+
+Currently, both Mochi and Hunyuan models support Lora finetuning through diffusers. To generate personalized videos from your own dataset, you'll need to follow three main steps: dataset preparation, finetuning, and inference.
+
+#### Dataset Preparation
+We provide scripts to better help you get started to train on your own characters!  
+You can run this to organize your dataset to get the videos2caption.json before preprocess. Specify your video folder and corresponding caption folder (caption files should be .txt files and have the same name with its video):
+```
+python scripts/dataset_preparation/prepare_json_file.py --video_dir data/input_videos/ --prompt_dir data/captions/ --output_path data/output_folder/videos2caption.json --verbose
+```
+Also, we provide script to resize your videos:
+```
+python scripts/data_preprocess/resize_videos.py 
+```
+#### Finetuning
+After basic dataset preparation and preprocess, you can start to finetune your model using Lora:
+```
+bash scripts/finetune/finetune_hunyuan_hf_lora.sh
+```
+#### Inference
+For inference with Lora checkpoint, you can run the following scripts with additional parameter `--lora_checkpoint_dir`:
+```
+bash scripts/inference/inference_hunyuan_hf.sh 
+```
+**We also provide scripts for Mochi in the same directory.**
+
+#### 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
+
+## 🤝 Contributing
+
+We welcome all contributions. Please run `bash format.sh --all` before submitting a pull request.
+
+## 🔧 Testing
+Run `pytest` to verify the data preprocessing, checkpoint saving, and sequence parallel pipelines. We recommend adding corresponding test cases in the `test` folder to support your contribution.
 
 ## 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.
+
+## Citation 
+If you use FastVideo for your research, please cite our paper:
+
+```bibtex
+@misc{zhang2025fastvideogenerationsliding,
+      title={Fast Video Generation with Sliding Tile Attention}, 
+      author={Peiyuan Zhang and Yongqi Chen and Runlong Su and Hangliang Ding and Ion Stoica and Zhenghong Liu and Hao Zhang},
+      year={2025},
+      eprint={2502.04507},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2502.04507}, 
+}
+@misc{ding2025efficientvditefficientvideodiffusion,
+      title={Efficient-vDiT: Efficient Video Diffusion Transformers With Attention Tile}, 
+      author={Hangliang Ding and Dacheng Li and Runlong Su and Peiyuan Zhang and Zhijie Deng and Ion Stoica and Hao Zhang},
+      year={2025},
+      eprint={2502.06155},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2502.06155}, 
+}
+```

assets/prompt.txt

@@ -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.

cog.yaml

@@ -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"

csrc/sliding_tile_attention/MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include tk *
+include config.py

csrc/sliding_tile_attention/README.md

@@ -0,0 +1,68 @@
+
+
+# Sliding Tile Atteniton Kernel
+
+
+## Installation
+We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
+First, install C++20 for ThunderKittens:
+
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+Install STA:
+```bash
+export CUDA_HOME=/usr/local/cuda-12.4
+export PATH=${CUDA_HOME}/bin:${PATH} 
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+git submodule update --init --recursive
+python setup.py install
+```
+
+## Usage
+
+```python
+from st_attn import sliding_tile_attention
+# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
+# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
+# a tile is a cube of size (6, 8, 8)
+# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
+# text_length: int ranging from 0 to 256
+# If your attention contains text token (Hunyuan)
+out = sliding_tile_attention(q, k, v, window_size, text_length)
+# If your attention does not contain text token (StepVideo)
+out = sliding_tile_attention(q, k, v, window_size, 0, False)
+
+```
+
+
+## Test
+```bash
+python test/test_sta.py
+```
+
+## How Does STA Work?
+We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 
+
+
+https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747
+
+## Why is STA Fast?
+2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. 
+
+STA removes mixed blocks.
+
+
+<div align="center">
+<img src=../../assets/sliding_tile_attn_map.png width="80%"/>
+</div>
+
+## Acknowledgement
+
+We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/sliding_tile_attention/config.py

@@ -0,0 +1,15 @@
+### ADD TO THIS TO REGISTER NEW KERNELS
+sources = {
+    'attn': {
+        'source_files': {
+            'h100': 'st_attn/st_attn_h100.cu'  # define these source files for each GPU target desired.
+        }
+    }
+}
+
+### WHICH KERNELS DO WE WANT TO BUILD?
+# (oftentimes during development work you don't need to redefine them all.)
+kernels = ['attn']
+
+### WHICH GPU TARGET DO WE WANT TO BUILD FOR?
+target = 'h100'

csrc/sliding_tile_attention/setup.py

@@ -0,0 +1,76 @@
+import os
+import subprocess
+
+from config import kernels, sources, target
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+target = target.lower()
+
+# Package metadata
+PACKAGE_NAME = "st_attn"
+VERSION = "0.0.2"
+AUTHOR = "Hao AI Lab"
+DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo"
+URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention"
+
+# Set environment variables
+tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/')))
+python_include = subprocess.check_output(['python', '-c',
+                                          "import sysconfig; print(sysconfig.get_path('include'))"]).decode().strip()
+torch_include = subprocess.check_output([
+    'python', '-c',
+    "import torch; from torch.utils.cpp_extension import include_paths; print(' '.join(['-I' + p for p in include_paths()]))"
+]).decode().strip()
+print('st_attn root:', tk_root)
+print('Python include:', python_include)
+print('Torch include directories:', torch_include)
+
+# CUDA flags
+cuda_flags = [
+    '-DNDEBUG', '-Xcompiler=-Wno-psabi', '-Xcompiler=-fno-strict-aliasing', '--expt-extended-lambda',
+    '--expt-relaxed-constexpr', '-forward-unknown-to-host-compiler', '--use_fast_math', '-std=c++20', '-O3',
+    '-Xnvlink=--verbose', '-Xptxas=--verbose', '-Xptxas=--warn-on-spills', f'-I{tk_root}/include',
+    f'-I{tk_root}/prototype', f'-I{python_include}', '-DTORCH_COMPILE'
+] + torch_include.split()
+cpp_flags = ['-std=c++20', '-O3']
+
+if target == 'h100':
+    cuda_flags.append('-DKITTENS_HOPPER')
+    cuda_flags.append('-arch=sm_90a')
+else:
+    raise ValueError(f'Target {target} not supported')
+
+source_files = ['st_attn.cpp']
+for k in kernels:
+    if target not in sources[k]['source_files']:
+        raise KeyError(f'Target {target} not found in source files for kernel {k}')
+    if isinstance(sources[k]['source_files'][target], list):
+        source_files.extend(sources[k]['source_files'][target])
+    else:
+        source_files.append(sources[k]['source_files'][target])
+    cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}')
+
+setup(name=PACKAGE_NAME,
+      version=VERSION,
+      author=AUTHOR,
+      description=DESCRIPTION,
+      url=URL,
+      packages=find_packages(),
+      ext_modules=[
+          CUDAExtension('st_attn_cuda',
+                        sources=source_files,
+                        extra_compile_args={
+                            'cxx': cpp_flags,
+                            'nvcc': cuda_flags
+                        },
+                        libraries=['cuda'])
+      ],
+      cmdclass={'build_ext': BuildExtension},
+      classifiers=[
+          "Programming Language :: Python :: 3",
+          "Environment :: GPU :: NVIDIA CUDA :: 12",
+          "License :: OSI Approved :: Apache Software License",
+      ],
+      python_requires='>=3.10',
+      install_requires=["torch>=2.5.0"])

csrc/sliding_tile_attention/st_attn.cpp

@@ -0,0 +1,24 @@
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+
+#include <vector>
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+
+#include <cuda_runtime.h>
+
+
+#ifdef TK_COMPILE_ATTN
+extern torch::Tensor sta_forward(
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text
+); 
+#endif
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "Sliding Block Attention Kernels"; // optional module docstring
+
+#ifdef TK_COMPILE_ATTN
+    m.def("sta_fwd",  torch::wrap_pybind_function(sta_forward), "sliding tile attention, assuming tile size is (6,8,8)");
+#endif
+}

csrc/sliding_tile_attention/st_attn/__init__.py

@@ -0,0 +1,35 @@
+import math
+
+import torch
+from st_attn_cuda import sta_fwd
+
+
+def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True):
+    seq_length = q_all.shape[2]
+    if has_text:
+        assert q_all.shape[
+            2] == 115456, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
+        assert q_all.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
+        target_size = math.ceil(seq_length / 384) * 384
+        pad_size = target_size - seq_length
+        if pad_size > 0:
+            q_all = torch.cat([q_all, q_all[:, :, -pad_size:]], dim=2)
+            k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
+            v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
+    else:
+        assert q_all.shape[2] == 82944
+
+    hidden_states = torch.empty_like(q_all)
+    # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
+    for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
+        for batch in range(q_all.shape[0]):
+            q_head, k_head, v_head, o_head = (q_all[batch:batch + 1, head_index:head_index + 1],
+                                              k_all[batch:batch + 1,
+                                                    head_index:head_index + 1], v_all[batch:batch + 1,
+                                                                                      head_index:head_index + 1],
+                                              hidden_states[batch:batch + 1, head_index:head_index + 1])
+
+            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text)
+    if has_text:
+        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True)
+    return hidden_states[:, :, :seq_length]

csrc/sliding_tile_attention/st_attn/st_attn_h100.cu

@@ -0,0 +1,687 @@
+// # Define TORCH_COMPILE macro
+
+#include "kittens.cuh"
+#include <cooperative_groups.h>
+#include <iostream>
+#include <stdio.h>
+
+#define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
+#define ABS(x) ((x) < 0 ? -(x) : (x))
+
+constexpr int CONSUMER_WARPGROUPS = (3); 
+constexpr int PRODUCER_WARPGROUPS = (1); 
+constexpr int NUM_WARPGROUPS      = (CONSUMER_WARPGROUPS+PRODUCER_WARPGROUPS); 
+constexpr int NUM_WORKERS         = (NUM_WARPGROUPS*kittens::WARPGROUP_WARPS); 
+
+using namespace kittens;
+namespace cg = cooperative_groups;
+
+template<int D> struct fwd_attend_ker_tile_dims {};
+template<> struct fwd_attend_ker_tile_dims<64> {
+    constexpr static int tile_width = (64);
+    constexpr static int qo_height  = (4*16);
+    constexpr static int kv_height  = (8*16);
+    constexpr static int stages     = (4); 
+};
+template<> struct fwd_attend_ker_tile_dims<128> {
+    constexpr static int tile_width = (128);
+    constexpr static int qo_height  = (4*16);
+    constexpr static int kv_height  = (8*16);
+    constexpr static int stages     = (2); 
+};
+
+template<int D> struct fwd_globals {
+    using q_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using k_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::kv_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using v_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::kv_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>>;
+    using o_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+
+    using q_gl = gl<bf16,  -1, -1, -1, -1, q_tile>;
+    using k_gl = gl<bf16,  -1, -1, -1, -1, k_tile>;
+    using v_gl = gl<bf16,  -1, -1, -1, -1, v_tile>;
+    using l_gl = gl<float, -1, -1, -1, -1, l_col_vec>;
+    using o_gl = gl<bf16,  -1, -1, -1, -1, o_tile>;
+
+    q_gl q;
+    k_gl k;
+    v_gl v;
+    l_gl l;
+    o_gl o;
+
+    const int N; 
+    const int text_L;
+    const int hr;
+};
+
+
+template<int D, bool is_causal, bool text_q, bool text_kv, int DT, int DH, int DW, int CT, int CH, int CW>
+__global__  __launch_bounds__((NUM_WORKERS)*kittens::WARP_THREADS, 1)
+void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
+    extern __shared__ int __shm[]; 
+    tma_swizzle_allocator al((int*)&__shm[0]);
+    int warpid = kittens::warpid(), warpgroupid = warpid/kittens::WARPGROUP_WARPS;
+
+    using K = fwd_attend_ker_tile_dims<D>;
+
+    using q_tile    =         st_bf<K::qo_height, K::tile_width>;
+    using k_tile    =         st_bf<K::kv_height, K::tile_width>;
+    using v_tile    =         st_bf<K::kv_height, K::tile_width>;
+    using l_col_vec = col_vec<st_fl<K::qo_height, K::tile_width>>;
+    using o_tile    =         st_bf<K::qo_height, K::tile_width>;
+    
+    q_tile    (&q_smem)[CONSUMER_WARPGROUPS] = al.allocate<q_tile, CONSUMER_WARPGROUPS>();
+    k_tile    (&k_smem)[K::stages]           = al.allocate<k_tile, K::stages          >();
+    v_tile    (&v_smem)[K::stages]           = al.allocate<v_tile, K::stages          >();
+    l_col_vec (&l_smem)[CONSUMER_WARPGROUPS] = al.allocate<l_col_vec, CONSUMER_WARPGROUPS>();
+    auto      (*o_smem)                      = reinterpret_cast<o_tile(*)>(q_smem);
+    int img_kv_blocks;
+    int kv_blocks   = g.N / (K::kv_height);
+    if constexpr (text_kv) {
+        img_kv_blocks = kv_blocks - 3;
+    } else {
+        img_kv_blocks = kv_blocks;
+    }
+    int kv_head_idx = blockIdx.y / g.hr;
+    int seq_idx;
+    if constexpr (text_q) {
+        seq_idx = CT * CH * CW * 6.0 + blockIdx.x * CONSUMER_WARPGROUPS;
+    } else {
+        seq_idx = blockIdx.x * CONSUMER_WARPGROUPS; 
+    }
+    __shared__ kittens::semaphore qsmem_semaphore, k_smem_arrived[K::stages], v_smem_arrived[K::stages], compute_done[K::stages];
+    if (threadIdx.x == 0) { 
+        init_semaphore(qsmem_semaphore, 0, 1); 
+        for(int j = 0; j < K::stages; j++) {
+            init_semaphore(k_smem_arrived[j], 0, 1); 
+            init_semaphore(v_smem_arrived[j], 0, 1); 
+            init_semaphore(compute_done[j], CONSUMER_WARPGROUPS, 0); 
+        }
+
+        tma::expect_bytes(qsmem_semaphore, sizeof(q_smem));
+
+        for (int wg = 0; wg < CONSUMER_WARPGROUPS; wg++) {
+            coord<q_tile> q_tile_idx = {blockIdx.z, blockIdx.y, (seq_idx) + wg, 0};
+            tma::load_async(q_smem[wg], g.q, q_tile_idx, qsmem_semaphore);
+        }
+
+        if constexpr (text_q){
+            for (int j = 0; j < K::stages - 1; j++) {
+                coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0};
+                tma::expect_bytes(k_smem_arrived[j], sizeof(k_tile));
+                tma::load_async(k_smem[j], g.k, kv_tile_idx, k_smem_arrived[j]);
+                tma::expect_bytes(v_smem_arrived[j], sizeof(v_tile));
+                tma::load_async(v_smem[j], g.v, kv_tile_idx, v_smem_arrived[j]);
+            }
+        } else {
+            int qt = seq_idx / 6 / (CH * CW);
+            int qh = (seq_idx / 6) % (CH * CW) / CW;
+            int qw = (seq_idx / 6) % CW;
+            qt = CLAMP(qt, DT, CT-DT-1);
+            qh = CLAMP(qh, DH, CH-DH-1);
+            qw = CLAMP(qw, DW, CW-DW-1);
+            int count = 0;
+            int j = 0;
+            while (count < K::stages - 1) {
+                int kt = j / 3 / (CH * CW);
+                int kh = (j / 3) % (CH * CW) / CW;
+                int kw = (j / 3) % CW;
+                bool mask = (ABS(qt - kt) <= DT) && (ABS(qh - kh) <= DH) && (ABS(qw - kw) <= DW);
+                if (mask){
+                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0};
+                    tma::expect_bytes(k_smem_arrived[count], sizeof(k_tile));
+                    tma::load_async(k_smem[count], g.k, kv_tile_idx, k_smem_arrived[count]);
+                    tma::expect_bytes(v_smem_arrived[count], sizeof(v_tile));
+                    tma::load_async(v_smem[count], g.v, kv_tile_idx, v_smem_arrived[count]);
+                    count += 1;
+                }
+                j += 1;
+            }
+        }
+    }
+    __syncthreads(); 
+
+    int pipe_idx = K::stages - 1; 
+    
+    if(warpgroupid == NUM_WARPGROUPS-1) {
+        warpgroup::decrease_registers<32>();      
+        
+        int kv_iters; 
+        if constexpr (is_causal) {
+            kv_iters = (seq_idx * (K::qo_height/kittens::TILE_ROW_DIM<bf16>)) - 1 + (CONSUMER_WARPGROUPS * (K::qo_height/kittens::TILE_ROW_DIM<bf16>)); 
+            kv_iters = ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM<bf16>)) == 0) ? (0) : ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM<bf16>)) - 1);
+        }
+        else { kv_iters = kv_blocks-2;}
+
+        if(warpid == NUM_WORKERS-4) {
+            if constexpr (text_q){
+                for (auto kv_idx = pipe_idx - 1; kv_idx <= kv_iters; kv_idx++) {
+                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, kv_idx + 1, 0};
+                    tma::expect_bytes(k_smem_arrived[(kv_idx+1)%K::stages], sizeof(k_tile));
+                    tma::load_async(k_smem[(kv_idx+1)%K::stages], g.k, kv_tile_idx, k_smem_arrived[(kv_idx+1)%K::stages]);
+                    tma::expect_bytes(v_smem_arrived[(kv_idx+1)%K::stages], sizeof(v_tile));
+                    tma::load_async(v_smem[(kv_idx+1)%K::stages], g.v, kv_tile_idx, v_smem_arrived[(kv_idx+1)%K::stages]);
+                    kittens::wait(compute_done[(kv_idx)%K::stages], (kv_idx/K::stages)%2);
+                }
+            } else {
+                int qt = seq_idx / 6 / (CH * CW);
+                int qh = (seq_idx / 6) % (CH * CW) / CW;
+                int qw = (seq_idx / 6) % CW;
+                qt = CLAMP(qt, DT, CT-DT-1);
+                qh = CLAMP(qh, DH, CH-DH-1);
+                qw = CLAMP(qw, DW, CW-DW-1);
+                int k_t_min = CLAMP(qt-DT, 0, CT-1);
+                int k_t_max = CLAMP(qt+DT, 0, CT-1);
+                int k_h_min = CLAMP(qh-DH, 0, CH-1);
+                int k_h_max = CLAMP(qh+DH, 0, CH-1);
+                int k_w_min = CLAMP(qw-DW, 0, CW-1);
+                int k_w_max = CLAMP(qw+DW, 0, CW-1);
+                int count = 0;
+                for (int kt = k_t_min; kt <= k_t_max; kt++) {
+                    for (int kh = k_h_min; kh <= k_h_max; kh++) {
+                        for (int kw = k_w_min; kw <= k_w_max; kw++) {
+                            for (int j = 0; j <= 2; j++){
+                                if (count >= K::stages - 1) {
+                                    int index = ((kt * (CH * CW)) + (kh * CW) + kw) * 3 + j;
+                                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, index, 0};
+                                    tma::expect_bytes(k_smem_arrived[count%K::stages], sizeof(k_tile));
+                                    tma::load_async(k_smem[count%K::stages], g.k, kv_tile_idx, k_smem_arrived[count%K::stages]);
+                                    tma::expect_bytes(v_smem_arrived[count%K::stages], sizeof(v_tile));
+                                    tma::load_async(v_smem[count%K::stages], g.v, kv_tile_idx, v_smem_arrived[count%K::stages]);
+                                    kittens::wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2);
+                                    count += 1;
+                                } else {
+                                    count += 1;
+                                }
+                            }
+                        }
+                    }
+                }
+                // for text 
+                for (int index = img_kv_blocks; index < kv_blocks; index++) {
+                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, index, 0};
+                    tma::expect_bytes(k_smem_arrived[count%K::stages], sizeof(k_tile));
+                    tma::load_async(k_smem[count%K::stages], g.k, kv_tile_idx, k_smem_arrived[count%K::stages]);
+                    tma::expect_bytes(v_smem_arrived[count%K::stages], sizeof(v_tile));
+                    tma::load_async(v_smem[count%K::stages], g.v, kv_tile_idx, v_smem_arrived[count%K::stages]);
+                    kittens::wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2);
+                    count += 1;
+                }
+            }
+
+
+        }
+    }
+    else {
+        warpgroup::increase_registers<160>();
+
+        rt_fl<16, K::kv_height>  att_block;
+        rt_bf<16, K::kv_height>  att_block_mma;
+        rt_fl<16, K::tile_width> o_reg;
+        
+        col_vec<rt_fl<16, K::kv_height>> max_vec, norm_vec, max_vec_last_scaled, max_vec_scaled;
+        
+        neg_infty(max_vec);
+        zero(norm_vec);
+        zero(o_reg);
+
+        int kv_iters; 
+        if constexpr (is_causal) {
+            kv_iters = (seq_idx * 4) - 1 + (CONSUMER_WARPGROUPS * 4);
+            kv_iters = (kv_iters/8);
+        }
+        else if constexpr (text_q){ 
+            // the last three kv blocks are for text, we process them separately
+            kv_iters = img_kv_blocks - 1;
+        } else {
+            kv_iters = CLAMP(DT*2+1, 1, CT) * CLAMP(DH*2+1, 1, CH) * CLAMP(DW*2+1, 1, CW) * 3 - 1 ; 
+        }
+
+        kittens::wait(qsmem_semaphore, 0);
+        for (auto kv_idx = 0; kv_idx <= kv_iters; kv_idx++) {
+
+            kittens::wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2);
+            warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]);
+            
+            copy(max_vec_last_scaled, max_vec);
+            if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); }
+            else                   { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); }
+            
+            warpgroup::mma_async_wait();
+
+            row_max(max_vec, att_block, max_vec);
+            
+            if constexpr (D == 64) { 
+                mul(att_block, att_block,    1.44269504089f*0.125f); 
+                mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f);
+            }
+            else                   { 
+                mul(att_block, att_block,    1.44269504089f*0.08838834764f); 
+                mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f);
+            }
+
+            sub_row(att_block, att_block, max_vec_scaled);
+            exp2(att_block, att_block);
+            sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled);
+            exp2(max_vec_last_scaled,       max_vec_last_scaled);
+            mul(norm_vec,            norm_vec,     max_vec_last_scaled);
+            row_sum(norm_vec,  att_block, norm_vec);
+            add(att_block, att_block, 0.f);
+            copy(att_block_mma, att_block); 
+            mul_row(o_reg, o_reg, max_vec_last_scaled); 
+
+            kittens::wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); 
+
+            warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]);
+            warpgroup::mma_async_wait();
+
+            if(warpgroup::laneid() == 0) arrive(compute_done[(kv_idx)%K::stages], 1);
+        }
+        // the last three kv blocks are for text, we process them separately
+        if constexpr(text_kv) {
+            for (auto kv_idx = kv_iters + 1; kv_idx <= kv_iters + 3; kv_idx++) {
+
+                kittens::wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2);
+                warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]);
+                
+                copy(max_vec_last_scaled, max_vec);
+                if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); }
+                else                   { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); }
+                
+                warpgroup::mma_async_wait();
+                // apply non-pad mask
+                int offset = g.text_L - (kv_idx - (kv_iters + 1)) * K::kv_height;
+                // printf("k_idx_start: %d, k_idx_end: %d, text_end: %d, offset: %d\n", k_idx_start, k_idx_end, text_end, offset);
+                right_fill(att_block, att_block, offset, base_types::constants<float>::neg_infty());
+
+
+                row_max(max_vec, att_block, max_vec);
+                
+                if constexpr (D == 64) { 
+                    mul(att_block, att_block,    1.44269504089f*0.125f); 
+                    mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f);
+                }
+                else                   { 
+                    mul(att_block, att_block,    1.44269504089f*0.08838834764f); 
+                    mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f);
+                }
+
+                sub_row(att_block, att_block, max_vec_scaled);
+                exp2(att_block, att_block);
+                sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled);
+                exp2(max_vec_last_scaled,       max_vec_last_scaled);
+                mul(norm_vec,            norm_vec,     max_vec_last_scaled);
+                row_sum(norm_vec,  att_block, norm_vec);
+                add(att_block, att_block, 0.f);
+                copy(att_block_mma, att_block); 
+                mul_row(o_reg, o_reg, max_vec_last_scaled); 
+
+                kittens::wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); 
+
+                warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]);
+                warpgroup::mma_async_wait();
+
+                if(warpgroup::laneid() == 0) arrive(compute_done[(kv_idx)%K::stages], 1);
+            }
+        }
+
+        div_row(o_reg, o_reg, norm_vec);
+        warpgroup::store(o_smem[warpgroupid], o_reg); 
+        warpgroup::sync(warpgroupid+4);
+
+        if (warpid % 4 == 0) {
+            coord<o_tile> o_tile_idx = {blockIdx.z, blockIdx.y, (seq_idx) + warpgroupid, 0};
+            tma::store_async(g.o, o_smem[warpgroupid], o_tile_idx);
+        }
+
+        mul(max_vec_scaled,   max_vec_scaled, 0.69314718056f);
+        log(norm_vec, norm_vec);
+        add(norm_vec, norm_vec, max_vec_scaled);
+
+        if constexpr (D == 64) { mul(norm_vec, norm_vec, -8.0f); }
+        else                   { mul(norm_vec, norm_vec, -11.313708499f); }
+    
+        warpgroup::store(l_smem[warpgroupid], norm_vec);
+        warpgroup::sync(warpgroupid+4);
+
+        if (warpid % 4 == 0) {
+            coord<l_col_vec> tile_idx = {blockIdx.z, blockIdx.y, 0, (seq_idx) + warpgroupid};
+            tma::store_async(g.l, l_smem[warpgroupid], tile_idx);
+        }
+        tma::store_async_wait();
+    }
+}
+
+
+
+#include "pyutils/torch_helpers.cuh"
+#include <ATen/cuda/CUDAContext.h>
+#include <iostream>
+
+torch::Tensor 
+sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text)
+{
+    CHECK_INPUT(q);
+    CHECK_INPUT(k);
+    CHECK_INPUT(v);
+
+    auto batch    = q.size(0);
+    auto seq_len  = q.size(2); 
+    auto head_dim = q.size(3);  
+    auto qo_heads = q.size(1);
+    auto kv_heads = k.size(1);
+
+    // check to see that these dimensions match for all inputs
+    TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs");
+    TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs");
+    TORCH_CHECK(v.size(0) == batch, "V batch dimension - idx 0 - must match for all inputs");
+
+    TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs");
+    TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs");
+    TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs");
+
+    TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs");
+    TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs");
+    TORCH_CHECK(v.size(3) == head_dim, "V head dimension - idx 3 - must match for all non-vector inputs");
+
+    TORCH_CHECK(qo_heads >= kv_heads, "QO heads must be greater than or equal to KV heads");
+    TORCH_CHECK(qo_heads % kv_heads == 0, "QO heads must be divisible by KV heads");
+    TORCH_CHECK(q.size(1) == qo_heads, "QO head dimension - idx 1 - must match for all inputs");
+    TORCH_CHECK(k.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs");
+    TORCH_CHECK(v.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs");  
+
+    auto hr = qo_heads / kv_heads;
+
+    c10::BFloat16* q_ptr = q.data_ptr<c10::BFloat16>();
+    c10::BFloat16* k_ptr = k.data_ptr<c10::BFloat16>();
+    c10::BFloat16* v_ptr = v.data_ptr<c10::BFloat16>();
+
+    bf16*  d_q = reinterpret_cast<bf16*>(q_ptr);
+    bf16*  d_k = reinterpret_cast<bf16*>(k_ptr);
+    bf16*  d_v = reinterpret_cast<bf16*>(v_ptr);
+    
+
+    
+    torch::Tensor l_vec = torch::empty({static_cast<const uint>(batch), 
+                                        static_cast<const uint>(qo_heads), 
+                                        static_cast<const uint>(seq_len), 
+                                        static_cast<const uint>(1)}, 
+                                        torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous));
+        
+
+    bf16*  o_ptr = reinterpret_cast<bf16*>(o.data_ptr<c10::BFloat16>());
+    bf16*  d_o   = reinterpret_cast<bf16*>(o_ptr);
+
+    float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
+    float* d_l   = reinterpret_cast<float*>(l_ptr);
+
+    cudaDeviceSynchronize();
+    auto stream = at::cuda::getCurrentCUDAStream().stream(); 
+
+
+    if (head_dim == 128) {
+        using q_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using k_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using v_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>>;
+        using o_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+
+        using q_global = gl<bf16,  -1, -1, -1, -1, q_tile>;
+        using k_global = gl<bf16,  -1, -1, -1, -1, k_tile>;
+        using v_global = gl<bf16,  -1, -1, -1, -1, v_tile>;
+        using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
+        using o_global = gl<bf16,  -1, -1, -1, -1, o_tile>;
+
+        using globals      = fwd_globals<128>;
+
+        q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+        k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,   static_cast<unsigned int>(seq_len)};
+        o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+
+        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len),  static_cast<int>(text_length), static_cast<int>(hr)};
+
+        auto mem_size = kittens::MAX_SHARED_MEMORY;
+        auto threads  = NUM_WORKERS * kittens::WARP_THREADS;
+        if (has_text) {
+            // TORCH_CHECK(seq_len % (CONSUMER_WARPGROUPS*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 192");
+            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4-2), qo_heads, batch);
+            dim3 grid_text(2, qo_heads, batch);
+            if (!process_text) {
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 1, 1, 1, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 1, 1, 1, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 1, 1, 2, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true,1, 1, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 5 && kernel_h_size == 3 && kernel_w_size == 3) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 1, 1, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 1, 1, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                }else if (kernel_t_size ==3 && kernel_h_size == 5 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 1, 2, 2, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 1, 2, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size ==5 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 3, 0, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 3, 0, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size ==5 && kernel_h_size == 3 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 1, 2, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 1, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 2, 2, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 2, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 2, 3, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 2, 3, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 3, 5, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 3, 5, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 0, 0, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 0, 0, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 0, 3, 5, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 0, 3, 5, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 0, 5, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true,2, 0, 5, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else {
+                    // print error
+                    std::cout << "Invalid kernel size" << std::endl;
+                    //print kernel size
+                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                }
+            } else {
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, true, true, 1, 1, 1, 5, 6, 10>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, true, true, 1, 1, 1, 5, 6, 10><<<grid_text, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            }
+
+        } else {
+            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
+
+            if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+            }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+            } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+            } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+            }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+            } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+            } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else {
+                // print error
+                std::cout << "Invalid kernel size" << std::endl;
+                //print kernel size
+                std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+            }
+
+        }
+        CHECK_CUDA_ERROR(cudaGetLastError());
+        cudaStreamSynchronize(stream);
+    }
+
+    return o;
+    cudaDeviceSynchronize();
+}

csrc/sliding_tile_attention/test/bench.py

@@ -0,0 +1,151 @@
+import os
+from collections import defaultdict
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from st_attn import sliding_tile_attention
+
+
+def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
+    assert mode in ["fwd", "bwd", "fwd_bwd"]
+    f = 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1)
+    return f if mode == "fwd" else (2.5 * f if mode == "bwd" else 3.5 * f)
+
+
+def efficiency(flop, time):
+    flop = flop / 1e12
+    time = time / 1e6
+    return flop / time
+
+
+def benchmark_attention(configurations):
+    results = {'fwd': defaultdict(list), 'bwd': defaultdict(list)}
+
+    for B, H, N, D, causal in configurations:
+        print("=" * 60)
+        print(f"Timing forward and backward pass for B={B}, H={H}, N={N}, D={D}, causal={causal}")
+
+        q = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
+        k = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
+        v = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
+
+        grad_output = torch.randn_like(q, requires_grad=False).contiguous()
+
+        qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
+        kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
+        vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
+
+        # Prepare for timing forward pass
+        start_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+        end_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+        # Warmup for forward pass
+        for _ in range(10):
+            o = sliding_tile_attention(q, k, v, [[6, 6, 6]] * 24, 0, False)
+
+        # Time the forward pass
+        for i in range(10):
+            start_events_fwd[i].record()
+            o = sliding_tile_attention(q, k, v, [[6, 6, 6]] * 24, 0, False)
+            end_events_fwd[i].record()
+
+        torch.cuda.synchronize()
+        times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
+        time_us_fwd = np.mean(times_fwd) * 1000
+
+        tflops_fwd = efficiency(flops(B, N, H, D, causal, 'fwd'), time_us_fwd)
+        results['fwd'][(D, causal)].append((N, tflops_fwd))
+
+        print(f"Average time for forward pass in us: {time_us_fwd:.2f}")
+        print(f"Average efficiency for forward pass in TFLOPS: {tflops_fwd}")
+        print("-" * 60)
+
+        # torch.cuda.empty_cache()
+        # torch.cuda.synchronize()
+
+        # # Prepare for timing backward pass
+        # start_events_bwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+        # end_events_bwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+
+        # # Warmup for backward pass
+        # for _ in range(10):
+        #     qg, kg, vg = tk.mha_backward(q, k, v, o, l_vec, grad_output, causal)
+
+        # # Time the backward pass
+        # for i in range(10):
+        #     start_events_bwd[i].record()
+        #     qg, kg, vg = tk.mha_backward(q, k, v, o, l_vec, grad_output, causal)
+        #     end_events_bwd[i].record()
+
+        # torch.cuda.synchronize()
+        # times_bwd = [s.elapsed_time(e) for s, e in zip(start_events_bwd, end_events_bwd)]
+        # time_us_bwd = np.mean(times_bwd) * 1000
+
+        # tflops_bwd = efficiency(flops(B, N, H, D, causal, 'bwd'), time_us_bwd)
+        # results['bwd'][(D, causal)].append((N, tflops_bwd))
+
+        # print(f"Average time for backward pass in us: {time_us_bwd:.2f}")
+        # print(f"Average efficiency for backward pass in TFLOPS: {tflops_bwd}")
+        print("=" * 60)
+
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+    return results
+
+
+def plot_results(results):
+    os.makedirs('benchmark_results', exist_ok=True)
+    for mode in ['fwd', 'bwd']:
+        for (D, causal), values in results[mode].items():
+            seq_lens = [x[0] for x in values]
+            tflops = [x[1] for x in values]
+
+            plt.figure(figsize=(10, 6))
+            bars = plt.bar(range(len(seq_lens)), tflops, tick_label=seq_lens)
+            plt.xlabel('Sequence Length')
+            plt.ylabel('TFLOPS')
+            plt.title(f'{mode.upper()} Pass - Head Dim: {D}, Causal: {causal}')
+            plt.grid(True)
+
+            # Adding the numerical y value on top of each bar
+            for bar in bars:
+                yval = bar.get_height()
+                plt.text(bar.get_x() + bar.get_width() / 2, yval, round(yval, 2), ha='center', va='bottom')
+
+            filename = f'benchmark_results/{mode}_D{D}_causal{causal}.png'
+            plt.savefig(filename)
+            plt.close()
+
+
+# Example list of configurations to test
+configurations = [
+    (2, 24, 82944, 128, False),
+    # (16, 16, 768*16,  128, False),
+    # (16, 16, 768*2,  128, False),
+    # (16, 16, 768*4,  128, False),
+    # (16, 16, 768*8,  128, False),
+    # (16, 16, 768*16, 128, False),
+    # (16, 16, 768,    128, True),
+    # (16, 16, 768*2,  128, True),
+    # (16, 16, 768*4,  128, True),
+    # (16, 16, 768*8,  128, True),
+    # (16, 16, 768*16, 128, True),
+    # (16, 32, 768,    64,  False),
+    # (16, 32, 768*2,  64,  False),
+    # (16, 32, 768*4,  64,  False),
+    # (16, 32, 768*8,  64,  False),
+    # (16, 32, 768*16, 64,  False),
+    # (16, 32, 768,    64,  True),
+    # (16, 32, 768*2,  64,  True),
+    # (16, 32, 768*4,  64,  True),
+    # (16, 32, 768*8,  64,  True),
+    # (16, 32, 768*16, 64,  True),
+]
+
+results = benchmark_attention(configurations)
+# plot_results(results)

csrc/sliding_tile_attention/test/flex_sta_ref.py

@@ -0,0 +1,71 @@
+from typing import Tuple
+
+import torch
+from torch import BoolTensor, IntTensor
+from torch.nn.attention.flex_attention import create_block_mask
+
+# Peiyuan: This is neccesay. Dont know why. see https://github.com/pytorch/pytorch/issues/135028
+torch._inductor.config.realize_opcount_threshold = 100
+
+
+def generate_sta_mask(canvas_twh, kernel_twh, tile_twh, text_length):
+    """Generates a 3D NATTEN attention mask with a given kernel size.
+    
+    Args:
+        canvas_t: The time dimension of the canvas.
+        canvas_h: The height of the canvas.
+        canvas_w: The width of the canvas.
+        kernel_t: The time dimension of the kernel.
+        kernel_h: The height of the kernel.
+        kernel_w: The width of the kernel.
+    """
+    canvas_t, canvas_h, canvas_w = canvas_twh
+    kernel_t, kernel_h, kernel_w = kernel_twh
+    tile_t_size, tile_h_size, tile_w_size = tile_twh
+    total_tile_size = tile_t_size * tile_h_size * tile_w_size
+    canvas_tile_t, canvas_tile_h, canvas_tile_w = canvas_t // tile_t_size, canvas_h // tile_h_size, canvas_w // tile_w_size
+    img_seq_len = canvas_t * canvas_h * canvas_w
+
+    def get_tile_t_x_y(idx: IntTensor) -> Tuple[IntTensor, IntTensor, IntTensor]:
+        tile_id = idx // total_tile_size
+        tile_t = tile_id // (canvas_tile_h * canvas_tile_w)
+        tile_h = (tile_id % (canvas_tile_h * canvas_tile_w)) // canvas_tile_w
+        tile_w = tile_id % canvas_tile_w
+        return tile_t, tile_h, tile_w
+
+    def sta_mask_3d(
+        b: IntTensor,
+        h: IntTensor,
+        q_idx: IntTensor,
+        kv_idx: IntTensor,
+    ) -> BoolTensor:
+        q_t_tile, q_x_tile, q_y_tile = get_tile_t_x_y(q_idx)
+        kv_t_tile, kv_x_tile, kv_y_tile = get_tile_t_x_y(kv_idx)
+        # kernel nominally attempts to center itself on the query, but kernel center
+        # is clamped to a fixed distance (kernel half-length) from the canvas edge
+        kernel_center_t = q_t_tile.clamp(kernel_t // 2, (canvas_tile_t - 1) - kernel_t // 2)
+        kernel_center_x = q_x_tile.clamp(kernel_h // 2, (canvas_tile_h - 1) - kernel_h // 2)
+        kernel_center_y = q_y_tile.clamp(kernel_w // 2, (canvas_tile_w - 1) - kernel_w // 2)
+        time_mask = (kernel_center_t - kv_t_tile).abs() <= kernel_t // 2
+        hori_mask = (kernel_center_x - kv_x_tile).abs() <= kernel_h // 2
+        vert_mask = (kernel_center_y - kv_y_tile).abs() <= kernel_w // 2
+        image_mask = (q_idx < img_seq_len) & (kv_idx < img_seq_len)
+        image_to_text_mask = (q_idx < img_seq_len) & (kv_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length)
+        text_to_all_mask = (q_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length)
+        return (image_mask & time_mask & hori_mask & vert_mask) | image_to_text_mask | text_to_all_mask
+
+    sta_mask_3d.__name__ = f"natten_3d_c{canvas_t}x{canvas_w}x{canvas_h}_k{kernel_t}x{kernel_w}x{kernel_h}"
+    return sta_mask_3d
+
+
+def get_sliding_tile_attention_mask(kernel_size, tile_size, img_size, text_length, device, text_max_len=256):
+    img_seq_len = img_size[0] * img_size[1] * img_size[2]
+    image_mask = generate_sta_mask(img_size, kernel_size, tile_size, text_length)
+    mask = create_block_mask(image_mask,
+                             B=None,
+                             H=None,
+                             Q_LEN=img_seq_len + text_max_len,
+                             KV_LEN=img_seq_len + text_max_len,
+                             device=device,
+                             _compile=True)
+    return mask

csrc/sliding_tile_attention/test/test_sta.py

@@ -0,0 +1,96 @@
+import torch
+from flex_sta_ref import get_sliding_tile_attention_mask
+from st_attn import sliding_tile_attention
+from torch.nn.attention.flex_attention import flex_attention
+# from flash_attn_interface import flash_attn_func
+from tqdm import tqdm
+
+flex_attention = torch.compile(flex_attention, dynamic=False)
+
+
+def flex_test(Q, K, V, kernel_size):
+    mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (36, 48, 48), 39, 'cuda', 0)
+    output = flex_attention(Q, K, V, block_mask=mask)
+
+    return output
+
+
+def h100_fwd_kernel_test(Q, K, V, kernel_size):
+    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 39, False)
+    return o
+
+
+def generate_tensor(shape, mean, std, dtype, device):
+    tensor = torch.randn(shape, dtype=dtype, device=device)
+
+    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
+    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
+
+    return scaled_tensor.contiguous()
+
+
+def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mode='all'):
+    results = {
+        'TK vs FLEX': {
+            'sum_diff': 0,
+            'sum_abs': 0,
+            'max_diff': 0
+        },
+    }
+    kernel_size_ls = [(6, 1, 6), (6, 6, 1)]
+    from tqdm import tqdm
+    for kernel_size in tqdm(kernel_size_ls):
+        for _ in range(num_iterations):
+            torch.manual_seed(0)
+
+            Q = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+            K = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+            V = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+            tk_o = h100_fwd_kernel_test(Q, K, V, kernel_size)
+            pt_o = flex_test(Q, K, V, kernel_size)
+
+            diff = pt_o - tk_o
+            abs_diff = torch.abs(diff)
+            results['TK vs FLEX']['sum_diff'] += torch.sum(abs_diff).item()
+            results['TK vs FLEX']['max_diff'] = max(results['TK vs FLEX']['max_diff'], torch.max(abs_diff).item())
+
+            torch.cuda.empty_cache()
+        print("kernel_size", kernel_size)
+        print("max_diff", torch.max(abs_diff).item())
+        print(
+            "avg_diff",
+            torch.sum(abs_diff).item() / (b * h * n * d *
+                                          (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4)))
+
+    total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else
+                                                       3 if error_mode == 'backward' else 4) * len(kernel_size_ls)
+    for name, data in results.items():
+        avg_diff = data['sum_diff'] / total_elements
+        max_diff = data['max_diff']
+        results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff}
+
+    return results
+
+
+def generate_error_graphs(b, h, d, causal, mean, std, error_mode='all'):
+    seq_lengths = [82944]
+
+    tk_avg_errors, tk_max_errors = [], []
+
+    for n in tqdm(seq_lengths, desc="Generating error data"):
+        results = check_correctness(b, h, n, d, causal, mean, std, error_mode=error_mode)
+
+        tk_avg_errors.append(results['TK vs FLEX']['avg_diff'])
+        tk_max_errors.append(results['TK vs FLEX']['max_diff'])
+
+
+# Example usage
+b, h, d = 2, 24, 128
+causal = False
+mean = 1e-1
+std = 10
+
+for mode in ['output']:
+    generate_error_graphs(b, h, d, causal, mean, std, error_mode=mode)
+
+print("Error graphs generated and saved for all modes.")

demo/gradio_web_demo.py

@@ -1,78 +1,85 @@
+import argparse
+import os
+import tempfile
+
 import gradio as gr
 import torch
-from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
-from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
-from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.asymm_models_joint import AsymmDiTJoint
 from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import export_to_video
+
 from fastvideo.distill.solver import PCMFMScheduler
-import tempfile
-import os
-import argparse
-from safetensors.torch import load_file
+from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
+from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
+
 
 def init_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--prompts", nargs='+', default=[])
-    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument("--prompts", nargs="+", default=[])
+    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()
 
+
 def load_model(args):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     if args.scheduler_type == "euler":
         scheduler = FlowMatchEulerDiscreteScheduler()
     else:
-        scheduler = PCMFMScheduler(1000, args.shift, args.num_euler_timesteps, False, args.linear_threshold, args.linear_range)
-    
-    mochi_genmo = True
-    if mochi_genmo:
-        model_path = "/root/fastmochi_genmo/dit.safetensors"
-        state_dcit = load_file(model_path)
-        transformer = AsymmDiTJoint()
-        transformer.load_state_dict(state_dcit)
-        # from IPython import embed
-        # embed()
-        transformer.config.in_channels = 12
-        print("load gennmo mochi successfully")
+        linear_quadratic = True if "linear_quadratic" in args.scheduler_type else False
+        scheduler = PCMFMScheduler(
+            1000,
+            args.shift,
+            args.num_euler_timesteps,
+            linear_quadratic,
+            args.linear_threshold,
+            args.linear_range,
+        )
+
+    if args.transformer_path:
+        transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path)
     else:
-        if args.transformer_path:
-            transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path)
-        else:
-            transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder='transformer/')
-    
+        transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder="transformer/")
+
     pipe = MochiPipeline.from_pretrained(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
 
-def generate_video(prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, 
-                  num_frames, height, width, num_inference_steps, randomize_seed=False):
+
+def generate_video(
+    prompt,
+    negative_prompt,
+    use_negative_prompt,
+    seed,
+    guidance_scale,
+    num_frames,
+    height,
+    width,
+    num_inference_steps,
+    randomize_seed=False,
+):
     if randomize_seed:
-        seed = torch.randint(0, 1000000, (1,)).item()
-    
-    pipe = load_model(args)
-    print("load model successfully")
+        seed = torch.randint(0, 1000000, (1, )).item()
+
     generator = torch.Generator(device="cuda").manual_seed(seed)
-    
+
     if not use_negative_prompt:
         negative_prompt = None
-    
+
     with torch.autocast("cuda", dtype=torch.bfloat16):
         output = pipe(
             prompt=[prompt],
@@ -84,22 +91,24 @@ def generate_video(prompt, negative_prompt, use_negative_prompt, seed, guidance_
             guidance_scale=guidance_scale,
             generator=generator,
         ).frames[0]
-    
+
     output_path = os.path.join(tempfile.mkdtemp(), "output.mp4")
     export_to_video(output, output_path, fps=30)
     return output_path, seed
 
+
 examples = [
     "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."
+    "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.",
 ]
 
 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():
             prompt = gr.Text(
@@ -111,33 +120,54 @@ with gr.Blocks() as demo:
             )
             run_button = gr.Button("Run", scale=0)
         result = gr.Video(label="Result", show_label=False)
-    
+
     with gr.Accordion("Advanced options", open=False):
         with gr.Group():
             with gr.Row():
-                height = gr.Slider(label="Height", minimum=256, maximum=1024, step=32, value=args.height)
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=1024,
+                    step=32,
+                    value=args.height,
+                )
                 width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=args.width)
-            
+
             with gr.Row():
-                num_frames = gr.Slider(label="Number of Frames", minimum=8, maximum=256, value=args.num_frames)
-                guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=20, value=args.guidance_scale)
-                num_inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=args.num_inference_steps)
-            
+                num_frames = gr.Slider(
+                    label="Number of Frames",
+                    minimum=21,
+                    maximum=163,
+                    value=args.num_frames,
+                )
+                guidance_scale = gr.Slider(
+                    label="Guidance Scale",
+                    minimum=1,
+                    maximum=12,
+                    value=args.guidance_scale,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Inference Steps",
+                    minimum=4,
+                    maximum=100,
+                    value=args.num_inference_steps,
+                )
+
             with gr.Row():
                 use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
                 placeholder="Enter a negative prompt",
-                visible=False
+                visible=False,
             )
-            
+
             seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, step=1, value=args.seed)
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             seed_output = gr.Number(label="Used Seed")
-    
+
     gr.Examples(examples=examples, inputs=prompt)
-    
+
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
@@ -146,10 +176,20 @@ with gr.Blocks() as demo:
 
     run_button.click(
         fn=generate_video,
-        inputs=[prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, 
-                num_frames, height, width, num_inference_steps, randomize_seed],
-        outputs=[result, seed_output]
+        inputs=[
+            prompt,
+            negative_prompt,
+            use_negative_prompt,
+            seed,
+            guidance_scale,
+            num_frames,
+            height,
+            width,
+            num_inference_steps,
+            randomize_seed,
+        ],
+        outputs=[result, seed_output],
     )
 
 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)

demo/qualitative_results.md

@@ -0,0 +1,15 @@
+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
+
+Fast-Mochi comparison with original Mochi, achieving an 8X diffusion speed boost with the FastVideo framework.
+
+https://github.com/user-attachments/assets/5fbc4596-56d6-43aa-98e0-da472cf8e26c
+
+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

docs/data_preprocess.md

@@ -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.

env_setup.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+# install torch
+pip install torch==2.5.0 torchvision --index-url https://download.pytorch.org/whl/cu124
+
+# install FA2 and diffusers
+pip install packaging ninja && pip install flash-attn==2.7.0.post2 --no-build-isolation 
+
+pip install -r requirements-lint.txt
+
+pip install -r requirements.txt
+
+# install fastvideo
+pip install -e .

fastvideo/data_preprocess/preprocess_text_embeddings.py

@@ -1,19 +1,23 @@
 import argparse
-import torch
-from accelerate.logging import get_logger
-from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
-from diffusers.utils import export_to_video
 import json
 import os
+
+import torch
 import torch.distributed as dist
+from accelerate.logging import get_logger
+from diffusers.utils import export_to_video
+from diffusers.video_processor import VideoProcessor
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+
+from fastvideo.utils.load import load_text_encoder, load_vae
 
 logger = get_logger(__name__)
-from torch.utils.data import Dataset
-from torch.utils.data.distributed import DistributedSampler
-from torch.utils.data import DataLoader
 
 
 class T5dataset(Dataset):
+
     def __init__(
         self,
         json_path,
@@ -31,9 +35,7 @@ class T5dataset(Dataset):
         length = self.train_dataset[idx]["length"]
         if self.vae_debug:
             latents = torch.load(
-                os.path.join(
-                    args.output_dir, "latent", self.train_dataset[idx]["latent_path"]
-                ),
+                os.path.join(args.output_dir, "latent", self.train_dataset[idx]["latent_path"]),
                 map_location="cpu",
             )
         else:
@@ -53,12 +55,9 @@ def main(args):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     torch.cuda.set_device(local_rank)
     if not dist.is_initialized():
-        dist.init_process_group(
-            backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
-        )
+        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
 
-    pipe = MochiPipeline.from_pretrained(args.model_path).to(device)
-    pipe.vae.enable_tiling()
+    videoprocessor = VideoProcessor(vae_scale_factor=8)
     os.makedirs(args.output_dir, exist_ok=True)
     os.makedirs(os.path.join(args.output_dir, "video"), exist_ok=True)
     os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
@@ -67,9 +66,10 @@ def main(args):
 
     latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
     train_dataset = T5dataset(latents_json_path, args.vae_debug)
-    sampler = DistributedSampler(
-        train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True
-    )
+    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)
+    vae.enable_tiling()
+    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
     train_dataloader = DataLoader(
         train_dataset,
         sampler=sampler,
@@ -78,32 +78,25 @@ def main(args):
     )
 
     json_data = []
-    for _, data in enumerate(train_dataloader):
+    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
         with torch.inference_mode():
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                prompt_embeds, prompt_attention_mask, _, _ = pipe.encode_prompt(
-                    prompt=data["caption"],
-                )
+            with torch.autocast("cuda", dtype=autocast_type):
+                prompt_embeds, prompt_attention_mask = text_encoder.encode_prompt(prompt=data["caption"], )
                 if args.vae_debug:
                     latents = data["latents"]
-                    video = pipe.vae.decode(latents.to(device), return_dict=False)[0]
-                    video = pipe.video_processor.postprocess_video(video)
+                    video = vae.decode(latents.to(device), return_dict=False)[0]
+                    video = videoprocessor.postprocess_video(video)
                 for idx, video_name in enumerate(data["filename"]):
-                    prompt_embed_path = os.path.join(
-                        args.output_dir, "prompt_embed", video_name + ".pt"
-                    )
-                    video_path = os.path.join(
-                        args.output_dir, "video", video_name + ".mp4"
-                    )
-                    prompt_attention_mask_path = os.path.join(
-                        args.output_dir, "prompt_attention_mask", video_name + ".pt"
-                    )
+                    prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt")
+                    video_path = os.path.join(args.output_dir, "video", video_name + ".mp4")
+                    prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask",
+                                                              video_name + ".pt")
                     # save latent
                     torch.save(prompt_embeds[idx], prompt_embed_path)
                     torch.save(prompt_attention_mask[idx], prompt_attention_mask_path)
                     print(f"sample {video_name} saved")
                     if args.vae_debug:
-                        export_to_video(video[idx], video_path, fps=30)
+                        export_to_video(video[idx], video_path, fps=fps)
                     item = {}
                     item["length"] = int(data["length"][idx])
                     item["latent_path"] = video_name + ".pt"
@@ -126,6 +119,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     # dataset & dataloader
     parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--model_type", type=str, default="mochi")
     # text encoder & vae & diffusion model
     parser.add_argument(
         "--dataloader_num_workers",

fastvideo/data_preprocess/preprocess_vae_latents.py

@@ -1,16 +1,16 @@
-from fastvideo.dataset import getdataset
-from torch.utils.data import DataLoader
-from fastvideo.utils.dataset_utils import Collate
 import argparse
-import torch
-from accelerate import Accelerator
-from accelerate.logging import get_logger
-from accelerate.utils import ProjectConfiguration
 import json
 import os
-from diffusers import AutoencoderKLMochi
+
+import torch
 import torch.distributed as dist
+from accelerate.logging import get_logger
+from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+
+from fastvideo.dataset import getdataset
+from fastvideo.utils.load import load_vae
 
 logger = get_logger(__name__)
 
@@ -19,25 +19,8 @@ def main(args):
     local_rank = int(os.getenv("RANK", 0))
     world_size = int(os.getenv("WORLD_SIZE", 1))
     print("world_size", world_size, "local rank", local_rank)
-    args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = 4, 8, 8
-    args.ae_stride = args.ae_stride_h
-    patch_size_t, patch_size_h, patch_size_w = 1, 2, 2
-    args.patch_size = patch_size_h
-    args.patch_size_t, args.patch_size_h, args.patch_size_w = (
-        patch_size_t,
-        patch_size_h,
-        patch_size_w,
-    )
-    accelerator_project_config = ProjectConfiguration(
-        project_dir=args.output_dir, logging_dir=args.logging_dir
-    )
-    accelerator = Accelerator(
-        project_config=accelerator_project_config,
-    )
     train_dataset = getdataset(args)
-    sampler = DistributedSampler(
-        train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True
-    )
+    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
     train_dataloader = DataLoader(
         train_dataset,
         sampler=sampler,
@@ -45,38 +28,30 @@ def main(args):
         num_workers=args.dataloader_num_workers,
     )
 
-    encoder_device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
+    encoder_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     torch.cuda.set_device(local_rank)
     if not dist.is_initialized():
-        dist.init_process_group(
-            backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
-        )
-    vae = AutoencoderKLMochi.from_pretrained(args.model_path, subfolder="vae").to(
-        "cuda"
-    )
+        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+    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)
 
     json_data = []
-    for _, data in enumerate(train_dataloader):
+    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
         with torch.inference_mode():
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                latents = vae.encode(data["pixel_values"].to(encoder_device))[
-                    "latent_dist"
-                ].sample()
-                for idx, video_path in enumerate(data["path"]):
-                    video_name = os.path.basename(video_path).split(".")[0]
-                    latent_path = os.path.join(
-                        args.output_dir, "latent", video_name + ".pt"
-                    )
-                    torch.save(latents[idx].to(torch.bfloat16), latent_path)
-                    item = {}
-                    item["length"] = latents[idx].shape[1]
-                    item["latent_path"] = video_name + ".pt"
-                    item["caption"] = data["text"][idx]
-                    json_data.append(item)
-                    print(f"{video_name} processed")
+            with torch.autocast("cuda", dtype=autocast_type):
+                latents = vae.encode(data["pixel_values"].to(encoder_device))["latent_dist"].sample()
+            for idx, video_path in enumerate(data["path"]):
+                video_name = os.path.basename(video_path).split(".")[0]
+                latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt")
+                torch.save(latents[idx].to(torch.bfloat16), latent_path)
+                item = {}
+                item["length"] = latents[idx].shape[1]
+                item["latent_path"] = video_name + ".pt"
+                item["caption"] = data["text"][idx]
+                json_data.append(item)
+                print(f"{video_name} processed")
     dist.barrier()
     local_data = json_data
     gathered_data = [None] * world_size
@@ -91,6 +66,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     # dataset & dataloader
     parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--model_type", type=str, default="mochi")
     parser.add_argument("--data_merge_path", type=str, required=True)
     parser.add_argument("--num_frames", type=int, default=163)
     parser.add_argument(
@@ -105,9 +81,7 @@ if __name__ == "__main__":
         default=16,
         help="Batch size (per device) for the training dataloader.",
     )
-    parser.add_argument(
-        "--num_latent_t", type=int, default=28, help="Number of latent timesteps."
-    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
     parser.add_argument("--max_height", type=int, default=480)
     parser.add_argument("--max_width", type=int, default=848)
     parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
@@ -133,10 +107,8 @@ if __name__ == "__main__":
         "--logging_dir",
         type=str,
         default="logs",
-        help=(
-            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
-        ),
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
     )
 
     args = parser.parse_args()

fastvideo/data_preprocess/preprocess_validation_text_embeddings.py

@@ -0,0 +1,67 @@
+import argparse
+import os
+
+import torch
+import torch.distributed as dist
+from accelerate.logging import get_logger
+
+from fastvideo.utils.load import load_text_encoder
+
+logger = get_logger(__name__)
+
+
+def main(args):
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    print("world_size", world_size, "local rank", local_rank)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.cuda.set_device(local_rank)
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+
+    text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
+    autocast_type = torch.float16 if args.model_type == "hunyuan" else torch.bfloat16
+    # output_dir/validation/prompt_attention_mask
+    # output_dir/validation/prompt_embed
+    os.makedirs(os.path.join(args.output_dir, "validation"), exist_ok=True)
+    os.makedirs(
+        os.path.join(args.output_dir, "validation", "prompt_attention_mask"),
+        exist_ok=True,
+    )
+    os.makedirs(os.path.join(args.output_dir, "validation", "prompt_embed"), exist_ok=True)
+
+    with open(args.validation_prompt_txt, "r", encoding="utf-8") as file:
+        lines = file.readlines()
+    prompts = [line.strip() for line in lines]
+    for prompt in prompts:
+        with torch.inference_mode():
+            with torch.autocast("cuda", dtype=autocast_type):
+                prompt_embeds, prompt_attention_mask = text_encoder.encode_prompt(prompt)
+                file_name = prompt.split(".")[0]
+                prompt_embed_path = os.path.join(args.output_dir, "validation", "prompt_embed", f"{file_name}.pt")
+                prompt_attention_mask_path = os.path.join(
+                    args.output_dir,
+                    "validation",
+                    "prompt_attention_mask",
+                    f"{file_name}.pt",
+                )
+                torch.save(prompt_embeds[0], prompt_embed_path)
+                torch.save(prompt_attention_mask[0], prompt_attention_mask_path)
+                print(f"sample {file_name} saved")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--model_type", type=str, default="mochi")
+    parser.add_argument("--validation_prompt_txt", type=str)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    args = parser.parse_args()
+    main(args)

fastvideo/dataset/__init__.py

@@ -1,14 +1,9 @@
-from transformers import AutoTokenizer
-
 from torchvision import transforms
 from torchvision.transforms import Lambda
+from transformers import AutoTokenizer
+
 from fastvideo.dataset.t2v_datasets import T2V_dataset
-from fastvideo.dataset.latent_datasets import LatentDataset
-from fastvideo.dataset.transform import (
-    Normalize255,
-    TemporalRandomCrop,
-    CenterCropResizeVideo,
-)
+from fastvideo.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop
 
 
 def getdataset(args):
@@ -20,26 +15,17 @@ def getdataset(args):
     resize = [
         CenterCropResizeVideo((args.max_height, args.max_width)),
     ]
-    transform = transforms.Compose(
-        [
-            # Normalize255(),
-            *resize,
-            # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
-            # norm_fun
-        ]
-    )
-    transform_topcrop = transforms.Compose(
-        [
-            Normalize255(),
-            *resize_topcrop,
-            # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
-            norm_fun,
-        ]
-    )
+    transform = transforms.Compose([
+        # Normalize255(),
+        *resize,
+    ])
+    transform_topcrop = transforms.Compose([
+        Normalize255(),
+        *resize_topcrop,
+        norm_fun,
+    ])
     # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
-    tokenizer = AutoTokenizer.from_pretrained(
-        args.text_encoder_name, cache_dir=args.cache_dir
-    )
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
     if args.dataset == "t2v":
         return T2V_dataset(
             args,
@@ -53,11 +39,13 @@ def getdataset(args):
 
 
 if __name__ == "__main__":
-    from accelerate import Accelerator
-    from fastvideo.dataset.t2v_datasets import dataset_prog
     import random
+
+    from accelerate import Accelerator
     from tqdm import tqdm
 
+    from fastvideo.dataset.t2v_datasets import dataset_prog
+
     args = type(
         "args",
         (),
@@ -92,9 +80,7 @@ if __name__ == "__main__":
     zero = 0
     for idx in tqdm(range(num)):
         image_data = dataset_prog.img_cap_list[idx]
-        caps = [
-            i["cap"] if isinstance(i["cap"], list) else [i["cap"]] for i in image_data
-        ]
+        caps = [i["cap"] if isinstance(i["cap"], list) else [i["cap"]] for i in image_data]
         try:
             caps = [[random.choice(i)] for i in caps]
         except Exception as e:

fastvideo/dataset/latent_datasets.py

@@ -1,11 +1,13 @@
-import torch
-from torch.utils.data import Dataset
 import json
 import os
 import random
 
+import torch
+from torch.utils.data import Dataset
+
 
 class LatentDataset(Dataset):
+
     def __init__(
         self,
         json_path,
@@ -19,9 +21,7 @@ class LatentDataset(Dataset):
         self.video_dir = os.path.join(self.datase_dir_path, "video")
         self.latent_dir = os.path.join(self.datase_dir_path, "latent")
         self.prompt_embed_dir = os.path.join(self.datase_dir_path, "prompt_embed")
-        self.prompt_attention_mask_dir = os.path.join(
-            self.datase_dir_path, "prompt_attention_mask"
-        )
+        self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask")
         with open(self.json_path, "r") as f:
             self.data_anno = json.load(f)
         # json.load(f) already keeps the order
@@ -31,10 +31,7 @@ class LatentDataset(Dataset):
         self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32)
         # 256 zeros
         self.uncond_prompt_mask = torch.zeros(256).bool()
-        self.lengths = [
-            data_item["length"] if "length" in data_item else 1
-            for data_item in self.data_anno
-        ]
+        self.lengths = [data_item["length"] if "length" in data_item else 1 for data_item in self.data_anno]
 
     def __getitem__(self, idx):
         latent_file = self.data_anno[idx]["latent_path"]
@@ -46,7 +43,7 @@ class LatentDataset(Dataset):
             map_location="cpu",
             weights_only=True,
         )
-        latent = latent.squeeze(0)[:, -self.num_latent_t :]
+        latent = latent.squeeze(0)[:, -self.num_latent_t:]
         if random.random() < self.cfg_rate:
             prompt_embed = self.uncond_prompt_embed
             prompt_attention_mask = self.uncond_prompt_mask
@@ -57,9 +54,7 @@ class LatentDataset(Dataset):
                 weights_only=True,
             )
             prompt_attention_mask = torch.load(
-                os.path.join(
-                    self.prompt_attention_mask_dir, prompt_attention_mask_file
-                ),
+                os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file),
                 map_location="cpu",
                 weights_only=True,
             )
@@ -92,16 +87,15 @@ def latent_collate_function(batch):
                 0,
                 max_w - latent.shape[3],
             ),
-        )
-        for latent in latents
+        ) for latent in latents
     ]
     # attn mask
     latent_attn_mask = torch.ones(len(latents), max_t, max_h, max_w)
     # set to 0 if padding
     for i, latent in enumerate(latents):
-        latent_attn_mask[i, latent.shape[1] :, :, :] = 0
-        latent_attn_mask[i, :, latent.shape[2] :, :] = 0
-        latent_attn_mask[i, :, :, latent.shape[3] :] = 0
+        latent_attn_mask[i, latent.shape[1]:, :, :] = 0
+        latent_attn_mask[i, :, latent.shape[2]:, :] = 0
+        latent_attn_mask[i, :, :, latent.shape[3]:] = 0
 
     prompt_embeds = torch.stack(prompt_embeds, dim=0)
     prompt_attention_masks = torch.stack(prompt_attention_masks, dim=0)
@@ -111,9 +105,7 @@ def latent_collate_function(batch):
 
 if __name__ == "__main__":
     dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt", num_latent_t=28)
-    dataloader = torch.utils.data.DataLoader(
-        dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function
-    )
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function)
     for latent, prompt_embed, latent_attn_mask, prompt_attention_mask in dataloader:
         print(
             latent.shape,

fastvideo/dataset/t2v_datasets.py

@@ -1,21 +1,19 @@
 import json
-import os, io, csv, math, random
-import numpy as np
-from einops import rearrange
-from decord import VideoReader
-from os.path import join as opj
+import math
+import os
+import random
 from collections import Counter
+from os.path import join as opj
 
+import numpy as np
 import torch
-from torch.utils.data.dataset import Dataset
-from torch.utils.data import DataLoader, Dataset, get_worker_info
-from tqdm import tqdm
-from PIL import Image
-from accelerate.logging import get_logger
-from fastvideo.utils.dataset_utils import DecordInit
 import torchvision
+from einops import rearrange
+from PIL import Image
+from torch.utils.data import Dataset
 
-logger = get_logger(__name__)
+from fastvideo.utils.dataset_utils import DecordInit
+from fastvideo.utils.logging_ import main_print
 
 
 class SingletonMeta(type):
@@ -29,6 +27,7 @@ class SingletonMeta(type):
 
 
 class DataSetProg(metaclass=SingletonMeta):
+
     def __init__(self):
         self.cap_list = []
         self.elements = []
@@ -58,9 +57,7 @@ class DataSetProg(metaclass=SingletonMeta):
         else:
             worker_id = work_info.id
 
-        idx = self.worker_elements[worker_id][
-            self.n_used_elements[worker_id] % len(self.worker_elements[worker_id])
-        ]
+        idx = self.worker_elements[worker_id][self.n_used_elements[worker_id] % len(self.worker_elements[worker_id])]
         self.n_used_elements[worker_id] += 1
         return idx
 
@@ -75,6 +72,7 @@ def filter_resolution(h, w, max_h_div_w_ratio=17 / 16, min_h_div_w_ratio=8 / 16)
 
 
 class T2V_dataset(Dataset):
+
     def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
         self.data = args.data_merge_path
         self.num_frames = args.num_frames
@@ -94,7 +92,7 @@ class T2V_dataset(Dataset):
         self.v_decoder = DecordInit()
         self.video_length_tolerance_range = args.video_length_tolerance_range
         self.support_Chinese = True
-        if not ("mt5" in args.text_encoder_name):
+        if "mt5" not in args.text_encoder_name:
             self.support_Chinese = False
 
         cap_list = self.get_cap_list()
@@ -116,14 +114,9 @@ class T2V_dataset(Dataset):
         return dataset_prog.n_elements
 
     def __getitem__(self, idx):
-        try:
-            data = self.get_data(idx)
-            return data
-        except Exception as e:
-            logger.info(f"Error with {e}")
-            if idx in dataset_prog.cap_list:
-                logger.info(f"Caught an exception! {dataset_prog.cap_list[idx]}")
-            return self.__getitem__(random.randint(0, self.__len__() - 1))
+
+        data = self.get_data(idx)
+        return data
 
     def get_data(self, idx):
         path = dataset_prog.cap_list[idx]["path"]
@@ -136,9 +129,7 @@ class T2V_dataset(Dataset):
         video_path = dataset_prog.cap_list[idx]["path"]
         assert os.path.exists(video_path), f"file {video_path} do not exist!"
         frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
-        torchvision_video, _, metadata = torchvision.io.read_video(
-            video_path, output_format="TCHW"
-        )
+        torchvision_video, _, metadata = torchvision.io.read_video(video_path, output_format="TCHW")
         video = torchvision_video[frame_indices]
         video = self.transform(video)
         video = rearrange(video, "t c h w -> c t h w")
@@ -187,24 +178,17 @@ class T2V_dataset(Dataset):
         #     h, w = i.shape[-2:]
         #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
 
-        image = (
-            self.transform_topcrop(image)
-            if "human_images" in image_data["path"]
-            else self.transform(image)
-        )  #  [1 C H W] -> num_img [1 C H W]
+        image = (self.transform_topcrop(image) if "human_images" in image_data["path"] else self.transform(image)
+                 )  #  [1 C H W] -> num_img [1 C H W]
         image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
 
         image = image.float() / 127.5 - 1.0
 
-        caps = (
-            image_data["cap"]
-            if isinstance(image_data["cap"], list)
-            else [image_data["cap"]]
-        )
+        caps = (image_data["cap"] if isinstance(image_data["cap"], list) else [image_data["cap"]])
         caps = [random.choice(caps)]
         text = caps
         input_ids, cond_mask = [], []
-        text = text if random.random() > self.cfg else ""
+        text = text[0] if random.random() > self.cfg else ""
         text_tokens_and_mask = self.tokenizer(
             text,
             max_length=self.text_max_length,
@@ -254,10 +238,7 @@ class T2V_dataset(Dataset):
                     cnt_no_resolution += 1
                     continue
                 else:
-                    if (
-                        resolution.get("height", None) is None
-                        or resolution.get("width", None) is None
-                    ):
+                    if (resolution.get("height", None) is None or resolution.get("width", None) is None):
                         cnt_no_resolution += 1
                         continue
                     height, width = i["resolution"]["height"], i["resolution"]["width"]
@@ -277,26 +258,19 @@ class T2V_dataset(Dataset):
                 # import ipdb;ipdb.set_trace()
                 i["num_frames"] = math.ceil(fps * duration)
                 # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
-                if (
-                    i["num_frames"] / fps
-                    > self.video_length_tolerance_range
-                    * (self.num_frames / self.train_fps * self.speed_factor)
-                ):  # too long video is not suitable for this training stage (self.num_frames)
+                if i["num_frames"] / fps > self.video_length_tolerance_range * (
+                        self.num_frames / self.train_fps *
+                        self.speed_factor):  # too long video is not suitable for this training stage (self.num_frames)
                     cnt_too_long += 1
                     continue
 
                 # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
                 frame_interval = fps / self.train_fps
                 start_frame_idx = 0
-                frame_indices = np.arange(
-                    start_frame_idx, i["num_frames"], frame_interval
-                ).astype(int)
+                frame_indices = np.arange(start_frame_idx, i["num_frames"], frame_interval).astype(int)
 
                 # comment out it to enable dynamic frames training
-                if (
-                    len(frame_indices) < self.num_frames
-                    and random.random() < self.drop_short_ratio
-                ):
+                if (len(frame_indices) < self.num_frames and random.random() < self.drop_short_ratio):
                     cnt_too_short += 1
                     continue
 
@@ -307,9 +281,7 @@ class T2V_dataset(Dataset):
                     # frame_indices = frame_indices[:self.num_frames]  # head crop
                 i["sample_frame_index"] = frame_indices.tolist()
                 new_cap_list.append(i)
-                i["sample_num_frames"] = len(
-                    i["sample_frame_index"]
-                )  # will use in dataloader(group sampler)
+                i["sample_num_frames"] = len(i["sample_frame_index"])  # will use in dataloader(group sampler)
                 sample_num_frames.append(i["sample_num_frames"])
             elif path.endswith(".jpg"):  # image
                 cnt_img += 1
@@ -318,15 +290,13 @@ class T2V_dataset(Dataset):
                 sample_num_frames.append(i["sample_num_frames"])
             else:
                 raise NameError(
-                    f"Unknown file extention {path.split('.')[-1]}, only support .mp4 for video and .jpg for image"
-                )
+                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image")
         # import ipdb;ipdb.set_trace()
-        logger.info(
+        main_print(
             f"no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, "
             f"no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, "
             f"Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, "
-            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}"
-        )
+            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}")
         return new_cap_list, sample_num_frames
 
     def decord_read(self, path, frame_indices):
@@ -339,14 +309,11 @@ class T2V_dataset(Dataset):
     def read_jsons(self, data):
         cap_lists = []
         with open(data, "r") as f:
-            folder_anno = [
-                i.strip().split(",") for i in f.readlines() if len(i.strip()) > 0
-            ]
+            folder_anno = [i.strip().split(",") for i in f.readlines() if len(i.strip()) > 0]
         print(folder_anno)
         for folder, anno in folder_anno:
             with open(anno, "r") as f:
                 sub_list = json.load(f)
-            logger.info(f"Building {anno}...")
             for i in range(len(sub_list)):
                 sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
             cap_lists += sub_list

fastvideo/dataset/transform.py

@@ -1,7 +1,8 @@
-import torch
-import random
 import numbers
-from torchvision.transforms import RandomCrop, RandomResizedCrop
+import random
+
+import torch
+from PIL import Image
 
 
 def _is_tensor_video_clip(clip):
@@ -20,21 +21,15 @@ def center_crop_arr(pil_image, image_size):
     https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
     """
     while min(*pil_image.size) >= 2 * image_size:
-        pil_image = pil_image.resize(
-            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
-        )
+        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
 
     scale = image_size / min(*pil_image.size)
-    pil_image = pil_image.resize(
-        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
-    )
+    pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
 
     arr = np.array(pil_image)
     crop_y = (arr.shape[0] - image_size) // 2
     crop_x = (arr.shape[1] - image_size) // 2
-    return Image.fromarray(
-        arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
-    )
+    return Image.fromarray(arr[crop_y:crop_y + image_size, crop_x:crop_x + image_size])
 
 
 def crop(clip, i, j, h, w):
@@ -44,14 +39,12 @@ def crop(clip, i, j, h, w):
     """
     if len(clip.size()) != 4:
         raise ValueError("clip should be a 4D tensor")
-    return clip[..., i : i + h, j : j + w]
+    return clip[..., i:i + h, j:j + w]
 
 
 def resize(clip, target_size, interpolation_mode):
     if len(target_size) != 2:
-        raise ValueError(
-            f"target size should be tuple (height, width), instead got {target_size}"
-        )
+        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
     return torch.nn.functional.interpolate(
         clip,
         size=target_size,
@@ -63,9 +56,7 @@ def resize(clip, target_size, interpolation_mode):
 
 def resize_scale(clip, target_size, interpolation_mode):
     if len(target_size) != 2:
-        raise ValueError(
-            f"target size should be tuple (height, width), instead got {target_size}"
-        )
+        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
     H, W = clip.size(-2), clip.size(-1)
     scale_ = target_size[0] / min(H, W)
     return torch.nn.functional.interpolate(
@@ -153,16 +144,14 @@ def random_shift_crop(clip):
     h, w = clip.size(-2), clip.size(-1)
 
     if h <= w:
-        long_edge = w
         short_edge = h
     else:
-        long_edge = h
         short_edge = w
 
     th, tw = short_edge, short_edge
 
-    i = torch.randint(0, h - th + 1, size=(1,)).item()
-    j = torch.randint(0, w - tw + 1, size=(1,)).item()
+    i = torch.randint(0, h - th + 1, size=(1, )).item()
+    j = torch.randint(0, w - tw + 1, size=(1, )).item()
     return crop(clip, i, j, th, tw)
 
 
@@ -177,9 +166,7 @@ def normalize_video(clip):
     """
     _is_tensor_video_clip(clip)
     if not clip.dtype == torch.uint8:
-        raise TypeError(
-            "clip tensor should have data type uint8. Got %s" % str(clip.dtype)
-        )
+        raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype))
     # return clip.float().permute(3, 0, 1, 2) / 255.0
     return clip.float() / 255.0
 
@@ -217,6 +204,7 @@ def hflip(clip):
 
 
 class RandomCropVideo:
+
     def __init__(self, size):
         if isinstance(size, numbers.Number):
             self.size = (int(size), int(size))
@@ -239,15 +227,13 @@ class RandomCropVideo:
         th, tw = self.size
 
         if h < th or w < tw:
-            raise ValueError(
-                f"Required crop size {(th, tw)} is larger than input image size {(h, w)}"
-            )
+            raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}")
 
         if w == tw and h == th:
             return 0, 0, h, w
 
-        i = torch.randint(0, h - th + 1, size=(1,)).item()
-        j = torch.randint(0, w - tw + 1, size=(1,)).item()
+        i = torch.randint(0, h - th + 1, size=(1, )).item()
+        j = torch.randint(0, w - tw + 1, size=(1, )).item()
 
         return i, j, th, tw
 
@@ -256,6 +242,7 @@ class RandomCropVideo:
 
 
 class SpatialStrideCropVideo:
+
     def __init__(self, stride):
         self.stride = stride
 
@@ -314,9 +301,7 @@ class LongSideResizeVideo:
         else:
             h = int(h * self.size / w)
             w = self.size
-        resize_clip = resize(
-            clip, target_size=(h, w), interpolation_mode=self.interpolation_mode
-        )
+        resize_clip = resize(clip, target_size=(h, w), interpolation_mode=self.interpolation_mode)
         return resize_clip
 
     def __repr__(self) -> str:
@@ -336,9 +321,7 @@ class CenterCropResizeVideo:
         interpolation_mode="bilinear",
     ):
         if len(size) != 2:
-            raise ValueError(
-                f"size should be tuple (height, width), instead got {size}"
-            )
+            raise ValueError(f"size should be tuple (height, width), instead got {size}")
         self.size = size
         self.top_crop = top_crop
         self.interpolation_mode = interpolation_mode
@@ -352,9 +335,7 @@ class CenterCropResizeVideo:
                 size is (T, C, crop_size, crop_size)
         """
         # clip_center_crop = center_crop_using_short_edge(clip)
-        clip_center_crop = center_crop_th_tw(
-            clip, self.size[0], self.size[1], top_crop=self.top_crop
-        )
+        clip_center_crop = center_crop_th_tw(clip, self.size[0], self.size[1], top_crop=self.top_crop)
         # import ipdb;ipdb.set_trace()
         clip_center_crop_resize = resize(
             clip_center_crop,
@@ -380,9 +361,7 @@ class UCFCenterCropVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(
-                    f"size should be tuple (height, width), instead got {size}"
-                )
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -397,9 +376,7 @@ class UCFCenterCropVideo:
             torch.tensor: scale resized / center cropped video clip.
                 size is (T, C, crop_size, crop_size)
         """
-        clip_resize = resize_scale(
-            clip=clip, target_size=self.size, interpolation_mode=self.interpolation_mode
-        )
+        clip_resize = resize_scale(clip=clip, target_size=self.size, interpolation_mode=self.interpolation_mode)
         clip_center_crop = center_crop(clip_resize, self.size)
         return clip_center_crop
 
@@ -419,9 +396,7 @@ class KineticsRandomCropResizeVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(
-                    f"size should be tuple (height, width), instead got {size}"
-                )
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -435,6 +410,7 @@ class KineticsRandomCropResizeVideo:
 
 
 class CenterCropVideo:
+
     def __init__(
         self,
         size,
@@ -442,9 +418,7 @@ class CenterCropVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(
-                    f"size should be tuple (height, width), instead got {size}"
-                )
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -571,9 +545,7 @@ class DynamicSampleDuration(object):
     def __call__(self, t, h, w):
         if self.extra_1:
             t = t - 1
-        truncate_t_list = list(range(t + 1))[t // 2 :][
-            :: self.t_stride
-        ]  # need half at least
+        truncate_t_list = list(range(t + 1))[t // 2:][::self.t_stride]  # need half at least
         truncate_t = random.choice(truncate_t_list)
         if self.extra_1:
             truncate_t = truncate_t + 1
@@ -581,27 +553,22 @@ class DynamicSampleDuration(object):
 
 
 if __name__ == "__main__":
-    from torchvision import transforms
-    import torchvision.io as io
-    import numpy as np
-    from torchvision.utils import save_image
     import os
 
-    vframes, aframes, info = io.read_video(
-        filename="./v_Archery_g01_c03.avi", pts_unit="sec", output_format="TCHW"
-    )
+    import numpy as np
+    import torchvision.io as io
+    from torchvision import transforms
+    from torchvision.utils import save_image
 
-    trans = transforms.Compose(
-        [
-            Normalize255(),
-            RandomHorizontalFlipVideo(),
-            UCFCenterCropVideo(512),
-            # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
-            transforms.Normalize(
-                mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True
-            ),
-        ]
-    )
+    vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi", pts_unit="sec", output_format="TCHW")
+
+    trans = transforms.Compose([
+        Normalize255(),
+        RandomHorizontalFlipVideo(),
+        UCFCenterCropVideo(512),
+        # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+    ])
 
     target_video_len = 32
     frame_interval = 1
@@ -615,9 +582,7 @@ if __name__ == "__main__":
     # print(start_frame_ind)
     # print(end_frame_ind)
     assert end_frame_ind - start_frame_ind >= target_video_len
-    frame_indice = np.linspace(
-        start_frame_ind, end_frame_ind - 1, target_video_len, dtype=int
-    )
+    frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, target_video_len, dtype=int)
     print(frame_indice)
 
     select_vframes = vframes[frame_indice]
@@ -628,9 +593,7 @@ if __name__ == "__main__":
     print(select_vframes_trans.shape)
     print(select_vframes_trans.dtype)
 
-    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * 255).to(
-        dtype=torch.uint8
-    )
+    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * 255).to(dtype=torch.uint8)
     print(select_vframes_trans_int.dtype)
     print(select_vframes_trans_int.permute(0, 2, 3, 1).shape)
 

fastvideo/distill.py

@@ -1,86 +1,48 @@
+# !/bin/python3
+# isort: skip_file
 import argparse
 import math
 import os
-from fastvideo.utils.parallel_states import (
-    initialize_sequence_parallel_state,
-    destroy_sequence_parallel_group,
-    get_sequence_parallel_state,
-    nccl_info,
-)
-from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast
-from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_mochi_dit_input
-from fastvideo.utils.validation import log_validation
 import time
-from torch.utils.data import DataLoader
+from collections import deque
+from copy import deepcopy
+
 import torch
-from torch.distributed.fsdp import ShardingStrategy
-from torch.distributed.fsdp import (
-    FullyShardedDataParallel as FSDP,
-    StateDictType,
-    FullStateDictConfig,
-)
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
-import json
-from torch.utils.data.distributed import DistributedSampler
-from fastvideo.utils.dataset_utils import LengthGroupedSampler
+import torch.distributed as dist
 import wandb
 from accelerate.utils import set_seed
-from tqdm.auto import tqdm
-from fastvideo.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing
-from diffusers import (
-    FlowMatchEulerDiscreteScheduler,
-)
-from fastvideo.distill.solver import EulerSolver, extract_into_tensor
-from copy import deepcopy
+from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.optimization import get_scheduler
-from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
 from diffusers.utils import check_min_version
-from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_function
-import torch.distributed as dist
-from safetensors.torch import save_file
 from peft import LoraConfig
-from torch.distributed.fsdp import (
-    FullyShardedDataParallel as FSDP,
-)
-from fastvideo.utils.checkpoint import (
-    save_checkpoint,
-    save_lora_checkpoint,
-    resume_lora_optimizer,
-)
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import ShardingStrategy
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from tqdm.auto import tqdm
+
+from fastvideo.dataset.latent_datasets import (LatentDataset, latent_collate_function)
+from fastvideo.distill.solver import EulerSolver, extract_into_tensor
+from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input
+from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
+from fastvideo.utils.checkpoint import (resume_lora_optimizer, save_checkpoint, save_lora_checkpoint)
+from fastvideo.utils.communications import (broadcast, sp_parallel_dataloader_wrapper)
+from fastvideo.utils.dataset_utils import LengthGroupedSampler
+from fastvideo.utils.fsdp_util import (apply_fsdp_checkpointing, get_dit_fsdp_kwargs)
+from fastvideo.utils.load import load_transformer
+from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group, get_sequence_parallel_state,
+                                             initialize_sequence_parallel_state)
+from fastvideo.utils.validation import log_validation
 
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
 check_min_version("0.31.0")
-import time
-from collections import deque
+
 
 def main_print(content):
     if int(os.environ["LOCAL_RANK"]) <= 0:
         print(content)
 
 
-def save_checkpoint(transformer: MochiTransformer3DModel, rank, output_dir, step):
-    main_print(f"--> saving checkpoint at step {step}")
-    with FSDP.state_dict_type(
-        transformer,
-        StateDictType.FULL_STATE_DICT,
-        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
-    ):
-        cpu_state = transformer.state_dict()
-    # todo move to get_state_dict
-    if rank <= 0:
-        save_dir = os.path.join(output_dir, f"checkpoint-{step}")
-        os.makedirs(save_dir, exist_ok=True)
-        # save using safetensors
-        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
-        save_file(cpu_state, weight_path)
-        config_dict = dict(transformer.config)
-        config_path = os.path.join(save_dir, "config.json")
-        # save dict as json
-        with open(config_path, "w") as f:
-            json.dump(config_dict, f, indent=4)
-    main_print(f"--> checkpoint saved at step {step}")
-
-
 def reshard_fsdp(model):
     for m in FSDP.fsdp_modules(model):
         if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
@@ -89,24 +51,23 @@ def reshard_fsdp(model):
 
 def get_norm(model_pred, norms, gradient_accumulation_steps):
     fro_norm = (
-        torch.linalg.matrix_norm(model_pred, ord="fro") / gradient_accumulation_steps
-    )
-    largest_singular_value = (
-        torch.linalg.matrix_norm(model_pred, ord=2) / gradient_accumulation_steps
-    )
+        torch.linalg.matrix_norm(model_pred, ord="fro") /  # codespell:ignore
+        gradient_accumulation_steps)
+    largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) / gradient_accumulation_steps)
     absolute_mean = torch.mean(torch.abs(model_pred)) / gradient_accumulation_steps
     absolute_max = torch.max(torch.abs(model_pred)) / gradient_accumulation_steps
     dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG)
     dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG)
     dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG)
-    norms["fro"] += torch.mean(fro_norm).item()
+    norms["fro"] += torch.mean(fro_norm).item()  # codespell:ignore
     norms["largest singular value"] += torch.mean(largest_singular_value).item()
     norms["absolute mean"] += absolute_mean.item()
     norms["absolute max"] += absolute_max.item()
 
 
-def train_one_step_mochi(
+def distill_one_step(
     transformer,
+    model_type,
     teacher_transformer,
     ema_transformer,
     optimizer,
@@ -127,11 +88,12 @@ def train_one_step_mochi(
     ema_decay,
     pred_decay_weight,
     pred_decay_type,
+    hunyuan_teacher_disable_cfg,
 ):
     total_loss = 0.0
     optimizer.zero_grad()
     model_pred_norm = {
-        "fro": 0.0,
+        "fro": 0.0,  # codespell:ignore
         "largest singular value": 0.0,
         "absolute mean": 0.0,
         "absolute max": 0.0,
@@ -143,12 +105,10 @@ def train_one_step_mochi(
             latents_attention_mask,
             encoder_attention_mask,
         ) = next(loader)
-        model_input = normalize_mochi_dit_input(latents)
+        model_input = normalize_dit_input(model_type, latents)
         noise = torch.randn_like(model_input)
         bsz = model_input.shape[0]
-        index = torch.randint(
-            0, num_euler_timesteps, (bsz,), device=model_input.device
-        ).long()
+        index = torch.randint(0, num_euler_timesteps, (bsz, ), device=model_input.device).long()
         if sp_size > 1:
             broadcast(index)
         # Add noise according to flow matching.
@@ -159,25 +119,25 @@ def train_one_step_mochi(
         timesteps = (sigmas * noise_scheduler.config.num_train_timesteps).view(-1)
         # if squeeze to [], unsqueeze to [1]
 
-        timesteps_prev = (
-            sigmas_prev * noise_scheduler.config.num_train_timesteps
-        ).view(-1)
+        timesteps_prev = (sigmas_prev * noise_scheduler.config.num_train_timesteps).view(-1)
         noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
-
         # Predict the noise residual
         with torch.autocast("cuda", dtype=torch.bfloat16):
-            model_pred = transformer(
-                noisy_model_input,
-                encoder_hidden_states,
-                timesteps,
-                encoder_attention_mask,  # B, L
-                return_dict=False,
-            )[0]
+            teacher_kwargs = {
+                "hidden_states": noisy_model_input,
+                "encoder_hidden_states": encoder_hidden_states,
+                "timestep": timesteps,
+                "encoder_attention_mask": encoder_attention_mask,  # B, L
+                "return_dict": False,
+            }
+            if hunyuan_teacher_disable_cfg:
+                teacher_kwargs["guidance"] = torch.tensor([1000.0],
+                                                          device=noisy_model_input.device,
+                                                          dtype=torch.bfloat16)
+            model_pred = transformer(**teacher_kwargs)[0]
 
         # if accelerator.is_main_process:
-        model_pred, end_index = solver.euler_style_multiphase_pred(
-            noisy_model_input, model_pred, index, multiphase
-        )
+        model_pred, end_index = solver.euler_style_multiphase_pred(noisy_model_input, model_pred, index, multiphase)
         with torch.no_grad():
             w = distill_cfg
             with torch.autocast("cuda", dtype=torch.bfloat16):
@@ -200,9 +160,7 @@ def train_one_step_mochi(
                         uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
                         return_dict=False,
                     )[0].float()
-            teacher_output = cond_teacher_output + w * (
-                cond_teacher_output - uncond_teacher_output
-            )
+            teacher_output = uncond_teacher_output + w * (cond_teacher_output - uncond_teacher_output)
             x_prev = solver.euler_step(noisy_model_input, teacher_output, index)
 
         # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
@@ -225,65 +183,44 @@ def train_one_step_mochi(
                         return_dict=False,
                     )[0]
 
-            target, end_index = solver.euler_style_multiphase_pred(
-                x_prev, target_pred, index, multiphase, True
-            )
+            target, end_index = solver.euler_style_multiphase_pred(x_prev, target_pred, index, multiphase, True)
 
         huber_c = 0.001
         # loss = loss.mean()
-        loss = (
-            torch.mean(
-                torch.sqrt((model_pred.float() - target.float()) ** 2 + huber_c**2)
-                - huber_c
-            )
-            / gradient_accumulation_steps
-        )
+        loss = (torch.mean(torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) - huber_c) /
+                gradient_accumulation_steps)
         if pred_decay_weight > 0:
             if pred_decay_type == "l1":
-                pred_decay_loss = (
-                    torch.mean(torch.sqrt(model_pred.float() ** 2))
-                    * pred_decay_weight
-                    / gradient_accumulation_steps
-                )
+                pred_decay_loss = (torch.mean(torch.sqrt(model_pred.float()**2)) * pred_decay_weight /
+                                   gradient_accumulation_steps)
                 loss += pred_decay_loss
             elif pred_decay_type == "l2":
                 # essnetially k2?
-                pred_decay_loss = (
-                    torch.mean(model_pred.float() ** 2)
-                    * pred_decay_weight
-                    / gradient_accumulation_steps
-                )
+                pred_decay_loss = (torch.mean(model_pred.float()**2) * pred_decay_weight / gradient_accumulation_steps)
                 loss += pred_decay_loss
             else:
                 assert NotImplementedError("pred_decay_type is not implemented")
 
         # calculate model_pred norm and mean
-        get_norm(
-            model_pred.detach().float(), model_pred_norm, gradient_accumulation_steps
-        )
+        get_norm(model_pred.detach().float(), model_pred_norm, gradient_accumulation_steps)
         loss.backward()
 
         avg_loss = loss.detach().clone()
         dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
-        dist.all_reduce(pred_decay_loss.detach(), op=dist.ReduceOp.AVG)
         total_loss += avg_loss.item()
 
     # update ema
     if ema_transformer is not None:
         reshard_fsdp(ema_transformer)
-        for p_averaged, p_model in zip(
-            ema_transformer.parameters(), transformer.parameters()
-        ):
+        for p_averaged, p_model in zip(ema_transformer.parameters(), transformer.parameters()):
             with torch.no_grad():
-                p_averaged.copy_(
-                    torch.lerp(p_averaged.detach(), p_model.detach(), 1 - ema_decay)
-                )
+                p_averaged.copy_(torch.lerp(p_averaged.detach(), p_model.detach(), 1 - ema_decay))
 
     grad_norm = transformer.clip_grad_norm_(max_grad_norm)
     optimizer.step()
     lr_scheduler.step()
 
-    return total_loss, grad_norm.item(), model_pred_norm, pred_decay_loss.item()
+    return total_loss, grad_norm.item(), model_pred_norm
 
 
 def main(args):
@@ -308,26 +245,20 @@ def main(args):
     if rank <= 0 and args.output_dir is not None:
         os.makedirs(args.output_dir, exist_ok=True)
 
-    # For mixed precision training we cast all non-trainable weigths to half-precision
+    # For mixed precision training we cast all non-trainable weights to half-precision
     # as these weights are only used for inference, keeping weights in full precision is not required.
 
     # Create model:
 
     main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
-    # keep the master weight to float32
-    if args.dit_model_name_or_path:
-        transformer = transformer = MochiTransformer3DModel.from_pretrained(
-            args.dit_model_name_or_path,
-            torch_dtype=torch.float32,
-            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
-        )
-    else:
-        transformer = MochiTransformer3DModel.from_pretrained(
-            args.pretrained_model_name_or_path,
-            subfolder="transformer",
-            torch_dtype=torch.float32,
-            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
-        )
+
+    transformer = load_transformer(
+        args.model_type,
+        args.dit_model_name_or_path,
+        args.pretrained_model_name_or_path,
+        torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16,
+    )
+
     teacher_transformer = deepcopy(transformer)
     if args.use_ema:
         ema_transformer = deepcopy(transformer)
@@ -335,6 +266,7 @@ def main(args):
         ema_transformer = None
 
     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,
@@ -345,12 +277,10 @@ def main(args):
         transformer.add_adapter(transformer_lora_config)
 
     main_print(
-        f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M"
-    )
-    main_print(
-        f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}"
-    )
-    fsdp_kwargs = get_dit_fsdp_kwargs(
+        f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M")
+    main_print(f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}")
+    fsdp_kwargs, no_split_modules = get_dit_fsdp_kwargs(
+        transformer,
         args.fsdp_sharding_startegy,
         args.use_lora,
         args.use_cpu_offload,
@@ -361,7 +291,7 @@ 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_modules
         fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)
 
     transformer = FSDP(
@@ -377,13 +307,13 @@ def main(args):
             ema_transformer,
             **fsdp_kwargs,
         )
-    main_print(f"--> model loaded")
+    main_print("--> model loaded")
 
     if args.gradient_checkpointing:
-        apply_fsdp_checkpointing(transformer, args.selective_checkpointing)
-        apply_fsdp_checkpointing(teacher_transformer, args.selective_checkpointing)
+        apply_fsdp_checkpointing(transformer, no_split_modules, args.selective_checkpointing)
+        apply_fsdp_checkpointing(teacher_transformer, no_split_modules, args.selective_checkpointing)
         if args.use_ema:
-            apply_fsdp_checkpointing(ema_transformer, args.selective_checkpointing)
+            apply_fsdp_checkpointing(ema_transformer, no_split_modules, args.selective_checkpointing)
     # Set model as trainable.
     transformer.train()
     teacher_transformer.requires_grad_(False)
@@ -391,9 +321,7 @@ def main(args):
         ema_transformer.requires_grad_(False)
     noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=args.shift)
     if args.scheduler_type == "pcm_linear_quadratic":
-        linear_steps = int(
-            noise_scheduler.config.num_train_timesteps * args.linear_range
-        )
+        linear_steps = int(noise_scheduler.config.num_train_timesteps * args.linear_range)
         sigmas = linear_quadratic_schedule(
             noise_scheduler.config.num_train_timesteps,
             args.linear_quadratic_threshold,
@@ -421,9 +349,8 @@ def main(args):
 
     init_steps = 0
     if args.resume_from_lora_checkpoint:
-        transformer, optimizer, init_steps = resume_lora_optimizer(
-            transformer, args.resume_from_lora_checkpoint, optimizer
-        )
+        transformer, optimizer, init_steps = resume_lora_optimizer(transformer, args.resume_from_lora_checkpoint,
+                                                                   optimizer)
     main_print(f"optimizer: {optimizer}")
 
     # todo add lr scheduler
@@ -440,20 +367,15 @@ def main(args):
     train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
     uncond_prompt_embed = train_dataset.uncond_prompt_embed
     uncond_prompt_mask = train_dataset.uncond_prompt_mask
-    sampler = (
-        LengthGroupedSampler(
-            args.train_batch_size,
-            rank=rank,
-            world_size=world_size,
-            lengths=train_dataset.lengths,
-            group_frame=args.group_frame,
-            group_resolution=args.group_resolution,
-        )
-        if (args.group_frame or args.group_resolution)
-        else DistributedSampler(
-            train_dataset, rank=rank, num_replicas=world_size, shuffle=False
-        )
-    )
+    sampler = (LengthGroupedSampler(
+        args.train_batch_size,
+        rank=rank,
+        world_size=world_size,
+        lengths=train_dataset.lengths,
+        group_frame=args.group_frame,
+        group_resolution=args.group_resolution,
+    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
+        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
 
     train_dataloader = DataLoader(
         train_dataset,
@@ -466,11 +388,7 @@ def main(args):
     )
 
     num_update_steps_per_epoch = math.ceil(
-        len(train_dataloader)
-        / args.gradient_accumulation_steps
-        * args.sp_size
-        / args.train_sp_batch_size
-    )
+        len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
     args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
     if rank <= 0:
@@ -478,22 +396,14 @@ def main(args):
         wandb.init(project=project, config=args)
 
     # Train!
-    total_batch_size = (
-        args.train_batch_size
-        * world_size
-        * args.gradient_accumulation_steps
-        / args.sp_size
-        * args.train_sp_batch_size
-    )
+    total_batch_size = (world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size)
     main_print("***** Running training *****")
     main_print(f"  Num examples = {len(train_dataset)}")
     main_print(f"  Dataloader size = {len(train_dataloader)}")
     main_print(f"  Num Epochs = {args.num_train_epochs}")
     main_print(f"  Resume training from step {init_steps}")
     main_print(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    main_print(
-        f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
-    )
+    main_print(f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
     main_print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
     main_print(f"  Total optimization steps = {args.max_train_steps}")
     main_print(
@@ -546,8 +456,9 @@ def main(args):
         assert args.multi_phased_distill_schedule is not None
         num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
 
-        loss, grad_norm, pred_norm, aux_loss = train_one_step_mochi(
+        loss, grad_norm, pred_norm = distill_one_step(
             transformer,
+            args.model_type,
             teacher_transformer,
             ema_transformer,
             optimizer,
@@ -568,20 +479,19 @@ def main(args):
             args.ema_decay,
             args.pred_decay_weight,
             args.pred_decay_type,
+            args.hunyuan_teacher_disable_cfg,
         )
 
         step_time = time.time() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 
-        progress_bar.set_postfix(
-            {
-                "loss": f"{loss:.4f}",
-                "step_time": f"{step_time:.2f}s",
-                "grad_norm": grad_norm,
-                "phases": num_phases,
-            }
-        )
+        progress_bar.set_postfix({
+            "loss": f"{loss:.4f}",
+            "step_time": f"{step_time:.2f}s",
+            "grad_norm": grad_norm,
+            "phases": num_phases,
+        })
         progress_bar.update(1)
         if rank <= 0:
             wandb.log(
@@ -591,20 +501,17 @@ def main(args):
                     "step_time": step_time,
                     "avg_step_time": avg_step_time,
                     "grad_norm": grad_norm,
-                    "pred_fro_norm": pred_norm["fro"],
+                    "pred_fro_norm": pred_norm["fro"],  # codespell:ignore
                     "pred_largest_singular_value": pred_norm["largest singular value"],
                     "pred_absolute_mean": pred_norm["absolute mean"],
                     "pred_absolute_max": pred_norm["absolute max"],
-                    "aux_loss": aux_loss,
                 },
                 step=step,
             )
         if step % args.checkpointing_steps == 0:
             if args.use_lora:
                 # Save LoRA weights
-                save_lora_checkpoint(
-                    transformer, optimizer, rank, args.output_dir, step
-                )
+                save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
             else:
                 # Your existing checkpoint saving code
                 if args.use_ema:
@@ -642,9 +549,7 @@ def main(args):
                 )
 
     if args.use_lora:
-        save_lora_checkpoint(
-            transformer, optimizer, rank, args.output_dir, args.max_train_steps
-        )
+        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
     else:
         save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
 
@@ -655,8 +560,12 @@ 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_height", type=int, default=480)
+    parser.add_argument("--num_width", type=int, default=848)
     parser.add_argument("--num_frames", type=int, default=163)
     parser.add_argument(
         "--dataloader_num_workers",
@@ -670,9 +579,7 @@ if __name__ == "__main__":
         default=16,
         help="Batch size (per device) for the training dataloader.",
     )
-    parser.add_argument(
-        "--num_latent_t", type=int, default=28, help="Number of latent timesteps."
-    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
     parser.add_argument("--group_frame", action="store_true")  # TODO
     parser.add_argument("--group_resolution", action="store_true")  # TODO
 
@@ -694,9 +601,7 @@ if __name__ == "__main__":
     parser.add_argument("--validation_steps", type=float, default=64)
     parser.add_argument("--log_validation", action="store_true")
     parser.add_argument("--tracker_project_name", type=str, default=None)
-    parser.add_argument(
-        "--seed", type=int, default=None, help="A seed for reproducible training."
-    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
     parser.add_argument(
         "--output_dir",
         type=str,
@@ -713,39 +618,31 @@ if __name__ == "__main__":
         "--checkpointing_steps",
         type=int,
         default=500,
-        help=(
-            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
-            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
-            " training using `--resume_from_checkpoint`."
-        ),
+        help=("Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+              " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+              " training using `--resume_from_checkpoint`."),
     )
     parser.add_argument("--shift", type=float, default=1.0)
     parser.add_argument(
         "--resume_from_checkpoint",
         type=str,
         default=None,
-        help=(
-            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
-            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
-        ),
+        help=("Whether training should be resumed from a previous checkpoint. Use a path saved by"
+              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
     parser.add_argument(
         "--resume_from_lora_checkpoint",
         type=str,
         default=None,
-        help=(
-            "Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
-            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
-        ),
+        help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
+              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
     parser.add_argument(
         "--logging_dir",
         type=str,
         default="logs",
-        help=(
-            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
-        ),
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
     )
 
     # optimizer & scheduler & Training
@@ -780,9 +677,7 @@ if __name__ == "__main__":
         default=10,
         help="Number of steps for the warmup in the lr scheduler.",
     )
-    parser.add_argument(
-        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
-    )
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
     parser.add_argument(
         "--gradient_checkpointing",
         action="store_true",
@@ -792,10 +687,8 @@ if __name__ == "__main__":
     parser.add_argument(
         "--allow_tf32",
         action="store_true",
-        help=(
-            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
-            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
-        ),
+        help=("Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+              " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"),
     )
     parser.add_argument(
         "--mixed_precision",
@@ -805,8 +698,7 @@ if __name__ == "__main__":
         help=(
             "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
             " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
-            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
-        ),
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."),
     )
     parser.add_argument(
         "--use_cpu_offload",
@@ -828,12 +720,8 @@ if __name__ == "__main__":
         default=False,
         help="Whether to use LoRA for finetuning.",
     )
-    parser.add_argument(
-        "--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA."
-    )
-    parser.add_argument(
-        "--lora_rank", type=int, default=128, help="LoRA rank parameter. "
-    )
+    parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.")
+    parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ")
     parser.add_argument("--fsdp_sharding_startegy", default="full")
 
     # lr_scheduler
@@ -841,10 +729,8 @@ if __name__ == "__main__":
         "--lr_scheduler",
         type=str,
         default="constant",
-        help=(
-            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
-            ' "constant", "constant_with_warmup"]'
-        ),
+        help=('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+              ' "constant", "constant_with_warmup"]'),
     )
     parser.add_argument("--num_euler_timesteps", type=int, default=100)
     parser.add_argument(
@@ -864,13 +750,9 @@ if __name__ == "__main__":
         action="store_true",
         help="Whether to apply the cfg_solver.",
     )
-    parser.add_argument(
-        "--distill_cfg", type=float, default=3.0, help="Distillation coefficient."
-    )
+    parser.add_argument("--distill_cfg", type=float, default=3.0, help="Distillation coefficient.")
     # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
-    parser.add_argument(
-        "--scheduler_type", type=str, default="pcm", help="The scheduler type to use."
-    )
+    parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.")
     parser.add_argument(
         "--linear_quadratic_threshold",
         type=float,
@@ -883,14 +765,12 @@ if __name__ == "__main__":
         default=0.5,
         help="Range for linear quadratic scheduler.",
     )
-    parser.add_argument(
-        "--weight_decay", type=float, default=0.001, help="Weight decay to apply."
-    )
+    parser.add_argument("--weight_decay", type=float, default=0.001, help="Weight decay to apply.")
     parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA.")
     parser.add_argument("--multi_phased_distill_schedule", type=str, default=None)
-    parser.add_argument("--finetune_weight", type=float, default=0.0)
     parser.add_argument("--pred_decay_weight", type=float, default=0.0)
     parser.add_argument("--pred_decay_type", default="l1")
+    parser.add_argument("--hunyuan_teacher_disable_cfg", action="store_true")
     parser.add_argument(
         "--master_weight_type",
         type=str,

fastvideo/distill/discriminator.py

@@ -1,45 +1,23 @@
-from typing import Any, Dict, Optional, Union
-
-import torch
 import torch.nn as nn
-
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
-from diffusers.models.attention import JointTransformerBlock
-from diffusers.models.attention_processor import Attention, AttentionProcessor
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.normalization import AdaLayerNormContinuous
-from diffusers.utils import (
-    USE_PEFT_BACKEND,
-    is_torch_version,
-    logging,
-    scale_lora_layers,
-    unscale_lora_layers,
-)
-from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
-from diffusers.models.transformers.transformer_2d import Transformer2DModelOutput
-from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
+from diffusers.utils import logging
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
 class DiscriminatorHead(nn.Module):
+
     def __init__(self, input_channel, output_channel=1):
         super().__init__()
         inner_channel = 1024
         self.conv1 = nn.Sequential(
             nn.Conv2d(input_channel, inner_channel, 1, 1, 0),
             nn.GroupNorm(32, inner_channel),
-            nn.LeakyReLU(
-                inplace=True
-            ),  # use LeakyReLu instead of GELU shown in the paper to save memory
+            nn.LeakyReLU(inplace=True),  # use LeakyReLu instead of GELU shown in the paper to save memory
         )
         self.conv2 = nn.Sequential(
             nn.Conv2d(inner_channel, inner_channel, 1, 1, 0),
             nn.GroupNorm(32, inner_channel),
-            nn.LeakyReLU(
-                inplace=True
-            ),  # use LeakyReLu instead of GELU shown in the paper to save memory
+            nn.LeakyReLU(inplace=True),  # use LeakyReLu instead of GELU shown in the paper to save memory
         )
 
         self.conv_out = nn.Conv2d(inner_channel, output_channel, 1, 1, 0)
@@ -57,41 +35,37 @@ 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,
     ):
         super().__init__()
-        adapter_channel_dims = adapter_channel_dims * (48 // stride)
+        adapter_channel_dims = adapter_channel_dims * (total_layers // stride)
         self.stride = stride
         self.num_h_per_head = num_h_per_head
         self.head_num = len(adapter_channel_dims)
-        self.heads = nn.ModuleList(
-            [
-                nn.ModuleList(
-                    [
-                        DiscriminatorHead(adapter_channel)
-                        for _ in range(self.num_h_per_head)
-                    ]
-                )
-                for adapter_channel in adapter_channel_dims
-            ]
-        )
+        self.heads = nn.ModuleList([
+            nn.ModuleList([DiscriminatorHead(adapter_channel) for _ in range(self.num_h_per_head)])
+            for adapter_channel in adapter_channel_dims
+        ])
 
     def forward(self, features):
         outputs = []
 
         def create_custom_forward(module):
+
             def custom_forward(*inputs):
                 return module(*inputs)
 
             return custom_forward
 
-        assert len(features) // self.stride == len(self.heads)
-        for i in range(0, len(features), self.stride):
-            for h in self.heads[i // self.stride]:
+        assert len(features) == len(self.heads)
+        for i in range(0, len(features)):
+            for h in self.heads[i]:
                 # out = torch.utils.checkpoint.checkpoint(
                 #     create_custom_forward(h),
                 #     features[i],

fastvideo/distill/solver.py

@@ -3,11 +3,10 @@ from typing import Optional, Tuple, Union
 
 import numpy as np
 import torch
-
 from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.utils import BaseOutput, logging
-from diffusers.utils.torch_utils import randn_tensor
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput, logging
+
 from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -21,7 +20,7 @@ class PCMFMSchedulerOutput(BaseOutput):
 def extract_into_tensor(a, t, x_shape):
     b, *_ = t.shape
     out = a.gather(-1, t)
-    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+    return out.reshape(b, *((1, ) * (len(x_shape) - 1)))
 
 
 class PCMFMScheduler(SchedulerMixin, ConfigMixin):
@@ -40,20 +39,15 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
     ):
         if linear_quadratic:
             linear_steps = int(num_train_timesteps * linear_range)
-            sigmas = linear_quadratic_schedule(
-                num_train_timesteps, linear_quadratic_threshold, linear_steps
-            )
+            sigmas = linear_quadratic_schedule(num_train_timesteps, linear_quadratic_threshold, linear_steps)
             sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
         else:
-            timesteps = np.linspace(
-                1, num_train_timesteps, num_train_timesteps, dtype=np.float32
-            )[::-1].copy()
+            timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
             timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
             sigmas = timesteps / num_train_timesteps
             sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
-        self.euler_timesteps = (
-            np.arange(1, pcm_timesteps + 1) * (num_train_timesteps // pcm_timesteps)
-        ).round().astype(np.int64) - 1
+        self.euler_timesteps = (np.arange(1, pcm_timesteps + 1) *
+                                (num_train_timesteps // pcm_timesteps)).round().astype(np.int64) - 1
         self.sigmas = sigmas.numpy()[::-1][self.euler_timesteps]
         self.sigmas = torch.from_numpy((self.sigmas[::-1].copy()))
         self.timesteps = self.sigmas * num_train_timesteps
@@ -118,9 +112,7 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
     def _sigma_to_t(self, sigma):
         return sigma * self.config.num_train_timesteps
 
-    def set_timesteps(
-        self, num_inference_steps: int, device: Union[str, torch.device] = None
-    ):
+    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
@@ -131,18 +123,14 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
                 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
         """
         self.num_inference_steps = num_inference_steps
-        inference_indices = np.linspace(
-            0, self.config.pcm_timesteps, num=num_inference_steps, endpoint=False
-        )
+        inference_indices = np.linspace(0, self.config.pcm_timesteps, num=num_inference_steps, endpoint=False)
         inference_indices = np.floor(inference_indices).astype(np.int64)
         inference_indices = torch.from_numpy(inference_indices).long()
 
         self.sigmas_ = self.sigmas[inference_indices]
         timesteps = self.sigmas_ * self.config.num_train_timesteps
         self.timesteps = timesteps.to(device=device)
-        self.sigmas_ = torch.cat(
-            [self.sigmas_, torch.zeros(1, device=self.sigmas_.device)]
-        )
+        self.sigmas_ = torch.cat([self.sigmas_, torch.zeros(1, device=self.sigmas_.device)])
         self._step_index = None
         self._begin_index = None
 
@@ -204,18 +192,11 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
                 returned, otherwise a tuple is returned where the first element is the sample tensor.
         """
 
-        if (
-            isinstance(timestep, int)
-            or isinstance(timestep, torch.IntTensor)
-            or isinstance(timestep, torch.LongTensor)
-        ):
-            raise ValueError(
-                (
-                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-                    " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-                    " one of the `scheduler.timesteps` as a timestep."
-                ),
-            )
+        if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor)
+                or isinstance(timestep, torch.LongTensor)):
+            raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                              " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                              " one of the `scheduler.timesteps` as a timestep."), )
 
         if self.step_index is None:
             self._init_step_index(timestep)
@@ -233,7 +214,7 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
         self._step_index += 1
 
         if not return_dict:
-            return (prev_sample,)
+            return (prev_sample, )
 
         return PCMFMSchedulerOutput(prev_sample=prev_sample)
 
@@ -242,16 +223,14 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
 
 
 class EulerSolver:
+
     def __init__(self, sigmas, timesteps=1000, euler_timesteps=50):
         self.step_ratio = timesteps // euler_timesteps
-        self.euler_timesteps = (
-            np.arange(1, euler_timesteps + 1) * self.step_ratio
-        ).round().astype(np.int64) - 1
+        self.euler_timesteps = (np.arange(1, euler_timesteps + 1) * self.step_ratio).round().astype(np.int64) - 1
         self.euler_timesteps_prev = np.asarray([0] + self.euler_timesteps[:-1].tolist())
         self.sigmas = sigmas[self.euler_timesteps]
-        self.sigmas_prev = np.asarray(
-            [sigmas[0]] + sigmas[self.euler_timesteps[:-1]].tolist()
-        )  # either use sigma0 or 0
+        self.sigmas_prev = np.asarray([sigmas[0]] +
+                                      sigmas[self.euler_timesteps[:-1]].tolist())  # either use sigma0 or 0
 
         self.euler_timesteps = torch.from_numpy(self.euler_timesteps).long()
         self.euler_timesteps_prev = torch.from_numpy(self.euler_timesteps_prev).long()
@@ -268,9 +247,7 @@ class EulerSolver:
 
     def euler_step(self, sample, model_pred, timestep_index):
         sigma = extract_into_tensor(self.sigmas, timestep_index, model_pred.shape)
-        sigma_prev = extract_into_tensor(
-            self.sigmas_prev, timestep_index, model_pred.shape
-        )
+        sigma_prev = extract_into_tensor(self.sigmas_prev, timestep_index, model_pred.shape)
         x_prev = sample + (sigma_prev - sigma) * model_pred
         return x_prev
 
@@ -282,16 +259,10 @@ class EulerSolver:
         multiphase,
         is_target=False,
     ):
-        inference_indices = np.linspace(
-            0, len(self.euler_timesteps), num=multiphase, endpoint=False
-        )
+        inference_indices = np.linspace(0, len(self.euler_timesteps), num=multiphase, endpoint=False)
         inference_indices = np.floor(inference_indices).astype(np.int64)
-        inference_indices = (
-            torch.from_numpy(inference_indices).long().to(self.euler_timesteps.device)
-        )
-        expanded_timestep_index = timestep_index.unsqueeze(1).expand(
-            -1, inference_indices.size(0)
-        )
+        inference_indices = (torch.from_numpy(inference_indices).long().to(self.euler_timesteps.device))
+        expanded_timestep_index = timestep_index.unsqueeze(1).expand(-1, inference_indices.size(0))
         valid_indices_mask = expanded_timestep_index >= inference_indices
         last_valid_index = valid_indices_mask.flip(dims=[1]).long().argmax(dim=1)
         last_valid_index = inference_indices.size(0) - 1 - last_valid_index
@@ -301,9 +272,7 @@ class EulerSolver:
             sigma = extract_into_tensor(self.sigmas_prev, timestep_index, sample.shape)
         else:
             sigma = extract_into_tensor(self.sigmas, timestep_index, sample.shape)
-        sigma_prev = extract_into_tensor(
-            self.sigmas_prev, timestep_index_end, sample.shape
-        )
+        sigma_prev = extract_into_tensor(self.sigmas_prev, timestep_index_end, sample.shape)
         x_prev = sample + (sigma_prev - sigma) * model_pred
 
         return x_prev, timestep_index_end

fastvideo/distill_adv.py

@@ -1,70 +1,43 @@
+# !/bin/python3
+# isort: skip_file
 import argparse
-from email.policy import strict
-import logging
 import math
 import os
-import shutil
-from pathlib import Path
-from fastvideo.utils.parallel_states import (
-    initialize_sequence_parallel_state,
-    destroy_sequence_parallel_group,
-    get_sequence_parallel_state,
-    nccl_info,
-)
-from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast
-from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_mochi_dit_input
-from fastvideo.utils.validation import log_validation
 import time
-from torch.utils.data import DataLoader
-import torch
-from torch.distributed.fsdp import (
-    FullyShardedDataParallel as FSDP,
-    StateDictType,
-    FullStateDictConfig,
-)
+from collections import deque
+from copy import deepcopy
 
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
-import json
-from torch.utils.data.distributed import DistributedSampler
-from fastvideo.utils.dataset_utils import LengthGroupedSampler
+import torch
+import torch.distributed as dist
 import wandb
 from accelerate.utils import set_seed
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.optimization import get_scheduler
+from diffusers.utils import check_min_version
+from peft import LoraConfig
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
 from tqdm.auto import tqdm
-from fastvideo.fsdp_util import (
-    get_dit_fsdp_kwargs,
-    apply_fsdp_checkpointing,
-    get_discriminator_fsdp_kwargs,
-)
-import diffusers
-from diffusers import (
-    FlowMatchEulerDiscreteScheduler,
-)
+
+from fastvideo.dataset.latent_datasets import (LatentDataset, latent_collate_function)
 from fastvideo.distill.discriminator import Discriminator
 from fastvideo.distill.solver import EulerSolver, extract_into_tensor
-from copy import deepcopy
-from diffusers.optimization import get_scheduler
-from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
-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 fastvideo.utils.checkpoint import (
-    save_checkpoint,
-    save_lora_checkpoint,
-    resume_lora_optimizer,
-    resume_training,
-    save_checkpoint_generator_discriminator,
-    resume_training_generator_discriminator,
-)
-from fastvideo.utils.logging import main_print
+from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input
+from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
+from fastvideo.utils.checkpoint import (resume_lora_optimizer, resume_training_generator_discriminator, save_checkpoint,
+                                        save_lora_checkpoint)
+from fastvideo.utils.communications import (broadcast, sp_parallel_dataloader_wrapper)
+from fastvideo.utils.dataset_utils import LengthGroupedSampler
+from fastvideo.utils.fsdp_util import (apply_fsdp_checkpointing, get_discriminator_fsdp_kwargs, get_dit_fsdp_kwargs)
+from fastvideo.utils.load import load_transformer
+from fastvideo.utils.logging_ import main_print
+from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group, get_sequence_parallel_state,
+                                             initialize_sequence_parallel_state)
+from fastvideo.utils.validation import log_validation
 
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
 check_min_version("0.31.0")
-import time
-from collections import deque
 
 
 def gan_d_loss(
@@ -76,6 +49,7 @@ def gan_d_loss(
     encoder_hidden_states,
     encoder_attention_mask,
     weight,
+    discriminator_head_stride,
 ):
     loss = 0.0
     # collate sample_fake and sample_real
@@ -85,7 +59,8 @@ def gan_d_loss(
             encoder_hidden_states,
             timestep,
             encoder_attention_mask,
-            output_attn=True,
+            output_features=True,
+            output_features_stride=discriminator_head_stride,
             return_dict=False,
         )[1]
         real_features = teacher_transformer(
@@ -93,17 +68,16 @@ def gan_d_loss(
             encoder_hidden_states,
             timestep,
             encoder_attention_mask,
-            output_attn=True,
+            output_features=True,
+            output_features_stride=discriminator_head_stride,
             return_dict=False,
         )[1]
 
     fake_outputs = discriminator(fake_features)
     real_outputs = discriminator(real_features)
     for fake_output, real_output in zip(fake_outputs, real_outputs):
-        loss += (
-            torch.mean(weight * torch.relu(fake_output.float() + 1))
-            + torch.mean(weight * torch.relu(1 - real_output.float()))
-        ) / (discriminator.head_num * discriminator.num_h_per_head)
+        loss += (torch.mean(weight * torch.relu(fake_output.float() + 1)) + torch.mean(
+            weight * torch.relu(1 - real_output.float()))) / (discriminator.head_num * discriminator.num_h_per_head)
     return loss
 
 
@@ -115,6 +89,7 @@ def gan_g_loss(
     encoder_hidden_states,
     encoder_attention_mask,
     weight,
+    discriminator_head_stride,
 ):
     loss = 0.0
     features = teacher_transformer(
@@ -122,33 +97,30 @@ def gan_g_loss(
         encoder_hidden_states,
         timestep,
         encoder_attention_mask,
-        output_attn=True,
+        output_features=True,
+        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
-        )
+        loss += torch.mean(
+            weight * torch.relu(1 - fake_output.float())) / (discriminator.head_num * discriminator.num_h_per_head)
     return loss
 
 
-def train_one_step_mochi(
+def distill_one_step_adv(
     transformer,
+    model_type,
     teacher_transformer,
     optimizer,
     discriminator,
     discriminator_optimizer,
-    global_step,
     lr_scheduler,
     loader,
     noise_scheduler,
     solver,
     noise_random_generator,
     sp_size,
-    precondition_outputs,
     max_grad_norm,
     uncond_prompt_embed,
     uncond_prompt_mask,
@@ -157,6 +129,7 @@ def train_one_step_mochi(
     not_apply_cfg_solver,
     distill_cfg,
     adv_weight,
+    discriminator_head_stride,
 ):
     optimizer.zero_grad()
     discriminator_optimizer.zero_grad()
@@ -167,12 +140,10 @@ def train_one_step_mochi(
         latents_attention_mask,
         encoder_attention_mask,
     ) = next(loader)
-    model_input = normalize_mochi_dit_input(latents)
+    model_input = normalize_dit_input(model_type, latents)
     noise = torch.randn_like(model_input)
     bsz = model_input.shape[0]
-    index = torch.randint(
-        0, num_euler_timesteps, (bsz,), device=model_input.device
-    ).long()
+    index = torch.randint(0, num_euler_timesteps, (bsz, ), device=model_input.device).long()
     if sp_size > 1:
         broadcast(index)
     # Add noise according to flow matching.
@@ -197,11 +168,8 @@ def train_one_step_mochi(
         )[0]
 
     # if accelerator.is_main_process:
-    model_pred, end_index = solver.euler_style_multiphase_pred(
-        noisy_model_input, model_pred, index, multiphase
-    )
+    model_pred, end_index = solver.euler_style_multiphase_pred(noisy_model_input, model_pred, index, multiphase)
 
-    weighting = 1.0
     # # simplified flow matching aka 0-rectified flow matching loss
     # # target = model_input - noise
     # target = model_input
@@ -210,14 +178,13 @@ def train_one_step_mochi(
         adv_index[i] = torch.randint(
             end_index[i].item(),
             end_index[i].item() + num_euler_timesteps // multiphase,
-            (1,),
+            (1, ),
             dtype=end_index.dtype,
             device=end_index.device,
         )
 
     sigmas_end = extract_into_tensor(solver.sigmas_prev, end_index, model_input.shape)
     sigmas_adv = extract_into_tensor(solver.sigmas_prev, adv_index, model_input.shape)
-    timesteps_end = (sigmas_end * noise_scheduler.config.num_train_timesteps).view(-1)
     timesteps_adv = (sigmas_adv * noise_scheduler.config.num_train_timesteps).view(-1)
 
     with torch.no_grad():
@@ -242,9 +209,7 @@ def train_one_step_mochi(
                     uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
                     return_dict=False,
                 )[0].float()
-        teacher_output = cond_teacher_output + w * (
-            cond_teacher_output - uncond_teacher_output
-        )
+        teacher_output = cond_teacher_output + w * (cond_teacher_output - uncond_teacher_output)
         x_prev = solver.euler_step(noisy_model_input, teacher_output, index)
 
     # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
@@ -258,22 +223,14 @@ def train_one_step_mochi(
                 return_dict=False,
             )[0]
 
-        target, end_index = solver.euler_style_multiphase_pred(
-            x_prev, target_pred, index, multiphase, True
-        )
+        target, end_index = solver.euler_style_multiphase_pred(x_prev, target_pred, index, multiphase, True)
 
-    real_adv = (
-        (1 - sigmas_adv) * target + (sigmas_adv - sigmas_end) * torch.randn_like(target)
-    ) / (1 - sigmas_end)
-    fake_adv = (
-        (1 - sigmas_adv) * model_pred
-        + (sigmas_adv - sigmas_end) * torch.randn_like(model_pred)
-    ) / (1 - sigmas_end)
+    real_adv = ((1 - sigmas_adv) * target + (sigmas_adv - sigmas_end) * torch.randn_like(target)) / (1 - sigmas_end)
+    fake_adv = ((1 - sigmas_adv) * model_pred +
+                (sigmas_adv - sigmas_end) * torch.randn_like(model_pred)) / (1 - sigmas_end)
 
     huber_c = 0.001
-    g_loss = torch.mean(
-        torch.sqrt((model_pred.float() - target.float()) ** 2 + huber_c**2) - huber_c
-    )
+    g_loss = torch.mean(torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) - huber_c)
     discriminator.requires_grad_(False)
     with torch.autocast("cuda", dtype=torch.bfloat16):
         g_gan_loss = adv_weight * gan_g_loss(
@@ -284,6 +241,7 @@ def train_one_step_mochi(
             encoder_hidden_states.float(),
             encoder_attention_mask,
             1.0,
+            discriminator_head_stride,
         )
     g_loss += g_gan_loss
     g_loss.backward()
@@ -308,6 +266,7 @@ def train_one_step_mochi(
             encoder_hidden_states,
             encoder_attention_mask,
             1.0,
+            discriminator_head_stride,
         )
 
     d_loss.backward()
@@ -340,28 +299,24 @@ def main(args):
     if rank <= 0 and args.output_dir is not None:
         os.makedirs(args.output_dir, exist_ok=True)
 
-    # For mixed precision training we cast all non-trainable weigths to half-precision
+    # For mixed precision training we cast all non-trainable weights to half-precision
     # as these weights are only used for inference, keeping weights in full precision is not required.
 
     # Create model:
 
     main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
     # keep the master weight to float32
-    if args.dit_model_name_or_path:
-        transformer = transformer = MochiTransformer3DModel.from_pretrained(
-            args.dit_model_name_or_path,
-            torch_dtype=torch.float32,
-            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
-        )
-    else:
-        transformer = MochiTransformer3DModel.from_pretrained(
-            args.pretrained_model_name_or_path,
-            subfolder="transformer",
-            torch_dtype=torch.float32,
-            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
-        )
+    transformer = load_transformer(
+        args.model_type,
+        args.dit_model_name_or_path,
+        args.pretrained_model_name_or_path,
+        torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16,
+    )
     teacher_transformer = deepcopy(transformer)
-    discriminator = Discriminator(args.discriminator_head_stride)
+    discriminator = Discriminator(
+        args.discriminator_head_stride,
+        total_layers=48 if args.model_type == "mochi" else 40,
+    )
 
     if args.use_lora:
         transformer.requires_grad_(False)
@@ -380,18 +335,21 @@ def main(args):
     main_print(
         f"  Total discriminator parameters = {sum(p.numel() for p in discriminator.parameters() if p.requires_grad) / 1e6} M"
     )
-    main_print(
-        f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}"
-    )
-    fsdp_kwargs = get_dit_fsdp_kwargs(
-        args.fsdp_sharding_startegy, args.use_lora, args.use_cpu_offload
+    main_print(f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}")
+    fsdp_kwargs, no_split_modules = get_dit_fsdp_kwargs(
+        transformer,
+        args.fsdp_sharding_startegy,
+        args.use_lora,
+        args.use_cpu_offload,
+        args.master_weight_type,
     )
     discriminator_fsdp_kwargs = get_discriminator_fsdp_kwargs(args.master_weight_type)
     if args.use_lora:
+        assert args.model_type == "mochi", "LoRA is only supported for Mochi model."
         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_modules
         fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)
 
     transformer = FSDP(
@@ -406,19 +364,17 @@ def main(args):
         discriminator,
         **discriminator_fsdp_kwargs,
     )
-    main_print(f"--> model loaded")
+    main_print("--> model loaded")
 
     if args.gradient_checkpointing:
-        apply_fsdp_checkpointing(transformer, args.selective_checkpointing)
-        apply_fsdp_checkpointing(teacher_transformer, args.selective_checkpointing)
+        apply_fsdp_checkpointing(transformer, no_split_modules, args.selective_checkpointing)
+        apply_fsdp_checkpointing(teacher_transformer, no_split_modules, args.selective_checkpointing)
     # Set model as trainable.
     transformer.train()
     teacher_transformer.requires_grad_(False)
     noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=args.shift)
     if args.scheduler_type == "pcm_linear_quadratic":
-        sigmas = linear_quadratic_schedule(
-            noise_scheduler.config.num_train_timesteps, args.linear_quadratic_threshold
-        )
+        sigmas = linear_quadratic_schedule(noise_scheduler.config.num_train_timesteps, args.linear_quadratic_threshold)
         sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
     else:
         sigmas = noise_scheduler.sigmas
@@ -435,7 +391,7 @@ def main(args):
         params_to_optimize,
         lr=args.learning_rate,
         betas=(0.9, 0.999),
-        weight_decay=1e-3,
+        weight_decay=args.weight_decay,
         eps=1e-8,
     )
 
@@ -443,15 +399,14 @@ def main(args):
         discriminator.parameters(),
         lr=args.discriminator_learning_rate,
         betas=(0, 0.999),
-        weight_decay=1e-3,
+        weight_decay=args.weight_decay,
         eps=1e-8,
     )
 
     init_steps = 0
     if args.resume_from_lora_checkpoint:
-        transformer, optimizer, init_steps = resume_lora_optimizer(
-            transformer, args.resume_from_lora_checkpoint, optimizer
-        )
+        transformer, optimizer, init_steps = resume_lora_optimizer(transformer, args.resume_from_lora_checkpoint,
+                                                                   optimizer)
     elif args.resume_from_checkpoint:
         (
             transformer,
@@ -483,20 +438,15 @@ def main(args):
     train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
     uncond_prompt_embed = train_dataset.uncond_prompt_embed
     uncond_prompt_mask = train_dataset.uncond_prompt_mask
-    sampler = (
-        LengthGroupedSampler(
-            args.train_batch_size,
-            rank=rank,
-            world_size=world_size,
-            lengths=train_dataset.lengths,
-            group_frame=args.group_frame,
-            group_resolution=args.group_resolution,
-        )
-        if (args.group_frame or args.group_resolution)
-        else DistributedSampler(
-            train_dataset, rank=rank, num_replicas=world_size, shuffle=False
-        )
-    )
+    sampler = (LengthGroupedSampler(
+        args.train_batch_size,
+        rank=rank,
+        world_size=world_size,
+        lengths=train_dataset.lengths,
+        group_frame=args.group_frame,
+        group_resolution=args.group_resolution,
+    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
+        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
 
     train_dataloader = DataLoader(
         train_dataset,
@@ -509,11 +459,7 @@ def main(args):
     )
     assert args.gradient_accumulation_steps == 1
     num_update_steps_per_epoch = math.ceil(
-        len(train_dataloader)
-        / args.gradient_accumulation_steps
-        * args.sp_size
-        / args.train_sp_batch_size
-    )
+        len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
     args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
     if rank <= 0:
@@ -521,22 +467,14 @@ def main(args):
         wandb.init(project=project, config=args)
 
     # Train!
-    total_batch_size = (
-        args.train_batch_size
-        * world_size
-        * args.gradient_accumulation_steps
-        / args.sp_size
-        * args.train_sp_batch_size
-    )
+    total_batch_size = (world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size)
     main_print("***** Running training *****")
     main_print(f"  Num examples = {len(train_dataset)}")
     main_print(f"  Dataloader size = {len(train_dataloader)}")
     main_print(f"  Num Epochs = {args.num_train_epochs}")
     main_print(f"  Resume training from step {init_steps}")
     main_print(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    main_print(
-        f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
-    )
+    main_print(f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
     main_print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
     main_print(f"  Total optimization steps = {args.max_train_steps}")
     main_print(
@@ -561,55 +499,65 @@ def main(args):
     )
 
     step_times = deque(maxlen=100)
+
     # log_validation(args, transformer, device,
-    #         torch.bfloat16, init_steps, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, ema=False)
+    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
+    def get_num_phases(multi_phased_distill_schedule, step):
+        # step-phase,step-phase
+        multi_phases = multi_phased_distill_schedule.split(",")
+        phase = multi_phases[-1].split("-")[-1]
+        for step_phases in multi_phases:
+            phase_step, phase = step_phases.split("-")
+            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):
+        assert args.multi_phased_distill_schedule is not None
+        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
         start_time = time.time()
         (
             generator_loss,
             generator_grad_norm,
             discriminator_loss,
             discriminator_grad_norm,
-        ) = train_one_step_mochi(
+        ) = distill_one_step_adv(
             transformer,
+            args.model_type,
             teacher_transformer,
             optimizer,
             discriminator,
             discriminator_optimizer,
-            step,
             lr_scheduler,
             loader,
             noise_scheduler,
             solver,
             noise_random_generator,
             args.sp_size,
-            args.precondition_outputs,
             args.max_grad_norm,
             uncond_prompt_embed,
             uncond_prompt_mask,
             args.num_euler_timesteps,
-            args.validation_sampling_steps,
+            num_phases,
             args.not_apply_cfg_solver,
             args.distill_cfg,
             args.adv_weight,
+            args.discriminator_head_stride,
         )
 
         step_time = time.time() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 
-        progress_bar.set_postfix(
-            {
-                "g_loss": f"{generator_loss:.4f}",
-                "d_loss": f"{discriminator_loss:.4f}",
-                "g_grad_norm": generator_grad_norm,
-                "d_grad_norm": discriminator_grad_norm,
-                "step_time": f"{step_time:.2f}s",
-            }
-        )
+        progress_bar.set_postfix({
+            "g_loss": f"{generator_loss:.4f}",
+            "d_loss": f"{discriminator_loss:.4f}",
+            "g_grad_norm": generator_grad_norm,
+            "d_grad_norm": discriminator_grad_norm,
+            "step_time": f"{step_time:.2f}s",
+        })
         progress_bar.update(1)
         if rank <= 0:
             wandb.log(
@@ -628,20 +576,20 @@ def main(args):
             main_print(f"--> saving checkpoint at step {step}")
             if args.use_lora:
                 # Save LoRA weights
-                save_lora_checkpoint(
-                    transformer, optimizer, rank, args.output_dir, step
-                )
+                save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
             else:
                 # Your existing checkpoint saving code
-                save_checkpoint_generator_discriminator(
-                    transformer,
-                    optimizer,
-                    discriminator,
-                    discriminator_optimizer,
-                    rank,
-                    args.output_dir,
-                    step,
-                )
+                # TODO
+                # save_checkpoint_generator_discriminator(
+                #     transformer,
+                #     optimizer,
+                #     discriminator,
+                #     discriminator_optimizer,
+                #     rank,
+                #     args.output_dir,
+                #     step,
+                # )
+                save_checkpoint(transformer, rank, args.output_dir, step)
             main_print(f"--> checkpoint saved at step {step}")
             dist.barrier()
         if args.log_validation and step % args.validation_steps == 0:
@@ -655,25 +603,14 @@ def main(args):
                 shift=args.shift,
                 num_euler_timesteps=args.num_euler_timesteps,
                 linear_quadratic_threshold=args.linear_quadratic_threshold,
+                linear_range=args.linear_range,
                 ema=False,
             )
 
     if args.use_lora:
-        save_lora_checkpoint(
-            transformer, optimizer, rank, args.output_dir, args.max_train_steps
-        )
+        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
     else:
-        save_checkpoint(
-            transformer, optimizer, rank, args.output_dir, args.max_train_steps
-        )
-        save_checkpoint(
-            discriminator,
-            discriminator_optimizer,
-            rank,
-            args.output_dir,
-            step,
-            discriminator=True,
-        )
+        save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
 
     if get_sequence_parallel_state():
         destroy_sequence_parallel_group()
@@ -682,8 +619,11 @@ 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_height", type=int, default=480)
+    parser.add_argument("--num_width", type=int, default=848)
     parser.add_argument("--num_frames", type=int, default=163)
     parser.add_argument(
         "--dataloader_num_workers",
@@ -697,9 +637,7 @@ if __name__ == "__main__":
         default=16,
         help="Batch size (per device) for the training dataloader.",
     )
-    parser.add_argument(
-        "--num_latent_t", type=int, default=28, help="Number of latent timesteps."
-    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
     parser.add_argument("--group_frame", action="store_true")  # TODO
     parser.add_argument("--group_resolution", action="store_true")  # TODO
 
@@ -712,22 +650,13 @@ if __name__ == "__main__":
     parser.add_argument("--ema_decay", type=float, default=0.999)
     parser.add_argument("--ema_start_step", type=int, default=0)
     parser.add_argument("--cfg", type=float, default=0.1)
-    parser.add_argument(
-        "--precondition_outputs",
-        action="store_true",
-        help="Whether to precondition the outputs of the model.",
-    )
-
     # validation & logs
-    parser.add_argument("--validation_prompt_dir", type=str)
-    parser.add_argument("--validation_sampling_steps", type=int, default=64)
-    parser.add_argument("--validation_guidance_scale", type=float, default=4.5)
+    parser.add_argument("--validation_sampling_steps", type=str, default="64")
+    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
     parser.add_argument("--validation_steps", type=float, default=64)
     parser.add_argument("--log_validation", action="store_true")
     parser.add_argument("--tracker_project_name", type=str, default=None)
-    parser.add_argument(
-        "--seed", type=int, default=None, help="A seed for reproducible training."
-    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
     parser.add_argument(
         "--output_dir",
         type=str,
@@ -744,39 +673,32 @@ if __name__ == "__main__":
         "--checkpointing_steps",
         type=int,
         default=500,
-        help=(
-            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
-            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
-            " training using `--resume_from_checkpoint`."
-        ),
+        help=("Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+              " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+              " training using `--resume_from_checkpoint`."),
     )
+    parser.add_argument("--validation_prompt_dir", type=str)
     parser.add_argument("--shift", type=float, default=1.0)
     parser.add_argument(
         "--resume_from_checkpoint",
         type=str,
         default=None,
-        help=(
-            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
-            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
-        ),
+        help=("Whether training should be resumed from a previous checkpoint. Use a path saved by"
+              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
     parser.add_argument(
         "--resume_from_lora_checkpoint",
         type=str,
         default=None,
-        help=(
-            "Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
-            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
-        ),
+        help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
+              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
     parser.add_argument(
         "--logging_dir",
         type=str,
         default="logs",
-        help=(
-            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
-        ),
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
     )
 
     # optimizer & scheduler & Training
@@ -811,9 +733,7 @@ if __name__ == "__main__":
         default=10,
         help="Number of steps for the warmup in the lr scheduler.",
     )
-    parser.add_argument(
-        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
-    )
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
     parser.add_argument(
         "--gradient_checkpointing",
         action="store_true",
@@ -823,10 +743,8 @@ if __name__ == "__main__":
     parser.add_argument(
         "--allow_tf32",
         action="store_true",
-        help=(
-            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
-            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
-        ),
+        help=("Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+              " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"),
     )
     parser.add_argument(
         "--mixed_precision",
@@ -836,8 +754,7 @@ if __name__ == "__main__":
         help=(
             "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
             " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
-            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
-        ),
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."),
     )
     parser.add_argument(
         "--use_cpu_offload",
@@ -859,13 +776,10 @@ if __name__ == "__main__":
         default=False,
         help="Whether to use LoRA for finetuning.",
     )
-    parser.add_argument(
-        "--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA."
-    )
-    parser.add_argument(
-        "--lora_rank", type=int, default=128, help="LoRA rank parameter. "
-    )
+    parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.")
+    parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ")
     parser.add_argument("--fsdp_sharding_startegy", default="full")
+    parser.add_argument("--multi_phased_distill_schedule", type=str, default=None)
     parser.add_argument(
         "--gradient_accumulation_steps",
         type=int,
@@ -878,10 +792,8 @@ if __name__ == "__main__":
         "--lr_scheduler",
         type=str,
         default="constant",
-        help=(
-            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
-            ' "constant", "constant_with_warmup"]'
-        ),
+        help=('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+              ' "constant", "constant_with_warmup"]'),
     )
     parser.add_argument("--num_euler_timesteps", type=int, default=100)
     parser.add_argument(
@@ -901,13 +813,9 @@ if __name__ == "__main__":
         action="store_true",
         help="Whether to apply the cfg_solver.",
     )
-    parser.add_argument(
-        "--distill_cfg", type=float, default=3.0, help="Distillation coefficient."
-    )
+    parser.add_argument("--distill_cfg", type=float, default=3.0, help="Distillation coefficient.")
     # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
-    parser.add_argument(
-        "--scheduler_type", type=str, default="pcm", help="The scheduler type to use."
-    )
+    parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.")
     parser.add_argument(
         "--adv_weight",
         type=float,
@@ -920,6 +828,13 @@ if __name__ == "__main__":
         default=2,
         help="The stride of the discriminator head.",
     )
+    parser.add_argument(
+        "--linear_range",
+        type=float,
+        default=0.5,
+        help="Range for linear quadratic scheduler.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.001, help="Weight decay to apply.")
     parser.add_argument(
         "--linear_quadratic_threshold",
         type=float,

fastvideo/models/flash_attn_no_pad.py

@@ -0,0 +1,28 @@
+from einops import rearrange
+from flash_attn import flash_attn_varlen_qkvpacked_func
+from flash_attn.bert_padding import pad_input, unpad_input
+
+
+def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None):
+    # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27
+    batch_size = qkv.shape[0]
+    seqlen = qkv.shape[1]
+    nheads = qkv.shape[-2]
+    x = rearrange(qkv, "b s three h d -> b s (three h d)")
+    x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask)
+
+    x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads)
+    output_unpad = flash_attn_varlen_qkvpacked_func(
+        x_unpad,
+        cu_seqlens,
+        max_s,
+        dropout_p,
+        softmax_scale=softmax_scale,
+        causal=causal,
+    )
+    output = rearrange(
+        pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen),
+        "b s (h d) -> b s h d",
+        h=nheads,
+    )
+    return output

fastvideo/models/hunyuan/constants.py

@@ -0,0 +1,89 @@
+import os
+
+import torch
+
+__all__ = [
+    "C_SCALE",
+    "PROMPT_TEMPLATE",
+    "MODEL_BASE",
+    "PRECISIONS",
+    "NORMALIZATION_TYPE",
+    "ACTIVATION_TYPE",
+    "VAE_PATH",
+    "TEXT_ENCODER_PATH",
+    "TOKENIZER_PATH",
+    "TEXT_PROJECTION",
+    "DATA_TYPE",
+    "NEGATIVE_PROMPT",
+]
+
+PRECISION_TO_TYPE = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+# =================== Constant Values =====================
+# Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid
+# overflow error when tensorboard logging values.
+C_SCALE = 1_000_000_000_000_000
+
+# When using decoder-only models, we must provide a prompt template to instruct the text encoder
+# on how to generate the text.
+# --------------------------------------------------------------------
+PROMPT_TEMPLATE_ENCODE = (
+    "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, "
+    "quantity, text, spatial relationships of the objects and background:<|eot_id|>"
+    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>")
+PROMPT_TEMPLATE_ENCODE_VIDEO = (
+    "<|start_header_id|>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 objects."
+    "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
+    "4. background environment, light, style and atmosphere."
+    "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
+    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>")
+
+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-video": {
+        "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
+        "crop_start": 95,
+    },
+}
+
+# ======================= Model ======================
+PRECISIONS = {"fp32", "fp16", "bf16"}
+NORMALIZATION_TYPE = {"layer", "rms"}
+ACTIVATION_TYPE = {"relu", "silu", "gelu", "gelu_tanh"}
+
+# =================== Model Path =====================
+MODEL_BASE = os.getenv("MODEL_BASE", "./data/hunyuan")
+
+# =================== Data =======================
+DATA_TYPE = {"image", "video", "image_video"}
+
+# 3D VAE
+VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"}
+
+# Text Encoder
+TEXT_ENCODER_PATH = {
+    "clipL": f"{MODEL_BASE}/text_encoder_2",
+    "llm": f"{MODEL_BASE}/text_encoder",
+}
+
+# Tokenizer
+TOKENIZER_PATH = {
+    "clipL": f"{MODEL_BASE}/text_encoder_2",
+    "llm": f"{MODEL_BASE}/text_encoder",
+}
+
+TEXT_PROJECTION = {
+    "linear",  # Default, an nn.Linear() layer
+    "single_refiner",  # Single TokenRefiner. Refer to LI-DiT
+}

fastvideo/models/hunyuan/diffusion/__init__.py

@@ -0,0 +1,3 @@
+# ruff: noqa: F401
+from .pipelines import HunyuanVideoPipeline
+from .schedulers import FlowMatchDiscreteScheduler

fastvideo/models/hunyuan/diffusion/pipelines/__init__.py

@@ -0,0 +1,2 @@
+# ruff: noqa: F401
+from .pipeline_hunyuan_video import HunyuanVideoPipeline

fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py

@@ -0,0 +1,931 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.configuration_utils import FrozenDict
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, deprecate, logging, replace_example_docstring,
+                             scale_lora_layers)
+from diffusers.utils.torch_utils import randn_tensor
+from einops import rearrange
+
+from fastvideo.utils.communications import all_gather
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+from ...constants import PRECISION_TO_TYPE
+from ...modules import HYVideoDiffusionTransformer
+from ...text_encoder import TextEncoder
+from ...vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """"""
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = (guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg)
+    return noise_cfg
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler.")
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler.")
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+@dataclass
+class HunyuanVideoPipelineOutput(BaseOutput):
+    videos: Union[torch.Tensor, np.ndarray]
+
+
+class HunyuanVideoPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-video generation using HunyuanVideo.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
+        text_encoder ([`TextEncoder`]):
+            Frozen text-encoder.
+        text_encoder_2 ([`TextEncoder`]):
+            Frozen text-encoder_2.
+        transformer ([`HYVideoDiffusionTransformer`]):
+            A `HYVideoDiffusionTransformer` to denoise the encoded video latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+    _optional_components = ["text_encoder_2"]
+    _exclude_from_cpu_offload = ["transformer"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: TextEncoder,
+        transformer: HYVideoDiffusionTransformer,
+        scheduler: KarrasDiffusionSchedulers,
+        text_encoder_2: Optional[TextEncoder] = None,
+        progress_bar_config: Dict[str, Any] = None,
+        args=None,
+    ):
+        super().__init__()
+
+        # ==========================================================================================
+        if progress_bar_config is None:
+            progress_bar_config = {}
+        if not hasattr(self, "_progress_bar_config"):
+            self._progress_bar_config = {}
+        self._progress_bar_config.update(progress_bar_config)
+
+        self.args = args
+        # ==========================================================================================
+
+        if (hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1):
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file")
+            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if (hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True):
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file")
+            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            transformer=transformer,
+            scheduler=scheduler,
+            text_encoder_2=text_encoder_2,
+        )
+        self.vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    def encode_prompt(
+        self,
+        prompt,
+        device,
+        num_videos_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_attention_mask: Optional[torch.Tensor] = None,
+        lora_scale: Optional[float] = None,
+        clip_skip: Optional[int] = None,
+        text_encoder: Optional[TextEncoder] = None,
+        data_type: Optional[str] = "image",
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_videos_per_prompt (`int`):
+                number of videos that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the video generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            attention_mask (`torch.Tensor`, *optional*):
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            negative_attention_mask (`torch.Tensor`, *optional*):
+            lora_scale (`float`, *optional*):
+                A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            text_encoder (TextEncoder, *optional*):
+            data_type (`str`, *optional*):
+        """
+        if text_encoder is None:
+            text_encoder = self.text_encoder
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if not USE_PEFT_BACKEND:
+                adjust_lora_scale_text_encoder(text_encoder.model, lora_scale)
+            else:
+                scale_lora_layers(text_encoder.model, lora_scale)
+
+        if prompt_embeds is None:
+            # textual inversion: process multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, text_encoder.tokenizer)
+
+            text_inputs = text_encoder.text2tokens(prompt, data_type=data_type)
+            if clip_skip is None:
+                prompt_outputs = text_encoder.encode(text_inputs, data_type=data_type, device=device)
+                prompt_embeds = prompt_outputs.hidden_state
+            else:
+                prompt_outputs = text_encoder.encode(
+                    text_inputs,
+                    output_hidden_states=True,
+                    data_type=data_type,
+                    device=device,
+                )
+                # Access the `hidden_states` first, that contains a tuple of
+                # all the hidden states from the encoder layers. Then index into
+                # the tuple to access the hidden states from the desired layer.
+                prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)]
+                # We also need to apply the final LayerNorm here to not mess with the
+                # representations. The `last_hidden_states` that we typically use for
+                # obtaining the final prompt representations passes through the LayerNorm
+                # layer.
+                prompt_embeds = text_encoder.model.text_model.final_layer_norm(prompt_embeds)
+
+            attention_mask = prompt_outputs.attention_mask
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(device)
+                bs_embed, seq_len = attention_mask.shape
+                attention_mask = attention_mask.repeat(1, num_videos_per_prompt)
+                attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len)
+
+        if text_encoder is not None:
+            prompt_embeds_dtype = text_encoder.dtype
+        elif self.transformer is not None:
+            prompt_embeds_dtype = self.transformer.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        if prompt_embeds.ndim == 2:
+            bs_embed, _ = prompt_embeds.shape
+            # duplicate text embeddings for each generation per prompt, using mps friendly method
+            prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt)
+            prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, -1)
+        else:
+            bs_embed, seq_len, _ = prompt_embeds.shape
+            # duplicate text embeddings for each generation per prompt, using mps friendly method
+            prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+            prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1)
+
+        return (
+            prompt_embeds,
+            negative_prompt_embeds,
+            attention_mask,
+            negative_attention_mask,
+        )
+
+    def decode_latents(self, latents, enable_tiling=True):
+        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
+        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
+
+        latents = 1 / self.vae.config.scaling_factor * latents
+        if enable_tiling:
+            self.vae.enable_tiling()
+        image = self.vae.decode(latents, return_dict=False)[0]
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        if image.ndim == 4:
+            image = image.cpu().permute(0, 2, 3, 1).float()
+        else:
+            image = image.cpu().float()
+        return image
+
+    def prepare_extra_func_kwargs(self, func, kwargs):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+        extra_step_kwargs = {}
+
+        for k, v in kwargs.items():
+            accepts = k in set(inspect.signature(func).parameters.keys())
+            if accepts:
+                extra_step_kwargs[k] = v
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        video_length,
+        callback_steps,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        vae_ver="88-4c-sd",
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if video_length is not None:
+            if "884" in vae_ver:
+                if video_length != 1 and (video_length - 1) % 4 != 0:
+                    raise ValueError(f"`video_length` has to be 1 or a multiple of 4 but is {video_length}.")
+            elif "888" in vae_ver:
+                if video_length != 1 and (video_length - 1) % 8 != 0:
+                    raise ValueError(f"`video_length` has to be 1 or a multiple of 8 but is {video_length}.")
+
+        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
+            raise ValueError(f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                             f" {type(callback_steps)}.")
+        if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs
+                                                                      for k in callback_on_step_end_tensor_inputs):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two.")
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.")
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                             f" {negative_prompt_embeds}. Please make sure to only forward one of the two.")
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}.")
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        video_length,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            video_length,
+            int(height) // self.vae_scale_factor,
+            int(width) // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler
+        if hasattr(self.scheduler, "init_noise_sigma"):
+            # scale the initial noise by the standard deviation required by the scheduler
+            latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
+    def get_guidance_scale_embedding(
+        self,
+        w: torch.Tensor,
+        embedding_dim: int = 512,
+        dtype: torch.dtype = torch.float32,
+    ) -> torch.Tensor:
+        """
+        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+
+        Args:
+            w (`torch.Tensor`):
+                Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
+            embedding_dim (`int`, *optional*, defaults to 512):
+                Dimension of the embeddings to generate.
+            dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
+                Data type of the generated embeddings.
+
+        Returns:
+            `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
+        """
+        assert len(w.shape) == 1
+        w = w * 1000.0
+
+        half_dim = embedding_dim // 2
+        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+        emb = w.to(dtype)[:, None] * emb[None, :]
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+        if embedding_dim % 2 == 1:  # zero pad
+            emb = torch.nn.functional.pad(emb, (0, 1))
+        assert emb.shape == (w.shape[0], embedding_dim)
+        return emb
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def guidance_rescale(self):
+        return self._guidance_rescale
+
+    @property
+    def clip_skip(self):
+        return self._clip_skip
+
+    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+    # corresponds to doing no classifier free guidance.
+    @property
+    def do_classifier_free_guidance(self):
+        # return self._guidance_scale > 1 and self.transformer.config.time_cond_proj_dim is None
+        return self._guidance_scale > 1
+
+    @property
+    def cross_attention_kwargs(self):
+        return self._cross_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        height: int,
+        width: int,
+        video_length: int,
+        data_type: str = "video",
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        sigmas: List[float] = None,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_videos_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_attention_mask: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback,
+                                             MultiPipelineCallbacks, ]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        vae_ver: str = "88-4c-sd",
+        enable_tiling: bool = False,
+        enable_vae_sp: bool = False,
+        n_tokens: Optional[int] = None,
+        embedded_guidance_scale: Optional[float] = None,
+        mask_strategy: Optional[Dict[str, list]] = None,
+        **kwargs,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`):
+                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+            height (`int`):
+                The height in pixels of the generated image.
+            width (`int`):
+                The width in pixels of the generated image.
+            video_length (`int`):
+                The number of frames in the generated video.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
+                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
+
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a
+                plain tuple.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.0):
+                Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
+                using zero terminal SNR.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
+                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
+                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
+                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+
+        Examples:
+
+        Returns:
+            [`~HunyuanVideoPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is a list with the generated images and the
+                second element is a list of `bool`s indicating whether the corresponding generated image contains
+                "not-safe-for-work" (nsfw) content.
+        """
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        # 0. Default height and width to unet
+        # height = height or self.transformer.config.sample_size * self.vae_scale_factor
+        # width = width or self.transformer.config.sample_size * self.vae_scale_factor
+        # to deal with lora scaling and other possible forward hooks
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            video_length,
+            callback_steps,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            callback_on_step_end_tensor_inputs,
+            vae_ver=vae_ver,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = (torch.device(f"cuda:{dist.get_rank()}") if dist.is_initialized() else self._execution_device)
+
+        # 3. Encode input prompt
+        lora_scale = (self.cross_attention_kwargs.get("scale", None)
+                      if self.cross_attention_kwargs is not None else None)
+
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_mask,
+            negative_prompt_mask,
+        ) = self.encode_prompt(
+            prompt,
+            device,
+            num_videos_per_prompt,
+            self.do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            attention_mask=attention_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_attention_mask=negative_attention_mask,
+            lora_scale=lora_scale,
+            clip_skip=self.clip_skip,
+            data_type=data_type,
+        )
+        if self.text_encoder_2 is not None:
+            (
+                prompt_embeds_2,
+                negative_prompt_embeds_2,
+                prompt_mask_2,
+                negative_prompt_mask_2,
+            ) = self.encode_prompt(
+                prompt,
+                device,
+                num_videos_per_prompt,
+                self.do_classifier_free_guidance,
+                negative_prompt,
+                prompt_embeds=None,
+                attention_mask=None,
+                negative_prompt_embeds=None,
+                negative_attention_mask=None,
+                lora_scale=lora_scale,
+                clip_skip=self.clip_skip,
+                text_encoder=self.text_encoder_2,
+                data_type=data_type,
+            )
+        else:
+            prompt_embeds_2 = None
+            negative_prompt_embeds_2 = None
+            prompt_mask_2 = None
+            negative_prompt_mask_2 = None
+
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+            if prompt_mask is not None:
+                prompt_mask = torch.cat([negative_prompt_mask, prompt_mask])
+            if prompt_embeds_2 is not None:
+                prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2])
+            if prompt_mask_2 is not None:
+                prompt_mask_2 = torch.cat([negative_prompt_mask_2, prompt_mask_2])
+
+        # 4. Prepare timesteps
+        extra_set_timesteps_kwargs = self.prepare_extra_func_kwargs(self.scheduler.set_timesteps,
+                                                                    {"n_tokens": n_tokens})
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            **extra_set_timesteps_kwargs,
+        )
+        if "884" in vae_ver:
+            video_length = (video_length - 1) // 4 + 1
+        elif "888" in vae_ver:
+            video_length = (video_length - 1) // 8 + 1
+        else:
+            video_length = video_length
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            video_length,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
+        if get_sequence_parallel_state():
+            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
+            latents = latents[:, :, rank, :, :, :]
+
+        # 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
+            },
+        )
+
+        target_dtype = PRECISION_TO_TYPE[self.args.precision]
+        autocast_enabled = (target_dtype != torch.float32) and not self.args.disable_autocast
+        vae_dtype = PRECISION_TO_TYPE[self.args.vae_precision]
+        vae_autocast_enabled = (vae_dtype != torch.float32) and not self.args.disable_autocast
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        def dict_to_3d_list(mask_strategy, t_max=50, l_max=60, h_max=24):
+            result = [[[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max)]
+            if mask_strategy is None:
+                return result
+            for key, value in mask_strategy.items():
+                t, l, h = map(int, key.split('_'))
+                result[t][l][h] = value
+            return result
+
+        mask_strategy = dict_to_3d_list(mask_strategy)
+        # if is_progress_bar:
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents)
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                t_expand = t.repeat(latent_model_input.shape[0])
+                guidance_expand = (torch.tensor(
+                    [embedded_guidance_scale] * latent_model_input.shape[0],
+                    dtype=torch.float32,
+                    device=device,
+                ).to(target_dtype) * 1000.0 if embedded_guidance_scale is not None else None)
+                # predict the noise residual
+                with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+                    # concat prompt_embeds_2 and prompt_embeds. Mismatch fill with zeros
+                    if prompt_embeds_2.shape[-1] != prompt_embeds.shape[-1]:
+                        prompt_embeds_2 = F.pad(
+                            prompt_embeds_2,
+                            (0, prompt_embeds.shape[2] - prompt_embeds_2.shape[1]),
+                            value=0,
+                        ).unsqueeze(1)
+                    encoder_hidden_states = torch.cat([prompt_embeds_2, prompt_embeds], dim=1)
+                    noise_pred = self.transformer(  # For an input image (129, 192, 336) (1, 256, 256)
+                        latent_model_input,
+                        encoder_hidden_states,
+                        t_expand,
+                        prompt_mask,
+                        mask_strategy=mask_strategy[i],
+                        guidance=guidance_expand,
+                        return_dict=False,
+                    )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(
+                        noise_pred,
+                        noise_pred_text,
+                        guidance_rescale=self.guidance_rescale,
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    if progress_bar is not None:
+                        progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+        if get_sequence_parallel_state():
+            latents = all_gather(latents, dim=2)
+
+        if not output_type == "latent":
+            expand_temporal_dim = False
+            if len(latents.shape) == 4:
+                if isinstance(self.vae, AutoencoderKLCausal3D):
+                    latents = latents.unsqueeze(2)
+                    expand_temporal_dim = True
+            elif len(latents.shape) == 5:
+                pass
+            else:
+                raise ValueError(
+                    f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}.")
+
+            if (hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor):
+                latents = (latents / self.vae.config.scaling_factor + self.vae.config.shift_factor)
+            else:
+                latents = latents / self.vae.config.scaling_factor
+
+            with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
+                if enable_tiling:
+                    self.vae.enable_tiling()
+                if enable_vae_sp:
+                    self.vae.enable_parallel()
+                image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
+
+            if expand_temporal_dim or image.shape[2] == 1:
+                image = image.squeeze(2)
+
+        else:
+            image = latents
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        image = image.cpu().float()
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return image
+
+        return HunyuanVideoPipelineOutput(videos=image)

fastvideo/models/hunyuan/diffusion/schedulers/__init__.py

@@ -0,0 +1,2 @@
+# ruff: noqa: F401
+from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler

fastvideo/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py

@@ -0,0 +1,239 @@
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput, logging
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class FlowMatchDiscreteSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Euler scheduler.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
+        reverse (`bool`, defaults to `True`):
+            Whether to reverse the timestep schedule.
+    """
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        reverse: bool = True,
+        solver: str = "euler",
+        n_tokens: Optional[int] = None,
+    ):
+        sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
+
+        if not reverse:
+            sigmas = sigmas.flip(0)
+
+        self.sigmas = sigmas
+        # the value fed to model
+        self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32)
+
+        self._step_index = None
+        self._begin_index = None
+
+        self.supported_solver = ["euler"]
+        if solver not in self.supported_solver:
+            raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}")
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Union[str, torch.device] = None,
+        n_tokens: int = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            n_tokens (`int`, *optional*):
+                Number of tokens in the input sequence.
+        """
+        self.num_inference_steps = num_inference_steps
+
+        sigmas = torch.linspace(1, 0, num_inference_steps + 1)
+        sigmas = self.sd3_time_shift(sigmas)
+
+        if not self.config.reverse:
+            sigmas = 1 - sigmas
+
+        self.sigmas = sigmas
+        self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device)
+
+        # Reset step index
+        self._step_index = None
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+        return sample
+
+    def sd3_time_shift(self, t: torch.Tensor):
+        return (self.config.shift * t) / (1 + (self.config.shift - 1) * t)
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        return_dict: bool = True,
+    ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            n_tokens (`int`, *optional*):
+                Number of tokens in the input sequence.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
+                tuple.
+
+        Returns:
+            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
+                returned, otherwise a tuple is returned where the first element is the sample tensor.
+        """
+
+        if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor)
+                or isinstance(timestep, torch.LongTensor)):
+            raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                              " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                              " one of the `scheduler.timesteps` as a timestep."), )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+
+        dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
+
+        if self.config.solver == "euler":
+            prev_sample = sample + model_output.to(torch.float32) * dt
+        else:
+            raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}")
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample, )
+
+        return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    def __len__(self):
+        return self.config.num_train_timesteps

fastvideo/models/hunyuan/idle_config.py

@@ -0,0 +1,380 @@
+# ruff: noqa: F405, F403
+import argparse
+import re
+
+from .constants import *
+from .modules.models import HUNYUAN_VIDEO_CONFIG
+
+
+def parse_args(namespace=None):
+    parser = argparse.ArgumentParser(description="HunyuanVideo inference script")
+
+    parser = add_network_args(parser)
+    parser = add_extra_models_args(parser)
+    parser = add_denoise_schedule_args(parser)
+    parser = add_inference_args(parser)
+    parser = add_parallel_args(parser)
+
+    args = parser.parse_args(namespace=namespace)
+    args = sanity_check_args(args)
+
+    return args
+
+
+def add_network_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="HunyuanVideo network args")
+
+    # Main model
+    group.add_argument(
+        "--model",
+        type=str,
+        choices=list(HUNYUAN_VIDEO_CONFIG.keys()),
+        default="HYVideo-T/2-cfgdistill",
+    )
+    group.add_argument(
+        "--latent-channels",
+        type=str,
+        default=16,
+        help="Number of latent channels of DiT. If None, it will be determined by `vae`. If provided, "
+        "it still needs to match the latent channels of the VAE model.",
+    )
+    group.add_argument(
+        "--precision",
+        type=str,
+        default="bf16",
+        choices=PRECISIONS,
+        help="Precision mode. Options: fp32, fp16, bf16. Applied to the backbone model and optimizer.",
+    )
+
+    # RoPE
+    group.add_argument("--rope-theta", type=int, default=256, help="Theta used in RoPE.")
+    return parser
+
+
+def add_extra_models_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="Extra models args, including vae, text encoders and tokenizers)")
+
+    # - VAE
+    group.add_argument(
+        "--vae",
+        type=str,
+        default="884-16c-hy",
+        choices=list(VAE_PATH),
+        help="Name of the VAE model.",
+    )
+    group.add_argument(
+        "--vae-precision",
+        type=str,
+        default="fp16",
+        choices=PRECISIONS,
+        help="Precision mode for the VAE model.",
+    )
+    group.add_argument(
+        "--vae-tiling",
+        action="store_true",
+        help="Enable tiling for the VAE model to save GPU memory.",
+    )
+    group.set_defaults(vae_tiling=True)
+
+    group.add_argument(
+        "--text-encoder",
+        type=str,
+        default="llm",
+        choices=list(TEXT_ENCODER_PATH),
+        help="Name of the text encoder model.",
+    )
+    group.add_argument(
+        "--text-encoder-precision",
+        type=str,
+        default="fp16",
+        choices=PRECISIONS,
+        help="Precision mode for the text encoder model.",
+    )
+    group.add_argument(
+        "--text-states-dim",
+        type=int,
+        default=4096,
+        help="Dimension of the text encoder hidden states.",
+    )
+    group.add_argument("--text-len", type=int, default=256, help="Maximum length of the text input.")
+    group.add_argument(
+        "--tokenizer",
+        type=str,
+        default="llm",
+        choices=list(TOKENIZER_PATH),
+        help="Name of the tokenizer model.",
+    )
+    group.add_argument(
+        "--prompt-template",
+        type=str,
+        default="dit-llm-encode",
+        choices=PROMPT_TEMPLATE,
+        help="Image prompt template for the decoder-only text encoder model.",
+    )
+    group.add_argument(
+        "--prompt-template-video",
+        type=str,
+        default="dit-llm-encode-video",
+        choices=PROMPT_TEMPLATE,
+        help="Video prompt template for the decoder-only text encoder model.",
+    )
+    group.add_argument(
+        "--hidden-state-skip-layer",
+        type=int,
+        default=2,
+        help="Skip layer for hidden states.",
+    )
+    group.add_argument(
+        "--apply-final-norm",
+        action="store_true",
+        help="Apply final normalization to the used text encoder hidden states.",
+    )
+
+    # - CLIP
+    group.add_argument(
+        "--text-encoder-2",
+        type=str,
+        default="clipL",
+        choices=list(TEXT_ENCODER_PATH),
+        help="Name of the second text encoder model.",
+    )
+    group.add_argument(
+        "--text-encoder-precision-2",
+        type=str,
+        default="fp16",
+        choices=PRECISIONS,
+        help="Precision mode for the second text encoder model.",
+    )
+    group.add_argument(
+        "--text-states-dim-2",
+        type=int,
+        default=768,
+        help="Dimension of the second text encoder hidden states.",
+    )
+    group.add_argument(
+        "--tokenizer-2",
+        type=str,
+        default="clipL",
+        choices=list(TOKENIZER_PATH),
+        help="Name of the second tokenizer model.",
+    )
+    group.add_argument(
+        "--text-len-2",
+        type=int,
+        default=77,
+        help="Maximum length of the second text input.",
+    )
+
+    return parser
+
+
+def add_denoise_schedule_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="Denoise schedule args")
+
+    group.add_argument(
+        "--denoise-type",
+        type=str,
+        default="flow",
+        help="Denoise type for noised inputs.",
+    )
+
+    # Flow Matching
+    group.add_argument(
+        "--flow-shift",
+        type=float,
+        default=7.0,
+        help="Shift factor for flow matching schedulers.",
+    )
+    group.add_argument(
+        "--flow-reverse",
+        action="store_true",
+        help="If reverse, learning/sampling from t=1 -> t=0.",
+    )
+    group.add_argument(
+        "--flow-solver",
+        type=str,
+        default="euler",
+        help="Solver for flow matching.",
+    )
+    group.add_argument(
+        "--use-linear-quadratic-schedule",
+        action="store_true",
+        help="Use linear quadratic schedule for flow matching."
+        "Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)",
+    )
+    group.add_argument(
+        "--linear-schedule-end",
+        type=int,
+        default=25,
+        help="End step for linear quadratic schedule for flow matching.",
+    )
+
+    return parser
+
+
+def add_inference_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="Inference args")
+
+    # ======================== Model loads ========================
+    group.add_argument(
+        "--model-base",
+        type=str,
+        default="ckpts",
+        help="Root path of all the models, including t2v models and extra models.",
+    )
+    group.add_argument(
+        "--dit-weight",
+        type=str,
+        default="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
+        help="Path to the HunyuanVideo model. If None, search the model in the args.model_root."
+        "1. If it is a file, load the model directly."
+        "2. If it is a directory, search the model in the directory. Support two types of models: "
+        "1) named `pytorch_model_*.pt`"
+        "2) named `*_model_states.pt`, where * can be `mp_rank_00`.",
+    )
+    group.add_argument(
+        "--model-resolution",
+        type=str,
+        default="540p",
+        choices=["540p", "720p"],
+        help="Root path of all the models, including t2v models and extra models.",
+    )
+    group.add_argument(
+        "--load-key",
+        type=str,
+        default="module",
+        help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.",
+    )
+    group.add_argument(
+        "--use-cpu-offload",
+        action="store_true",
+        help="Use CPU offload for the model load.",
+    )
+
+    # ======================== Inference general setting ========================
+    group.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        help="Batch size for inference and evaluation.",
+    )
+    group.add_argument(
+        "--infer-steps",
+        type=int,
+        default=50,
+        help="Number of denoising steps for inference.",
+    )
+    group.add_argument(
+        "--disable-autocast",
+        action="store_true",
+        help="Disable autocast for denoising loop and vae decoding in pipeline sampling.",
+    )
+    group.add_argument(
+        "--save-path",
+        type=str,
+        default="./results",
+        help="Path to save the generated samples.",
+    )
+    group.add_argument(
+        "--save-path-suffix",
+        type=str,
+        default="",
+        help="Suffix for the directory of saved samples.",
+    )
+    group.add_argument(
+        "--name-suffix",
+        type=str,
+        default="",
+        help="Suffix for the names of saved samples.",
+    )
+    group.add_argument(
+        "--num-videos",
+        type=int,
+        default=1,
+        help="Number of videos to generate for each prompt.",
+    )
+    # ---sample size---
+    group.add_argument(
+        "--video-size",
+        type=int,
+        nargs="+",
+        default=(720, 1280),
+        help="Video size for training. If a single value is provided, it will be used for both height "
+        "and width. If two values are provided, they will be used for height and width "
+        "respectively.",
+    )
+    group.add_argument(
+        "--video-length",
+        type=int,
+        default=129,
+        help="How many frames to sample from a video. if using 3d vae, the number should be 4n+1",
+    )
+    # --- prompt ---
+    group.add_argument(
+        "--prompt",
+        type=str,
+        default=None,
+        help="Prompt for sampling during evaluation.",
+    )
+    group.add_argument(
+        "--seed-type",
+        type=str,
+        default="auto",
+        choices=["file", "random", "fixed", "auto"],
+        help="Seed type for evaluation. If file, use the seed from the CSV file. If random, generate a "
+        "random seed. If fixed, use the fixed seed given by `--seed`. If auto, `csv` will use the "
+        "seed column if available, otherwise use the fixed `seed` value. `prompt` will use the "
+        "fixed `seed` value.",
+    )
+    group.add_argument("--seed", type=int, default=None, help="Seed for evaluation.")
+
+    # Classifier-Free Guidance
+    group.add_argument("--neg-prompt", type=str, default=None, help="Negative prompt for sampling.")
+    group.add_argument("--cfg-scale", type=float, default=1.0, help="Classifier free guidance scale.")
+    group.add_argument(
+        "--embedded-cfg-scale",
+        type=float,
+        default=6.0,
+        help="Embedded classifier free guidance scale.",
+    )
+
+    group.add_argument(
+        "--reproduce",
+        action="store_true",
+        help="Enable reproducibility by setting random seeds and deterministic algorithms.",
+    )
+
+    return parser
+
+
+def add_parallel_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="Parallel args")
+
+    # ======================== Model loads ========================
+    group.add_argument(
+        "--ulysses-degree",
+        type=int,
+        default=1,
+        help="Ulysses degree.",
+    )
+    group.add_argument(
+        "--ring-degree",
+        type=int,
+        default=1,
+        help="Ulysses degree.",
+    )
+
+    return parser
+
+
+def sanity_check_args(args):
+    # VAE channels
+    vae_pattern = r"\d{2,3}-\d{1,2}c-\w+"
+    if not re.match(vae_pattern, args.vae):
+        raise ValueError(f"Invalid VAE model: {args.vae}. Must be in the format of '{vae_pattern}'.")
+    vae_channels = int(args.vae.split("-")[1][:-1])
+    if args.latent_channels is None:
+        args.latent_channels = vae_channels
+    if vae_channels != args.latent_channels:
+        raise ValueError(f"Latent channels ({args.latent_channels}) must match the VAE channels ({vae_channels}).")
+    return args

fastvideo/models/hunyuan/inference.py

@@ -0,0 +1,482 @@
+import os
+import random
+import time
+from pathlib import Path
+
+import torch
+from loguru import logger
+from safetensors.torch import load_file as safetensors_load_file
+
+from fastvideo.models.hunyuan.constants import NEGATIVE_PROMPT, PRECISION_TO_TYPE, PROMPT_TEMPLATE
+from fastvideo.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline
+from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler
+from fastvideo.models.hunyuan.modules import load_model
+from fastvideo.models.hunyuan.text_encoder import TextEncoder
+from fastvideo.models.hunyuan.utils.data_utils import align_to
+from fastvideo.models.hunyuan.vae import load_vae
+from fastvideo.utils.parallel_states import nccl_info
+
+
+class Inference(object):
+
+    def __init__(
+        self,
+        args,
+        vae,
+        vae_kwargs,
+        text_encoder,
+        model,
+        text_encoder_2=None,
+        pipeline=None,
+        use_cpu_offload=False,
+        device=None,
+        logger=None,
+        parallel_args=None,
+    ):
+        self.vae = vae
+        self.vae_kwargs = vae_kwargs
+
+        self.text_encoder = text_encoder
+        self.text_encoder_2 = text_encoder_2
+
+        self.model = model
+        self.pipeline = pipeline
+        self.use_cpu_offload = use_cpu_offload
+
+        self.args = args
+        self.device = (device if device is not None else "cuda" if torch.cuda.is_available() else "cpu")
+        self.logger = logger
+        self.parallel_args = parallel_args
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_path, args, device=None, **kwargs):
+        """
+        Initialize the Inference pipeline.
+
+        Args:
+            pretrained_model_path (str or pathlib.Path): The model path, including t2v, text encoder and vae checkpoints.
+            args (argparse.Namespace): The arguments for the pipeline.
+            device (int): The device for inference. Default is 0.
+        """
+        # ========================================================================
+        logger.info(f"Got text-to-video model root path: {pretrained_model_path}")
+
+        # ==================== Initialize Distributed Environment ================
+        if nccl_info.sp_size > 1:
+            device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}")
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        parallel_args = None  # {"ulysses_degree": args.ulysses_degree, "ring_degree": args.ring_degree}
+
+        # ======================== Get the args path =============================
+
+        # Disable gradient
+        torch.set_grad_enabled(False)
+
+        # =========================== Build main model ===========================
+        logger.info("Building model...")
+        factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]}
+        in_channels = args.latent_channels
+        out_channels = args.latent_channels
+
+        model = load_model(
+            args,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            factor_kwargs=factor_kwargs,
+        )
+        model = model.to(device)
+        model = Inference.load_state_dict(args, model, pretrained_model_path)
+        if args.enable_torch_compile:
+            model = torch.compile(model)
+        model.eval()
+
+        # ============================= Build extra models ========================
+        # VAE
+        vae, _, s_ratio, t_ratio = load_vae(
+            args.vae,
+            args.vae_precision,
+            logger=logger,
+            device=device if not args.use_cpu_offload else "cpu",
+        )
+        vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio}
+
+        # Text encoder
+        if args.prompt_template_video is not None:
+            crop_start = PROMPT_TEMPLATE[args.prompt_template_video].get("crop_start", 0)
+        elif args.prompt_template is not None:
+            crop_start = PROMPT_TEMPLATE[args.prompt_template].get("crop_start", 0)
+        else:
+            crop_start = 0
+        max_length = args.text_len + crop_start
+
+        # prompt_template
+        prompt_template = (PROMPT_TEMPLATE[args.prompt_template] if args.prompt_template is not None else None)
+
+        # prompt_template_video
+        prompt_template_video = (PROMPT_TEMPLATE[args.prompt_template_video]
+                                 if args.prompt_template_video is not None else None)
+
+        text_encoder = TextEncoder(
+            text_encoder_type=args.text_encoder,
+            max_length=max_length,
+            text_encoder_precision=args.text_encoder_precision,
+            tokenizer_type=args.tokenizer,
+            prompt_template=prompt_template,
+            prompt_template_video=prompt_template_video,
+            hidden_state_skip_layer=args.hidden_state_skip_layer,
+            apply_final_norm=args.apply_final_norm,
+            reproduce=args.reproduce,
+            logger=logger,
+            device=device if not args.use_cpu_offload else "cpu",
+        )
+        text_encoder_2 = None
+        if args.text_encoder_2 is not None:
+            text_encoder_2 = TextEncoder(
+                text_encoder_type=args.text_encoder_2,
+                max_length=args.text_len_2,
+                text_encoder_precision=args.text_encoder_precision_2,
+                tokenizer_type=args.tokenizer_2,
+                reproduce=args.reproduce,
+                logger=logger,
+                device=device if not args.use_cpu_offload else "cpu",
+            )
+
+        return cls(
+            args=args,
+            vae=vae,
+            vae_kwargs=vae_kwargs,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            model=model,
+            use_cpu_offload=args.use_cpu_offload,
+            device=device,
+            logger=logger,
+            parallel_args=parallel_args,
+        )
+
+    @staticmethod
+    def load_state_dict(args, model, pretrained_model_path):
+        load_key = args.load_key
+        dit_weight = Path(args.dit_weight)
+
+        if dit_weight is None:
+            model_dir = pretrained_model_path / f"t2v_{args.model_resolution}"
+            files = list(model_dir.glob("*.pt"))
+            if len(files) == 0:
+                raise ValueError(f"No model weights found in {model_dir}")
+            if str(files[0]).startswith("pytorch_model_"):
+                model_path = dit_weight / f"pytorch_model_{load_key}.pt"
+                bare_model = True
+            elif any(str(f).endswith("_model_states.pt") for f in files):
+                files = [f for f in files if str(f).endswith("_model_states.pt")]
+                model_path = files[0]
+                if len(files) > 1:
+                    logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}")
+                bare_model = False
+            else:
+                raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: "
+                                 f"{list(map(str, files))}. When given a directory as --dit-weight, only "
+                                 f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and "
+                                 f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
+                                 f"specific weight file, please provide the full path to the file.")
+        else:
+            if dit_weight.is_dir():
+                files = list(dit_weight.glob("*.pt"))
+                if len(files) == 0:
+                    raise ValueError(f"No model weights found in {dit_weight}")
+                if str(files[0]).startswith("pytorch_model_"):
+                    model_path = dit_weight / f"pytorch_model_{load_key}.pt"
+                    bare_model = True
+                elif any(str(f).endswith("_model_states.pt") for f in files):
+                    files = [f for f in files if str(f).endswith("_model_states.pt")]
+                    model_path = files[0]
+                    if len(files) > 1:
+                        logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}")
+                    bare_model = False
+                else:
+                    raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: "
+                                     f"{list(map(str, files))}. When given a directory as --dit-weight, only "
+                                     f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and "
+                                     f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
+                                     f"specific weight file, please provide the full path to the file.")
+            elif dit_weight.is_file():
+                model_path = dit_weight
+                bare_model = "unknown"
+            else:
+                raise ValueError(f"Invalid model path: {dit_weight}")
+
+        if not model_path.exists():
+            raise ValueError(f"model_path not exists: {model_path}")
+        logger.info(f"Loading torch model {model_path}...")
+        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
+        if bare_model is False:
+            if load_key in state_dict:
+                state_dict = state_dict[load_key]
+            else:
+                raise KeyError(f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint "
+                               f"are: {list(state_dict.keys())}.")
+        model.load_state_dict(state_dict, strict=True)
+        return model
+
+    @staticmethod
+    def parse_size(size):
+        if isinstance(size, int):
+            size = [size]
+        if not isinstance(size, (list, tuple)):
+            raise ValueError(f"Size must be an integer or (height, width), got {size}.")
+        if len(size) == 1:
+            size = [size[0], size[0]]
+        if len(size) != 2:
+            raise ValueError(f"Size must be an integer or (height, width), got {size}.")
+        return size
+
+
+class HunyuanVideoSampler(Inference):
+
+    def __init__(
+        self,
+        args,
+        vae,
+        vae_kwargs,
+        text_encoder,
+        model,
+        text_encoder_2=None,
+        pipeline=None,
+        use_cpu_offload=False,
+        device=0,
+        logger=None,
+        parallel_args=None,
+    ):
+        super().__init__(
+            args,
+            vae,
+            vae_kwargs,
+            text_encoder,
+            model,
+            text_encoder_2=text_encoder_2,
+            pipeline=pipeline,
+            use_cpu_offload=use_cpu_offload,
+            device=device,
+            logger=logger,
+            parallel_args=parallel_args,
+        )
+
+        self.pipeline = self.load_diffusion_pipeline(
+            args=args,
+            vae=self.vae,
+            text_encoder=self.text_encoder,
+            text_encoder_2=self.text_encoder_2,
+            model=self.model,
+            device=self.device,
+        )
+
+        self.default_negative_prompt = NEGATIVE_PROMPT
+
+    def load_diffusion_pipeline(
+        self,
+        args,
+        vae,
+        text_encoder,
+        text_encoder_2,
+        model,
+        scheduler=None,
+        device=None,
+        progress_bar_config=None,
+        data_type="video",
+    ):
+        """Load the denoising scheduler for inference."""
+        if scheduler is None:
+            if args.denoise_type == "flow":
+                scheduler = FlowMatchDiscreteScheduler(
+                    shift=args.flow_shift,
+                    reverse=args.flow_reverse,
+                    solver=args.flow_solver,
+                )
+            else:
+                raise ValueError(f"Invalid denoise type {args.denoise_type}")
+
+        pipeline = HunyuanVideoPipeline(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            transformer=model,
+            scheduler=scheduler,
+            progress_bar_config=progress_bar_config,
+            args=args,
+        )
+        if self.use_cpu_offload:
+            pipeline.enable_sequential_cpu_offload()
+        else:
+            pipeline = pipeline.to(device)
+
+        return pipeline
+
+    @torch.no_grad()
+    def predict(
+        self,
+        prompt,
+        height=192,
+        width=336,
+        video_length=129,
+        seed=None,
+        negative_prompt=None,
+        infer_steps=50,
+        guidance_scale=6,
+        flow_shift=5.0,
+        embedded_guidance_scale=None,
+        batch_size=1,
+        num_videos_per_prompt=1,
+        mask_strategy=None,
+        **kwargs,
+    ):
+        """
+        Predict the image/video from the given text.
+
+        Args:
+            prompt (str or List[str]): The input text.
+            kwargs:
+                height (int): The height of the output video. Default is 192.
+                width (int): The width of the output video. Default is 336.
+                video_length (int): The frame number of the output video. Default is 129.
+                seed (int or List[str]): The random seed for the generation. Default is a random integer.
+                negative_prompt (str or List[str]): The negative text prompt. Default is an empty string.
+                guidance_scale (float): The guidance scale for the generation. Default is 6.0.
+                num_images_per_prompt (int): The number of images per prompt. Default is 1.
+                infer_steps (int): The number of inference steps. Default is 100.
+        """
+
+        out_dict = dict()
+
+        # ========================================================================
+        # Arguments: seed
+        # ========================================================================
+        if isinstance(seed, torch.Tensor):
+            seed = seed.tolist()
+        if seed is None:
+            seeds = [random.randint(0, 1_000_000) for _ in range(batch_size * num_videos_per_prompt)]
+        elif isinstance(seed, int):
+            seeds = [seed + i for _ in range(batch_size) for i in range(num_videos_per_prompt)]
+        elif isinstance(seed, (list, tuple)):
+            if len(seed) == batch_size:
+                seeds = [int(seed[i]) + j for i in range(batch_size) for j in range(num_videos_per_prompt)]
+            elif len(seed) == batch_size * num_videos_per_prompt:
+                seeds = [int(s) for s in seed]
+            else:
+                raise ValueError(
+                    f"Length of seed must be equal to number of prompt(batch_size) or "
+                    f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}.")
+        else:
+            raise ValueError(f"Seed must be an integer, a list of integers, or None, got {seed}.")
+        # 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
+
+        # ========================================================================
+        # Arguments: target_width, target_height, target_video_length
+        # ========================================================================
+        if width <= 0 or height <= 0 or video_length <= 0:
+            raise ValueError(
+                f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}"
+            )
+        if (video_length - 1) % 4 != 0:
+            raise ValueError(f"`video_length-1` must be a multiple of 4, got {video_length}")
+
+        logger.info(f"Input (height, width, video_length) = ({height}, {width}, {video_length})")
+
+        target_height = align_to(height, 16)
+        target_width = align_to(width, 16)
+        target_video_length = video_length
+
+        out_dict["size"] = (target_height, target_width, target_video_length)
+
+        # ========================================================================
+        # Arguments: prompt, new_prompt, negative_prompt
+        # ========================================================================
+        if not isinstance(prompt, str):
+            raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
+        prompt = [prompt.strip()]
+
+        # negative prompt
+        if negative_prompt is None or negative_prompt == "":
+            negative_prompt = self.default_negative_prompt
+        if not isinstance(negative_prompt, str):
+            raise TypeError(f"`negative_prompt` must be a string, but got {type(negative_prompt)}")
+        negative_prompt = [negative_prompt.strip()]
+
+        # ========================================================================
+        # Scheduler
+        # ========================================================================
+        scheduler = FlowMatchDiscreteScheduler(
+            shift=flow_shift,
+            reverse=self.args.flow_reverse,
+            solver=self.args.flow_solver,
+        )
+        self.pipeline.scheduler = scheduler
+
+        if "884" in self.args.vae:
+            latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8]
+        elif "888" in self.args.vae:
+            latents_size = [(video_length - 1) // 8 + 1, height // 8, width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+
+        # ========================================================================
+        # Print infer args
+        # ========================================================================
+        debug_str = f"""
+                        height: {target_height}
+                         width: {target_width}
+                  video_length: {target_video_length}
+                        prompt: {prompt}
+                    neg_prompt: {negative_prompt}
+                          seed: {seed}
+                   infer_steps: {infer_steps}
+         num_videos_per_prompt: {num_videos_per_prompt}
+                guidance_scale: {guidance_scale}
+                      n_tokens: {n_tokens}
+                    flow_shift: {flow_shift}
+       embedded_guidance_scale: {embedded_guidance_scale}"""
+        logger.debug(debug_str)
+
+        # ========================================================================
+        # Pipeline inference
+        # ========================================================================
+        start_time = time.time()
+        samples = self.pipeline(
+            prompt=prompt,
+            height=target_height,
+            width=target_width,
+            video_length=target_video_length,
+            num_inference_steps=infer_steps,
+            guidance_scale=guidance_scale,
+            negative_prompt=negative_prompt,
+            num_videos_per_prompt=num_videos_per_prompt,
+            generator=generator,
+            output_type="pil",
+            n_tokens=n_tokens,
+            embedded_guidance_scale=embedded_guidance_scale,
+            data_type="video" if target_video_length > 1 else "image",
+            is_progress_bar=True,
+            vae_ver=self.args.vae,
+            enable_tiling=self.args.vae_tiling,
+            enable_vae_sp=self.args.vae_sp,
+            mask_strategy=mask_strategy,
+        )[0]
+        out_dict["samples"] = samples
+        out_dict["prompts"] = prompt
+
+        gen_time = time.time() - start_time
+        logger.info(f"Success, time: {gen_time}")
+
+        return out_dict

fastvideo/models/hunyuan/modules/__init__.py

@@ -0,0 +1,25 @@
+from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer
+
+
+def load_model(args, in_channels, out_channels, factor_kwargs):
+    """load hunyuan video model
+
+    Args:
+        args (dict): model args
+        in_channels (int): input channels number
+        out_channels (int): output channels number
+        factor_kwargs (dict): factor kwargs
+
+    Returns:
+        model (nn.Module): The hunyuan video model
+    """
+    if args.model in HUNYUAN_VIDEO_CONFIG.keys():
+        model = HYVideoDiffusionTransformer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            **HUNYUAN_VIDEO_CONFIG[args.model],
+            **factor_kwargs,
+        )
+        return model
+    else:
+        raise NotImplementedError()

fastvideo/models/hunyuan/modules/activation_layers.py

@@ -0,0 +1,23 @@
+import torch.nn as nn
+
+
+def get_activation_layer(act_type):
+    """get activation layer
+
+    Args:
+        act_type (str): the activation type
+
+    Returns:
+        torch.nn.functional: the activation layer
+    """
+    if act_type == "gelu":
+        return lambda: nn.GELU()
+    elif act_type == "gelu_tanh":
+        # Approximate `tanh` requires torch >= 1.13
+        return lambda: nn.GELU(approximate="tanh")
+    elif act_type == "relu":
+        return nn.ReLU
+    elif act_type == "silu":
+        return nn.SiLU
+    else:
+        raise ValueError(f"Unknown activation type: {act_type}")

fastvideo/models/hunyuan/modules/attenion.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+try:
+    from st_attn import sliding_tile_attention
+except ImportError:
+    print("Could not load Sliding Tile Attention.")
+    sliding_tile_attention = None
+
+from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
+from fastvideo.utils.communications import all_gather, all_to_all_4D
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+
+def attention(
+    q,
+    k,
+    v,
+    drop_rate=0,
+    attn_mask=None,
+    causal=False,
+):
+
+    qkv = torch.stack([q, k, v], dim=2)
+
+    if attn_mask is not None and attn_mask.dtype != torch.bool:
+        attn_mask = attn_mask.bool()
+
+    x = flash_attn_no_pad(qkv, attn_mask, causal=causal, dropout_p=drop_rate, softmax_scale=None)
+
+    b, s, a, d = x.shape
+    out = x.reshape(b, s, -1)
+    return out
+
+
+def tile(x, sp_size):
+    x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80)
+    return rearrange(x,
+                     "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
+                     n_t=5,
+                     n_h=6,
+                     n_w=10,
+                     ts_t=6,
+                     ts_h=8,
+                     ts_w=8)
+
+
+def untile(x, sp_size):
+    x = rearrange(x,
+                  "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
+                  n_t=5,
+                  n_h=6,
+                  n_w=10,
+                  ts_t=6,
+                  ts_h=8,
+                  ts_w=8)
+    return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80)
+
+
+def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=None):
+    query, encoder_query = q
+    key, encoder_key = k
+    value, encoder_value = v
+    text_length = text_mask.sum()
+
+    if get_sequence_parallel_state():
+        # batch_size, seq_len, attn_heads, head_dim
+        query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
+        key = all_to_all_4D(key, scatter_dim=2, gather_dim=1)
+        value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
+
+        def shrink_head(encoder_state, dim):
+            local_heads = encoder_state.shape[dim] // nccl_info.sp_size
+            return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
+
+        encoder_query = shrink_head(encoder_query, dim=2)
+        encoder_key = shrink_head(encoder_key, dim=2)
+        encoder_value = shrink_head(encoder_value, dim=2)
+        # [b, s, h, d]
+
+    sequence_length = query.size(1)
+    encoder_sequence_length = encoder_query.size(1)
+
+    if mask_strategy[0] is not None:
+        query = torch.cat([tile(query, nccl_info.sp_size), encoder_query], dim=1).transpose(1, 2)
+        key = torch.cat([tile(key, nccl_info.sp_size), encoder_key], dim=1).transpose(1, 2)
+        value = torch.cat([tile(value, nccl_info.sp_size), encoder_value], dim=1).transpose(1, 2)
+
+        head_num = query.size(1)
+        current_rank = nccl_info.rank_within_group
+        start_head = current_rank * head_num
+        windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)]
+
+        hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2)
+    else:
+        query = torch.cat([query, encoder_query], dim=1)
+        key = torch.cat([key, encoder_key], dim=1)
+        value = torch.cat([value, encoder_value], dim=1)
+        # B, S, 3, H, D
+        qkv = torch.stack([query, key, value], dim=2)
+
+        attn_mask = F.pad(text_mask, (sequence_length, 0), value=True)
+        hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None)
+
+    hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length),
+                                                                          dim=1)
+
+    if mask_strategy[0] is not None:
+        hidden_states = untile(hidden_states, nccl_info.sp_size)
+
+    if get_sequence_parallel_state():
+        hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)
+        encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous()
+
+    hidden_states = hidden_states.to(query.dtype)
+    encoder_hidden_states = encoder_hidden_states.to(query.dtype)
+
+    attn = torch.cat([hidden_states, encoder_hidden_states], dim=1)
+
+    b, s, a, d = attn.shape
+    attn = attn.reshape(b, s, -1)
+
+    return attn

fastvideo/models/hunyuan/modules/embed_layers.py

@@ -0,0 +1,150 @@
+import math
+
+import torch
+import torch.nn as nn
+
+from ..utils.helpers import to_2tuple
+
+
+class PatchEmbed(nn.Module):
+    """2D Image to Patch Embedding
+
+    Image to Patch Embedding using Conv2d
+
+    A convolution based approach to patchifying a 2D image w/ embedding projection.
+
+    Based on the impl in https://github.com/google-research/vision_transformer
+
+    Hacked together by / Copyright 2020 Ross Wightman
+
+    Remove the _assert function in forward function to be compatible with multi-resolution images.
+    """
+
+    def __init__(
+        self,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        norm_layer=None,
+        flatten=True,
+        bias=True,
+        dtype=None,
+        device=None,
+    ):
+        factory_kwargs = {"dtype": dtype, "device": device}
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+        self.patch_size = patch_size
+        self.flatten = flatten
+
+        self.proj = nn.Conv3d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=bias,
+            **factory_kwargs,
+        )
+        nn.init.xavier_uniform_(self.proj.weight.view(self.proj.weight.size(0), -1))
+        if bias:
+            nn.init.zeros_(self.proj.bias)
+
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+class TextProjection(nn.Module):
+    """
+    Projects text embeddings. Also handles dropout for classifier-free guidance.
+
+    Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
+    """
+
+    def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None):
+        factory_kwargs = {"dtype": dtype, "device": device}
+        super().__init__()
+        self.linear_1 = nn.Linear(
+            in_features=in_channels,
+            out_features=hidden_size,
+            bias=True,
+            **factory_kwargs,
+        )
+        self.act_1 = act_layer()
+        self.linear_2 = nn.Linear(
+            in_features=hidden_size,
+            out_features=hidden_size,
+            bias=True,
+            **factory_kwargs,
+        )
+
+    def forward(self, caption):
+        hidden_states = self.linear_1(caption)
+        hidden_states = self.act_1(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+
+def timestep_embedding(t, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+
+    Args:
+        t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.
+        dim (int): the dimension of the output.
+        max_period (int): controls the minimum frequency of the embeddings.
+
+    Returns:
+        embedding (torch.Tensor): An (N, D) Tensor of positional embeddings.
+
+    .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+    """
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) /
+                      half).to(device=t.device)
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        act_layer,
+        frequency_embedding_size=256,
+        max_period=10000,
+        out_size=None,
+        dtype=None,
+        device=None,
+    ):
+        factory_kwargs = {"dtype": dtype, "device": device}
+        super().__init__()
+        self.frequency_embedding_size = frequency_embedding_size
+        self.max_period = max_period
+        if out_size is None:
+            out_size = hidden_size
+
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs),
+            act_layer(),
+            nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs),
+        )
+        nn.init.normal_(self.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.mlp[2].weight, std=0.02)
+
+    def forward(self, t):
+        t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb

fastvideo/models/hunyuan/modules/mlp_layers.py

@@ -0,0 +1,107 @@
+# Modified from timm library:
+# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from ..utils.helpers import to_2tuple
+from .modulate_layers import modulate
+
+
+class MLP(nn.Module):
+    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        norm_layer=None,
+        bias=True,
+        drop=0.0,
+        use_conv=False,
+        device=None,
+        dtype=None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        out_features = out_features or in_channels
+        hidden_channels = hidden_channels or in_channels
+        bias = to_2tuple(bias)
+        drop_probs = to_2tuple(drop)
+        linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
+
+        self.fc1 = linear_layer(in_channels, hidden_channels, bias=bias[0], **factory_kwargs)
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop_probs[0])
+        self.norm = (norm_layer(hidden_channels, **factory_kwargs) if norm_layer is not None else nn.Identity())
+        self.fc2 = linear_layer(hidden_channels, out_features, bias=bias[1], **factory_kwargs)
+        self.drop2 = nn.Dropout(drop_probs[1])
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.norm(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+
+
+#
+class MLPEmbedder(nn.Module):
+    """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py"""
+
+    def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class FinalLayer(nn.Module):
+    """The final layer of DiT."""
+
+    def __init__(self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        # Just use LayerNorm for the final layer
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+        if isinstance(patch_size, int):
+            self.linear = nn.Linear(
+                hidden_size,
+                patch_size * patch_size * out_channels,
+                bias=True,
+                **factory_kwargs,
+            )
+        else:
+            self.linear = nn.Linear(
+                hidden_size,
+                patch_size[0] * patch_size[1] * patch_size[2] * out_channels,
+                bias=True,
+            )
+        nn.init.zeros_(self.linear.weight)
+        nn.init.zeros_(self.linear.bias)
+
+        # Here we don't distinguish between the modulate types. Just use the simple one.
+        self.adaLN_modulation = nn.Sequential(
+            act_layer(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
+        )
+        # Zero-initialize the modulation
+        nn.init.zeros_(self.adaLN_modulation[1].weight)
+        nn.init.zeros_(self.adaLN_modulation[1].bias)
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift=shift, scale=scale)
+        x = self.linear(x)
+        return x

fastvideo/models/hunyuan/modules/models.py

@@ -0,0 +1,666 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models import ModelMixin
+from einops import rearrange
+
+from fastvideo.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_embed
+from fastvideo.utils.parallel_states import nccl_info
+
+from .activation_layers import get_activation_layer
+from .attenion import parallel_attention
+from .embed_layers import PatchEmbed, TextProjection, TimestepEmbedder
+from .mlp_layers import MLP, FinalLayer, MLPEmbedder
+from .modulate_layers import ModulateDiT, apply_gate, modulate
+from .norm_layers import get_norm_layer
+from .posemb_layers import apply_rotary_emb
+from .token_refiner import SingleTokenRefiner
+
+
+class MMDoubleStreamBlock(nn.Module):
+    """
+    A multimodal dit block with separate modulation for
+    text and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206
+                                     (Flux.1): https://github.com/black-forest-labs/flux
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        heads_num: int,
+        mlp_width_ratio: float,
+        mlp_act_type: str = "gelu_tanh",
+        qk_norm: bool = True,
+        qk_norm_type: str = "rms",
+        qkv_bias: bool = False,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        self.deterministic = False
+        self.heads_num = heads_num
+        head_dim = hidden_size // heads_num
+        mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
+
+        self.img_mod = ModulateDiT(
+            hidden_size,
+            factor=6,
+            act_layer=get_activation_layer("silu"),
+            **factory_kwargs,
+        )
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+        self.img_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs)
+        qk_norm_layer = get_norm_layer(qk_norm_type)
+        self.img_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                                if qk_norm else nn.Identity())
+        self.img_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                                if qk_norm else nn.Identity())
+        self.img_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs)
+
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+        self.img_mlp = MLP(
+            hidden_size,
+            mlp_hidden_dim,
+            act_layer=get_activation_layer(mlp_act_type),
+            bias=True,
+            **factory_kwargs,
+        )
+
+        self.txt_mod = ModulateDiT(
+            hidden_size,
+            factor=6,
+            act_layer=get_activation_layer("silu"),
+            **factory_kwargs,
+        )
+        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+        self.txt_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs)
+        self.txt_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                                if qk_norm else nn.Identity())
+        self.txt_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                                if qk_norm else nn.Identity())
+        self.txt_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs)
+
+        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+        self.txt_mlp = MLP(
+            hidden_size,
+            mlp_hidden_dim,
+            act_layer=get_activation_layer(mlp_act_type),
+            bias=True,
+            **factory_kwargs,
+        )
+        self.hybrid_seq_parallel_attn = None
+
+    def enable_deterministic(self):
+        self.deterministic = True
+
+    def disable_deterministic(self):
+        self.deterministic = False
+
+    def forward(
+        self,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        freqs_cis: tuple = None,
+        text_mask: torch.Tensor = None,
+        mask_strategy=None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        (
+            img_mod1_shift,
+            img_mod1_scale,
+            img_mod1_gate,
+            img_mod2_shift,
+            img_mod2_scale,
+            img_mod2_gate,
+        ) = self.img_mod(vec).chunk(6, dim=-1)
+        (
+            txt_mod1_shift,
+            txt_mod1_scale,
+            txt_mod1_gate,
+            txt_mod2_shift,
+            txt_mod2_scale,
+            txt_mod2_gate,
+        ) = self.txt_mod(vec).chunk(6, dim=-1)
+
+        # Prepare image for attention.
+        img_modulated = self.img_norm1(img)
+        img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale)
+        img_qkv = self.img_attn_qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
+        # Apply QK-Norm if needed
+        img_q = self.img_attn_q_norm(img_q).to(img_v)
+        img_k = self.img_attn_k_norm(img_k).to(img_v)
+
+        # Apply RoPE if needed.
+        if freqs_cis is not None:
+
+            def shrink_head(encoder_state, dim):
+                local_heads = encoder_state.shape[dim] // nccl_info.sp_size
+                return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
+
+            freqs_cis = (
+                shrink_head(freqs_cis[0], dim=0),
+                shrink_head(freqs_cis[1], dim=0),
+            )
+
+            img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False)
+            assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape
+                    ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}"
+            img_q, img_k = img_qq, img_kk
+
+        # Prepare txt for attention.
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale)
+        txt_qkv = self.txt_attn_qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
+        # Apply QK-Norm if needed.
+        txt_q = self.txt_attn_q_norm(txt_q).to(txt_v)
+        txt_k = self.txt_attn_k_norm(txt_k).to(txt_v)
+
+        attn = parallel_attention(
+            (img_q, txt_q),
+            (img_k, txt_k),
+            (img_v, txt_v),
+            img_q_len=img_q.shape[1],
+            img_kv_len=img_k.shape[1],
+            text_mask=text_mask,
+            mask_strategy=mask_strategy,
+        )
+
+        # attention computation end
+
+        img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:]
+
+        # Calculate the img blocks.
+        img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate)
+        img = img + apply_gate(
+            self.img_mlp(modulate(self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale)),
+            gate=img_mod2_gate,
+        )
+
+        # Calculate the txt blocks.
+        txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate)
+        txt = txt + apply_gate(
+            self.txt_mlp(modulate(self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale)),
+            gate=txt_mod2_gate,
+        )
+        return img, txt
+
+
+class MMSingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    Also refer to (SD3): https://arxiv.org/abs/2403.03206
+                  (Flux.1): https://github.com/black-forest-labs/flux
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        heads_num: int,
+        mlp_width_ratio: float = 4.0,
+        mlp_act_type: str = "gelu_tanh",
+        qk_norm: bool = True,
+        qk_norm_type: str = "rms",
+        qk_scale: float = None,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        self.deterministic = False
+        self.hidden_size = hidden_size
+        self.heads_num = heads_num
+        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
+
+        # qkv and mlp_in
+        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + mlp_hidden_dim, **factory_kwargs)
+        # proj and mlp_out
+        self.linear2 = nn.Linear(hidden_size + mlp_hidden_dim, hidden_size, **factory_kwargs)
+
+        qk_norm_layer = get_norm_layer(qk_norm_type)
+        self.q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                       if qk_norm else nn.Identity())
+        self.k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                       if qk_norm else nn.Identity())
+
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+        self.mlp_act = get_activation_layer(mlp_act_type)()
+        self.modulation = ModulateDiT(
+            hidden_size,
+            factor=3,
+            act_layer=get_activation_layer("silu"),
+            **factory_kwargs,
+        )
+        self.hybrid_seq_parallel_attn = None
+
+    def enable_deterministic(self):
+        self.deterministic = True
+
+    def disable_deterministic(self):
+        self.deterministic = False
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        vec: torch.Tensor,
+        txt_len: int,
+        freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
+        text_mask: torch.Tensor = None,
+        mask_strategy=None,
+    ) -> torch.Tensor:
+        mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
+        x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale)
+        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
+
+        # Apply QK-Norm if needed.
+        q = self.q_norm(q).to(v)
+        k = self.k_norm(k).to(v)
+
+        def shrink_head(encoder_state, dim):
+            local_heads = encoder_state.shape[dim] // nccl_info.sp_size
+            return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
+
+        freqs_cis = (
+            shrink_head(freqs_cis[0], dim=0),
+            shrink_head(freqs_cis[1], dim=0),
+        )
+
+        img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :]
+        img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :]
+        img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :]
+        img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False)
+        assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape
+                ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}"
+        img_q, img_k = img_qq, img_kk
+
+        attn = parallel_attention(
+            (img_q, txt_q),
+            (img_k, txt_k),
+            (img_v, txt_v),
+            img_q_len=img_q.shape[1],
+            img_kv_len=img_k.shape[1],
+            text_mask=text_mask,
+            mask_strategy=mask_strategy,
+        )
+
+        # attention computation end
+
+        # Compute activation in mlp stream, cat again and run second linear layer.
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        return x + apply_gate(output, gate=mod_gate)
+
+
+class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
+    """
+    HunyuanVideo Transformer backbone
+
+    Inherited from ModelMixin and ConfigMixin for compatibility with diffusers' sampler StableDiffusionPipeline.
+
+    Reference:
+    [1] Flux.1: https://github.com/black-forest-labs/flux
+    [2] MMDiT: http://arxiv.org/abs/2403.03206
+
+    Parameters
+    ----------
+    args: argparse.Namespace
+        The arguments parsed by argparse.
+    patch_size: list
+        The size of the patch.
+    in_channels: int
+        The number of input channels.
+    out_channels: int
+        The number of output channels.
+    hidden_size: int
+        The hidden size of the transformer backbone.
+    heads_num: int
+        The number of attention heads.
+    mlp_width_ratio: float
+        The ratio of the hidden size of the MLP in the transformer block.
+    mlp_act_type: str
+        The activation function of the MLP in the transformer block.
+    depth_double_blocks: int
+        The number of transformer blocks in the double blocks.
+    depth_single_blocks: int
+        The number of transformer blocks in the single blocks.
+    rope_dim_list: list
+        The dimension of the rotary embedding for t, h, w.
+    qkv_bias: bool
+        Whether to use bias in the qkv linear layer.
+    qk_norm: bool
+        Whether to use qk norm.
+    qk_norm_type: str
+        The type of qk norm.
+    guidance_embed: bool
+        Whether to use guidance embedding for distillation.
+    text_projection: str
+        The type of the text projection, default is single_refiner.
+    use_attention_mask: bool
+        Whether to use attention mask for text encoder.
+    dtype: torch.dtype
+        The dtype of the model.
+    device: torch.device
+        The device of the model.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: list = [1, 2, 2],
+        in_channels: int = 4,  # Should be VAE.config.latent_channels.
+        out_channels: int = None,
+        hidden_size: int = 3072,
+        heads_num: int = 24,
+        mlp_width_ratio: float = 4.0,
+        mlp_act_type: str = "gelu_tanh",
+        mm_double_blocks_depth: int = 20,
+        mm_single_blocks_depth: int = 40,
+        rope_dim_list: List[int] = [16, 56, 56],
+        qkv_bias: bool = True,
+        qk_norm: bool = True,
+        qk_norm_type: str = "rms",
+        guidance_embed: bool = False,  # For modulation.
+        text_projection: str = "single_refiner",
+        use_attention_mask: bool = True,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        text_states_dim: int = 4096,
+        text_states_dim_2: int = 768,
+        rope_theta: int = 256,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.unpatchify_channels = self.out_channels
+        self.guidance_embed = guidance_embed
+        self.rope_dim_list = rope_dim_list
+        self.rope_theta = rope_theta
+        # Text projection. Default to linear projection.
+        # Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831
+        self.use_attention_mask = use_attention_mask
+        self.text_projection = text_projection
+
+        if hidden_size % heads_num != 0:
+            raise ValueError(f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}")
+        pe_dim = hidden_size // heads_num
+        if sum(rope_dim_list) != pe_dim:
+            raise ValueError(f"Got {rope_dim_list} but expected positional dim {pe_dim}")
+        self.hidden_size = hidden_size
+        self.heads_num = heads_num
+
+        # image projection
+        self.img_in = PatchEmbed(self.patch_size, self.in_channels, self.hidden_size, **factory_kwargs)
+
+        # text projection
+        if self.text_projection == "linear":
+            self.txt_in = TextProjection(
+                self.config.text_states_dim,
+                self.hidden_size,
+                get_activation_layer("silu"),
+                **factory_kwargs,
+            )
+        elif self.text_projection == "single_refiner":
+            self.txt_in = SingleTokenRefiner(
+                self.config.text_states_dim,
+                hidden_size,
+                heads_num,
+                depth=2,
+                **factory_kwargs,
+            )
+        else:
+            raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}")
+
+        # time modulation
+        self.time_in = TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs)
+
+        # text modulation
+        self.vector_in = MLPEmbedder(self.config.text_states_dim_2, self.hidden_size, **factory_kwargs)
+
+        # guidance modulation
+        self.guidance_in = (TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs)
+                            if guidance_embed else None)
+
+        # double blocks
+        self.double_blocks = nn.ModuleList([
+            MMDoubleStreamBlock(
+                self.hidden_size,
+                self.heads_num,
+                mlp_width_ratio=mlp_width_ratio,
+                mlp_act_type=mlp_act_type,
+                qk_norm=qk_norm,
+                qk_norm_type=qk_norm_type,
+                qkv_bias=qkv_bias,
+                **factory_kwargs,
+            ) for _ in range(mm_double_blocks_depth)
+        ])
+
+        # single blocks
+        self.single_blocks = nn.ModuleList([
+            MMSingleStreamBlock(
+                self.hidden_size,
+                self.heads_num,
+                mlp_width_ratio=mlp_width_ratio,
+                mlp_act_type=mlp_act_type,
+                qk_norm=qk_norm,
+                qk_norm_type=qk_norm_type,
+                **factory_kwargs,
+            ) for _ in range(mm_single_blocks_depth)
+        ])
+
+        self.final_layer = FinalLayer(
+            self.hidden_size,
+            self.patch_size,
+            self.out_channels,
+            get_activation_layer("silu"),
+            **factory_kwargs,
+        )
+
+    def enable_deterministic(self):
+        for block in self.double_blocks:
+            block.enable_deterministic()
+        for block in self.single_blocks:
+            block.enable_deterministic()
+
+    def disable_deterministic(self):
+        for block in self.double_blocks:
+            block.disable_deterministic()
+        for block in self.single_blocks:
+            block.disable_deterministic()
+
+    def get_rotary_pos_embed(self, rope_sizes):
+        target_ndim = 3
+
+        head_dim = self.hidden_size // self.heads_num
+        rope_dim_list = self.rope_dim_list
+        if rope_dim_list is None:
+            rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
+        assert (sum(rope_dim_list) == head_dim), "sum(rope_dim_list) should equal to head_dim of attention layer"
+        freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
+            rope_dim_list,
+            rope_sizes,
+            theta=self.rope_theta,
+            use_real=True,
+            theta_rescale_factor=1,
+        )
+        return freqs_cos, freqs_sin
+        # x: torch.Tensor,
+        # t: torch.Tensor,  # Should be in range(0, 1000).
+        # text_states: torch.Tensor = None,
+        # text_mask: torch.Tensor = None,  # Now we don't use it.
+        # text_states_2: Optional[torch.Tensor] = None,  # Text embedding for modulation.
+        # guidance: torch.Tensor = None,  # Guidance for modulation, should be cfg_scale x 1000.
+        # return_dict: bool = True,
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_attention_mask: torch.Tensor,
+        mask_strategy=None,
+        output_features=False,
+        output_features_stride=8,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = False,
+        guidance=None,
+    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+        if guidance is None:
+            guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16)
+        if mask_strategy is None:
+            mask_strategy = [[None] * self.heads_num for _ in range(len(self.double_blocks) + len(self.single_blocks))]
+        img = x = hidden_states
+        text_mask = encoder_attention_mask
+        t = timestep
+        txt = encoder_hidden_states[:, 1:]
+        text_states_2 = encoder_hidden_states[:, 0, :self.config.text_states_dim_2]
+        _, _, ot, oh, ow = x.shape  # codespell:ignore
+        tt, th, tw = (
+            ot // self.patch_size[0],  # codespell:ignore
+            oh // self.patch_size[1],  # codespell:ignore
+            ow // self.patch_size[2],  # codespell:ignore
+        )
+        original_tt = nccl_info.sp_size * tt
+        freqs_cos, freqs_sin = self.get_rotary_pos_embed((original_tt, th, tw))
+        # Prepare modulation vectors.
+        vec = self.time_in(t)
+
+        # text modulation
+        vec = vec + self.vector_in(text_states_2)
+
+        # guidance modulation
+        if self.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+
+            # our timestep_embedding is merged into guidance_in(TimestepEmbedder)
+            vec = vec + self.guidance_in(guidance)
+
+        # Embed image and text.
+        img = self.img_in(img)
+        if self.text_projection == "linear":
+            txt = self.txt_in(txt)
+        elif self.text_projection == "single_refiner":
+            txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None)
+        else:
+            raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}")
+
+        txt_seq_len = txt.shape[1]
+        img_seq_len = img.shape[1]
+
+        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
+        # --------------------- Pass through DiT blocks ------------------------
+
+        for index, block in enumerate(self.double_blocks):
+            double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]]
+            img, txt = block(*double_block_args)
+        # Merge txt and img to pass through single stream blocks.
+        x = torch.cat((img, txt), 1)
+        if output_features:
+            features_list = []
+        if len(self.single_blocks) > 0:
+            for index, block in enumerate(self.single_blocks):
+                single_block_args = [
+                    x,
+                    vec,
+                    txt_seq_len,
+                    (freqs_cos, freqs_sin),
+                    text_mask,
+                    mask_strategy[index + len(self.double_blocks)],
+                ]
+                x = block(*single_block_args)
+                if output_features and _ % output_features_stride == 0:
+                    features_list.append(x[:, :img_seq_len, ...])
+
+        img = x[:, :img_seq_len, ...]
+
+        # ---------------------------- Final layer ------------------------------
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+
+        img = self.unpatchify(img, tt, th, tw)
+        assert not return_dict, "return_dict is not supported."
+        if output_features:
+            features_list = torch.stack(features_list, dim=0)
+        else:
+            features_list = None
+        return (img, features_list)
+
+    def unpatchify(self, x, t, h, w):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.unpatchify_channels
+        pt, ph, pw = self.patch_size
+        assert t * h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
+        x = torch.einsum("nthwcopq->nctohpwq", x)
+        imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
+
+        return imgs
+
+    def params_count(self):
+        counts = {
+            "double":
+            sum([
+                sum(p.numel()
+                    for p in block.img_attn_qkv.parameters()) + sum(p.numel()
+                                                                    for p in block.img_attn_proj.parameters()) +
+                sum(p.numel() for p in block.img_mlp.parameters()) + sum(p.numel()
+                                                                         for p in block.txt_attn_qkv.parameters()) +
+                sum(p.numel() for p in block.txt_attn_proj.parameters()) + sum(p.numel()
+                                                                               for p in block.txt_mlp.parameters())
+                for block in self.double_blocks
+            ]),
+            "single":
+            sum([
+                sum(p.numel() for p in block.linear1.parameters()) + sum(p.numel() for p in block.linear2.parameters())
+                for block in self.single_blocks
+            ]),
+            "total":
+            sum(p.numel() for p in self.parameters()),
+        }
+        counts["attn+mlp"] = counts["double"] + counts["single"]
+        return counts
+
+
+#################################################################################
+#                             HunyuanVideo Configs                              #
+#################################################################################
+
+HUNYUAN_VIDEO_CONFIG = {
+    "HYVideo-T/2": {
+        "mm_double_blocks_depth": 20,
+        "mm_single_blocks_depth": 40,
+        "rope_dim_list": [16, 56, 56],
+        "hidden_size": 3072,
+        "heads_num": 24,
+        "mlp_width_ratio": 4,
+    },
+    "HYVideo-T/2-cfgdistill": {
+        "mm_double_blocks_depth": 20,
+        "mm_single_blocks_depth": 40,
+        "rope_dim_list": [16, 56, 56],
+        "hidden_size": 3072,
+        "heads_num": 24,
+        "mlp_width_ratio": 4,
+        "guidance_embed": True,
+    },
+}

fastvideo/models/hunyuan/modules/modulate_layers.py

@@ -0,0 +1,152 @@
+from typing import Callable
+
+import torch
+import torch.nn as nn
+
+
+class ModulateDiT(nn.Module):
+    """Modulation layer for DiT."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        factor: int,
+        act_layer: Callable,
+        dtype=None,
+        device=None,
+    ):
+        factory_kwargs = {"dtype": dtype, "device": device}
+        super().__init__()
+        self.act = act_layer()
+        self.linear = nn.Linear(hidden_size, factor * hidden_size, bias=True, **factory_kwargs)
+        # Zero-initialize the modulation
+        nn.init.zeros_(self.linear.weight)
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(self.act(x))
+
+
+def modulate(x, shift=None, scale=None):
+    """modulate by shift and scale
+
+    Args:
+        x (torch.Tensor): input tensor.
+        shift (torch.Tensor, optional): shift tensor. Defaults to None.
+        scale (torch.Tensor, optional): scale tensor. Defaults to None.
+
+    Returns:
+        torch.Tensor: the output tensor after modulate.
+    """
+    if scale is None and shift is None:
+        return x
+    elif shift is None:
+        return x * (1 + scale.unsqueeze(1))
+    elif scale is None:
+        return x + shift.unsqueeze(1)
+    else:
+        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def apply_gate(x, gate=None, tanh=False):
+    """AI is creating summary for apply_gate
+
+    Args:
+        x (torch.Tensor): input tensor.
+        gate (torch.Tensor, optional): gate tensor. Defaults to None.
+        tanh (bool, optional): whether to use tanh function. Defaults to False.
+
+    Returns:
+        torch.Tensor: the output tensor after apply gate.
+    """
+    if gate is None:
+        return x
+    if tanh:
+        return x * gate.unsqueeze(1).tanh()
+    else:
+        return x * gate.unsqueeze(1)
+
+
+def ckpt_wrapper(module):
+
+    def ckpt_forward(*inputs):
+        outputs = module(*inputs)
+        return outputs
+
+    return ckpt_forward
+
+
+class RMSNorm(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        elementwise_affine=True,
+        eps: float = 1e-6,
+        device=None,
+        dtype=None,
+    ):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.eps = eps
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        if hasattr(self, "weight"):
+            output = output * self.weight
+        return output
+
+
+def get_norm_layer(norm_layer):
+    """
+    Get the normalization layer.
+
+    Args:
+        norm_layer (str): The type of normalization layer.
+
+    Returns:
+        norm_layer (nn.Module): The normalization layer.
+    """
+    if norm_layer == "layer":
+        return nn.LayerNorm
+    elif norm_layer == "rms":
+        return RMSNorm
+    else:
+        raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")

fastvideo/models/hunyuan/modules/norm_layers.py

@@ -0,0 +1,78 @@
+import torch
+import torch.nn as nn
+
+
+class RMSNorm(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        elementwise_affine=True,
+        eps: float = 1e-6,
+        device=None,
+        dtype=None,
+    ):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.eps = eps
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        if hasattr(self, "weight"):
+            output = output * self.weight
+        return output
+
+
+def get_norm_layer(norm_layer):
+    """
+    Get the normalization layer.
+
+    Args:
+        norm_layer (str): The type of normalization layer.
+
+    Returns:
+        norm_layer (nn.Module): The normalization layer.
+    """
+    if norm_layer == "layer":
+        return nn.LayerNorm
+    elif norm_layer == "rms":
+        return RMSNorm
+    else:
+        raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")

fastvideo/models/hunyuan/modules/posemb_layers.py

@@ -0,0 +1,289 @@
+from typing import List, Tuple, Union
+
+import torch
+
+
+def _to_tuple(x, dim=2):
+    if isinstance(x, int):
+        return (x, ) * dim
+    elif len(x) == dim:
+        return x
+    else:
+        raise ValueError(f"Expected length {dim} or int, but got {x}")
+
+
+def get_meshgrid_nd(start, *args, dim=2):
+    """
+    Get n-D meshgrid with start, stop and num.
+
+    Args:
+        start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop,
+            step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num
+            should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in
+            n-tuples.
+        *args: See above.
+        dim (int): Dimension of the meshgrid. Defaults to 2.
+
+    Returns:
+        grid (np.ndarray): [dim, ...]
+    """
+    if len(args) == 0:
+        # start is grid_size
+        num = _to_tuple(start, dim=dim)
+        start = (0, ) * dim
+        stop = num
+    elif len(args) == 1:
+        # start is start, args[0] is stop, step is 1
+        start = _to_tuple(start, dim=dim)
+        stop = _to_tuple(args[0], dim=dim)
+        num = [stop[i] - start[i] for i in range(dim)]
+    elif len(args) == 2:
+        # start is start, args[0] is stop, args[1] is num
+        start = _to_tuple(start, dim=dim)  # Left-Top       eg: 12,0
+        stop = _to_tuple(args[0], dim=dim)  # Right-Bottom   eg: 20,32
+        num = _to_tuple(args[1], dim=dim)  # Target Size    eg: 32,124
+    else:
+        raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}")
+
+    # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False)
+    axis_grid = []
+    for i in range(dim):
+        a, b, n = start[i], stop[i], num[i]
+        g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n]
+        axis_grid.append(g)
+    grid = torch.meshgrid(*axis_grid, indexing="ij")  # dim x [W, H, D]
+    grid = torch.stack(grid, dim=0)  # [dim, W, H, D]
+
+    return grid
+
+
+#################################################################################
+#                   Rotary Positional Embedding Functions                       #
+#################################################################################
+# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80
+
+
+def reshape_for_broadcast(
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+    x: torch.Tensor,
+    head_first=False,
+):
+    """
+    Reshape frequency tensor for broadcasting it with another tensor.
+
+    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
+    for the purpose of broadcasting the frequency tensor during element-wise operations.
+
+    Notes:
+        When using FlashMHAModified, head_first should be False.
+        When using Attention, head_first should be True.
+
+    Args:
+        freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped.
+        x (torch.Tensor): Target tensor for broadcasting compatibility.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        torch.Tensor: Reshaped frequency tensor.
+
+    Raises:
+        AssertionError: If the frequency tensor doesn't match the expected shape.
+        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
+    """
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+
+    if isinstance(freqs_cis, tuple):
+        # freqs_cis: (cos, sin) in real space
+        if head_first:
+            assert freqs_cis[0].shape == (
+                x.shape[-2],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
+            shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        else:
+            assert freqs_cis[0].shape == (
+                x.shape[1],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
+            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
+    else:
+        # freqs_cis: values in complex space
+        if head_first:
+            assert freqs_cis.shape == (
+                x.shape[-2],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
+            shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        else:
+            assert freqs_cis.shape == (
+                x.shape[1],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
+            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        return freqs_cis.view(*shape)
+
+
+def rotate_half(x):
+    x_real, x_imag = (x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1))  # [B, S, H, D//2]
+    return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+    head_first: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor.
+
+    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
+    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
+    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
+    returned as real tensors.
+
+    Args:
+        xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D]
+        xk (torch.Tensor): Key tensor to apply rotary embeddings.   [B, S, H, D]
+        freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+
+    """
+    xk_out = None
+    if isinstance(freqs_cis, tuple):
+        cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first)  # [S, D]
+        cos, sin = cos.to(xq.device), sin.to(xq.device)
+        # real * cos - imag * sin
+        # imag * cos + real * sin
+        xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
+        xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk)
+    else:
+        # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex)
+        xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))  # [B, S, H, D//2]
+        freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device)  # [S, D//2] --> [1, S, 1, D//2]
+        # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin)
+        # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real)
+        xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq)
+        xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))  # [B, S, H, D//2]
+        xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk)
+
+    return xq_out, xk_out
+
+
+def get_nd_rotary_pos_embed(
+    rope_dim_list,
+    start,
+    *args,
+    theta=10000.0,
+    use_real=False,
+    theta_rescale_factor: Union[float, List[float]] = 1.0,
+    interpolation_factor: Union[float, List[float]] = 1.0,
+):
+    """
+    This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
+
+    Args:
+        rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n.
+            sum(rope_dim_list) should equal to head_dim of attention layer.
+        start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start,
+            args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num.
+        *args: See above.
+        theta (float): Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+            Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real
+            part and an imaginary part separately.
+        theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0.
+
+    Returns:
+        pos_embed (torch.Tensor): [HW, D/2]
+    """
+
+    grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list))  # [3, W, H, D] / [2, W, H]
+
+    if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float):
+        theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
+    elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
+        theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list)
+    assert len(theta_rescale_factor) == len(
+        rope_dim_list), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
+
+    if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float):
+        interpolation_factor = [interpolation_factor] * len(rope_dim_list)
+    elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
+        interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list)
+    assert len(interpolation_factor) == len(
+        rope_dim_list), "len(interpolation_factor) should equal to len(rope_dim_list)"
+
+    # use 1/ndim of dimensions to encode grid_axis
+    embs = []
+    for i in range(len(rope_dim_list)):
+        emb = get_1d_rotary_pos_embed(
+            rope_dim_list[i],
+            grid[i].reshape(-1),
+            theta,
+            use_real=use_real,
+            theta_rescale_factor=theta_rescale_factor[i],
+            interpolation_factor=interpolation_factor[i],
+        )  # 2 x [WHD, rope_dim_list[i]]
+        embs.append(emb)
+
+    if use_real:
+        cos = torch.cat([emb[0] for emb in embs], dim=1)  # (WHD, D/2)
+        sin = torch.cat([emb[1] for emb in embs], dim=1)  # (WHD, D/2)
+        return cos, sin
+    else:
+        emb = torch.cat(embs, dim=1)  # (WHD, D/2)
+        return emb
+
+
+def get_1d_rotary_pos_embed(
+    dim: int,
+    pos: Union[torch.FloatTensor, int],
+    theta: float = 10000.0,
+    use_real: bool = False,
+    theta_rescale_factor: float = 1.0,
+    interpolation_factor: float = 1.0,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Precompute the frequency tensor for complex exponential (cis) with given dimensions.
+    (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
+
+    This function calculates a frequency tensor with complex exponential using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+    The returned tensor contains complex values in complex64 data type.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (bool, optional): If True, return real part and imaginary part separately.
+                                   Otherwise, return complex numbers.
+        theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
+
+    Returns:
+        freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2]
+        freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
+    """
+    if isinstance(pos, int):
+        pos = torch.arange(pos).float()
+
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    if theta_rescale_factor != 1.0:
+        theta *= theta_rescale_factor**(dim / (dim - 2))
+
+    freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))  # [D/2]
+    # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}"
+    freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
+    if use_real:
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=1)  # [S, D]
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=1)  # [S, D]
+        return freqs_cos, freqs_sin
+    else:
+        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64     # [S, D/2]
+        return freqs_cis

fastvideo/models/hunyuan/modules/token_refiner.py

@@ -0,0 +1,202 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from .activation_layers import get_activation_layer
+from .attenion import attention
+from .embed_layers import TextProjection, TimestepEmbedder
+from .mlp_layers import MLP
+from .modulate_layers import apply_gate
+from .norm_layers import get_norm_layer
+
+
+class IndividualTokenRefinerBlock(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size,
+        heads_num,
+        mlp_width_ratio: str = 4.0,
+        mlp_drop_rate: float = 0.0,
+        act_type: str = "silu",
+        qk_norm: bool = False,
+        qk_norm_type: str = "layer",
+        qkv_bias: bool = True,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.heads_num = heads_num
+        head_dim = hidden_size // heads_num
+        mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
+
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+        self.self_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs)
+        qk_norm_layer = get_norm_layer(qk_norm_type)
+        self.self_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                                 if qk_norm else nn.Identity())
+        self.self_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                                 if qk_norm else nn.Identity())
+        self.self_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs)
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+        act_layer = get_activation_layer(act_type)
+        self.mlp = MLP(
+            in_channels=hidden_size,
+            hidden_channels=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=mlp_drop_rate,
+            **factory_kwargs,
+        )
+
+        self.adaLN_modulation = nn.Sequential(
+            act_layer(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
+        )
+        # Zero-initialize the modulation
+        nn.init.zeros_(self.adaLN_modulation[1].weight)
+        nn.init.zeros_(self.adaLN_modulation[1].bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        c: torch.Tensor,  # timestep_aware_representations + context_aware_representations
+        attn_mask: torch.Tensor = None,
+    ):
+        gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1)
+
+        norm_x = self.norm1(x)
+        qkv = self.self_attn_qkv(norm_x)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
+        # Apply QK-Norm if needed
+        q = self.self_attn_q_norm(q).to(v)
+        k = self.self_attn_k_norm(k).to(v)
+
+        # Self-Attention
+        attn = attention(q, k, v, attn_mask=attn_mask)
+
+        x = x + apply_gate(self.self_attn_proj(attn), gate_msa)
+
+        # FFN Layer
+        x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp)
+
+        return x
+
+
+class IndividualTokenRefiner(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size,
+        heads_num,
+        depth,
+        mlp_width_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        act_type: str = "silu",
+        qk_norm: bool = False,
+        qk_norm_type: str = "layer",
+        qkv_bias: bool = True,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.blocks = nn.ModuleList([
+            IndividualTokenRefinerBlock(
+                hidden_size=hidden_size,
+                heads_num=heads_num,
+                mlp_width_ratio=mlp_width_ratio,
+                mlp_drop_rate=mlp_drop_rate,
+                act_type=act_type,
+                qk_norm=qk_norm,
+                qk_norm_type=qk_norm_type,
+                qkv_bias=qkv_bias,
+                **factory_kwargs,
+            ) for _ in range(depth)
+        ])
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        c: torch.LongTensor,
+        mask: Optional[torch.Tensor] = None,
+    ):
+        mask = mask.clone().bool()
+        # avoid attention weight become NaN
+        mask[:, 0] = True
+        for block in self.blocks:
+            x = block(x, c, mask)
+        return x
+
+
+class SingleTokenRefiner(nn.Module):
+    """
+    A single token refiner block for llm text embedding refine.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_size,
+        heads_num,
+        depth,
+        mlp_width_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        act_type: str = "silu",
+        qk_norm: bool = False,
+        qk_norm_type: str = "layer",
+        qkv_bias: bool = True,
+        attn_mode: str = "torch",
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.attn_mode = attn_mode
+        assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner."
+
+        self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True, **factory_kwargs)
+
+        act_layer = get_activation_layer(act_type)
+        # Build timestep embedding layer
+        self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs)
+        # Build context embedding layer
+        self.c_embedder = TextProjection(in_channels, hidden_size, act_layer, **factory_kwargs)
+
+        self.individual_token_refiner = IndividualTokenRefiner(
+            hidden_size=hidden_size,
+            heads_num=heads_num,
+            depth=depth,
+            mlp_width_ratio=mlp_width_ratio,
+            mlp_drop_rate=mlp_drop_rate,
+            act_type=act_type,
+            qk_norm=qk_norm,
+            qk_norm_type=qk_norm_type,
+            qkv_bias=qkv_bias,
+            **factory_kwargs,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        t: torch.LongTensor,
+        mask: Optional[torch.LongTensor] = None,
+    ):
+        timestep_aware_representations = self.t_embedder(t)
+
+        if mask is None:
+            context_aware_representations = x.mean(dim=1)
+        else:
+            mask_float = mask.float().unsqueeze(-1)  # [b, s1, 1]
+            context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+        context_aware_representations = self.c_embedder(context_aware_representations)
+        c = timestep_aware_representations + context_aware_representations
+
+        x = self.input_embedder(x)
+
+        x = self.individual_token_refiner(x, c, mask)
+
+        return x

fastvideo/models/hunyuan/prompt_rewrite.py

@@ -0,0 +1,52 @@
+normal_mode_prompt = """Normal mode - Video Recaption Task:
+
+You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description.
+
+0. Preserve ALL information, including style words and technical terms.
+
+1. If the input is in Chinese, translate the entire description to English. 
+
+2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences.
+
+3. If the input does not include style, lighting, atmosphere, you can make reasonable associations.
+
+4. Output ALL must be in English.
+
+Given Input:
+input: "{input}"
+"""
+
+master_mode_prompt = """Master mode - Video Recaption Task:
+
+You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description.
+
+0. Preserve ALL information, including style words and technical terms.
+
+1. If the input is in Chinese, translate the entire description to English. 
+
+2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences.
+
+3. If the input does not include style, lighting, atmosphere, you can make reasonable associations.
+
+4. Output ALL must be in English.
+
+Given Input:
+input: "{input}"
+"""
+
+
+def get_rewrite_prompt(ori_prompt, mode="Normal"):
+    if mode == "Normal":
+        prompt = normal_mode_prompt.format(input=ori_prompt)
+    elif mode == "Master":
+        prompt = master_mode_prompt.format(input=ori_prompt)
+    else:
+        raise Exception("Only supports Normal and Normal", mode)
+    return prompt
+
+
+ori_prompt = "一只小狗在草地上奔跑。"
+normal_prompt = get_rewrite_prompt(ori_prompt, mode="Normal")
+master_prompt = get_rewrite_prompt(ori_prompt, mode="Master")
+
+# Then you can use the normal_prompt or master_prompt to access the hunyuan-large rewrite model to get the final prompt.

fastvideo/models/hunyuan/text_encoder/__init__.py

@@ -0,0 +1,323 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer, CLIPTextModel, CLIPTokenizer
+from transformers.utils import ModelOutput
+
+from ..constants import PRECISION_TO_TYPE, TEXT_ENCODER_PATH, TOKENIZER_PATH
+
+
+def use_default(value, default):
+    return value if value is not None else default
+
+
+def load_text_encoder(
+    text_encoder_type,
+    text_encoder_precision=None,
+    text_encoder_path=None,
+    logger=None,
+    device=None,
+):
+    if text_encoder_path is None:
+        text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type]
+    if logger is not None:
+        logger.info(f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}")
+
+    if text_encoder_type == "clipL":
+        text_encoder = CLIPTextModel.from_pretrained(text_encoder_path)
+        text_encoder.final_layer_norm = text_encoder.text_model.final_layer_norm
+    elif text_encoder_type == "llm":
+        text_encoder = AutoModel.from_pretrained(text_encoder_path, low_cpu_mem_usage=True)
+        text_encoder.final_layer_norm = text_encoder.norm
+    else:
+        raise ValueError(f"Unsupported text encoder type: {text_encoder_type}")
+    # from_pretrained will ensure that the model is in eval mode.
+
+    if text_encoder_precision is not None:
+        text_encoder = text_encoder.to(dtype=PRECISION_TO_TYPE[text_encoder_precision])
+
+    text_encoder.requires_grad_(False)
+
+    if logger is not None:
+        logger.info(f"Text encoder to dtype: {text_encoder.dtype}")
+
+    if device is not None:
+        text_encoder = text_encoder.to(device)
+
+    return text_encoder, text_encoder_path
+
+
+def load_tokenizer(tokenizer_type, tokenizer_path=None, padding_side="right", logger=None):
+    if tokenizer_path is None:
+        tokenizer_path = TOKENIZER_PATH[tokenizer_type]
+    if logger is not None:
+        logger.info(f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}")
+
+    if tokenizer_type == "clipL":
+        tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path, max_length=77)
+    elif tokenizer_type == "llm":
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, padding_side=padding_side)
+    else:
+        raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}")
+
+    return tokenizer, tokenizer_path
+
+
+@dataclass
+class TextEncoderModelOutput(ModelOutput):
+    """
+    Base class for model's outputs that also contains a pooling of the last hidden states.
+
+    Args:
+        hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
+        hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        text_outputs (`list`, *optional*, returned when `return_texts=True` is passed):
+            List of decoded texts.
+    """
+
+    hidden_state: torch.FloatTensor = None
+    attention_mask: Optional[torch.LongTensor] = None
+    hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None
+    text_outputs: Optional[list] = None
+
+
+class TextEncoder(nn.Module):
+
+    def __init__(
+        self,
+        text_encoder_type: str,
+        max_length: int,
+        text_encoder_precision: Optional[str] = None,
+        text_encoder_path: Optional[str] = None,
+        tokenizer_type: Optional[str] = None,
+        tokenizer_path: Optional[str] = None,
+        output_key: Optional[str] = None,
+        use_attention_mask: bool = True,
+        input_max_length: Optional[int] = None,
+        prompt_template: Optional[dict] = None,
+        prompt_template_video: Optional[dict] = None,
+        hidden_state_skip_layer: Optional[int] = None,
+        apply_final_norm: bool = False,
+        reproduce: bool = False,
+        logger=None,
+        device=None,
+    ):
+        super().__init__()
+        self.text_encoder_type = text_encoder_type
+        self.max_length = max_length
+        self.precision = text_encoder_precision
+        self.model_path = text_encoder_path
+        self.tokenizer_type = (tokenizer_type if tokenizer_type is not None else text_encoder_type)
+        self.tokenizer_path = (tokenizer_path if tokenizer_path is not None else text_encoder_path)
+        self.use_attention_mask = use_attention_mask
+        if prompt_template_video is not None:
+            assert (use_attention_mask is True), "Attention mask is True required when training videos."
+        self.input_max_length = (input_max_length if input_max_length is not None else max_length)
+        self.prompt_template = prompt_template
+        self.prompt_template_video = prompt_template_video
+        self.hidden_state_skip_layer = hidden_state_skip_layer
+        self.apply_final_norm = apply_final_norm
+        self.reproduce = reproduce
+        self.logger = logger
+
+        self.use_template = self.prompt_template is not None
+        if self.use_template:
+            assert (isinstance(self.prompt_template, dict) and "template" in self.prompt_template
+                    ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}"
+            assert "{}" in str(self.prompt_template["template"]), (
+                "`prompt_template['template']` must contain a placeholder `{}` for the input text, "
+                f"got {self.prompt_template['template']}")
+
+        self.use_video_template = self.prompt_template_video is not None
+        if self.use_video_template:
+            if self.prompt_template_video is not None:
+                assert (
+                    isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video
+                ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}"
+            assert "{}" in str(self.prompt_template_video["template"]), (
+                "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, "
+                f"got {self.prompt_template_video['template']}")
+
+        if "t5" in text_encoder_type:
+            self.output_key = output_key or "last_hidden_state"
+        elif "clip" in text_encoder_type:
+            self.output_key = output_key or "pooler_output"
+        elif "llm" in text_encoder_type or "glm" in text_encoder_type:
+            self.output_key = output_key or "last_hidden_state"
+        else:
+            raise ValueError(f"Unsupported text encoder type: {text_encoder_type}")
+
+        self.model, self.model_path = load_text_encoder(
+            text_encoder_type=self.text_encoder_type,
+            text_encoder_precision=self.precision,
+            text_encoder_path=self.model_path,
+            logger=self.logger,
+            device=device,
+        )
+        self.dtype = self.model.dtype
+        self.device = self.model.device
+
+        self.tokenizer, self.tokenizer_path = load_tokenizer(
+            tokenizer_type=self.tokenizer_type,
+            tokenizer_path=self.tokenizer_path,
+            padding_side="right",
+            logger=self.logger,
+        )
+
+    def __repr__(self):
+        return f"{self.text_encoder_type} ({self.precision} - {self.model_path})"
+
+    @staticmethod
+    def apply_text_to_template(text, template, prevent_empty_text=True):
+        """
+        Apply text to template.
+
+        Args:
+            text (str): Input text.
+            template (str or list): Template string or list of chat conversation.
+            prevent_empty_text (bool): If True, we will prevent the user text from being empty
+                by adding a space. Defaults to True.
+        """
+        if isinstance(template, str):
+            # Will send string to tokenizer. Used for llm
+            return template.format(text)
+        else:
+            raise TypeError(f"Unsupported template type: {type(template)}")
+
+    def text2tokens(self, text, data_type="image"):
+        """
+        Tokenize the input text.
+
+        Args:
+            text (str or list): Input text.
+        """
+        tokenize_input_type = "str"
+        if self.use_template:
+            if data_type == "image":
+                prompt_template = self.prompt_template["template"]
+            elif data_type == "video":
+                prompt_template = self.prompt_template_video["template"]
+            else:
+                raise ValueError(f"Unsupported data type: {data_type}")
+            if isinstance(text, (list, tuple)):
+                text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text]
+                if isinstance(text[0], list):
+                    tokenize_input_type = "list"
+            elif isinstance(text, str):
+                text = self.apply_text_to_template(text, prompt_template)
+                if isinstance(text, list):
+                    tokenize_input_type = "list"
+            else:
+                raise TypeError(f"Unsupported text type: {type(text)}")
+
+        kwargs = dict(
+            truncation=True,
+            max_length=self.max_length,
+            padding="max_length",
+            return_tensors="pt",
+        )
+        if tokenize_input_type == "str":
+            return self.tokenizer(
+                text,
+                return_length=False,
+                return_overflowing_tokens=False,
+                return_attention_mask=True,
+                **kwargs,
+            )
+        elif tokenize_input_type == "list":
+            return self.tokenizer.apply_chat_template(
+                text,
+                add_generation_prompt=True,
+                tokenize=True,
+                return_dict=True,
+                **kwargs,
+            )
+        else:
+            raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}")
+
+    def encode(
+        self,
+        batch_encoding,
+        use_attention_mask=None,
+        output_hidden_states=False,
+        do_sample=None,
+        hidden_state_skip_layer=None,
+        return_texts=False,
+        data_type="image",
+        device=None,
+    ):
+        """
+        Args:
+            batch_encoding (dict): Batch encoding from tokenizer.
+            use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask.
+                Defaults to None.
+            output_hidden_states (bool): Whether to output hidden states. If False, return the value of
+                self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer,
+                output_hidden_states will be set True. Defaults to False.
+            do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None.
+                When self.produce is False, do_sample is set to True by default.
+            hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer.
+                If None, self.output_key will be used. Defaults to None.
+            return_texts (bool): Whether to return the decoded texts. Defaults to False.
+        """
+        device = self.model.device if device is None else device
+        use_attention_mask = use_default(use_attention_mask, self.use_attention_mask)
+        hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer)
+        do_sample = use_default(do_sample, not self.reproduce)
+        attention_mask = (batch_encoding["attention_mask"].to(device) if use_attention_mask else None)
+        outputs = self.model(
+            input_ids=batch_encoding["input_ids"].to(device),
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None,
+        )
+        if hidden_state_skip_layer is not None:
+            last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)]
+            # Real last hidden state already has layer norm applied. So here we only apply it
+            # for intermediate layers.
+            if hidden_state_skip_layer > 0 and self.apply_final_norm:
+                last_hidden_state = self.model.final_layer_norm(last_hidden_state)
+        else:
+            last_hidden_state = outputs[self.output_key]
+
+        # Remove hidden states of instruction tokens, only keep prompt tokens.
+        if self.use_template:
+            if data_type == "image":
+                crop_start = self.prompt_template.get("crop_start", -1)
+            elif data_type == "video":
+                crop_start = self.prompt_template_video.get("crop_start", -1)
+            else:
+                raise ValueError(f"Unsupported data type: {data_type}")
+            if crop_start > 0:
+                last_hidden_state = last_hidden_state[:, crop_start:]
+                attention_mask = (attention_mask[:, crop_start:] if use_attention_mask else None)
+
+        if output_hidden_states:
+            return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states)
+        return TextEncoderModelOutput(last_hidden_state, attention_mask)
+
+    def forward(
+        self,
+        text,
+        use_attention_mask=None,
+        output_hidden_states=False,
+        do_sample=False,
+        hidden_state_skip_layer=None,
+        return_texts=False,
+    ):
+        batch_encoding = self.text2tokens(text)
+        return self.encode(
+            batch_encoding,
+            use_attention_mask=use_attention_mask,
+            output_hidden_states=output_hidden_states,
+            do_sample=do_sample,
+            hidden_state_skip_layer=hidden_state_skip_layer,
+            return_texts=return_texts,
+        )

fastvideo/models/hunyuan/utils/data_utils.py

@@ -0,0 +1,14 @@
+import math
+
+
+def align_to(value, alignment):
+    """align height, width according to alignment
+
+    Args:
+        value (int): height or width
+        alignment (int): target alignment factor
+
+    Returns:
+        int: the aligned value
+    """
+    return int(math.ceil(value / alignment) * alignment)

fastvideo/models/hunyuan/utils/file_utils.py

@@ -0,0 +1,71 @@
+import os
+from pathlib import Path
+
+import imageio
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+
+CODE_SUFFIXES = {
+    ".py",  # Python codes
+    ".sh",  # Shell scripts
+    ".yaml",
+    ".yml",  # Configuration files
+}
+
+
+def safe_dir(path):
+    """
+    Create a directory (or the parent directory of a file) if it does not exist.
+
+    Args:
+        path (str or Path): Path to the directory.
+
+    Returns:
+        path (Path): Path object of the directory.
+    """
+    path = Path(path)
+    path.mkdir(exist_ok=True, parents=True)
+    return path
+
+
+def safe_file(path):
+    """
+    Create the parent directory of a file if it does not exist.
+
+    Args:
+        path (str or Path): Path to the file.
+
+    Returns:
+        path (Path): Path object of the file.
+    """
+    path = Path(path)
+    path.parent.mkdir(exist_ok=True, parents=True)
+    return path
+
+
+def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=1, fps=24):
+    """save videos by video tensor
+       copy from https://github.com/guoyww/AnimateDiff/blob/e92bd5671ba62c0d774a32951453e328018b7c5b/animatediff/utils/util.py#L61
+
+    Args:
+        videos (torch.Tensor): video tensor predicted by the model
+        path (str): path to save video
+        rescale (bool, optional): rescale the video tensor from [-1, 1] to  . Defaults to False.
+        n_rows (int, optional): Defaults to 1.
+        fps (int, optional): video save fps. Defaults to 8.
+    """
+    videos = rearrange(videos, "b c t h w -> t b c h w")
+    outputs = []
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=n_rows)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+        if rescale:
+            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
+        x = torch.clamp(x, 0, 1)
+        x = (x * 255).numpy().astype(np.uint8)
+        outputs.append(x)
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    imageio.mimsave(path, outputs, fps=fps)

fastvideo/models/hunyuan/utils/helpers.py

@@ -0,0 +1,41 @@
+import collections.abc
+from itertools import repeat
+
+
+def _ntuple(n):
+
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            x = tuple(x)
+            if len(x) == 1:
+                x = tuple(repeat(x[0], n))
+            return x
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+
+
+def as_tuple(x):
+    if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+        return tuple(x)
+    if x is None or isinstance(x, (int, float, str)):
+        return (x, )
+    else:
+        raise ValueError(f"Unknown type {type(x)}")
+
+
+def as_list_of_2tuple(x):
+    x = as_tuple(x)
+    if len(x) == 1:
+        x = (x[0], x[0])
+    assert len(x) % 2 == 0, f"Expect even length, got {len(x)}."
+    lst = []
+    for i in range(0, len(x), 2):
+        lst.append((x[i], x[i + 1]))
+    return lst

fastvideo/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py

@@ -0,0 +1,41 @@
+import argparse
+
+import torch
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+
+
+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,
+    ).to(0)
+
+    model.language_model.save_pretrained(f"{args.output_dir}")
+    processor.tokenizer.save_pretrained(f"{args.output_dir}")
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="The path to the llava-llama-3-8b-v1_1-transformers.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="",
+        help="The output path of the llava-llama-3-8b-text-encoder-tokenizer."
+        "if '', the parent dir of output will be the same as input dir.",
+    )
+    args = parser.parse_args()
+
+    if len(args.output_dir) == 0:
+        args.output_dir = "/".join(args.input_dir.split("/")[:-1])
+
+    preprocess_text_encoder_tokenizer(args)

fastvideo/models/hunyuan/vae/__init__.py

@@ -0,0 +1,64 @@
+from pathlib import Path
+
+import torch
+
+from ..constants import PRECISION_TO_TYPE, VAE_PATH
+from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D
+
+
+def load_vae(
+    vae_type: str = "884-16c-hy",
+    vae_precision: str = None,
+    sample_size: tuple = None,
+    vae_path: str = None,
+    logger=None,
+    device=None,
+):
+    """the function to load the 3D VAE model
+
+    Args:
+        vae_type (str): the type of the 3D VAE model. Defaults to "884-16c-hy".
+        vae_precision (str, optional): the precision to load vae. Defaults to None.
+        sample_size (tuple, optional): the tiling size. Defaults to None.
+        vae_path (str, optional): the path to vae. Defaults to None.
+        logger (_type_, optional): logger. Defaults to None.
+        device (_type_, optional): device to load vae. Defaults to None.
+    """
+    if vae_path is None:
+        vae_path = VAE_PATH[vae_type]
+
+    if logger is not None:
+        logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}")
+    config = AutoencoderKLCausal3D.load_config(vae_path)
+    if sample_size:
+        vae = AutoencoderKLCausal3D.from_config(config, sample_size=sample_size)
+    else:
+        vae = AutoencoderKLCausal3D.from_config(config)
+
+    vae_ckpt = Path(vae_path) / "pytorch_model.pt"
+    assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}"
+
+    ckpt = torch.load(vae_ckpt, map_location=vae.device)
+    if "state_dict" in ckpt:
+        ckpt = ckpt["state_dict"]
+    if any(k.startswith("vae.") for k in ckpt.keys()):
+        ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")}
+    vae.load_state_dict(ckpt)
+
+    spatial_compression_ratio = vae.config.spatial_compression_ratio
+    time_compression_ratio = vae.config.time_compression_ratio
+
+    if vae_precision is not None:
+        vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision])
+
+    vae.requires_grad_(False)
+
+    if logger is not None:
+        logger.info(f"VAE to dtype: {vae.dtype}")
+
+    if device is not None:
+        vae = vae.to(device)
+
+    vae.eval()
+
+    return vae, vae_path, spatial_compression_ratio, time_compression_ratio

fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py

@@ -0,0 +1,764 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+from dataclasses import dataclass
+from math import prod
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+from fastvideo.utils.parallel_states import nccl_info
+
+try:
+    # This diffusers is modified and packed in the mirror.
+    from diffusers.loaders import FromOriginalVAEMixin
+except ImportError:
+    # Use this to be compatible with the original diffusers.
+    from diffusers.loaders.single_file_model import (
+        FromOriginalModelMixin as FromOriginalVAEMixin, )
+
+from diffusers.models.attention_processor import (ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, Attention,
+                                                  AttentionProcessor, AttnAddedKVProcessor, AttnProcessor)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+
+from .vae import BaseOutput, DecoderCausal3D, DecoderOutput, DiagonalGaussianDistribution, EncoderCausal3D
+
+
+@dataclass
+class DecoderOutput2(BaseOutput):
+    sample: torch.FloatTensor
+    posterior: Optional[DiagonalGaussianDistribution] = None
+
+
+class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin):
+    r"""
+    A VAE model with KL loss for encoding images/videos into latents and decoding latent representations into images/videos.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: Tuple[str] = ("DownEncoderBlockCausal3D", ),
+        up_block_types: Tuple[str] = ("UpDecoderBlockCausal3D", ),
+        block_out_channels: Tuple[int] = (64, ),
+        layers_per_block: int = 1,
+        act_fn: str = "silu",
+        latent_channels: int = 4,
+        norm_num_groups: int = 32,
+        sample_size: int = 32,
+        sample_tsize: int = 64,
+        scaling_factor: float = 0.18215,
+        force_upcast: float = True,
+        spatial_compression_ratio: int = 8,
+        time_compression_ratio: int = 4,
+        mid_block_add_attention: bool = True,
+    ):
+        super().__init__()
+
+        self.time_compression_ratio = time_compression_ratio
+
+        self.encoder = EncoderCausal3D(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+            time_compression_ratio=time_compression_ratio,
+            spatial_compression_ratio=spatial_compression_ratio,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.decoder = DecoderCausal3D(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            time_compression_ratio=time_compression_ratio,
+            spatial_compression_ratio=spatial_compression_ratio,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1)
+        self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1)
+
+        self.use_slicing = False
+        self.use_spatial_tiling = False
+        self.use_temporal_tiling = False
+        self.use_parallel = False
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_tsize = sample_tsize
+        self.tile_latent_min_tsize = sample_tsize // time_compression_ratio
+
+        self.tile_sample_min_size = self.config.sample_size
+        sample_size = (self.config.sample_size[0] if isinstance(self.config.sample_size,
+                                                                (list, tuple)) else self.config.sample_size)
+        self.tile_latent_min_size = int(sample_size / (2**(len(self.config.block_out_channels) - 1)))
+        self.tile_overlap_factor = 0.25
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (EncoderCausal3D, DecoderCausal3D)):
+            module.gradient_checkpointing = value
+
+    def enable_temporal_tiling(self, use_tiling: bool = True):
+        self.use_temporal_tiling = use_tiling
+
+    def disable_temporal_tiling(self):
+        self.enable_temporal_tiling(False)
+
+    def enable_spatial_tiling(self, use_tiling: bool = True):
+        self.use_spatial_tiling = use_tiling
+
+    def disable_spatial_tiling(self):
+        self.enable_spatial_tiling(False)
+
+    def enable_tiling(self, use_tiling: bool = True):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger videos.
+        """
+        self.enable_spatial_tiling(use_tiling)
+        self.enable_temporal_tiling(use_tiling)
+
+    def disable_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.disable_spatial_tiling()
+        self.disable_temporal_tiling()
+
+    def enable_parallel(self):
+        r"""
+        Enable sequence parallelism for the model. This will allow the vae to decode (with tiling) in parallel.
+        """
+        self.use_parallel = True
+
+    def enable_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.use_slicing = True
+
+    def disable_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_slicing = False
+
+    @property
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(
+        self,
+        processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]],
+        _remove_lora=False,
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes.")
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor, _remove_lora=_remove_lora)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnAddedKVProcessor()
+        elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor, _remove_lora=True)
+
+    @apply_forward_hook
+    def encode(self,
+               x: torch.FloatTensor,
+               return_dict: bool = True) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        """
+        Encode a batch of images/videos into latents.
+
+        Args:
+            x (`torch.FloatTensor`): Input batch of images/videos.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images/videos. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        assert len(x.shape) == 5, "The input tensor should have 5 dimensions."
+
+        if self.use_temporal_tiling and x.shape[2] > self.tile_sample_min_tsize:
+            return self.temporal_tiled_encode(x, return_dict=return_dict)
+
+        if self.use_spatial_tiling and (x.shape[-1] > self.tile_sample_min_size
+                                        or x.shape[-2] > self.tile_sample_min_size):
+            return self.spatial_tiled_encode(x, return_dict=return_dict)
+
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self.encoder(x)
+
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior, )
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        assert len(z.shape) == 5, "The input tensor should have 5 dimensions."
+
+        if self.use_parallel:
+            return self.parallel_tiled_decode(z, return_dict=return_dict)
+
+        if self.use_temporal_tiling and z.shape[2] > self.tile_latent_min_tsize:
+            return self.temporal_tiled_decode(z, return_dict=return_dict)
+
+        if self.use_spatial_tiling and (z.shape[-1] > self.tile_latent_min_size
+                                        or z.shape[-2] > self.tile_latent_min_size):
+            return self.spatial_tiled_decode(z, return_dict=return_dict)
+
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+
+        if not return_dict:
+            return (dec, )
+
+        return DecoderOutput(sample=dec)
+
+    @apply_forward_hook
+    def decode(self,
+               z: torch.FloatTensor,
+               return_dict: bool = True,
+               generator=None) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Decode a batch of images/videos.
+
+        Args:
+            z (`torch.FloatTensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+
+        """
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        if not return_dict:
+            return (decoded, )
+
+        return DecoderOutput(sample=decoded)
+
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :,
+              y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent)
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :,
+              x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent)
+        return b
+
+    def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :,
+              x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :,
+                                                                                       x, :, :] * (x / blend_extent)
+        return b
+
+    def spatial_tiled_encode(
+        self,
+        x: torch.FloatTensor,
+        return_dict: bool = True,
+        return_moments: bool = False,
+    ) -> AutoencoderKLOutput:
+        r"""Encode a batch of images/videos using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image/videos size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.FloatTensor`): Input batch of images/videos.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+                `tuple` is returned.
+        """
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split video into tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[-2], overlap_size):
+            row = []
+            for j in range(0, x.shape[-1], overlap_size):
+                tile = x[:, :, :, i:i + self.tile_sample_min_size, j:j + self.tile_sample_min_size, ]
+                tile = self.encoder(tile)
+                tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        moments = torch.cat(result_rows, dim=-2)
+        if return_moments:
+            return moments
+
+        posterior = DiagonalGaussianDistribution(moments)
+        if not return_dict:
+            return (posterior, )
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def spatial_tiled_decode(self,
+                             z: torch.FloatTensor,
+                             return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        r"""
+        Decode a batch of images/videos using a tiled decoder.
+
+        Args:
+            z (`torch.FloatTensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[-2], overlap_size):
+            row = []
+            for j in range(0, z.shape[-1], overlap_size):
+                tile = z[:, :, :, i:i + self.tile_latent_min_size, j:j + self.tile_latent_min_size, ]
+                tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=-2)
+        if not return_dict:
+            return (dec, )
+
+        return DecoderOutput(sample=dec)
+
+    def temporal_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput:
+
+        B, C, T, H, W = x.shape
+        overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor)
+        t_limit = self.tile_latent_min_tsize - blend_extent
+
+        # Split the video into tiles and encode them separately.
+        row = []
+        for i in range(0, T, overlap_size):
+            tile = x[:, :, i:i + self.tile_sample_min_tsize + 1, :, :]
+            if self.use_spatial_tiling and (tile.shape[-1] > self.tile_sample_min_size
+                                            or tile.shape[-2] > self.tile_sample_min_size):
+                tile = self.spatial_tiled_encode(tile, return_moments=True)
+            else:
+                tile = self.encoder(tile)
+                tile = self.quant_conv(tile)
+            if i > 0:
+                tile = tile[:, :, 1:, :, :]
+            row.append(tile)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :t_limit, :, :])
+            else:
+                result_row.append(tile[:, :, :t_limit + 1, :, :])
+
+        moments = torch.cat(result_row, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior, )
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def temporal_tiled_decode(self,
+                              z: torch.FloatTensor,
+                              return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        # Split z into overlapping tiles and decode them separately.
+
+        B, C, T, H, W = z.shape
+        overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor)
+        t_limit = self.tile_sample_min_tsize - blend_extent
+
+        row = []
+        for i in range(0, T, overlap_size):
+            tile = z[:, :, i:i + self.tile_latent_min_tsize + 1, :, :]
+            if self.use_spatial_tiling and (tile.shape[-1] > self.tile_latent_min_size
+                                            or tile.shape[-2] > self.tile_latent_min_size):
+                decoded = self.spatial_tiled_decode(tile, return_dict=True).sample
+            else:
+                tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+            if i > 0:
+                decoded = decoded[:, :, 1:, :, :]
+            row.append(decoded)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :t_limit, :, :])
+            else:
+                result_row.append(tile[:, :, :t_limit + 1, :, :])
+
+        dec = torch.cat(result_row, dim=2)
+        if not return_dict:
+            return (dec, )
+
+        return DecoderOutput(sample=dec)
+
+    def _parallel_data_generator(self, gathered_results, gathered_dim_metadata):
+        global_idx = 0
+        for i, per_rank_metadata in enumerate(gathered_dim_metadata):
+            _start_shape = 0
+            for shape in per_rank_metadata:
+                mul_shape = prod(shape)
+                yield (gathered_results[i, _start_shape:_start_shape + mul_shape].reshape(shape), global_idx)
+                _start_shape += mul_shape
+                global_idx += 1
+
+    def parallel_tiled_decode(self,
+                              z: torch.FloatTensor,
+                              return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs
+        """
+        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
+        B, C, T, H, W = z.shape
+
+        # Calculate parameters
+        t_overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor))
+        t_blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor)
+        t_limit = self.tile_sample_min_tsize - t_blend_extent
+
+        s_overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        s_blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        s_row_limit = self.tile_sample_min_size - s_blend_extent
+
+        # Calculate tile dimensions
+        num_t_tiles = (T + t_overlap_size - 1) // t_overlap_size
+        num_h_tiles = (H + s_overlap_size - 1) // s_overlap_size
+        num_w_tiles = (W + s_overlap_size - 1) // s_overlap_size
+        total_spatial_tiles = num_h_tiles * num_w_tiles
+        total_tiles = num_t_tiles * total_spatial_tiles
+
+        # Calculate tiles per rank and padding
+        tiles_per_rank = (total_tiles + world_size - 1) // world_size
+        start_tile_idx = rank * tiles_per_rank
+        end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles)
+
+        local_results = []
+        local_dim_metadata = []
+        # Process assigned tiles
+        for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)):
+            # Convert flat index to 3D indices
+            t_idx = global_idx // total_spatial_tiles
+            spatial_idx = global_idx % total_spatial_tiles
+            h_idx = spatial_idx // num_w_tiles
+            w_idx = spatial_idx % num_w_tiles
+
+            # Calculate positions
+            t_start = t_idx * t_overlap_size
+            h_start = h_idx * s_overlap_size
+            w_start = w_idx * s_overlap_size
+
+            # Extract and process tile
+            tile = z[:, :, t_start:t_start + self.tile_latent_min_tsize + 1,
+                     h_start:h_start + self.tile_latent_min_size, w_start:w_start + self.tile_latent_min_size]
+
+            # Process tile
+            tile = self.post_quant_conv(tile)
+            decoded = self.decoder(tile)
+
+            if t_start > 0:
+                decoded = decoded[:, :, 1:, :, :]
+
+            # Store metadata
+            shape = decoded.shape
+            # Store decoded data (flattened)
+            decoded_flat = decoded.reshape(-1)
+            local_results.append(decoded_flat)
+            local_dim_metadata.append(shape)
+
+        results = torch.cat(local_results, dim=0).contiguous()
+        del local_results
+        torch.cuda.empty_cache()
+        # first gather size to pad the results
+        local_size = torch.tensor([results.size(0)], device=results.device, dtype=torch.int64)
+        all_sizes = [torch.zeros(1, device=results.device, dtype=torch.int64) for _ in range(world_size)]
+        dist.all_gather(all_sizes, local_size)
+        max_size = max(size.item() for size in all_sizes)
+        padded_results = torch.zeros(max_size, device=results.device)
+        padded_results[:results.size(0)] = results
+        del results
+        torch.cuda.empty_cache()
+        # Gather all results
+        gathered_dim_metadata = [None] * world_size
+        gathered_results = torch.zeros_like(padded_results).repeat(
+            world_size, *[1] * len(padded_results.shape)).contiguous(
+            )  # use contiguous to make sure it won't copy data in the following operations
+        dist.all_gather_into_tensor(gathered_results, padded_results)
+        dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
+        # Process gathered results
+        data = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] for _ in range(num_t_tiles)]
+        for current_data, global_idx in self._parallel_data_generator(gathered_results, gathered_dim_metadata):
+            t_idx = global_idx // total_spatial_tiles
+            spatial_idx = global_idx % total_spatial_tiles
+            h_idx = spatial_idx // num_w_tiles
+            w_idx = spatial_idx % num_w_tiles
+            data[t_idx][h_idx][w_idx] = current_data
+        # Merge results
+        result_slices = []
+        last_slice_data = None
+        for i, tem_data in enumerate(data):
+            slice_data = self._merge_spatial_tiles(tem_data, s_blend_extent, s_row_limit)
+            if i > 0:
+                slice_data = self.blend_t(last_slice_data, slice_data, t_blend_extent)
+                result_slices.append(slice_data[:, :, :t_limit, :, :])
+            else:
+                result_slices.append(slice_data[:, :, :t_limit + 1, :, :])
+            last_slice_data = slice_data
+        dec = torch.cat(result_slices, dim=2)
+
+        if not return_dict:
+            return (dec, )
+        return DecoderOutput(sample=dec)
+
+    def _merge_spatial_tiles(self, spatial_rows, blend_extent, row_limit):
+        """Helper function to merge spatial tiles with blending"""
+        result_rows = []
+        for i, row in enumerate(spatial_rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_v(spatial_rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        return_posterior: bool = False,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput2, torch.FloatTensor]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        if not return_dict:
+            if return_posterior:
+                return (dec, posterior)
+            else:
+                return (dec, )
+        if return_posterior:
+            return DecoderOutput2(sample=dec, posterior=posterior)
+        else:
+            return DecoderOutput2(sample=dec)
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query,
+        key, value) are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)

fastvideo/models/hunyuan/vae/unet_causal_3d_blocks.py

@@ -0,0 +1,760 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import Attention, SpatialNorm
+from diffusers.models.normalization import AdaGroupNorm, RMSNorm
+from diffusers.utils import logging
+from einops import rearrange
+from torch import nn
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None):
+    seq_len = n_frame * n_hw
+    mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
+    for i in range(seq_len):
+        i_frame = i // n_hw
+        mask[i, :(i_frame + 1) * n_hw] = 0
+    if batch_size is not None:
+        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+    return mask
+
+
+class CausalConv3d(nn.Module):
+    """
+    Implements a causal 3D convolution layer where each position only depends on previous timesteps and current spatial locations.
+    This maintains temporal causality in video generation tasks.
+    """
+
+    def __init__(
+        self,
+        chan_in,
+        chan_out,
+        kernel_size: Union[int, Tuple[int, int, int]],
+        stride: Union[int, Tuple[int, int, int]] = 1,
+        dilation: Union[int, Tuple[int, int, int]] = 1,
+        pad_mode="replicate",
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.pad_mode = pad_mode
+        padding = (
+            kernel_size // 2,
+            kernel_size // 2,
+            kernel_size // 2,
+            kernel_size // 2,
+            kernel_size - 1,
+            0,
+        )  # W, H, T
+        self.time_causal_padding = padding
+
+        self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
+
+    def forward(self, x):
+        x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
+        return self.conv(x)
+
+
+class UpsampleCausal3D(nn.Module):
+    """
+    A 3D upsampling layer with an optional convolution.
+    """
+
+    def __init__(
+            self,
+            channels: int,
+            use_conv: bool = False,
+            use_conv_transpose: bool = False,
+            out_channels: Optional[int] = None,
+            name: str = "conv",
+            kernel_size: Optional[int] = None,
+            padding=1,
+            norm_type=None,
+            eps=None,
+            elementwise_affine=None,
+            bias=True,
+            interpolate=True,
+            upsample_factor=(2, 2, 2),
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+        self.interpolate = interpolate
+        self.upsample_factor = upsample_factor
+
+        if norm_type == "ln_norm":
+            self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(channels, eps, elementwise_affine)
+        elif norm_type is None:
+            self.norm = None
+        else:
+            raise ValueError(f"unknown norm_type: {norm_type}")
+
+        conv = None
+        if use_conv_transpose:
+            raise NotImplementedError
+        elif use_conv:
+            if kernel_size is None:
+                kernel_size = 3
+            conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, bias=bias)
+
+        if name == "conv":
+            self.conv = conv
+        else:
+            self.Conv2d_0 = conv
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        output_size: Optional[int] = None,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        assert hidden_states.shape[1] == self.channels
+
+        if self.norm is not None:
+            raise NotImplementedError
+
+        if self.use_conv_transpose:
+            return self.conv(hidden_states)
+
+        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.float32)
+
+        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+        if hidden_states.shape[0] >= 64:
+            hidden_states = hidden_states.contiguous()
+
+        # if `output_size` is passed we force the interpolation output
+        # size and do not make use of `scale_factor=2`
+        if self.interpolate:
+            B, C, T, H, W = hidden_states.shape
+            first_h, other_h = hidden_states.split((1, T - 1), dim=2)
+            if output_size is None:
+                if T > 1:
+                    other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest")
+
+                first_h = first_h.squeeze(2)
+                first_h = F.interpolate(first_h, scale_factor=self.upsample_factor[1:], mode="nearest")
+                first_h = first_h.unsqueeze(2)
+            else:
+                raise NotImplementedError
+
+            if T > 1:
+                hidden_states = torch.cat((first_h, other_h), dim=2)
+            else:
+                hidden_states = first_h
+
+        # If the input is bfloat16, we cast back to bfloat16
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+
+        if self.use_conv:
+            if self.name == "conv":
+                hidden_states = self.conv(hidden_states)
+            else:
+                hidden_states = self.Conv2d_0(hidden_states)
+
+        return hidden_states
+
+
+class DownsampleCausal3D(nn.Module):
+    """
+    A 3D downsampling layer with an optional convolution.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        use_conv: bool = False,
+        out_channels: Optional[int] = None,
+        padding: int = 1,
+        name: str = "conv",
+        kernel_size=3,
+        norm_type=None,
+        eps=None,
+        elementwise_affine=None,
+        bias=True,
+        stride=2,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.padding = padding
+        stride = stride
+        self.name = name
+
+        if norm_type == "ln_norm":
+            self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(channels, eps, elementwise_affine)
+        elif norm_type is None:
+            self.norm = None
+        else:
+            raise ValueError(f"unknown norm_type: {norm_type}")
+
+        if use_conv:
+            conv = CausalConv3d(
+                self.channels,
+                self.out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                bias=bias,
+            )
+        else:
+            raise NotImplementedError
+
+        if name == "conv":
+            self.Conv2d_0 = conv
+            self.conv = conv
+        elif name == "Conv2d_0":
+            self.conv = conv
+        else:
+            self.conv = conv
+
+    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
+        assert hidden_states.shape[1] == self.channels
+
+        if self.norm is not None:
+            hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+        assert hidden_states.shape[1] == self.channels
+
+        hidden_states = self.conv(hidden_states)
+
+        return hidden_states
+
+
+class ResnetBlockCausal3D(nn.Module):
+    r"""
+    A Resnet block.
+    """
+
+    def __init__(
+        self,
+        *,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 512,
+        groups: int = 32,
+        groups_out: Optional[int] = None,
+        pre_norm: bool = True,
+        eps: float = 1e-6,
+        non_linearity: str = "swish",
+        skip_time_act: bool = False,
+        # default, scale_shift, ada_group, spatial
+        time_embedding_norm: str = "default",
+        kernel: Optional[torch.FloatTensor] = None,
+        output_scale_factor: float = 1.0,
+        use_in_shortcut: Optional[bool] = None,
+        up: bool = False,
+        down: bool = False,
+        conv_shortcut_bias: bool = True,
+        conv_3d_out_channels: Optional[int] = None,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.pre_norm = True
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.up = up
+        self.down = down
+        self.output_scale_factor = output_scale_factor
+        self.time_embedding_norm = time_embedding_norm
+        self.skip_time_act = skip_time_act
+
+        linear_cls = nn.Linear
+
+        if groups_out is None:
+            groups_out = groups
+
+        if self.time_embedding_norm == "ada_group":
+            self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps)
+        elif self.time_embedding_norm == "spatial":
+            self.norm1 = SpatialNorm(in_channels, temb_channels)
+        else:
+            self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+
+        self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, stride=1)
+
+        if temb_channels is not None:
+            if self.time_embedding_norm == "default":
+                self.time_emb_proj = linear_cls(temb_channels, out_channels)
+            elif self.time_embedding_norm == "scale_shift":
+                self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels)
+            elif (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"):
+                self.time_emb_proj = None
+            else:
+                raise ValueError(f"Unknown time_embedding_norm : {self.time_embedding_norm} ")
+        else:
+            self.time_emb_proj = None
+
+        if self.time_embedding_norm == "ada_group":
+            self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps)
+        elif self.time_embedding_norm == "spatial":
+            self.norm2 = SpatialNorm(out_channels, temb_channels)
+        else:
+            self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+
+        self.dropout = torch.nn.Dropout(dropout)
+        conv_3d_out_channels = conv_3d_out_channels or out_channels
+        self.conv2 = CausalConv3d(out_channels, conv_3d_out_channels, kernel_size=3, stride=1)
+
+        self.nonlinearity = get_activation(non_linearity)
+
+        self.upsample = self.downsample = None
+        if self.up:
+            self.upsample = UpsampleCausal3D(in_channels, use_conv=False)
+        elif self.down:
+            self.downsample = DownsampleCausal3D(in_channels, use_conv=False, name="op")
+
+        self.use_in_shortcut = (self.in_channels != conv_3d_out_channels
+                                if use_in_shortcut is None else use_in_shortcut)
+
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = CausalConv3d(
+                in_channels,
+                conv_3d_out_channels,
+                kernel_size=1,
+                stride=1,
+                bias=conv_shortcut_bias,
+            )
+
+    def forward(
+        self,
+        input_tensor: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        hidden_states = input_tensor
+
+        if (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"):
+            hidden_states = self.norm1(hidden_states, temb)
+        else:
+            hidden_states = self.norm1(hidden_states)
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        if self.upsample is not None:
+            # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+            if hidden_states.shape[0] >= 64:
+                input_tensor = input_tensor.contiguous()
+                hidden_states = hidden_states.contiguous()
+            input_tensor = self.upsample(input_tensor, scale=scale)
+            hidden_states = self.upsample(hidden_states, scale=scale)
+        elif self.downsample is not None:
+            input_tensor = self.downsample(input_tensor, scale=scale)
+            hidden_states = self.downsample(hidden_states, scale=scale)
+
+        hidden_states = self.conv1(hidden_states)
+
+        if self.time_emb_proj is not None:
+            if not self.skip_time_act:
+                temb = self.nonlinearity(temb)
+            temb = self.time_emb_proj(temb, scale)[:, :, None, None]
+
+        if temb is not None and self.time_embedding_norm == "default":
+            hidden_states = hidden_states + temb
+
+        if (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"):
+            hidden_states = self.norm2(hidden_states, temb)
+        else:
+            hidden_states = self.norm2(hidden_states)
+
+        if temb is not None and self.time_embedding_norm == "scale_shift":
+            scale, shift = torch.chunk(temb, 2, dim=1)
+            hidden_states = hidden_states * (1 + scale) + shift
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+        return output_tensor
+
+
+def get_down_block3d(
+    down_block_type: str,
+    num_layers: int,
+    in_channels: int,
+    out_channels: int,
+    temb_channels: int,
+    add_downsample: bool,
+    downsample_stride: int,
+    resnet_eps: float,
+    resnet_act_fn: str,
+    transformer_layers_per_block: int = 1,
+    num_attention_heads: Optional[int] = None,
+    resnet_groups: Optional[int] = None,
+    cross_attention_dim: Optional[int] = None,
+    downsample_padding: Optional[int] = None,
+    dual_cross_attention: bool = False,
+    use_linear_projection: bool = False,
+    only_cross_attention: bool = False,
+    upcast_attention: bool = False,
+    resnet_time_scale_shift: str = "default",
+    attention_type: str = "default",
+    resnet_skip_time_act: bool = False,
+    resnet_out_scale_factor: float = 1.0,
+    cross_attention_norm: Optional[str] = None,
+    attention_head_dim: Optional[int] = None,
+    downsample_type: Optional[str] = None,
+    dropout: float = 0.0,
+):
+    # If attn head dim is not defined, we default it to the number of heads
+    if attention_head_dim is None:
+        logger.warn(
+            f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+        )
+        attention_head_dim = num_attention_heads
+
+    down_block_type = (down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type)
+    if down_block_type == "DownEncoderBlockCausal3D":
+        return DownEncoderBlockCausal3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            dropout=dropout,
+            add_downsample=add_downsample,
+            downsample_stride=downsample_stride,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block3d(
+    up_block_type: str,
+    num_layers: int,
+    in_channels: int,
+    out_channels: int,
+    prev_output_channel: int,
+    temb_channels: int,
+    add_upsample: bool,
+    upsample_scale_factor: Tuple,
+    resnet_eps: float,
+    resnet_act_fn: str,
+    resolution_idx: Optional[int] = None,
+    transformer_layers_per_block: int = 1,
+    num_attention_heads: Optional[int] = None,
+    resnet_groups: Optional[int] = None,
+    cross_attention_dim: Optional[int] = None,
+    dual_cross_attention: bool = False,
+    use_linear_projection: bool = False,
+    only_cross_attention: bool = False,
+    upcast_attention: bool = False,
+    resnet_time_scale_shift: str = "default",
+    attention_type: str = "default",
+    resnet_skip_time_act: bool = False,
+    resnet_out_scale_factor: float = 1.0,
+    cross_attention_norm: Optional[str] = None,
+    attention_head_dim: Optional[int] = None,
+    upsample_type: Optional[str] = None,
+    dropout: float = 0.0,
+) -> nn.Module:
+    # If attn head dim is not defined, we default it to the number of heads
+    if attention_head_dim is None:
+        logger.warn(
+            f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+        )
+        attention_head_dim = num_attention_heads
+
+    up_block_type = (up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type)
+    if up_block_type == "UpDecoderBlockCausal3D":
+        return UpDecoderBlockCausal3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            resolution_idx=resolution_idx,
+            dropout=dropout,
+            add_upsample=add_upsample,
+            upsample_scale_factor=upsample_scale_factor,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            temb_channels=temb_channels,
+        )
+    raise ValueError(f"{up_block_type} does not exist.")
+
+
+class UNetMidBlockCausal3D(nn.Module):
+    """
+    A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",  # default, spatial
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        attn_groups: Optional[int] = None,
+        resnet_pre_norm: bool = True,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+        output_scale_factor: float = 1.0,
+    ):
+        super().__init__()
+        resnet_groups = (resnet_groups if resnet_groups is not None else min(in_channels // 4, 32))
+        self.add_attention = add_attention
+
+        if attn_groups is None:
+            attn_groups = (resnet_groups if resnet_time_scale_shift == "default" else None)
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlockCausal3D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                temb_channels=temb_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                time_embedding_norm=resnet_time_scale_shift,
+                non_linearity=resnet_act_fn,
+                output_scale_factor=output_scale_factor,
+                pre_norm=resnet_pre_norm,
+            )
+        ]
+        attentions = []
+
+        if attention_head_dim is None:
+            logger.warn(
+                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
+            )
+            attention_head_dim = in_channels
+
+        for _ in range(num_layers):
+            if self.add_attention:
+                attentions.append(
+                    Attention(
+                        in_channels,
+                        heads=in_channels // attention_head_dim,
+                        dim_head=attention_head_dim,
+                        rescale_output_factor=output_scale_factor,
+                        eps=resnet_eps,
+                        norm_num_groups=attn_groups,
+                        spatial_norm_dim=(temb_channels if resnet_time_scale_shift == "spatial" else None),
+                        residual_connection=True,
+                        bias=True,
+                        upcast_softmax=True,
+                        _from_deprecated_attn_block=True,
+                    ))
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                ResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                ))
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+    def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            if attn is not None:
+                B, C, T, H, W = hidden_states.shape
+                hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c")
+                attention_mask = prepare_causal_attention_mask(T,
+                                                               H * W,
+                                                               hidden_states.dtype,
+                                                               hidden_states.device,
+                                                               batch_size=B)
+                hidden_states = attn(hidden_states, temb=temb, attention_mask=attention_mask)
+                hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W)
+            hidden_states = resnet(hidden_states, temb)
+
+        return hidden_states
+
+
+class DownEncoderBlockCausal3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor: float = 1.0,
+        add_downsample: bool = True,
+        downsample_stride: int = 2,
+        downsample_padding: int = 1,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=None,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                ))
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList([
+                DownsampleCausal3D(
+                    out_channels,
+                    use_conv=True,
+                    out_channels=out_channels,
+                    padding=downsample_padding,
+                    name="op",
+                    stride=downsample_stride,
+                )
+            ])
+        else:
+            self.downsamplers = None
+
+    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, temb=None, scale=scale)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states, scale)
+
+        return hidden_states
+
+
+class UpDecoderBlockCausal3D(nn.Module):
+
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            resolution_idx: Optional[int] = None,
+            dropout: float = 0.0,
+            num_layers: int = 1,
+            resnet_eps: float = 1e-6,
+            resnet_time_scale_shift: str = "default",  # default, spatial
+            resnet_act_fn: str = "swish",
+            resnet_groups: int = 32,
+            resnet_pre_norm: bool = True,
+            output_scale_factor: float = 1.0,
+            add_upsample: bool = True,
+            upsample_scale_factor=(2, 2, 2),
+            temb_channels: Optional[int] = None,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlockCausal3D(
+                    in_channels=input_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                ))
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([
+                UpsampleCausal3D(
+                    out_channels,
+                    use_conv=True,
+                    out_channels=out_channels,
+                    upsample_factor=upsample_scale_factor,
+                )
+            ])
+        else:
+            self.upsamplers = None
+
+        self.resolution_idx = resolution_idx
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: Optional[torch.FloatTensor] = None,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, temb=temb, scale=scale)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states

fastvideo/models/hunyuan/vae/vae.py

@@ -0,0 +1,342 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from diffusers.models.attention_processor import SpatialNorm
+from diffusers.utils import BaseOutput, is_torch_version
+from diffusers.utils.torch_utils import randn_tensor
+
+from .unet_causal_3d_blocks import CausalConv3d, UNetMidBlockCausal3D, get_down_block3d, get_up_block3d
+
+
+@dataclass
+class DecoderOutput(BaseOutput):
+    r"""
+    Output of decoding method.
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            The decoded output sample from the last layer of the model.
+    """
+
+    sample: torch.FloatTensor
+
+
+class EncoderCausal3D(nn.Module):
+    r"""
+    The `EncoderCausal3D` layer of a variational autoencoder that encodes its input into a latent representation.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D", ),
+        block_out_channels: Tuple[int, ...] = (64, ),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        double_z: bool = True,
+        mid_block_add_attention=True,
+        time_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+
+        self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1)
+        self.mid_block = None
+        self.down_blocks = nn.ModuleList([])
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_downsample_layers = int(np.log2(time_compression_ratio))
+
+            if time_compression_ratio == 4:
+                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+                add_time_downsample = bool(i >= (len(block_out_channels) - 1 - num_time_downsample_layers)
+                                           and not is_final_block)
+            else:
+                raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.")
+
+            downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
+            downsample_stride_T = (2, ) if add_time_downsample else (1, )
+            downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
+            down_block = get_down_block3d(
+                down_block_type,
+                num_layers=self.layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                add_downsample=bool(add_spatial_downsample or add_time_downsample),
+                downsample_stride=downsample_stride,
+                resnet_eps=1e-6,
+                downsample_padding=0,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attention_head_dim=output_channel,
+                temb_channels=None,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetMidBlockCausal3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            output_scale_factor=1,
+            resnet_time_scale_shift="default",
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            temb_channels=None,
+            add_attention=mid_block_add_attention,
+        )
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+        self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3)
+
+    def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor:
+        r"""The forward method of the `EncoderCausal3D` class."""
+        assert len(sample.shape) == 5, "The input tensor should have 5 dimensions"
+
+        sample = self.conv_in(sample)
+
+        # down
+        for down_block in self.down_blocks:
+            sample = down_block(sample)
+
+        # middle
+        sample = self.mid_block(sample)
+
+        # post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        return sample
+
+
+class DecoderCausal3D(nn.Module):
+    r"""
+    The `DecoderCausal3D` layer of a variational autoencoder that decodes its latent representation into an output sample.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D", ),
+        block_out_channels: Tuple[int, ...] = (64, ),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        norm_type: str = "group",  # group, spatial
+        mid_block_add_attention=True,
+        time_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+
+        self.conv_in = CausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1)
+        self.mid_block = None
+        self.up_blocks = nn.ModuleList([])
+
+        temb_channels = in_channels if norm_type == "spatial" else None
+
+        # mid
+        self.mid_block = UNetMidBlockCausal3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            output_scale_factor=1,
+            resnet_time_scale_shift="default" if norm_type == "group" else norm_type,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            temb_channels=temb_channels,
+            add_attention=mid_block_add_attention,
+        )
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_upsample_layers = int(np.log2(time_compression_ratio))
+
+            if time_compression_ratio == 4:
+                add_spatial_upsample = bool(i < num_spatial_upsample_layers)
+                add_time_upsample = bool(i >= len(block_out_channels) - 1 - num_time_upsample_layers
+                                         and not is_final_block)
+            else:
+                raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.")
+
+            upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1)
+            upsample_scale_factor_T = (2, ) if add_time_upsample else (1, )
+            upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW)
+            up_block = get_up_block3d(
+                up_block_type,
+                num_layers=self.layers_per_block + 1,
+                in_channels=prev_output_channel,
+                out_channels=output_channel,
+                prev_output_channel=None,
+                add_upsample=bool(add_spatial_upsample or add_time_upsample),
+                upsample_scale_factor=upsample_scale_factor,
+                resnet_eps=1e-6,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attention_head_dim=output_channel,
+                temb_channels=temb_channels,
+                resnet_time_scale_shift=norm_type,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        if norm_type == "spatial":
+            self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels)
+        else:
+            self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+        self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        latent_embeds: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        r"""The forward method of the `DecoderCausal3D` class."""
+        assert len(sample.shape) == 5, "The input tensor should have 5 dimensions."
+
+        sample = self.conv_in(sample)
+
+        upscale_dtype = next(iter(self.up_blocks.parameters())).dtype
+        if self.training and self.gradient_checkpointing:
+
+            def create_custom_forward(module):
+
+                def custom_forward(*inputs):
+                    return module(*inputs)
+
+                return custom_forward
+
+            if is_torch_version(">=", "1.11.0"):
+                # middle
+                sample = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(self.mid_block),
+                    sample,
+                    latent_embeds,
+                    use_reentrant=False,
+                )
+                sample = sample.to(upscale_dtype)
+
+                # up
+                for up_block in self.up_blocks:
+                    sample = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(up_block),
+                        sample,
+                        latent_embeds,
+                        use_reentrant=False,
+                    )
+            else:
+                # middle
+                sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample, latent_embeds)
+                sample = sample.to(upscale_dtype)
+
+                # up
+                for up_block in self.up_blocks:
+                    sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds)
+        else:
+            # middle
+            sample = self.mid_block(sample, latent_embeds)
+            sample = sample.to(upscale_dtype)
+
+            # up
+            for up_block in self.up_blocks:
+                sample = up_block(sample, latent_embeds)
+
+        # post-process
+        if latent_embeds is None:
+            sample = self.conv_norm_out(sample)
+        else:
+            sample = self.conv_norm_out(sample, latent_embeds)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        return sample
+
+
+class DiagonalGaussianDistribution(object):
+
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+        if parameters.ndim == 3:
+            dim = 2  # (B, L, C)
+        elif parameters.ndim == 5 or parameters.ndim == 4:
+            dim = 1  # (B, C, T, H ,W) / (B, C, H, W)
+        else:
+            raise NotImplementedError
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean,
+                                                   device=self.parameters.device,
+                                                   dtype=self.parameters.dtype)
+
+    def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor:
+        # make sure sample is on the same device as the parameters and has same dtype
+        sample = randn_tensor(
+            self.mean.shape,
+            generator=generator,
+            device=self.parameters.device,
+            dtype=self.parameters.dtype,
+        )
+        x = self.mean + self.std * sample
+        return x
+
+    def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            reduce_dim = list(range(1, self.mean.ndim))
+            if other is None:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=reduce_dim,
+                )
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar +
+                    other.logvar,
+                    dim=reduce_dim,
+                )
+
+    def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self) -> torch.Tensor:
+        return self.mean

fastvideo/models/hunyuan_hf/modeling_hunyuan.py

@@ -0,0 +1,836 @@
+# Copyright 2024 The Hunyuan Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_processor import Attention, AttentionProcessor
+from diffusers.models.embeddings import (CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings,
+                                         get_1d_rotary_pos_embed)
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+
+from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
+from fastvideo.utils.communications import all_gather, all_to_all_4D
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def shrink_head(encoder_state, dim):
+    local_heads = encoder_state.shape[dim] // nccl_info.sp_size
+    return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
+
+
+class HunyuanVideoAttnProcessor2_0:
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "HunyuanVideoAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        sequence_length = hidden_states.size(1)
+        encoder_sequence_length = encoder_hidden_states.size(1)
+        if attn.add_q_proj is None and encoder_hidden_states is not None:
+            hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
+
+        # 1. QKV projections
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+        key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+        value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+
+        # 2. QK normalization
+        if attn.norm_q is not None:
+            query = attn.norm_q(query).to(value)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key).to(value)
+
+        image_rotary_emb = (
+            shrink_head(image_rotary_emb[0], dim=0),
+            shrink_head(image_rotary_emb[1], dim=0),
+        )
+
+        # 3. Rotational positional embeddings applied to latent stream
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            if attn.add_q_proj is None and encoder_hidden_states is not None:
+                query = torch.cat(
+                    [
+                        apply_rotary_emb(query[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb),
+                        query[:, :, -encoder_hidden_states.shape[1]:],
+                    ],
+                    dim=2,
+                )
+                key = torch.cat(
+                    [
+                        apply_rotary_emb(key[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb),
+                        key[:, :, -encoder_hidden_states.shape[1]:],
+                    ],
+                    dim=2,
+                )
+            else:
+                query = apply_rotary_emb(query, image_rotary_emb)
+                key = apply_rotary_emb(key, image_rotary_emb)
+
+        # 4. Encoder condition QKV projection and normalization
+        if attn.add_q_proj is not None and encoder_hidden_states is not None:
+            encoder_query = attn.add_q_proj(encoder_hidden_states)
+            encoder_key = attn.add_k_proj(encoder_hidden_states)
+            encoder_value = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_query = encoder_query.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+            encoder_key = encoder_key.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+            encoder_value = encoder_value.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_query = attn.norm_added_q(encoder_query).to(encoder_value)
+            if attn.norm_added_k is not None:
+                encoder_key = attn.norm_added_k(encoder_key).to(encoder_value)
+
+            query = torch.cat([query, encoder_query], dim=2)
+            key = torch.cat([key, encoder_key], dim=2)
+            value = torch.cat([value, encoder_value], dim=2)
+
+        if get_sequence_parallel_state():
+            query_img, query_txt = query[:, :, :sequence_length, :], query[:, :, sequence_length:, :]
+            key_img, key_txt = key[:, :, :sequence_length, :], key[:, :, sequence_length:, :]
+            value_img, value_txt = value[:, :, :sequence_length, :], value[:, :, sequence_length:, :]
+            query_img = all_to_all_4D(query_img, scatter_dim=1, gather_dim=2)  #
+            key_img = all_to_all_4D(key_img, scatter_dim=1, gather_dim=2)
+            value_img = all_to_all_4D(value_img, scatter_dim=1, gather_dim=2)
+
+            query_txt = shrink_head(query_txt, dim=1)
+            key_txt = shrink_head(key_txt, dim=1)
+            value_txt = shrink_head(value_txt, dim=1)
+            query = torch.cat([query_img, query_txt], dim=2)
+            key = torch.cat([key_img, key_txt], dim=2)
+            value = torch.cat([value_img, value_txt], dim=2)
+
+        query = query.unsqueeze(2)
+        key = key.unsqueeze(2)
+        value = value.unsqueeze(2)
+        qkv = torch.cat([query, key, value], dim=2)
+        qkv = qkv.transpose(1, 3)
+
+        # 5. Attention
+        attention_mask = attention_mask[:, 0, :]
+        seq_len = qkv.shape[1]
+        attn_len = attention_mask.shape[1]
+        attention_mask = F.pad(attention_mask, (seq_len - attn_len, 0), value=True)
+
+        hidden_states = flash_attn_no_pad(qkv, attention_mask, causal=False, dropout_p=0.0, softmax_scale=None)
+
+        if get_sequence_parallel_state():
+            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
+                (sequence_length * nccl_info.sp_size, encoder_sequence_length), dim=1)
+            hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)
+            encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous()
+            hidden_states = hidden_states.flatten(2, 3)
+            hidden_states = hidden_states.to(query.dtype)
+            encoder_hidden_states = encoder_hidden_states.flatten(2, 3)
+            encoder_hidden_states = encoder_hidden_states.to(query.dtype)
+        else:
+            hidden_states = hidden_states.flatten(2, 3)
+            hidden_states = hidden_states.to(query.dtype)
+
+            # 6. Output projection
+            if encoder_hidden_states is not None:
+                hidden_states, encoder_hidden_states = (
+                    hidden_states[:, :-encoder_hidden_states.shape[1]],
+                    hidden_states[:, -encoder_hidden_states.shape[1]:],
+                )
+
+        if encoder_hidden_states is not None:
+            if getattr(attn, "to_out", None) is not None:
+                hidden_states = attn.to_out[0](hidden_states)
+                hidden_states = attn.to_out[1](hidden_states)
+
+            if getattr(attn, "to_add_out", None) is not None:
+                encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+        return hidden_states, encoder_hidden_states
+
+
+class HunyuanVideoPatchEmbed(nn.Module):
+
+    def __init__(
+        self,
+        patch_size: Union[int, Tuple[int, int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+    ) -> None:
+        super().__init__()
+
+        patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.proj(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)  # BCFHW -> BNC
+        return hidden_states
+
+
+class HunyuanVideoAdaNorm(nn.Module):
+
+    def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
+        super().__init__()
+
+        out_features = out_features or 2 * in_features
+        self.linear = nn.Linear(in_features, out_features)
+        self.nonlinearity = nn.SiLU()
+
+    def forward(self,
+                temb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        temb = self.linear(self.nonlinearity(temb))
+        gate_msa, gate_mlp = temb.chunk(2, dim=1)
+        gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
+        return gate_msa, gate_mlp
+
+
+class HunyuanVideoIndividualTokenRefinerBlock(nn.Module):
+
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_width_ratio: str = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
+        self.attn = Attention(
+            query_dim=hidden_size,
+            cross_attention_dim=None,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            bias=attention_bias,
+        )
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
+        self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate)
+
+        self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        norm_hidden_states = self.norm1(hidden_states)
+
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=attention_mask,
+        )
+
+        gate_msa, gate_mlp = self.norm_out(temb)
+        hidden_states = hidden_states + attn_output * gate_msa
+
+        ff_output = self.ff(self.norm2(hidden_states))
+        hidden_states = hidden_states + ff_output * gate_mlp
+
+        return hidden_states
+
+
+class HunyuanVideoIndividualTokenRefiner(nn.Module):
+
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        num_layers: int,
+        mlp_width_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.refiner_blocks = nn.ModuleList([
+            HunyuanVideoIndividualTokenRefinerBlock(
+                num_attention_heads=num_attention_heads,
+                attention_head_dim=attention_head_dim,
+                mlp_width_ratio=mlp_width_ratio,
+                mlp_drop_rate=mlp_drop_rate,
+                attention_bias=attention_bias,
+            ) for _ in range(num_layers)
+        ])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> None:
+        self_attn_mask = None
+        if attention_mask is not None:
+            batch_size = attention_mask.shape[0]
+            seq_len = attention_mask.shape[1]
+            attention_mask = attention_mask.to(hidden_states.device).bool()
+            self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
+            self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
+            self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
+            self_attn_mask[:, :, :, 0] = True
+
+        for block in self.refiner_blocks:
+            hidden_states = block(hidden_states, temb, self_attn_mask)
+
+        return hidden_states
+
+
+class HunyuanVideoTokenRefiner(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        num_layers: int,
+        mlp_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.time_text_embed = CombinedTimestepTextProjEmbeddings(embedding_dim=hidden_size,
+                                                                  pooled_projection_dim=in_channels)
+        self.proj_in = nn.Linear(in_channels, hidden_size, bias=True)
+        self.token_refiner = HunyuanVideoIndividualTokenRefiner(
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            num_layers=num_layers,
+            mlp_width_ratio=mlp_ratio,
+            mlp_drop_rate=mlp_drop_rate,
+            attention_bias=attention_bias,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        if attention_mask is None:
+            pooled_projections = hidden_states.mean(dim=1)
+        else:
+            original_dtype = hidden_states.dtype
+            mask_float = attention_mask.float().unsqueeze(-1)
+            pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+            pooled_projections = pooled_projections.to(original_dtype)
+
+        temb = self.time_text_embed(timestep, pooled_projections)
+        hidden_states = self.proj_in(hidden_states)
+        hidden_states = self.token_refiner(hidden_states, temb, attention_mask)
+
+        return hidden_states
+
+
+class HunyuanVideoRotaryPosEmbed(nn.Module):
+
+    def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None:
+        super().__init__()
+
+        self.patch_size = patch_size
+        self.patch_size_t = patch_size_t
+        self.rope_dim = rope_dim
+        self.theta = theta
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        rope_sizes = [
+            num_frames * nccl_info.sp_size // self.patch_size_t, height // self.patch_size, width // self.patch_size
+        ]
+
+        axes_grids = []
+        for i in range(3):
+            # Note: The following line diverges from original behaviour. We create the grid on the device, whereas
+            # original implementation creates it on CPU and then moves it to device. This results in numerical
+            # differences in layerwise debugging outputs, but visually it is the same.
+            grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32)
+            axes_grids.append(grid)
+        grid = torch.meshgrid(*axes_grids, indexing="ij")  # [W, H, T]
+        grid = torch.stack(grid, dim=0)  # [3, W, H, T]
+
+        freqs = []
+        for i in range(3):
+            freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True)
+            freqs.append(freq)
+
+        freqs_cos = torch.cat([f[0] for f in freqs], dim=1)  # (W * H * T, D / 2)
+        freqs_sin = torch.cat([f[1] for f in freqs], dim=1)  # (W * H * T, D / 2)
+        return freqs_cos, freqs_sin
+
+
+class HunyuanVideoSingleTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float = 4.0,
+        qk_norm: str = "rms_norm",
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+        mlp_dim = int(hidden_size * mlp_ratio)
+
+        self.attn = Attention(
+            query_dim=hidden_size,
+            cross_attention_dim=None,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=hidden_size,
+            bias=True,
+            processor=HunyuanVideoAttnProcessor2_0(),
+            qk_norm=qk_norm,
+            eps=1e-6,
+            pre_only=True,
+        )
+
+        self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm")
+        self.proj_mlp = nn.Linear(hidden_size, mlp_dim)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.shape[1]
+        hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
+
+        residual = hidden_states
+
+        # 1. Input normalization
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+
+        norm_hidden_states, norm_encoder_hidden_states = (
+            norm_hidden_states[:, :-text_seq_length, :],
+            norm_hidden_states[:, -text_seq_length:, :],
+        )
+
+        # 2. Attention
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            attention_mask=attention_mask,
+            image_rotary_emb=image_rotary_emb,
+        )
+        attn_output = torch.cat([attn_output, context_attn_output], dim=1)
+
+        # 3. Modulation and residual connection
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        hidden_states = gate.unsqueeze(1) * self.proj_out(hidden_states)
+        hidden_states = hidden_states + residual
+
+        hidden_states, encoder_hidden_states = (
+            hidden_states[:, :-text_seq_length, :],
+            hidden_states[:, -text_seq_length:, :],
+        )
+        return hidden_states, encoder_hidden_states
+
+
+class HunyuanVideoTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float,
+        qk_norm: str = "rms_norm",
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
+        self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
+
+        self.attn = Attention(
+            query_dim=hidden_size,
+            cross_attention_dim=None,
+            added_kv_proj_dim=hidden_size,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=hidden_size,
+            context_pre_only=False,
+            bias=True,
+            processor=HunyuanVideoAttnProcessor2_0(),
+            qk_norm=qk_norm,
+            eps=1e-6,
+        )
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
+
+        self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Input normalization
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+            encoder_hidden_states, emb=temb)
+
+        # 2. Joint attention
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            attention_mask=attention_mask,
+            image_rotary_emb=freqs_cis,
+        )
+
+        # 3. Modulation and residual connection
+        hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1)
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+        # 4. Feed-forward
+        ff_output = self.ff(norm_hidden_states)
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+
+        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+
+        return hidden_states, encoder_hidden_states
+
+
+class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    r"""
+    A Transformer model for video-like data used in [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo).
+
+    Args:
+        in_channels (`int`, defaults to `16`):
+            The number of channels in the input.
+        out_channels (`int`, defaults to `16`):
+            The number of channels in the output.
+        num_attention_heads (`int`, defaults to `24`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, defaults to `128`):
+            The number of channels in each head.
+        num_layers (`int`, defaults to `20`):
+            The number of layers of dual-stream blocks to use.
+        num_single_layers (`int`, defaults to `40`):
+            The number of layers of single-stream blocks to use.
+        num_refiner_layers (`int`, defaults to `2`):
+            The number of layers of refiner blocks to use.
+        mlp_ratio (`float`, defaults to `4.0`):
+            The ratio of the hidden layer size to the input size in the feedforward network.
+        patch_size (`int`, defaults to `2`):
+            The size of the spatial patches to use in the patch embedding layer.
+        patch_size_t (`int`, defaults to `1`):
+            The size of the tmeporal patches to use in the patch embedding layer.
+        qk_norm (`str`, defaults to `rms_norm`):
+            The normalization to use for the query and key projections in the attention layers.
+        guidance_embeds (`bool`, defaults to `True`):
+            Whether to use guidance embeddings in the model.
+        text_embed_dim (`int`, defaults to `4096`):
+            Input dimension of text embeddings from the text encoder.
+        pooled_projection_dim (`int`, defaults to `768`):
+            The dimension of the pooled projection of the text embeddings.
+        rope_theta (`float`, defaults to `256.0`):
+            The value of theta to use in the RoPE layer.
+        rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`):
+            The dimensions of the axes to use in the RoPE layer.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+            self,
+            in_channels: int = 16,
+            out_channels: int = 16,
+            num_attention_heads: int = 24,
+            attention_head_dim: int = 128,
+            num_layers: int = 20,
+            num_single_layers: int = 40,
+            num_refiner_layers: int = 2,
+            mlp_ratio: float = 4.0,
+            patch_size: int = 2,
+            patch_size_t: int = 1,
+            qk_norm: str = "rms_norm",
+            guidance_embeds: bool = True,
+            text_embed_dim: int = 4096,
+            pooled_projection_dim: int = 768,
+            rope_theta: float = 256.0,
+            rope_axes_dim: Tuple[int] = (16, 56, 56),
+    ) -> None:
+        super().__init__()
+
+        inner_dim = num_attention_heads * attention_head_dim
+        out_channels = out_channels or in_channels
+
+        # 1. Latent and condition embedders
+        self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
+        self.context_embedder = HunyuanVideoTokenRefiner(text_embed_dim,
+                                                         num_attention_heads,
+                                                         attention_head_dim,
+                                                         num_layers=num_refiner_layers)
+        self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim)
+
+        # 2. RoPE
+        self.rope = HunyuanVideoRotaryPosEmbed(patch_size, patch_size_t, rope_axes_dim, rope_theta)
+
+        # 3. Dual stream transformer blocks
+        self.transformer_blocks = nn.ModuleList([
+            HunyuanVideoTransformerBlock(num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm)
+            for _ in range(num_layers)
+        ])
+
+        # 4. Single stream transformer blocks
+        self.single_transformer_blocks = nn.ModuleList([
+            HunyuanVideoSingleTransformerBlock(num_attention_heads,
+                                               attention_head_dim,
+                                               mlp_ratio=mlp_ratio,
+                                               qk_norm=qk_norm) for _ in range(num_single_layers)
+        ])
+
+        # 5. Output projection
+        self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels)
+
+        self.gradient_checkpointing = False
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes.")
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_attention_mask: torch.Tensor,
+        guidance: torch.Tensor = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+        if guidance is None:
+            guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16)
+
+        if attention_kwargs is not None:
+            attention_kwargs = attention_kwargs.copy()
+            lora_scale = attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+                logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.")
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p, p_t = self.config.patch_size, self.config.patch_size_t
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        pooled_projections = encoder_hidden_states[:, 0, :self.config.pooled_projection_dim]
+        encoder_hidden_states = encoder_hidden_states[:, 1:]
+
+        # 1. RoPE
+        image_rotary_emb = self.rope(hidden_states)
+
+        # 2. Conditional embeddings
+        temb = self.time_text_embed(timestep, guidance, pooled_projections)
+        hidden_states = self.x_embedder(hidden_states)
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask)
+
+        # 3. Attention mask preparation
+        latent_sequence_length = hidden_states.shape[1]
+        condition_sequence_length = encoder_hidden_states.shape[1]
+        sequence_length = latent_sequence_length + condition_sequence_length
+        attention_mask = torch.zeros(batch_size,
+                                     sequence_length,
+                                     sequence_length,
+                                     device=hidden_states.device,
+                                     dtype=torch.bool)  # [B, N, N]
+
+        effective_condition_sequence_length = encoder_attention_mask.sum(dim=1, dtype=torch.int)
+        effective_sequence_length = latent_sequence_length + effective_condition_sequence_length
+
+        for i in range(batch_size):
+            attention_mask[i, :effective_sequence_length[i], :effective_sequence_length[i]] = True
+
+        # 4. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+            def create_custom_forward(module, return_dict=None):
+
+                def custom_forward(*inputs):
+                    if return_dict is not None:
+                        return module(*inputs, return_dict=return_dict)
+                    else:
+                        return module(*inputs)
+
+                return custom_forward
+
+            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+
+            for block in self.transformer_blocks:
+                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    attention_mask,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            for block in self.single_transformer_blocks:
+                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    attention_mask,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+        else:
+            for block in self.transformer_blocks:
+                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask,
+                                                             image_rotary_emb)
+
+            for block in self.single_transformer_blocks:
+                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask,
+                                                             image_rotary_emb)
+
+        # 5. Output projection
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
+                                              -1, p_t, p, p)
+        hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
+        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (hidden_states, )
+
+        return Transformer2DModelOutput(sample=hidden_states)

fastvideo/models/hunyuan_hf/pipeline_hunyuan.py

@@ -0,0 +1,691 @@
+# Copyright 2024 The HunyuanVideo Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.loaders import HunyuanVideoLoraLoaderMixin
+from diffusers.models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel
+from diffusers.pipelines.hunyuan_video.pipeline_output import HunyuanVideoPipelineOutput
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from einops import rearrange
+from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast
+
+from fastvideo.utils.communications import all_gather
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```python
+        >>> import torch
+        >>> from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
+        >>> from diffusers.utils import export_to_video
+
+        >>> model_id = "tencent/HunyuanVideo"
+        >>> transformer = HunyuanVideoTransformer3DModel.from_pretrained(
+        ...     model_id, subfolder="transformer", torch_dtype=torch.bfloat16
+        ... )
+        >>> pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.float16)
+        >>> pipe.vae.enable_tiling()
+        >>> pipe.to("cuda")
+
+        >>> output = pipe(
+        ...     prompt="A cat walks on the grass, realistic",
+        ...     height=320,
+        ...     width=512,
+        ...     num_frames=61,
+        ...     num_inference_steps=30,
+        ... ).frames[0]
+        >>> export_to_video(output, "output.mp4", fps=15)
+        ```
+"""
+
+DEFAULT_PROMPT_TEMPLATE = {
+    "template": ("<|start_header_id|>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 objects."
+                 "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
+                 "4. background environment, light, style and atmosphere."
+                 "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
+                 "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"),
+    "crop_start":
+    95,
+}
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler.")
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler.")
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class HunyuanVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin):
+    r"""
+    Pipeline for text-to-video generation using HunyuanVideo.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    Args:
+        text_encoder ([`LlamaModel`]):
+            [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers).
+        tokenizer_2 (`LlamaTokenizer`):
+            Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers).
+        transformer ([`HunyuanVideoTransformer3DModel`]):
+            Conditional Transformer to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKLHunyuanVideo`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
+        text_encoder_2 ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        tokenizer_2 (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        text_encoder: LlamaModel,
+        tokenizer: LlamaTokenizerFast,
+        transformer: HunyuanVideoTransformer3DModel,
+        vae: AutoencoderKLHunyuanVideo,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        text_encoder_2: CLIPTextModel,
+        tokenizer_2: CLIPTokenizer,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            scheduler=scheduler,
+            text_encoder_2=text_encoder_2,
+            tokenizer_2=tokenizer_2,
+        )
+
+        self.vae_scale_factor_temporal = (self.vae.temporal_compression_ratio
+                                          if hasattr(self, "vae") and self.vae is not None else 4)
+        self.vae_scale_factor_spatial = (self.vae.spatial_compression_ratio
+                                         if hasattr(self, "vae") and self.vae is not None else 8)
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+
+    def _get_llama_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_template: Dict[str, Any],
+        num_videos_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        max_sequence_length: int = 256,
+        num_hidden_layers_to_skip: int = 2,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        prompt = [prompt_template["template"].format(p) for p in prompt]
+
+        crop_start = prompt_template.get("crop_start", None)
+        if crop_start is None:
+            prompt_template_input = self.tokenizer(
+                prompt_template["template"],
+                padding="max_length",
+                return_tensors="pt",
+                return_length=False,
+                return_overflowing_tokens=False,
+                return_attention_mask=False,
+            )
+            crop_start = prompt_template_input["input_ids"].shape[-1]
+            # Remove <|eot_id|> token and placeholder {}
+            crop_start -= 2
+
+        max_sequence_length += crop_start
+        text_inputs = self.tokenizer(
+            prompt,
+            max_length=max_sequence_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_attention_mask=True,
+        )
+        text_input_ids = text_inputs.input_ids.to(device=device)
+        prompt_attention_mask = text_inputs.attention_mask.to(device=device)
+
+        prompt_embeds = self.text_encoder(
+            input_ids=text_input_ids,
+            attention_mask=prompt_attention_mask,
+            output_hidden_states=True,
+        ).hidden_states[-(num_hidden_layers_to_skip + 1)]
+        prompt_embeds = prompt_embeds.to(dtype=dtype)
+
+        if crop_start is not None and crop_start > 0:
+            prompt_embeds = prompt_embeds[:, crop_start:]
+            prompt_attention_mask = prompt_attention_mask[:, crop_start:]
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+        prompt_attention_mask = prompt_attention_mask.repeat(1, num_videos_per_prompt)
+        prompt_attention_mask = prompt_attention_mask.view(batch_size * num_videos_per_prompt, seq_len)
+
+        return prompt_embeds, prompt_attention_mask
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_videos_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        max_sequence_length: int = 77,
+    ) -> torch.Tensor:
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder_2.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer_2(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1])
+            logger.warning("The following part of your input was truncated because CLIP can only handle sequences up to"
+                           f" {max_sequence_length} tokens: {removed_text}")
+
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt)
+        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, -1)
+
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]] = None,
+        prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+        num_videos_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        max_sequence_length: int = 256,
+    ):
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds(
+                prompt,
+                prompt_template,
+                num_videos_per_prompt,
+                device=device,
+                dtype=dtype,
+                max_sequence_length=max_sequence_length,
+            )
+
+        if pooled_prompt_embeds is None:
+            if prompt_2 is None and pooled_prompt_embeds is None:
+                prompt_2 = prompt
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                prompt,
+                num_videos_per_prompt,
+                device=device,
+                dtype=dtype,
+                max_sequence_length=77,
+            )
+
+        return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        height,
+        width,
+        prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        prompt_template=None,
+    ):
+        if height % 16 != 0 or width % 16 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
+
+        if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs
+                                                                      for k in callback_on_step_end_tensor_inputs):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two.")
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two.")
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.")
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+
+        if prompt_template is not None:
+            if not isinstance(prompt_template, dict):
+                raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}")
+            if "template" not in prompt_template:
+                raise ValueError(
+                    f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}")
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels_latents: 32,
+        height: int = 720,
+        width: int = 1280,
+        num_frames: int = 129,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        shape = (
+            batch_size,
+            num_channels_latents,
+            num_frames,
+            int(height) // self.vae_scale_factor_spatial,
+            int(width) // self.vae_scale_factor_spatial,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                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)
+        return latents
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Union[str, List[str]] = None,
+        height: int = 720,
+        width: int = 1280,
+        num_frames: int = 129,
+        num_inference_steps: int = 50,
+        sigmas: List[float] = None,
+        guidance_scale: float = 6.0,
+        num_videos_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback,
+                                             MultiPipelineCallbacks]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+        max_sequence_length: int = 256,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead.
+            height (`int`, defaults to `720`):
+                The height in pixels of the generated image.
+            width (`int`, defaults to `1280`):
+                The width in pixels of the generated image.
+            num_frames (`int`, defaults to `129`):
+                The number of frames in the generated video.
+            num_inference_steps (`int`, defaults to `50`):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, defaults to `6.0`):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality. Note that the only available HunyuanVideo model is
+                CFG-distilled, which means that traditional guidance between unconditional and conditional latent is
+                not applied.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
+                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
+                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
+                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+
+        Examples:
+
+        Returns:
+            [`~HunyuanVideoPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned, otherwise a `tuple` is returned
+                where the first element is a list with the generated images and the second element is a list of `bool`s
+                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
+        """
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds,
+            callback_on_step_end_tensor_inputs,
+            prompt_template,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+
+        device = self._execution_device
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # 3. Encode input prompt
+        prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt,
+            prompt_template=prompt_template,
+            num_videos_per_prompt=num_videos_per_prompt,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            device=device,
+            max_sequence_length=max_sequence_length,
+        )
+
+        transformer_dtype = self.transformer.dtype
+        prompt_embeds = prompt_embeds.to(transformer_dtype)
+        prompt_attention_mask = prompt_attention_mask.to(transformer_dtype)
+        if pooled_prompt_embeds is not None:
+            pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype)
+
+        # 4. Prepare timesteps
+        sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+        )
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            num_latent_frames,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+        # check sequence_parallel
+        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
+        if get_sequence_parallel_state():
+            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
+            latents = latents[:, :, rank, :, :, :]
+
+        # 6. Prepare guidance condition
+        guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                latent_model_input = latents.to(transformer_dtype)
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+                if pooled_prompt_embeds.shape[-1] != prompt_embeds.shape[-1]:
+                    pooled_prompt_embeds_padding = F.pad(
+                        pooled_prompt_embeds,
+                        (0, prompt_embeds.shape[2] - pooled_prompt_embeds.shape[1]),
+                        value=0,
+                    ).unsqueeze(1)
+                encoder_hidden_states = torch.cat([pooled_prompt_embeds_padding, prompt_embeds], dim=1)
+
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    encoder_hidden_states=encoder_hidden_states,  # [1, 257, 4096]
+                    timestep=timestep,
+                    encoder_attention_mask=prompt_attention_mask,
+                    guidance=guidance,
+                    attention_kwargs=attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        if get_sequence_parallel_state():
+            latents = all_gather(latents, dim=2)
+
+        if not output_type == "latent":
+            latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor
+            video = self.vae.decode(latents, return_dict=False)[0]
+            video = self.video_processor.postprocess_video(video, output_type=output_type)
+        else:
+            video = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (video, )
+
+        return HunyuanVideoPipelineOutput(frames=video)

fastvideo/models/mochi_genmo/lib/attn_imports.py

@@ -1,29 +0,0 @@
-from contextlib import contextmanager
-
-import torch
-
-
-try:
-    from flash_attn import flash_attn_varlen_func as flash_varlen_attn
-except ImportError:
-    flash_varlen_attn = None
-
-try:
-    from sageattention import sageattn as sage_attn
-except ImportError:
-    sage_attn = None
-
-from torch.nn.attention import SDPBackend, sdpa_kernel
-
-training_backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]
-eval_backends = list(training_backends)
-if torch.cuda.get_device_properties(0).major >= 9.0:
-    # Enable fast CuDNN attention on Hopper.
-    # This gives NaN on the backward pass for some reason,
-    # so only use it for evaluation.
-    eval_backends.append(SDPBackend.CUDNN_ATTENTION)
-
-@contextmanager
-def sdpa_attn_ctx(training: bool = False):
-    with sdpa_kernel(training_backends if training else eval_backends):
-        yield

fastvideo/models/mochi_genmo/lib/progress.py

@@ -1,87 +0,0 @@
-import contextlib
-from typing import Any, Iterable, Iterator, Optional
-
-try:
-    from tqdm import tqdm
-except ImportError:
-    tqdm = None
-
-try:
-    from ray.experimental.tqdm_ray import tqdm as ray_tqdm
-except:
-    ray_tqdm = None
-
-# Global state
-_current_progress_type = "none"
-_is_progress_bar_active = False
-
-
-class DummyProgressBar:
-    """A no-op progress bar that mimics tqdm interface"""
-
-    def __init__(self, iterable=None, **kwargs):
-        self.iterable = iterable
-
-    def __iter__(self):
-        return iter(self.iterable)
-
-    def update(self, n=1):
-        pass
-
-    def close(self):
-        pass
-
-    def set_description(self, desc):
-        pass
-
-
-def get_new_progress_bar(iterable: Optional[Iterable] = None, **kwargs) -> Any:
-    if not _is_progress_bar_active:
-        return DummyProgressBar(iterable=iterable, **kwargs)
-
-    if _current_progress_type == "tqdm":
-        if tqdm is None:
-            raise ImportError("tqdm is required but not installed. Please install tqdm to use the tqdm progress bar.")
-        return tqdm(iterable=iterable, **kwargs)
-    elif _current_progress_type == "ray_tqdm":
-        if ray_tqdm is None:
-            raise ImportError("ray is required but not installed. Please install ray to use the ray_tqdm progress bar.")
-        return ray_tqdm(iterable=iterable, **kwargs)
-    return DummyProgressBar(iterable=iterable, **kwargs)
-
-
-@contextlib.contextmanager
-def progress_bar(type: str = "none", enabled=True):
-    """
-    Context manager for setting progress bar type and options.
-
-    Args:
-        type: Type of progress bar ("none" or "tqdm")
-        **options: Options to pass to the progress bar (e.g., total, desc)
-
-    Raises:
-        ValueError: If progress bar type is invalid
-        RuntimeError: If progress bars are nested
-
-    Example:
-        with progress_bar(type="tqdm", total=100):
-            for i in get_new_progress_bar(range(100)):
-                process(i)
-    """
-    if type not in ("none", "tqdm", "ray_tqdm"):
-        raise ValueError("Progress bar type must be 'none' or 'tqdm' or 'ray_tqdm'")
-    if not enabled:
-        type = "none"
-    global _current_progress_type, _is_progress_bar_active
-
-    if _is_progress_bar_active:
-        raise RuntimeError("Nested progress bars are not supported")
-
-    _is_progress_bar_active = True
-    _current_progress_type = type
-
-    try:
-        yield
-    finally:
-        _is_progress_bar_active = False
-        _current_progress_type = "none"

fastvideo/models/mochi_genmo/lib/utils.py

@@ -1,67 +0,0 @@
-import os
-import subprocess
-import tempfile
-import time
-
-import numpy as np
-from moviepy.editor import ImageSequenceClip
-from PIL import Image
-
-from genmo.lib.progress import get_new_progress_bar
-
-
-class Timer:
-    def __init__(self):
-        self.times = {}  # Dictionary to store times per stage
-
-    def __call__(self, name):
-        print(f"Timing {name}")
-        return self.TimerContextManager(self, name)
-
-    def print_stats(self):
-        total_time = sum(self.times.values())
-        # Print table header
-        print("{:<20} {:>10} {:>10}".format("Stage", "Time(s)", "Percent"))
-        for name, t in self.times.items():
-            percent = (t / total_time) * 100 if total_time > 0 else 0
-            print("{:<20} {:>10.2f} {:>9.2f}%".format(name, t, percent))
-
-    class TimerContextManager:
-        def __init__(self, outer, name):
-            self.outer = outer  # Reference to the Timer instance
-            self.name = name
-            self.start_time = None
-
-        def __enter__(self):
-            self.start_time = time.perf_counter()
-            return self
-
-        def __exit__(self, exc_type, exc_value, traceback):
-            end_time = time.perf_counter()
-            elapsed = end_time - self.start_time
-            self.outer.times[self.name] = self.outer.times.get(self.name, 0) + elapsed
-
-
-def save_video(final_frames, output_path, fps=30):
-    assert final_frames.ndim == 4 and final_frames.shape[3] == 3, f"invalid shape: {final_frames} (need t h w c)"
-    if final_frames.dtype != np.uint8:
-        final_frames = (final_frames * 255).astype(np.uint8)
-    ImageSequenceClip(list(final_frames), fps=fps).write_videofile(output_path)
-
-
-def create_memory_tracker():
-    import torch
-
-    previous = [None]  # Use list for mutable closure state
-
-    def track(label="all2all"):
-        current = torch.cuda.memory_allocated() / 1e9
-        if previous[0] is not None:
-            diff = current - previous[0]
-            sign = "+" if diff >= 0 else ""
-            print(f"GPU memory ({label}): {current:.2f} GB ({sign}{diff:.2f} GB)")
-        else:
-            print(f"GPU memory ({label}): {current:.2f} GB")
-        previous[0] = current  # type: ignore
-
-    return track

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py

@@ -1,721 +0,0 @@
-import os
-from typing import Dict, List, Optional, Tuple, Any
-import warnings
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from torch.nn.attention import sdpa_kernel
-
-from fastvideo.models.mochi_genmo.lib.attn_imports import flash_varlen_attn, sage_attn, sdpa_attn_ctx
-from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.layers import (
-    FeedForward,
-    PatchEmbed,
-    RMSNorm,
-    TimestepEmbedder,
-)
-from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.lora import LoraLinear
-from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.mod_rmsnorm import modulated_rmsnorm
-from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.residual_tanh_gated_rmsnorm import (
-    residual_tanh_gated_rmsnorm,
-)
-from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.rope_mixed import (
-    compute_mixed_rotation,
-    create_position_matrix,
-)
-from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.temporal_rope import apply_rotary_emb_qk_real
-from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.utils import (
-    AttentionPool,
-    modulate,
-    pad_and_split_xy,
-)
-from fastvideo.models.mochi_genmo.mochi_preview.pipelines import compute_packed_indices
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.loaders import PeftAdapterMixin
-
-COMPILE_FINAL_LAYER = os.environ.get("COMPILE_DIT") == "1"
-COMPILE_MMDIT_BLOCK = os.environ.get("COMPILE_DIT") == "1"
-
-
-def ck(fn, *args, enabled=True, **kwargs) -> torch.Tensor:
-    if enabled:
-        return torch.utils.checkpoint.checkpoint(fn, *args, **kwargs, use_reentrant=False)
-
-    return fn(*args, **kwargs)
-
-
-class AsymmetricAttention(nn.Module):
-    def __init__(
-        self,
-        dim_x: int,
-        dim_y: int,
-        num_heads: int = 8,
-        qkv_bias: bool = False,
-        qk_norm: bool = True,
-        update_y: bool = True,
-        out_bias: bool = True,
-        attention_mode: str = "flash",
-        softmax_scale: Optional[float] = None,
-        device: Optional[torch.device] = None,
-        # Disable LoRA by default ...
-        qkv_proj_lora_rank: int = 0,
-        qkv_proj_lora_alpha: int = 0,
-        qkv_proj_lora_dropout: float = 0.0,
-        out_proj_lora_rank: int = 0,
-        out_proj_lora_alpha: int = 0,
-        out_proj_lora_dropout: float = 0.0,
-    ):
-        super().__init__()
-        self.attention_mode = attention_mode
-        self.dim_x = dim_x
-        self.dim_y = dim_y
-        self.num_heads = num_heads
-        self.head_dim = dim_x // num_heads
-        self.update_y = update_y
-        self.softmax_scale = softmax_scale
-        if dim_x % num_heads != 0:
-            raise ValueError(f"dim_x={dim_x} should be divisible by num_heads={num_heads}")
-
-        # Input layers.
-        self.qkv_bias = qkv_bias
-        qkv_lora_kwargs = dict(
-            bias=qkv_bias,
-            device=device,
-            r=qkv_proj_lora_rank,
-            lora_alpha=qkv_proj_lora_alpha,
-            lora_dropout=qkv_proj_lora_dropout,
-        )
-        self.qkv_x = LoraLinear(dim_x, 3 * dim_x, **qkv_lora_kwargs)
-        # Project text features to match visual features (dim_y -> dim_x)
-        self.qkv_y = LoraLinear(dim_y, 3 * dim_x, **qkv_lora_kwargs)
-
-        # Query and key normalization for stability.
-        assert qk_norm
-        self.q_norm_x = RMSNorm(self.head_dim, device=device)
-        self.k_norm_x = RMSNorm(self.head_dim, device=device)
-        self.q_norm_y = RMSNorm(self.head_dim, device=device)
-        self.k_norm_y = RMSNorm(self.head_dim, device=device)
-
-        # Output layers. y features go back down from dim_x -> dim_y.
-        proj_lora_kwargs = dict(
-            bias=out_bias,
-            device=device,
-            r=out_proj_lora_rank,
-            lora_alpha=out_proj_lora_alpha,
-            lora_dropout=out_proj_lora_dropout,
-        )
-        self.proj_x = LoraLinear(dim_x, dim_x, **proj_lora_kwargs)
-        self.proj_y = LoraLinear(dim_x, dim_y, **proj_lora_kwargs) if update_y else nn.Identity()
-
-    def run_qkv_y(self, y):
-        local_heads = self.num_heads
-
-        qkv_y = self.qkv_y(y)  # (B, L, 3 * dim)
-
-        qkv_y = qkv_y.view(qkv_y.size(0), qkv_y.size(1), 3, local_heads, self.head_dim)
-        q_y, k_y, v_y = qkv_y.unbind(2)
-
-        q_y = self.q_norm_y(q_y)
-        k_y = self.k_norm_y(k_y)
-        return q_y, k_y, v_y
-
-    def prepare_qkv(
-        self,
-        x: torch.Tensor,  # (B, M, dim_x)
-        y: torch.Tensor,  # (B, L, dim_y)
-        *,
-        scale_x: torch.Tensor,
-        scale_y: torch.Tensor,
-        rope_cos: torch.Tensor,
-        rope_sin: torch.Tensor,
-        valid_token_indices: torch.Tensor,
-        max_seqlen_in_batch: int,
-    ):
-        # Process visual features
-        x = modulated_rmsnorm(x, scale_x)  # (B, M, dim_x) where M = N 
-        qkv_x = self.qkv_x(x)  # (B, M, 3 * dim_x)
-        assert qkv_x.dtype == torch.bfloat16
-        B, M, _ = qkv_x.size()
-        qkv_x = qkv_x.view(B, M, 3, self.num_heads, -1)
-        qkv_x = qkv_x.permute(2, 0, 1, 3, 4)
-
-        # Split qkv_x into q, k, v
-        q_x, k_x, v_x = qkv_x.unbind(0)  # (B, N, local_h, head_dim)
-        q_x = self.q_norm_x(q_x)
-        q_x = apply_rotary_emb_qk_real(q_x, rope_cos, rope_sin)
-        k_x = self.k_norm_x(k_x)
-        k_x = apply_rotary_emb_qk_real(k_x, rope_cos, rope_sin)
-
-        # Concatenate streams
-        B, N, num_heads, head_dim = q_x.size()
-        D = num_heads * head_dim
-
-        # Process text features
-        if B == 1:
-            text_seqlen = max_seqlen_in_batch - N
-            if text_seqlen > 0:
-                y = y[:, :text_seqlen]  # Remove padding tokens.
-                y = modulated_rmsnorm(y, scale_y)  # (B, L, dim_y)
-                q_y, k_y, v_y = self.run_qkv_y(y)  # (B, L, local_heads, head_dim)
-
-                q = torch.cat([q_x, q_y], dim=1)
-                k = torch.cat([k_x, k_y], dim=1)
-                v = torch.cat([v_x, v_y], dim=1)
-            else:
-                q, k, v = q_x, k_x, v_x
-        else:
-            y = modulated_rmsnorm(y, scale_y)  # (B, L, dim_y)
-            q_y, k_y, v_y = self.run_qkv_y(y)  # (B, L, local_heads, head_dim)
-
-            indices = valid_token_indices[:, None].expand(-1, D)
-            q = torch.cat([q_x, q_y], dim=1).view(-1, D).gather(0, indices)  # (total, D)
-            k = torch.cat([k_x, k_y], dim=1).view(-1, D).gather(0, indices)  # (total, D)
-            v = torch.cat([v_x, v_y], dim=1).view(-1, D).gather(0, indices)  # (total, D)
-
-        q = q.view(-1, num_heads, head_dim)
-        k = k.view(-1, num_heads, head_dim)
-        v = v.view(-1, num_heads, head_dim)
-        return q, k, v
-
-    @torch.autocast("cuda", enabled=False)
-    def flash_attention(self, q, k, v, cu_seqlens, max_seqlen_in_batch, total, local_dim):
-        out: torch.Tensor = flash_varlen_attn(
-            q, k, v,
-            cu_seqlens_q=cu_seqlens,
-            cu_seqlens_k=cu_seqlens,
-            max_seqlen_q=max_seqlen_in_batch,
-            max_seqlen_k=max_seqlen_in_batch,
-            dropout_p=0.0,
-            softmax_scale=self.softmax_scale,
-        )  # (total, local_heads, head_dim)
-        return out.view(total, local_dim)
-
-    def sdpa_attention(self, q, k, v):
-        with sdpa_attn_ctx(training=self.training):
-            out = F.scaled_dot_product_attention(
-                q, k, v,
-                attn_mask=None,
-                dropout_p=0.0,
-                is_causal=False,
-            )
-            return out
-
-    @torch.autocast("cuda", enabled=False)
-    def sage_attention(self, q, k, v):
-        return sage_attn(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
-
-    def run_attention(
-        self,
-        q: torch.Tensor,  # (total <= B * (N + L), num_heads, head_dim)
-        k: torch.Tensor,  # (total <= B * (N + L), num_heads, head_dim)
-        v: torch.Tensor,  # (total <= B * (N + L), num_heads, head_dim)
-        *,
-        B: int,
-        cu_seqlens: Optional[torch.Tensor] = None,
-        max_seqlen_in_batch: Optional[int] = None,
-    ):
-        local_heads = self.num_heads
-        local_dim = local_heads * self.head_dim
-
-        # Check shapes
-        assert q.ndim == 3 and k.ndim == 3 and v.ndim == 3
-        total = q.size(0)
-        assert k.size(0) == total and v.size(0) == total
-
-        if self.attention_mode == "flash":
-            out = self.flash_attention(
-                q, k, v, cu_seqlens, max_seqlen_in_batch, total, local_dim)  # (total, local_dim)
-        else:
-            assert B == 1, \
-                f"Non-flash attention mode {self.attention_mode} only supports batch size 1, got {B}"
-
-            q = rearrange(q, "(b s) h d -> b h s d", b=B)
-            k = rearrange(k, "(b s) h d -> b h s d", b=B)
-            v = rearrange(v, "(b s) h d -> b h s d", b=B)
-
-            if self.attention_mode == "sdpa":
-                out = self.sdpa_attention(q, k, v)  # (B, local_heads, seq_len, head_dim)
-            elif self.attention_mode == "sage":
-                out = self.sage_attention(q, k, v)  # (B, local_heads, seq_len, head_dim)
-            else:
-                raise ValueError(f"Unknown attention mode: {self.attention_mode}")
-
-            out = rearrange(out, "b h s d -> (b s) (h d)")
-
-        return out
-
-    def post_attention(
-        self,
-        out: torch.Tensor,
-        B: int,
-        M: int,
-        L: int,
-        dtype: torch.dtype,
-        valid_token_indices: torch.Tensor,
-    ):
-        """
-        Args:
-            out: (total <= B * (N + L), local_dim)
-            valid_token_indices: (total <= B * (N + L),)
-            B: Batch size
-            M: Number of visual tokens per context parallel rank
-            L: Number of text tokens
-            dtype: Data type of the input and output tensors
-
-        Returns:
-            x: (B, N, dim_x) tensor of visual tokens where N = M 
-            y: (B, L, dim_y) tensor of text token features
-        """
-        local_heads = self.num_heads
-        local_dim = local_heads * self.head_dim
-        N = M
-
-        # Split sequence into visual and text tokens, adding back padding.
-        if B == 1:
-            out = out.view(B, -1, local_dim)
-            if out.size(1) > N:
-                x, y = torch.tensor_split(out, (N,), dim=1)  # (B, N, local_dim), (B, <= L, local_dim)
-                y = F.pad(y, (0, 0, 0, L - y.size(1)))  # (B, L, local_dim)
-            else:
-                # Empty prompt.
-                x, y = out, out.new_zeros(B, L, local_dim)
-        else:
-            x, y = pad_and_split_xy(out, valid_token_indices, B, N, L, dtype)
-        assert x.size() == (B, N, local_dim)
-        assert y.size() == (B, L, local_dim)
-
-        # Communicate across context parallel ranks.
-        x = x.view(B, N, local_heads, self.head_dim)
-        x = x.view(x.size(0), x.size(1), x.size(2) * x.size(3))  # (B, M, dim_x = num_heads * head_dim)
-
-        x = self.proj_x(x)
-        y = self.proj_y(y)
-        return x, y
-
-    def forward(
-        self,
-        x: torch.Tensor,  # (B, M, dim_x)
-        y: torch.Tensor,  # (B, L, dim_y)
-        *,
-        scale_x: torch.Tensor,  # (B, dim_x), modulation for pre-RMSNorm.
-        scale_y: torch.Tensor,  # (B, dim_y), modulation for pre-RMSNorm.
-        packed_indices: Dict[str, torch.Tensor] = None,
-        checkpoint_qkv: bool = False,
-        checkpoint_post_attn: bool = False,
-        **rope_rotation,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Forward pass of asymmetric multi-modal attention.
-
-        Args:
-            x: (B, M, dim_x) tensor of visual tokens
-            y: (B, L, dim_y) tensor of text token features
-            packed_indices: Dict with keys for Flash Attention
-            num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens
-
-        Returns:
-            x: (B, M, dim_x) tensor of visual tokens after multi-modal attention
-            y: (B, L, dim_y) tensor of text token features after multi-modal attention
-        """
-        B, L, _ = y.shape
-        _, M, _ = x.shape
-
-        # Predict a packed QKV tensor from visual and text features.
-        q, k, v = ck(self.prepare_qkv,
-            x=x,
-            y=y,
-            scale_x=scale_x,
-            scale_y=scale_y,
-            rope_cos=rope_rotation.get("rope_cos"),
-            rope_sin=rope_rotation.get("rope_sin"),
-            valid_token_indices=packed_indices["valid_token_indices_kv"],
-            max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"],
-            enabled=checkpoint_qkv,
-        )  # (total <= B * (N + L), 3, local_heads, head_dim)
-
-        # Self-attention is expensive, so don't checkpoint it.
-        out = self.run_attention(
-            q, k, v, B=B,
-            cu_seqlens=packed_indices["cu_seqlens_kv"],
-            max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"],
-        )
-
-        x, y = ck(self.post_attention,
-            out,
-            B=B, M=M, L=L,
-            dtype=v.dtype,
-            valid_token_indices=packed_indices["valid_token_indices_kv"],
-            enabled=checkpoint_post_attn,
-        )
-
-        return x, y
-
-
-@torch.compile(disable=not COMPILE_MMDIT_BLOCK)
-class AsymmetricJointBlock(nn.Module):
-    def __init__(
-        self,
-        hidden_size_x: int,
-        hidden_size_y: int,
-        num_heads: int,
-        *,
-        mlp_ratio_x: float = 8.0,  # Ratio of hidden size to d_model for MLP for visual tokens.
-        mlp_ratio_y: float = 4.0,  # Ratio of hidden size to d_model for MLP for text tokens.
-        update_y: bool = True,  # Whether to update text tokens in this block.
-        device: Optional[torch.device] = None,
-        **block_kwargs,
-    ):
-        super().__init__()
-        self.update_y = update_y
-        self.hidden_size_x = hidden_size_x
-        self.hidden_size_y = hidden_size_y
-        self.mod_x = nn.Linear(hidden_size_x, 4 * hidden_size_x, device=device)
-        if self.update_y:
-            self.mod_y = nn.Linear(hidden_size_x, 4 * hidden_size_y, device=device)
-        else:
-            self.mod_y = nn.Linear(hidden_size_x, hidden_size_y, device=device)
-
-        # Self-attention:
-        self.attn = AsymmetricAttention(
-            hidden_size_x,
-            hidden_size_y,
-            num_heads=num_heads,
-            update_y=update_y,
-            device=device,
-            **block_kwargs,
-        )
-
-        # MLP.
-        mlp_hidden_dim_x = int(hidden_size_x * mlp_ratio_x)
-        assert mlp_hidden_dim_x == int(1536 * 8)
-        self.mlp_x = FeedForward(
-            in_features=hidden_size_x,
-            hidden_size=mlp_hidden_dim_x,
-            multiple_of=256,
-            ffn_dim_multiplier=None,
-            device=device,
-        )
-
-        # MLP for text not needed in last block.
-        if self.update_y:
-            mlp_hidden_dim_y = int(hidden_size_y * mlp_ratio_y)
-            self.mlp_y = FeedForward(
-                in_features=hidden_size_y,
-                hidden_size=mlp_hidden_dim_y,
-                multiple_of=256,
-                ffn_dim_multiplier=None,
-                device=device,
-            )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        c: torch.Tensor,
-        y: torch.Tensor,
-        # TODO: These could probably just go into attn_kwargs
-        checkpoint_ff: bool = False,
-        checkpoint_qkv: bool = False,
-        checkpoint_post_attn: bool = False,
-        **attn_kwargs,
-    ):
-        """Forward pass of a block.
-
-        Args:
-            x: (B, N, dim) tensor of visual tokens
-            c: (B, dim) tensor of conditioned features
-            y: (B, L, dim) tensor of text tokens
-            num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens
-
-        Returns:
-            x: (B, N, dim) tensor of visual tokens after block
-            y: (B, L, dim) tensor of text tokens after block
-        """
-        N = x.size(1)
-
-        c = F.silu(c)
-        mod_x = self.mod_x(c)
-        scale_msa_x, gate_msa_x, scale_mlp_x, gate_mlp_x = mod_x.chunk(4, dim=1)
-        mod_y = self.mod_y(c)
-
-        if self.update_y:
-            scale_msa_y, gate_msa_y, scale_mlp_y, gate_mlp_y = mod_y.chunk(4, dim=1)
-        else:
-            scale_msa_y = mod_y
-
-        # Self-attention block.
-        x_attn, y_attn = self.attn(
-            x,
-            y,
-            scale_x=scale_msa_x,
-            scale_y=scale_msa_y,
-            checkpoint_qkv=checkpoint_qkv,
-            checkpoint_post_attn=checkpoint_post_attn,
-            **attn_kwargs,
-        )
-
-        assert x_attn.size(1) == N
-        x = residual_tanh_gated_rmsnorm(x, x_attn, gate_msa_x)
-
-        if self.update_y:
-            y = residual_tanh_gated_rmsnorm(y, y_attn, gate_msa_y)
-
-        # MLP block.
-        x = ck(self.ff_block_x, x, scale_mlp_x, gate_mlp_x, enabled=checkpoint_ff)
-        if self.update_y:
-            y = ck(self.ff_block_y, y, scale_mlp_y, gate_mlp_y, enabled=checkpoint_ff)  # type: ignore
-        return x, y
-
-    def ff_block_x(self, x, scale_x, gate_x):
-        x_mod = modulated_rmsnorm(x, scale_x)
-        x_res = self.mlp_x(x_mod)
-        x = residual_tanh_gated_rmsnorm(x, x_res, gate_x)  # Sandwich norm
-        return x
-
-    def ff_block_y(self, y, scale_y, gate_y):
-        y_mod = modulated_rmsnorm(y, scale_y)
-        y_res = self.mlp_y(y_mod)
-        y = residual_tanh_gated_rmsnorm(y, y_res, gate_y)  # Sandwich norm
-        return y
-
-
-@torch.compile(disable=not COMPILE_FINAL_LAYER)
-class FinalLayer(nn.Module):
-    """
-    The final layer of DiT.
-    """
-
-    def __init__(
-        self,
-        hidden_size,
-        patch_size,
-        out_channels,
-        device: Optional[torch.device] = None,
-    ):
-        super().__init__()
-        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, device=device)
-        self.mod = nn.Linear(hidden_size, 2 * hidden_size, device=device)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, device=device)
-
-    def forward(self, x, c):
-        c = F.silu(c)
-        shift, scale = self.mod(c).chunk(2, dim=1)
-        x = modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        return x
-
-
-class AsymmDiTJoint(ModelMixin, ConfigMixin, PeftAdapterMixin):
-    """
-    Diffusion model with a Transformer backbone.
-
-    Ingests text embeddings instead of a label.
-    """
-    @register_to_config
-    def __init__(
-        self,
-        *,
-        patch_size=2,
-        in_channels=12,
-        hidden_size_x=3072,
-        hidden_size_y=1536,
-        depth=48,
-        num_heads=24,
-        mlp_ratio_x=4.0,
-        mlp_ratio_y=4.0,
-        t5_feat_dim: int = 4096,
-        t5_token_length: int = 256,
-        patch_embed_bias: bool = True,
-        timestep_mlp_bias: bool = True,
-        timestep_scale: float = 1000.0,
-        use_extended_posenc: bool = False,
-        rope_theta: float = 10000.0,
-        device: Optional[torch.device] = None,
-        **block_kwargs,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.patch_size = patch_size
-        self.num_heads = num_heads
-        self.hidden_size_x = hidden_size_x
-        self.hidden_size_y = hidden_size_y
-        self.head_dim = hidden_size_x // num_heads  # Head dimension and count is determined by visual.
-        self.use_extended_posenc = use_extended_posenc
-        self.t5_token_length = t5_token_length
-        self.t5_feat_dim = t5_feat_dim
-        self.rope_theta = rope_theta  # Scaling factor for frequency computation for temporal RoPE.
-        self.timestep_scale = timestep_scale
-
-        self.x_embedder = PatchEmbed(
-            patch_size=patch_size,
-            in_chans=in_channels,
-            embed_dim=hidden_size_x,
-            bias=patch_embed_bias,
-            device=device,
-        )
-        # Conditionings
-        # Timestep
-        self.t_embedder = TimestepEmbedder(hidden_size_x, bias=timestep_mlp_bias, timestep_scale=timestep_scale)
-
-        # Caption Pooling (T5)
-        self.t5_y_embedder = AttentionPool(t5_feat_dim, num_heads=8, output_dim=hidden_size_x, device=device)
-
-        # Dense Embedding Projection (T5)
-        self.t5_yproj = nn.Linear(t5_feat_dim, hidden_size_y, bias=True, device=device)
-
-        # Initialize pos_frequencies as an empty parameter.
-        self.pos_frequencies = nn.Parameter(torch.empty(3, self.num_heads, self.head_dim // 2, device=device))
-
-        # for depth 48:
-        #  b =  0: AsymmetricJointBlock, update_y=True
-        #  b =  1: AsymmetricJointBlock, update_y=True
-        #  ...
-        #  b = 46: AsymmetricJointBlock, update_y=True
-        #  b = 47: AsymmetricJointBlock, update_y=False. No need to update text features.
-        blocks = []
-        for b in range(depth):
-            # Joint multi-modal block
-            update_y = b < depth - 1
-            block = AsymmetricJointBlock(
-                hidden_size_x,
-                hidden_size_y,
-                num_heads,
-                mlp_ratio_x=mlp_ratio_x,
-                mlp_ratio_y=mlp_ratio_y,
-                update_y=update_y,
-                device=device,
-                **block_kwargs,
-            )
-
-            blocks.append(block)
-        self.blocks = nn.ModuleList(blocks)
-
-        self.final_layer = FinalLayer(hidden_size_x, patch_size, self.out_channels, device=device)
-
-    def embed_x(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x: (B, C=12, T, H, W) tensor of visual tokens
-
-        Returns:
-            x: (B, C=3072, N) tensor of visual tokens with positional embedding.
-        """
-        return self.x_embedder(x)  # Convert BcTHW to BCN
-
-    @torch.compile(disable=not COMPILE_MMDIT_BLOCK)
-    def prepare(
-        self,
-        x: torch.Tensor,
-        sigma: torch.Tensor,
-        t5_feat: torch.Tensor,
-        t5_mask: torch.Tensor,
-    ):
-        """Prepare input and conditioning embeddings."""
-
-        # Visual patch embeddings with positional encoding.
-        T, H, W = x.shape[-3:]
-        pH, pW = H // self.patch_size, W // self.patch_size
-        x = self.embed_x(x)  # (B, N, D), where N = T * H * W / patch_size ** 2
-        assert x.ndim == 3
-        B = x.size(0)
-
-        # Construct position array of size [N, 3].
-        # pos[:, 0] is the frame index for each location,
-        # pos[:, 1] is the row index for each location, and
-        # pos[:, 2] is the column index for each location.
-        N = T * pH * pW
-        assert x.size(1) == N
-        pos = create_position_matrix(T, pH=pH, pW=pW, device=x.device, dtype=torch.float32)  # (N, 3)
-        rope_cos, rope_sin = compute_mixed_rotation(
-            freqs=self.pos_frequencies, pos=pos
-        )  # Each are (N, num_heads, dim // 2)
-
-        # Global vector embedding for conditionings.
-        c_t = self.t_embedder(1 - sigma)  # (B, D)
-
-        # Pool T5 tokens using attention pooler
-        # Note encoder_hidden_states[1] contains T5 token features.
-        assert (
-            t5_feat.size(1) == self.t5_token_length
-        ), f"Expected L={self.t5_token_length}, got {t5_feat.shape} for encoder_hidden_states."
-        t5_y_pool = self.t5_y_embedder(t5_feat, t5_mask)  # (B, D)
-        assert t5_y_pool.size(0) == B, f"Expected B={B}, got {t5_y_pool.shape} for t5_y_pool."
-
-        c = c_t + t5_y_pool
-
-        encoder_hidden_states = self.t5_yproj(t5_feat)  # (B, L, t5_feat_dim) --> (B, L, D)
-
-        return x, c, encoder_hidden_states, rope_cos, rope_sin
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor, # [1, 12, 7, 60, 106]
-        timestep: torch.Tensor, # [1] 
-        encoder_hidden_states: torch.Tensor, # [1, 256, 4096]
-        encoder_attention_mask: torch.Tensor, # [1, 256]
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = False,
-        rope_cos: torch.Tensor = None, 
-        rope_sin: torch.Tensor = None, 
-        num_ff_checkpoint: int = 48, # 48
-        num_qkv_checkpoint: int = 48, # 48
-        num_post_attn_checkpoint: int = 0, # 0
-    ):
-        """Forward pass of DiT.
-
-        Args:
-            hidden_states: (B, C, T, H, W) tensor of spatial inputs (images or latent representations of images)
-            sigma: (B,) tensor of noise standard deviations
-            encoder_hidden_states: List((B, L, encoder_hidden_states_dim) tensor of caption token features. For SDXL text encoders: L=77, encoder_hidden_states_dim=2048)
-            encoder_attention_mask: List((B, L) boolean tensor indicating which tokens are not padding)
-            packed_indices: Dict with keys for Flash Attention. Result of compute_packed_indices.
-            {'cu_seqlens_kv': tensor([    0, 11230], device='cuda:0', dtype=torch.int32),
-                'max_seqlen_in_batch_kv': 11230,
-                'valid_token_indices_kv': tensor([    0,     1,     2,  ..., 11227, 11228, 11229], device='cuda:0')}
-        """
-        sigma = timestep / self.timestep_scale
-        num_latent_toks = np.prod(hidden_states.shape[-3:])
-
-        packed_indices = compute_packed_indices(hidden_states.device, encoder_attention_mask, int(num_latent_toks))
-
-        _, _, T, H, W = hidden_states.shape
-        
-        if self.pos_frequencies.dtype != torch.float32:
-            warnings.warn(f"pos_frequencies dtype {self.pos_frequencies.dtype} != torch.float32")
-
-        # Use EFFICIENT_ATTENTION backend for T5 pooling, since we have a mask.
-        # Have to call sdpa_kernel outside of a torch.compile region.
-        with sdpa_kernel(torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION):
-            hidden_states, c, encoder_hidden_states, rope_cos, rope_sin = self.prepare(hidden_states, sigma, encoder_hidden_states, encoder_attention_mask) # [1, 11130, 3072], [1, 3072], [1, 256, 1536], [11130, 24, 64]
-        del encoder_attention_mask
-
-        for i, block in enumerate(self.blocks):
-            hidden_states, encoder_hidden_states = block( # [1, 11130, 3072], [1, 256, 1536]
-                hidden_states,
-                c,
-                encoder_hidden_states,
-                rope_cos=rope_cos,
-                rope_sin=rope_sin,
-                packed_indices=packed_indices,
-                checkpoint_ff=i < num_ff_checkpoint,
-                checkpoint_qkv=i < num_qkv_checkpoint,
-                checkpoint_post_attn=i < num_post_attn_checkpoint,
-            )  # (B, M, D), (B, L, D)
-        del encoder_hidden_states  # Final layers don't use dense text features.
-        
-        hidden_states = self.final_layer(hidden_states, c)  # (B, M, patch_size ** 2 * out_channels) [1, 11130, 48]
-
-        hidden_states = rearrange( # [1, 12, 7, 60, 106]
-            hidden_states,
-            "B (T hp wp) (p1 p2 c) -> B c T (hp p1) (wp p2)",
-            T=T,
-            hp=H // self.patch_size,
-            wp=W // self.patch_size,
-            p1=self.patch_size,
-            p2=self.patch_size,
-            c=self.out_channels,
-        )
-
-        attn_outputs_list = None
-        return (-hidden_states, attn_outputs_list)

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint_raw.py

@@ -1,737 +0,0 @@
-import os
-from typing import Dict, List, Optional, Tuple
-import warnings
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from torch.nn.attention import sdpa_kernel
-
-import genmo.mochi_preview.dit.joint_model.context_parallel as cp
-from genmo.lib.attn_imports import flash_varlen_attn, sage_attn, sdpa_attn_ctx
-from genmo.mochi_preview.dit.joint_model.layers import (
-    FeedForward,
-    PatchEmbed,
-    RMSNorm,
-    TimestepEmbedder,
-)
-from genmo.mochi_preview.dit.joint_model.lora import LoraLinear
-from genmo.mochi_preview.dit.joint_model.mod_rmsnorm import modulated_rmsnorm
-from genmo.mochi_preview.dit.joint_model.residual_tanh_gated_rmsnorm import (
-    residual_tanh_gated_rmsnorm,
-)
-from genmo.mochi_preview.dit.joint_model.rope_mixed import (
-    compute_mixed_rotation,
-    create_position_matrix,
-)
-from genmo.mochi_preview.dit.joint_model.temporal_rope import apply_rotary_emb_qk_real
-from genmo.mochi_preview.dit.joint_model.utils import (
-    AttentionPool,
-    modulate,
-    pad_and_split_xy,
-)
-
-COMPILE_FINAL_LAYER = os.environ.get("COMPILE_DIT") == "1"
-COMPILE_MMDIT_BLOCK = os.environ.get("COMPILE_DIT") == "1"
-
-
-def ck(fn, *args, enabled=True, **kwargs) -> torch.Tensor:
-    if enabled:
-        return torch.utils.checkpoint.checkpoint(fn, *args, **kwargs, use_reentrant=False)
-
-    return fn(*args, **kwargs)
-
-
-class AsymmetricAttention(nn.Module):
-    def __init__(
-        self,
-        dim_x: int,
-        dim_y: int,
-        num_heads: int = 8,
-        qkv_bias: bool = True,
-        qk_norm: bool = False,
-        update_y: bool = True,
-        out_bias: bool = True,
-        attention_mode: str = "flash",
-        softmax_scale: Optional[float] = None,
-        device: Optional[torch.device] = None,
-        # Disable LoRA by default ...
-        qkv_proj_lora_rank: int = 0,
-        qkv_proj_lora_alpha: int = 0,
-        qkv_proj_lora_dropout: float = 0.0,
-        out_proj_lora_rank: int = 0,
-        out_proj_lora_alpha: int = 0,
-        out_proj_lora_dropout: float = 0.0,
-    ):
-        super().__init__()
-        self.attention_mode = attention_mode
-        self.dim_x = dim_x
-        self.dim_y = dim_y
-        self.num_heads = num_heads
-        self.head_dim = dim_x // num_heads
-        self.update_y = update_y
-        self.softmax_scale = softmax_scale
-        if dim_x % num_heads != 0:
-            raise ValueError(f"dim_x={dim_x} should be divisible by num_heads={num_heads}")
-
-        # Input layers.
-        self.qkv_bias = qkv_bias
-        qkv_lora_kwargs = dict(
-            bias=qkv_bias,
-            device=device,
-            r=qkv_proj_lora_rank,
-            lora_alpha=qkv_proj_lora_alpha,
-            lora_dropout=qkv_proj_lora_dropout,
-        )
-        self.qkv_x = LoraLinear(dim_x, 3 * dim_x, **qkv_lora_kwargs)
-        # Project text features to match visual features (dim_y -> dim_x)
-        self.qkv_y = LoraLinear(dim_y, 3 * dim_x, **qkv_lora_kwargs)
-
-        # Query and key normalization for stability.
-        assert qk_norm
-        self.q_norm_x = RMSNorm(self.head_dim, device=device)
-        self.k_norm_x = RMSNorm(self.head_dim, device=device)
-        self.q_norm_y = RMSNorm(self.head_dim, device=device)
-        self.k_norm_y = RMSNorm(self.head_dim, device=device)
-
-        # Output layers. y features go back down from dim_x -> dim_y.
-        proj_lora_kwargs = dict(
-            bias=out_bias,
-            device=device,
-            r=out_proj_lora_rank,
-            lora_alpha=out_proj_lora_alpha,
-            lora_dropout=out_proj_lora_dropout,
-        )
-        self.proj_x = LoraLinear(dim_x, dim_x, **proj_lora_kwargs)
-        self.proj_y = LoraLinear(dim_x, dim_y, **proj_lora_kwargs) if update_y else nn.Identity()
-
-    def run_qkv_y(self, y):
-        cp_rank, cp_size = cp.get_cp_rank_size()
-        local_heads = self.num_heads // cp_size
-
-        if cp.is_cp_active():
-            # Only predict local heads.
-            assert not self.qkv_bias
-            W_qkv_y = self.qkv_y.weight.view(3, self.num_heads, self.head_dim, self.dim_y)
-            W_qkv_y = W_qkv_y.narrow(1, cp_rank * local_heads, local_heads)
-            W_qkv_y = W_qkv_y.reshape(3 * local_heads * self.head_dim, self.dim_y)
-            qkv_y = F.linear(y, W_qkv_y, None)  # (B, L, 3 * local_h * head_dim)
-        else:
-            qkv_y = self.qkv_y(y)  # (B, L, 3 * dim)
-
-        qkv_y = qkv_y.view(qkv_y.size(0), qkv_y.size(1), 3, local_heads, self.head_dim)
-        q_y, k_y, v_y = qkv_y.unbind(2)
-
-        q_y = self.q_norm_y(q_y)
-        k_y = self.k_norm_y(k_y)
-        return q_y, k_y, v_y
-
-    def prepare_qkv(
-        self,
-        x: torch.Tensor,  # (B, M, dim_x)
-        y: torch.Tensor,  # (B, L, dim_y)
-        *,
-        scale_x: torch.Tensor,
-        scale_y: torch.Tensor,
-        rope_cos: torch.Tensor,
-        rope_sin: torch.Tensor,
-        valid_token_indices: torch.Tensor,
-        max_seqlen_in_batch: int,
-    ):
-        # Process visual features
-        x = modulated_rmsnorm(x, scale_x)  # (B, M, dim_x) where M = N / cp_group_size
-        qkv_x = self.qkv_x(x)  # (B, M, 3 * dim_x)
-        assert qkv_x.dtype == torch.bfloat16
-
-        qkv_x = cp.all_to_all_collect_tokens(qkv_x, self.num_heads)  # (3, B, N, local_h, head_dim)
-
-        # Split qkv_x into q, k, v
-        q_x, k_x, v_x = qkv_x.unbind(0)  # (B, N, local_h, head_dim)
-        q_x = self.q_norm_x(q_x)
-        q_x = apply_rotary_emb_qk_real(q_x, rope_cos, rope_sin)
-        k_x = self.k_norm_x(k_x)
-        k_x = apply_rotary_emb_qk_real(k_x, rope_cos, rope_sin)
-
-        # Concatenate streams
-        B, N, num_heads, head_dim = q_x.size()
-        D = num_heads * head_dim
-
-        # Process text features
-        if B == 1:
-            text_seqlen = max_seqlen_in_batch - N
-            if text_seqlen > 0:
-                y = y[:, :text_seqlen]  # Remove padding tokens.
-                y = modulated_rmsnorm(y, scale_y)  # (B, L, dim_y)
-                q_y, k_y, v_y = self.run_qkv_y(y)  # (B, L, local_heads, head_dim)
-
-                q = torch.cat([q_x, q_y], dim=1)
-                k = torch.cat([k_x, k_y], dim=1)
-                v = torch.cat([v_x, v_y], dim=1)
-            else:
-                q, k, v = q_x, k_x, v_x
-        else:
-            y = modulated_rmsnorm(y, scale_y)  # (B, L, dim_y)
-            q_y, k_y, v_y = self.run_qkv_y(y)  # (B, L, local_heads, head_dim)
-
-            indices = valid_token_indices[:, None].expand(-1, D)
-            q = torch.cat([q_x, q_y], dim=1).view(-1, D).gather(0, indices)  # (total, D)
-            k = torch.cat([k_x, k_y], dim=1).view(-1, D).gather(0, indices)  # (total, D)
-            v = torch.cat([v_x, v_y], dim=1).view(-1, D).gather(0, indices)  # (total, D)
-
-        q = q.view(-1, num_heads, head_dim)
-        k = k.view(-1, num_heads, head_dim)
-        v = v.view(-1, num_heads, head_dim)
-        return q, k, v
-
-    @torch.autocast("cuda", enabled=False)
-    def flash_attention(self, q, k, v, cu_seqlens, max_seqlen_in_batch, total, local_dim):
-        out: torch.Tensor = flash_varlen_attn(
-            q, k, v,
-            cu_seqlens_q=cu_seqlens,
-            cu_seqlens_k=cu_seqlens,
-            max_seqlen_q=max_seqlen_in_batch,
-            max_seqlen_k=max_seqlen_in_batch,
-            dropout_p=0.0,
-            softmax_scale=self.softmax_scale,
-        )  # (total, local_heads, head_dim)
-        return out.view(total, local_dim)
-
-    def sdpa_attention(self, q, k, v):
-        with sdpa_attn_ctx(training=self.training):
-            out = F.scaled_dot_product_attention(
-                q, k, v,
-                attn_mask=None,
-                dropout_p=0.0,
-                is_causal=False,
-            )
-            return out
-
-    @torch.autocast("cuda", enabled=False)
-    def sage_attention(self, q, k, v):
-        return sage_attn(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
-
-    def run_attention(
-        self,
-        q: torch.Tensor,  # (total <= B * (N + L), num_heads, head_dim)
-        k: torch.Tensor,  # (total <= B * (N + L), num_heads, head_dim)
-        v: torch.Tensor,  # (total <= B * (N + L), num_heads, head_dim)
-        *,
-        B: int,
-        cu_seqlens: Optional[torch.Tensor] = None,
-        max_seqlen_in_batch: Optional[int] = None,
-    ):
-        _, cp_size = cp.get_cp_rank_size()
-        assert self.num_heads % cp_size == 0
-        local_heads = self.num_heads // cp_size
-        local_dim = local_heads * self.head_dim
-
-        # Check shapes
-        assert q.ndim == 3 and k.ndim == 3 and v.ndim == 3
-        total = q.size(0)
-        assert k.size(0) == total and v.size(0) == total
-
-        if self.attention_mode == "flash":
-            out = self.flash_attention(
-                q, k, v, cu_seqlens, max_seqlen_in_batch, total, local_dim)  # (total, local_dim)
-        else:
-            assert B == 1, \
-                f"Non-flash attention mode {self.attention_mode} only supports batch size 1, got {B}"
-
-            q = rearrange(q, "(b s) h d -> b h s d", b=B)
-            k = rearrange(k, "(b s) h d -> b h s d", b=B)
-            v = rearrange(v, "(b s) h d -> b h s d", b=B)
-
-            if self.attention_mode == "sdpa":
-                out = self.sdpa_attention(q, k, v)  # (B, local_heads, seq_len, head_dim)
-            elif self.attention_mode == "sage":
-                out = self.sage_attention(q, k, v)  # (B, local_heads, seq_len, head_dim)
-            else:
-                raise ValueError(f"Unknown attention mode: {self.attention_mode}")
-
-            out = rearrange(out, "b h s d -> (b s) (h d)")
-
-        return out
-
-    def post_attention(
-        self,
-        out: torch.Tensor,
-        B: int,
-        M: int,
-        L: int,
-        dtype: torch.dtype,
-        valid_token_indices: torch.Tensor,
-    ):
-        """
-        Args:
-            out: (total <= B * (N + L), local_dim)
-            valid_token_indices: (total <= B * (N + L),)
-            B: Batch size
-            M: Number of visual tokens per context parallel rank
-            L: Number of text tokens
-            dtype: Data type of the input and output tensors
-
-        Returns:
-            x: (B, N, dim_x) tensor of visual tokens where N = M * cp_size
-            y: (B, L, dim_y) tensor of text token features
-        """
-        _, cp_size = cp.get_cp_rank_size()
-        local_heads = self.num_heads // cp_size
-        local_dim = local_heads * self.head_dim
-        N = M * cp_size
-
-        # Split sequence into visual and text tokens, adding back padding.
-        if B == 1:
-            out = out.view(B, -1, local_dim)
-            if out.size(1) > N:
-                x, y = torch.tensor_split(out, (N,), dim=1)  # (B, N, local_dim), (B, <= L, local_dim)
-                y = F.pad(y, (0, 0, 0, L - y.size(1)))  # (B, L, local_dim)
-            else:
-                # Empty prompt.
-                x, y = out, out.new_zeros(B, L, local_dim)
-        else:
-            x, y = pad_and_split_xy(out, valid_token_indices, B, N, L, dtype)
-        assert x.size() == (B, N, local_dim)
-        assert y.size() == (B, L, local_dim)
-
-        # Communicate across context parallel ranks.
-        x = x.view(B, N, local_heads, self.head_dim)
-        x = cp.all_to_all_collect_heads(x)  # (B, M, dim_x = num_heads * head_dim)
-        if cp.is_cp_active():
-            y = cp.all_gather(y)  # (cp_size * B, L, local_heads * head_dim)
-            y = rearrange(y, "(G B) L D -> B L (G D)", G=cp_size, D=local_dim)  # (B, L, dim_x)
-
-        x = self.proj_x(x)
-        y = self.proj_y(y)
-        return x, y
-
-    def forward(
-        self,
-        x: torch.Tensor,  # (B, M, dim_x)
-        y: torch.Tensor,  # (B, L, dim_y)
-        *,
-        scale_x: torch.Tensor,  # (B, dim_x), modulation for pre-RMSNorm.
-        scale_y: torch.Tensor,  # (B, dim_y), modulation for pre-RMSNorm.
-        packed_indices: Dict[str, torch.Tensor] = None,
-        checkpoint_qkv: bool = False,
-        checkpoint_post_attn: bool = False,
-        **rope_rotation,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Forward pass of asymmetric multi-modal attention.
-
-        Args:
-            x: (B, M, dim_x) tensor of visual tokens
-            y: (B, L, dim_y) tensor of text token features
-            packed_indices: Dict with keys for Flash Attention
-            num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens
-
-        Returns:
-            x: (B, M, dim_x) tensor of visual tokens after multi-modal attention
-            y: (B, L, dim_y) tensor of text token features after multi-modal attention
-        """
-        B, L, _ = y.shape
-        _, M, _ = x.shape
-
-        # Predict a packed QKV tensor from visual and text features.
-        q, k, v = ck(self.prepare_qkv,
-            x=x,
-            y=y,
-            scale_x=scale_x,
-            scale_y=scale_y,
-            rope_cos=rope_rotation.get("rope_cos"),
-            rope_sin=rope_rotation.get("rope_sin"),
-            valid_token_indices=packed_indices["valid_token_indices_kv"],
-            max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"],
-            enabled=checkpoint_qkv,
-        )  # (total <= B * (N + L), 3, local_heads, head_dim)
-
-        # Self-attention is expensive, so don't checkpoint it.
-        out = self.run_attention(
-            q, k, v, B=B,
-            cu_seqlens=packed_indices["cu_seqlens_kv"],
-            max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"],
-        )
-
-        x, y = ck(self.post_attention,
-            out,
-            B=B, M=M, L=L,
-            dtype=v.dtype,
-            valid_token_indices=packed_indices["valid_token_indices_kv"],
-            enabled=checkpoint_post_attn,
-        )
-
-        return x, y
-
-
-@torch.compile(disable=not COMPILE_MMDIT_BLOCK)
-class AsymmetricJointBlock(nn.Module):
-    def __init__(
-        self,
-        hidden_size_x: int,
-        hidden_size_y: int,
-        num_heads: int,
-        *,
-        mlp_ratio_x: float = 8.0,  # Ratio of hidden size to d_model for MLP for visual tokens.
-        mlp_ratio_y: float = 4.0,  # Ratio of hidden size to d_model for MLP for text tokens.
-        update_y: bool = True,  # Whether to update text tokens in this block.
-        device: Optional[torch.device] = None,
-        **block_kwargs,
-    ):
-        super().__init__()
-        self.update_y = update_y
-        self.hidden_size_x = hidden_size_x
-        self.hidden_size_y = hidden_size_y
-        self.mod_x = nn.Linear(hidden_size_x, 4 * hidden_size_x, device=device)
-        if self.update_y:
-            self.mod_y = nn.Linear(hidden_size_x, 4 * hidden_size_y, device=device)
-        else:
-            self.mod_y = nn.Linear(hidden_size_x, hidden_size_y, device=device)
-
-        # Self-attention:
-        self.attn = AsymmetricAttention(
-            hidden_size_x,
-            hidden_size_y,
-            num_heads=num_heads,
-            update_y=update_y,
-            device=device,
-            **block_kwargs,
-        )
-
-        # MLP.
-        mlp_hidden_dim_x = int(hidden_size_x * mlp_ratio_x)
-        assert mlp_hidden_dim_x == int(1536 * 8)
-        self.mlp_x = FeedForward(
-            in_features=hidden_size_x,
-            hidden_size=mlp_hidden_dim_x,
-            multiple_of=256,
-            ffn_dim_multiplier=None,
-            device=device,
-        )
-
-        # MLP for text not needed in last block.
-        if self.update_y:
-            mlp_hidden_dim_y = int(hidden_size_y * mlp_ratio_y)
-            self.mlp_y = FeedForward(
-                in_features=hidden_size_y,
-                hidden_size=mlp_hidden_dim_y,
-                multiple_of=256,
-                ffn_dim_multiplier=None,
-                device=device,
-            )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        c: torch.Tensor,
-        y: torch.Tensor,
-        # TODO: These could probably just go into attn_kwargs
-        checkpoint_ff: bool = False,
-        checkpoint_qkv: bool = False,
-        checkpoint_post_attn: bool = False,
-        **attn_kwargs,
-    ):
-        """Forward pass of a block.
-
-        Args:
-            x: (B, N, dim) tensor of visual tokens
-            c: (B, dim) tensor of conditioned features
-            y: (B, L, dim) tensor of text tokens
-            num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens
-
-        Returns:
-            x: (B, N, dim) tensor of visual tokens after block
-            y: (B, L, dim) tensor of text tokens after block
-        """
-        N = x.size(1)
-
-        c = F.silu(c)
-        mod_x = self.mod_x(c)
-        scale_msa_x, gate_msa_x, scale_mlp_x, gate_mlp_x = mod_x.chunk(4, dim=1)
-        mod_y = self.mod_y(c)
-
-        if self.update_y:
-            scale_msa_y, gate_msa_y, scale_mlp_y, gate_mlp_y = mod_y.chunk(4, dim=1)
-        else:
-            scale_msa_y = mod_y
-
-        # Self-attention block.
-        x_attn, y_attn = self.attn(
-            x,
-            y,
-            scale_x=scale_msa_x,
-            scale_y=scale_msa_y,
-            checkpoint_qkv=checkpoint_qkv,
-            checkpoint_post_attn=checkpoint_post_attn,
-            **attn_kwargs,
-        )
-
-        assert x_attn.size(1) == N
-        x = residual_tanh_gated_rmsnorm(x, x_attn, gate_msa_x)
-
-        if self.update_y:
-            y = residual_tanh_gated_rmsnorm(y, y_attn, gate_msa_y)
-
-        # MLP block.
-        x = ck(self.ff_block_x, x, scale_mlp_x, gate_mlp_x, enabled=checkpoint_ff)
-        if self.update_y:
-            y = ck(self.ff_block_y, y, scale_mlp_y, gate_mlp_y, enabled=checkpoint_ff)  # type: ignore
-        return x, y
-
-    def ff_block_x(self, x, scale_x, gate_x):
-        x_mod = modulated_rmsnorm(x, scale_x)
-        x_res = self.mlp_x(x_mod)
-        x = residual_tanh_gated_rmsnorm(x, x_res, gate_x)  # Sandwich norm
-        return x
-
-    def ff_block_y(self, y, scale_y, gate_y):
-        y_mod = modulated_rmsnorm(y, scale_y)
-        y_res = self.mlp_y(y_mod)
-        y = residual_tanh_gated_rmsnorm(y, y_res, gate_y)  # Sandwich norm
-        return y
-
-
-@torch.compile(disable=not COMPILE_FINAL_LAYER)
-class FinalLayer(nn.Module):
-    """
-    The final layer of DiT.
-    """
-
-    def __init__(
-        self,
-        hidden_size,
-        patch_size,
-        out_channels,
-        device: Optional[torch.device] = None,
-    ):
-        super().__init__()
-        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, device=device)
-        self.mod = nn.Linear(hidden_size, 2 * hidden_size, device=device)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, device=device)
-
-    def forward(self, x, c):
-        c = F.silu(c)
-        shift, scale = self.mod(c).chunk(2, dim=1)
-        x = modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        return x
-
-
-class AsymmDiTJoint(nn.Module):
-    """
-    Diffusion model with a Transformer backbone.
-
-    Ingests text embeddings instead of a label.
-    """
-
-    def __init__(
-        self,
-        *,
-        patch_size=2,
-        in_channels=4,
-        hidden_size_x=1152,
-        hidden_size_y=1152,
-        depth=48,
-        num_heads=16,
-        mlp_ratio_x=8.0,
-        mlp_ratio_y=4.0,
-        t5_feat_dim: int = 4096,
-        t5_token_length: int = 256,
-        patch_embed_bias: bool = True,
-        timestep_mlp_bias: bool = True,
-        timestep_scale: Optional[float] = None,
-        use_extended_posenc: bool = False,
-        rope_theta: float = 10000.0,
-        device: Optional[torch.device] = None,
-        **block_kwargs,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.patch_size = patch_size
-        self.num_heads = num_heads
-        self.hidden_size_x = hidden_size_x
-        self.hidden_size_y = hidden_size_y
-        self.head_dim = hidden_size_x // num_heads  # Head dimension and count is determined by visual.
-        self.use_extended_posenc = use_extended_posenc
-        self.t5_token_length = t5_token_length
-        self.t5_feat_dim = t5_feat_dim
-        self.rope_theta = rope_theta  # Scaling factor for frequency computation for temporal RoPE.
-
-        self.x_embedder = PatchEmbed(
-            patch_size=patch_size,
-            in_chans=in_channels,
-            embed_dim=hidden_size_x,
-            bias=patch_embed_bias,
-            device=device,
-        )
-        # Conditionings
-        # Timestep
-        self.t_embedder = TimestepEmbedder(hidden_size_x, bias=timestep_mlp_bias, timestep_scale=timestep_scale)
-
-        # Caption Pooling (T5)
-        self.t5_y_embedder = AttentionPool(t5_feat_dim, num_heads=8, output_dim=hidden_size_x, device=device)
-
-        # Dense Embedding Projection (T5)
-        self.t5_yproj = nn.Linear(t5_feat_dim, hidden_size_y, bias=True, device=device)
-
-        # Initialize pos_frequencies as an empty parameter.
-        self.pos_frequencies = nn.Parameter(torch.empty(3, self.num_heads, self.head_dim // 2, device=device))
-
-        # for depth 48:
-        #  b =  0: AsymmetricJointBlock, update_y=True
-        #  b =  1: AsymmetricJointBlock, update_y=True
-        #  ...
-        #  b = 46: AsymmetricJointBlock, update_y=True
-        #  b = 47: AsymmetricJointBlock, update_y=False. No need to update text features.
-        blocks = []
-        for b in range(depth):
-            # Joint multi-modal block
-            update_y = b < depth - 1
-            block = AsymmetricJointBlock(
-                hidden_size_x,
-                hidden_size_y,
-                num_heads,
-                mlp_ratio_x=mlp_ratio_x,
-                mlp_ratio_y=mlp_ratio_y,
-                update_y=update_y,
-                device=device,
-                **block_kwargs,
-            )
-
-            blocks.append(block)
-        self.blocks = nn.ModuleList(blocks)
-
-        self.final_layer = FinalLayer(hidden_size_x, patch_size, self.out_channels, device=device)
-
-    def embed_x(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x: (B, C=12, T, H, W) tensor of visual tokens
-
-        Returns:
-            x: (B, C=3072, N) tensor of visual tokens with positional embedding.
-        """
-        return self.x_embedder(x)  # Convert BcTHW to BCN
-
-    @torch.compile(disable=not COMPILE_MMDIT_BLOCK)
-    def prepare(
-        self,
-        x: torch.Tensor,
-        sigma: torch.Tensor,
-        t5_feat: torch.Tensor,
-        t5_mask: torch.Tensor,
-    ):
-        """Prepare input and conditioning embeddings."""
-
-        # Visual patch embeddings with positional encoding.
-        T, H, W = x.shape[-3:]
-        pH, pW = H // self.patch_size, W // self.patch_size
-        x = self.embed_x(x)  # (B, N, D), where N = T * H * W / patch_size ** 2
-        assert x.ndim == 3
-        B = x.size(0)
-
-        # Construct position array of size [N, 3].
-        # pos[:, 0] is the frame index for each location,
-        # pos[:, 1] is the row index for each location, and
-        # pos[:, 2] is the column index for each location.
-        N = T * pH * pW
-        assert x.size(1) == N
-        pos = create_position_matrix(T, pH=pH, pW=pW, device=x.device, dtype=torch.float32)  # (N, 3)
-        rope_cos, rope_sin = compute_mixed_rotation(
-            freqs=self.pos_frequencies, pos=pos
-        )  # Each are (N, num_heads, dim // 2)
-
-        # Global vector embedding for conditionings.
-        c_t = self.t_embedder(1 - sigma)  # (B, D)
-
-        # Pool T5 tokens using attention pooler
-        # Note y_feat[1] contains T5 token features.
-        assert (
-            t5_feat.size(1) == self.t5_token_length
-        ), f"Expected L={self.t5_token_length}, got {t5_feat.shape} for y_feat."
-        t5_y_pool = self.t5_y_embedder(t5_feat, t5_mask)  # (B, D)
-        assert t5_y_pool.size(0) == B, f"Expected B={B}, got {t5_y_pool.shape} for t5_y_pool."
-
-        c = c_t + t5_y_pool
-
-        y_feat = self.t5_yproj(t5_feat)  # (B, L, t5_feat_dim) --> (B, L, D)
-
-        return x, c, y_feat, rope_cos, rope_sin
-
-    def forward(
-        self,
-        x: torch.Tensor, # [1, 12, 7, 60, 106]
-        sigma: torch.Tensor, # [1] 
-        y_feat: List[torch.Tensor], # [0][1, 256, 4096]
-        y_mask: List[torch.Tensor], # [0][1, 256]
-        packed_indices: Dict[str, torch.Tensor] = None,
-        rope_cos: torch.Tensor = None, 
-        rope_sin: torch.Tensor = None, 
-        num_ff_checkpoint: int = 0, # 48
-        num_qkv_checkpoint: int = 0, # 48
-        num_post_attn_checkpoint: int = 0, # 0
-    ):
-        """Forward pass of DiT.
-
-        Args:
-            x: (B, C, T, H, W) tensor of spatial inputs (images or latent representations of images)
-            sigma: (B,) tensor of noise standard deviations
-            y_feat: List((B, L, y_feat_dim) tensor of caption token features. For SDXL text encoders: L=77, y_feat_dim=2048)
-            y_mask: List((B, L) boolean tensor indicating which tokens are not padding)
-            packed_indices: Dict with keys for Flash Attention. Result of compute_packed_indices.
-        """
-        _, _, T, H, W = x.shape
-
-        if self.pos_frequencies.dtype != torch.float32:
-            warnings.warn(f"pos_frequencies dtype {self.pos_frequencies.dtype} != torch.float32")
-
-        # Use EFFICIENT_ATTENTION backend for T5 pooling, since we have a mask.
-        # Have to call sdpa_kernel outside of a torch.compile region.
-        with sdpa_kernel(torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION):
-            x, c, y_feat, rope_cos, rope_sin = self.prepare(x, sigma, y_feat[0], y_mask[0]) # [1, 11130, 3072], [1, 3072], [1, 256, 1536], [11130, 24, 64]
-        del y_mask
-
-        cp_rank, cp_size = cp.get_cp_rank_size()
-        N = x.size(1)
-        M = N // cp_size
-        assert N % cp_size == 0, f"Visual sequence length ({x.shape[1]}) must be divisible by cp_size ({cp_size})."
-
-        if cp_size > 1:
-            x = x.narrow(1, cp_rank * M, M)
-
-            assert self.num_heads % cp_size == 0
-            local_heads = self.num_heads // cp_size
-            rope_cos = rope_cos.narrow(1, cp_rank * local_heads, local_heads)
-            rope_sin = rope_sin.narrow(1, cp_rank * local_heads, local_heads)
-
-        for i, block in enumerate(self.blocks):
-            x, y_feat = block( # [1, 11130, 3072], [1, 256, 1536]
-                x,
-                c,
-                y_feat,
-                rope_cos=rope_cos,
-                rope_sin=rope_sin,
-                packed_indices=packed_indices,
-                checkpoint_ff=i < num_ff_checkpoint,
-                checkpoint_qkv=i < num_qkv_checkpoint,
-                checkpoint_post_attn=i < num_post_attn_checkpoint,
-            )  # (B, M, D), (B, L, D)
-        del y_feat  # Final layers don't use dense text features.
-
-        x = self.final_layer(x, c)  # (B, M, patch_size ** 2 * out_channels) [1, 11130, 48]
-
-        patch = x.size(2)
-        x = cp.all_gather(x)
-        x = rearrange(x, "(G B) M P -> B (G M) P", G=cp_size, P=patch) # [1, 11130, 48]
-        x = rearrange( # [1, 12, 7, 60, 106]
-            x,
-            "B (T hp wp) (p1 p2 c) -> B c T (hp p1) (wp p2)",
-            T=T,
-            hp=H // self.patch_size,
-            wp=W // self.patch_size,
-            p1=self.patch_size,
-            p2=self.patch_size,
-            c=self.out_channels,
-        )
-
-        return x

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/context_parallel.py

@@ -1,158 +0,0 @@
-from typing import Tuple
-
-import torch
-import torch.distributed as dist
-from einops import rearrange
-
-_CONTEXT_PARALLEL_GROUP = None
-_CONTEXT_PARALLEL_RANK = None
-_CONTEXT_PARALLEL_GROUP_SIZE = None
-_CONTEXT_PARALLEL_GROUP_RANKS = None
-
-
-def get_cp_rank_size() -> Tuple[int, int]:
-    if _CONTEXT_PARALLEL_GROUP:
-        assert isinstance(_CONTEXT_PARALLEL_RANK, int) and isinstance(_CONTEXT_PARALLEL_GROUP_SIZE, int)
-        return _CONTEXT_PARALLEL_RANK, _CONTEXT_PARALLEL_GROUP_SIZE
-    else:
-        return 0, 1
-
-
-def local_shard(x: torch.Tensor, dim: int = 2) -> torch.Tensor:
-    if not _CONTEXT_PARALLEL_GROUP:
-        return x
-
-    cp_rank, cp_size = get_cp_rank_size()
-    return x.tensor_split(cp_size, dim=dim)[cp_rank]
-
-
-def set_cp_group(cp_group, ranks, global_rank):
-    global _CONTEXT_PARALLEL_GROUP, _CONTEXT_PARALLEL_RANK, _CONTEXT_PARALLEL_GROUP_SIZE, _CONTEXT_PARALLEL_GROUP_RANKS
-    if _CONTEXT_PARALLEL_GROUP is not None:
-        raise RuntimeError("CP group already initialized.")
-    _CONTEXT_PARALLEL_GROUP = cp_group
-    _CONTEXT_PARALLEL_RANK = dist.get_rank(cp_group)
-    _CONTEXT_PARALLEL_GROUP_SIZE = dist.get_world_size(cp_group)
-    _CONTEXT_PARALLEL_GROUP_RANKS = ranks
-
-    assert _CONTEXT_PARALLEL_RANK == ranks.index(
-        global_rank
-    ), f"Rank mismatch: {global_rank} in {ranks} does not have position {_CONTEXT_PARALLEL_RANK} "
-    assert _CONTEXT_PARALLEL_GROUP_SIZE == len(
-        ranks
-    ), f"Group size mismatch: {_CONTEXT_PARALLEL_GROUP_SIZE} != len({ranks})"
-
-
-def get_cp_group():
-    if _CONTEXT_PARALLEL_GROUP is None:
-        raise RuntimeError("CP group not initialized")
-    return _CONTEXT_PARALLEL_GROUP
-
-
-def is_cp_active():
-    return _CONTEXT_PARALLEL_GROUP is not None
-
-
-class AllGatherIntoTensorFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x: torch.Tensor, reduce_dtype, group: dist.ProcessGroup):
-        ctx.reduce_dtype = reduce_dtype
-        ctx.group = group
-        ctx.batch_size = x.size(0)
-        group_size = dist.get_world_size(group)
-
-        x = x.contiguous()
-        output = torch.empty(group_size * x.size(0), *x.shape[1:], dtype=x.dtype, device=x.device)
-        dist.all_gather_into_tensor(output, x, group=group)
-        return output
-
-
-def all_gather(tensor: torch.Tensor) -> torch.Tensor:
-    if not _CONTEXT_PARALLEL_GROUP:
-        return tensor
-
-    return AllGatherIntoTensorFunction.apply(tensor, torch.float32, _CONTEXT_PARALLEL_GROUP)
-
-
-@torch.compiler.disable()
-def _all_to_all_single(output, input, group):
-    # Disable compilation since torch compile changes contiguity.
-    assert input.is_contiguous(), "Input tensor must be contiguous."
-    assert output.is_contiguous(), "Output tensor must be contiguous."
-    return dist.all_to_all_single(output, input, group=group)
-
-
-class CollectTokens(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, qkv: torch.Tensor, group: dist.ProcessGroup, num_heads: int):
-        """Redistribute heads and receive tokens.
-
-        Args:
-            qkv: query, key or value. Shape: [B, M, 3 * num_heads * head_dim]
-
-        Returns:
-            qkv: shape: [3, B, N, local_heads, head_dim]
-
-        where M is the number of local tokens,
-        N = cp_size * M is the number of global tokens,
-        local_heads = num_heads // cp_size is the number of local heads.
-        """
-        ctx.group = group
-        ctx.num_heads = num_heads
-        cp_size = dist.get_world_size(group)
-        assert num_heads % cp_size == 0
-        ctx.local_heads = num_heads // cp_size
-
-        qkv = rearrange(
-            qkv,
-            "B M (qkv G h d) -> G M h B (qkv d)",
-            qkv=3,
-            G=cp_size,
-            h=ctx.local_heads,
-        ).contiguous()
-
-        output_chunks = torch.empty_like(qkv)
-        _all_to_all_single(output_chunks, qkv, group=group)
-
-        return rearrange(output_chunks, "G M h B (qkv d) -> qkv B (G M) h d", qkv=3)
-
-
-def all_to_all_collect_tokens(x: torch.Tensor, num_heads: int) -> torch.Tensor:
-    if not _CONTEXT_PARALLEL_GROUP:
-        # Move QKV dimension to the front.
-        #   B M (3 H d) -> 3 B M H d
-        B, M, _ = x.size()
-        x = x.view(B, M, 3, num_heads, -1)
-        return x.permute(2, 0, 1, 3, 4)
-
-    return CollectTokens.apply(x, _CONTEXT_PARALLEL_GROUP, num_heads)
-
-
-class CollectHeads(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x: torch.Tensor, group: dist.ProcessGroup):
-        """Redistribute tokens and receive heads.
-
-        Args:
-            x: Output of attention. Shape: [B, N, local_heads, head_dim]
-
-        Returns:
-            Shape: [B, M, num_heads * head_dim]
-        """
-        ctx.group = group
-        ctx.local_heads = x.size(2)
-        ctx.head_dim = x.size(3)
-        group_size = dist.get_world_size(group)
-        x = rearrange(x, "B (G M) h D -> G h M B D", G=group_size).contiguous()
-        output = torch.empty_like(x)
-        _all_to_all_single(output, x, group=group)
-        del x
-        return rearrange(output, "G h M B D -> B M (G h D)")
-
-
-def all_to_all_collect_heads(x: torch.Tensor) -> torch.Tensor:
-    if not _CONTEXT_PARALLEL_GROUP:
-        # Merge heads.
-        return x.view(x.size(0), x.size(1), x.size(2) * x.size(3))
-
-    return CollectHeads.apply(x, _CONTEXT_PARALLEL_GROUP)

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/layers.py

@@ -1,179 +0,0 @@
-import collections.abc
-import math
-from itertools import repeat
-from typing import Callable, Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-
-
-# From PyTorch internals
-def _ntuple(n):
-    def parse(x):
-        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
-            return tuple(x)
-        return tuple(repeat(x, n))
-
-    return parse
-
-
-to_2tuple = _ntuple(2)
-
-
-class TimestepEmbedder(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        frequency_embedding_size: int = 256,
-        *,
-        bias: bool = True,
-        timestep_scale: Optional[float] = None,
-        device: Optional[torch.device] = None,
-    ):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency_embedding_size, hidden_size, bias=bias, device=device),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size, bias=bias, device=device),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-        self.timestep_scale = timestep_scale
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        half = dim // 2
-        freqs = torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
-        freqs.mul_(-math.log(max_period) / half).exp_()
-        args = t[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding
-
-    def forward(self, t):
-        if self.timestep_scale is not None:
-            t = t * self.timestep_scale
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-
-class PooledCaptionEmbedder(nn.Module):
-    def __init__(
-        self,
-        caption_feature_dim: int,
-        hidden_size: int,
-        *,
-        bias: bool = True,
-        device: Optional[torch.device] = None,
-    ):
-        super().__init__()
-        self.caption_feature_dim = caption_feature_dim
-        self.hidden_size = hidden_size
-        self.mlp = nn.Sequential(
-            nn.Linear(caption_feature_dim, hidden_size, bias=bias, device=device),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size, bias=bias, device=device),
-        )
-
-    def forward(self, x):
-        return self.mlp(x)
-
-
-class FeedForward(nn.Module):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_size: int,
-        multiple_of: int,
-        ffn_dim_multiplier: Optional[float],
-        device: Optional[torch.device] = None,
-    ):
-        super().__init__()
-        # keep parameter count and computation constant compared to standard FFN
-        hidden_size = int(2 * hidden_size / 3)
-        # custom dim factor multiplier
-        if ffn_dim_multiplier is not None:
-            hidden_size = int(ffn_dim_multiplier * hidden_size)
-        hidden_size = multiple_of * ((hidden_size + multiple_of - 1) // multiple_of)
-
-        self.hidden_dim = hidden_size
-        self.w1 = nn.Linear(in_features, 2 * hidden_size, bias=False, device=device)
-        self.w2 = nn.Linear(hidden_size, in_features, bias=False, device=device)
-
-    def forward(self, x):
-        # assert self.w1.weight.dtype == torch.bfloat16, f"FFN weight dtype {self.w1.weight.dtype} != bfloat16"
-        x, gate = self.w1(x).chunk(2, dim=-1)
-        x = self.w2(F.silu(x) * gate)
-        return x
-
-
-class PatchEmbed(nn.Module):
-    def __init__(
-        self,
-        patch_size: int = 16,
-        in_chans: int = 3,
-        embed_dim: int = 768,
-        norm_layer: Optional[Callable] = None,
-        flatten: bool = True,
-        bias: bool = True,
-        dynamic_img_pad: bool = False,
-        device: Optional[torch.device] = None,
-    ):
-        super().__init__()
-        self.patch_size = to_2tuple(patch_size)
-        self.flatten = flatten
-        self.dynamic_img_pad = dynamic_img_pad
-
-        self.proj = nn.Conv2d(
-            in_chans,
-            embed_dim,
-            kernel_size=patch_size,
-            stride=patch_size,
-            bias=bias,
-            device=device,
-        )
-        assert norm_layer is None
-        self.norm = norm_layer(embed_dim, device=device) if norm_layer else nn.Identity()
-
-    def forward(self, x):
-        B, _C, T, H, W = x.shape
-        if not self.dynamic_img_pad:
-            assert (
-                H % self.patch_size[0] == 0
-            ), f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]})."
-            assert (
-                W % self.patch_size[1] == 0
-            ), f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]})."
-        else:
-            pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0]
-            pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1]
-            x = F.pad(x, (0, pad_w, 0, pad_h))
-
-        x = rearrange(x, "B C T H W -> (B T) C H W", B=B, T=T)
-        x = self.proj(x)
-
-        # Flatten temporal and spatial dimensions.
-        if not self.flatten:
-            raise NotImplementedError("Must flatten output.")
-        x = rearrange(x, "(B T) C H W -> B (T H W) C", B=B, T=T)
-
-        x = self.norm(x)
-        return x
-
-
-class RMSNorm(torch.nn.Module):
-    def __init__(self, hidden_size, eps=1e-5, device=None):
-        super().__init__()
-        self.eps = eps
-        self.weight = torch.nn.Parameter(torch.empty(hidden_size, device=device))
-        self.register_parameter("bias", None)
-
-    def forward(self, x):
-        # assert self.weight.dtype == torch.float32, f"RMSNorm weight dtype {self.weight.dtype} != float32"
-
-        x_fp32 = x.float()
-        x_normed = x_fp32 * torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)
-        return (x_normed * self.weight).type_as(x)

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/lora.py

@@ -1,112 +0,0 @@
-#! /usr/bin/env python3
-import math
-from typing import Dict, List, Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class LoRALayer:
-    def __init__(
-        self,
-        r: int,
-        lora_alpha: int,
-        lora_dropout: float,
-        merge_weights: bool,
-    ):
-        self.r = r
-        self.lora_alpha = lora_alpha
-        if lora_dropout > 0.0:
-            self.lora_dropout = nn.Dropout(p=lora_dropout)
-        else:
-            self.lora_dropout = lambda x: x
-        self.merged = False
-        self.merge_weights = merge_weights
-
-
-def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None:
-    assert bias == "none", f"Only bias='none' is supported"
-    for n, p in model.named_parameters():
-        if "lora_" not in n:
-            p.requires_grad = False
-
-
-def lora_state_dict(model: nn.Module, bias: str = "none") -> Dict[str, torch.Tensor]:
-    assert bias == "none", f"Only bias='none' is supported"
-    my_state_dict = model.state_dict()
-    return {k: my_state_dict[k] for k in my_state_dict if "lora_" in k}
-
-
-class LoraLinear(nn.Linear, LoRALayer):
-    # LoRA implemented in a dense layer
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        r: int = 0,
-        lora_alpha: int = 1,
-        lora_dropout: float = 0.0,
-        fan_in_fan_out: bool = False,  # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
-        merge_weights: bool = True,
-        **kwargs,
-    ):
-        nn.Linear.__init__(self, in_features, out_features, **kwargs)
-        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
-
-        self.fan_in_fan_out = fan_in_fan_out
-        # Actual trainable parameters
-        if r > 0:
-            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)).to(torch.float32))
-            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)).to(torch.float32))
-            self.scaling = self.lora_alpha / self.r
-
-            # Freezing the pre-trained weight matrix
-            self.weight.requires_grad = False
-
-        self.reset_parameters()
-
-        if fan_in_fan_out:
-            self.weight.data = self.weight.data.transpose(0, 1)
-
-    def reset_parameters(self):
-        nn.Linear.reset_parameters(self)
-        if hasattr(self, "lora_A"):
-            # initialize B the same way as the default for nn.Linear and A to zero
-            # this is different than what is described in the paper but should not affect performance
-            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
-            nn.init.zeros_(self.lora_B)
-
-    def train(self, mode: bool = True):
-        def T(w):
-            return w.transpose(0, 1) if self.fan_in_fan_out else w
-
-        nn.Linear.train(self, mode)
-        if mode:
-            if self.merge_weights and self.merged:
-                # Make sure that the weights are not merged
-                if self.r > 0:
-                    self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
-                self.merged = False
-        else:
-            if self.merge_weights and not self.merged:
-                # Merge the weights and mark it
-                if self.r > 0:
-                    self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
-                self.merged = True
-
-    def forward(self, x: torch.Tensor):
-        def T(w):
-            return w.transpose(0, 1) if self.fan_in_fan_out else w
-
-        if self.r > 0 and not self.merged:
-            result = F.linear(x, T(self.weight), bias=self.bias)
-
-            x = self.lora_dropout(x)
-            x = x @ self.lora_A.transpose(0, 1)
-            x = x @ self.lora_B.transpose(0, 1)
-            x = x * self.scaling
-
-            return result + x
-        else:
-            return F.linear(x, T(self.weight), bias=self.bias)

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/mod_rmsnorm.py

@@ -1,15 +0,0 @@
-import torch
-
-
-def modulated_rmsnorm(x, scale, eps=1e-6):
-    dtype = x.dtype
-    x = x.float()
-
-    # Compute RMS
-    mean_square = x.pow(2).mean(-1, keepdim=True)
-    inv_rms = torch.rsqrt(mean_square + eps)
-
-    # Normalize and modulate
-    x_normed = x * inv_rms
-    x_modulated = x_normed * (1 + scale.unsqueeze(1).float())
-    return x_modulated.to(dtype)

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/residual_tanh_gated_rmsnorm.py

@@ -1,20 +0,0 @@
-import torch
-
-
-def residual_tanh_gated_rmsnorm(x, x_res, gate, eps=1e-6):
-    # Convert to fp32 for precision
-    x_res = x_res.float()
-
-    # Compute RMS
-    mean_square = x_res.pow(2).mean(-1, keepdim=True)
-    scale = torch.rsqrt(mean_square + eps)
-
-    # Apply tanh to gate
-    tanh_gate = torch.tanh(gate).unsqueeze(1)
-
-    # Normalize and apply gated scaling
-    x_normed = x_res * scale * tanh_gate
-
-    # Apply residual connection
-    output = x + x_normed.type_as(x)
-    return output

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/rope_mixed.py

@@ -1,88 +0,0 @@
-import functools
-import math
-
-import torch
-
-
-def centers(start: float, stop, num, dtype=None, device=None):
-    """linspace through bin centers.
-
-    Args:
-        start (float): Start of the range.
-        stop (float): End of the range.
-        num (int): Number of points.
-        dtype (torch.dtype): Data type of the points.
-        device (torch.device): Device of the points.
-
-    Returns:
-        centers (Tensor): Centers of the bins. Shape: (num,).
-    """
-    edges = torch.linspace(start, stop, num + 1, dtype=dtype, device=device)
-    return (edges[:-1] + edges[1:]) / 2
-
-
-@functools.lru_cache(maxsize=1)
-def create_position_matrix(
-    T: int,
-    pH: int,
-    pW: int,
-    device: torch.device,
-    dtype: torch.dtype,
-    *,
-    target_area: float = 36864,
-):
-    """
-    Args:
-        T: int - Temporal dimension
-        pH: int - Height dimension after patchify
-        pW: int - Width dimension after patchify
-
-    Returns:
-        pos: [T * pH * pW, 3] - position matrix
-    """
-    with torch.no_grad():
-        # Create 1D tensors for each dimension
-        t = torch.arange(T, dtype=dtype)
-
-        # Positionally interpolate to area 36864.
-        # (3072x3072 frame with 16x16 patches = 192x192 latents).
-        # This automatically scales rope positions when the resolution changes.
-        # We use a large target area so the model is more sensitive
-        # to changes in the learned pos_frequencies matrix.
-        scale = math.sqrt(target_area / (pW * pH))
-        w = centers(-pW * scale / 2, pW * scale / 2, pW)
-        h = centers(-pH * scale / 2, pH * scale / 2, pH)
-
-        # Use meshgrid to create 3D grids
-        grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
-
-        # Stack and reshape the grids.
-        pos = torch.stack([grid_t, grid_h, grid_w], dim=-1)  # [T, pH, pW, 3]
-        pos = pos.view(-1, 3)  # [T * pH * pW, 3]
-        pos = pos.to(dtype=dtype, device=device)
-
-    return pos
-
-
-def compute_mixed_rotation(
-    freqs: torch.Tensor,
-    pos: torch.Tensor,
-):
-    """
-    Project each 3-dim position into per-head, per-head-dim 1D frequencies.
-
-    Args:
-        freqs: [3, num_heads, num_freqs] - learned rotation frequency (for t, row, col) for each head position
-        pos: [N, 3] - position of each token
-        num_heads: int
-
-    Returns:
-        freqs_cos: [N, num_heads, num_freqs] - cosine components
-        freqs_sin: [N, num_heads, num_freqs] - sine components
-    """
-    with torch.autocast("cuda", enabled=False):
-        assert freqs.ndim == 3
-        freqs_sum = torch.einsum("Nd,dhf->Nhf", pos.to(freqs), freqs)
-        freqs_cos = torch.cos(freqs_sum)
-        freqs_sin = torch.sin(freqs_sum)
-    return freqs_cos, freqs_sin

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/temporal_rope.py

@@ -1,34 +0,0 @@
-# Based on Llama3 Implementation.
-import torch
-
-
-def apply_rotary_emb_qk_real(
-    xqk: torch.Tensor,
-    freqs_cos: torch.Tensor,
-    freqs_sin: torch.Tensor,
-) -> torch.Tensor:
-    """
-    Apply rotary embeddings to input tensors using the given frequency tensor without complex numbers.
-
-    Args:
-        xqk (torch.Tensor): Query and/or Key tensors to apply rotary embeddings. Shape: (B, S, *, num_heads, D)
-                            Can be either just query or just key, or both stacked along some batch or * dim.
-        freqs_cos (torch.Tensor): Precomputed cosine frequency tensor.
-        freqs_sin (torch.Tensor): Precomputed sine frequency tensor.
-
-    Returns:
-        torch.Tensor: The input tensor with rotary embeddings applied.
-    """
-    assert xqk.dtype == torch.bfloat16
-    # Split the last dimension into even and odd parts
-    xqk_even = xqk[..., 0::2]
-    xqk_odd = xqk[..., 1::2]
-
-    # Apply rotation
-    cos_part = (xqk_even * freqs_cos - xqk_odd * freqs_sin).type_as(xqk)
-    sin_part = (xqk_even * freqs_sin + xqk_odd * freqs_cos).type_as(xqk)
-
-    # Interleave the results back into the original shape
-    out = torch.stack([cos_part, sin_part], dim=-1).flatten(-2)
-    assert out.dtype == torch.bfloat16
-    return out

fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/utils.py

@@ -1,109 +0,0 @@
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-def modulate(x, shift, scale):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-def pool_tokens(x: torch.Tensor, mask: torch.Tensor, *, keepdim=False) -> torch.Tensor:
-    """
-    Pool tokens in x using mask.
-
-    NOTE: We assume x does not require gradients.
-
-    Args:
-        x: (B, L, D) tensor of tokens.
-        mask: (B, L) boolean tensor indicating which tokens are not padding.
-
-    Returns:
-        pooled: (B, D) tensor of pooled tokens.
-    """
-    assert x.size(1) == mask.size(1)  # Expected mask to have same length as tokens.
-    assert x.size(0) == mask.size(0)  # Expected mask to have same batch size as tokens.
-    mask = mask[:, :, None].to(dtype=x.dtype)
-    mask = mask / mask.sum(dim=1, keepdim=True).clamp(min=1)
-    pooled = (x * mask).sum(dim=1, keepdim=keepdim)
-    return pooled
-
-
-class AttentionPool(nn.Module):
-    def __init__(
-        self,
-        embed_dim: int,
-        num_heads: int,
-        output_dim: int = None,
-        device: Optional[torch.device] = None,
-    ):
-        """
-        Args:
-            spatial_dim (int): Number of tokens in sequence length.
-            embed_dim (int): Dimensionality of input tokens.
-            num_heads (int): Number of attention heads.
-            output_dim (int): Dimensionality of output tokens. Defaults to embed_dim.
-        """
-        super().__init__()
-        self.num_heads = num_heads
-        self.to_kv = nn.Linear(embed_dim, 2 * embed_dim, device=device)
-        self.to_q = nn.Linear(embed_dim, embed_dim, device=device)
-        self.to_out = nn.Linear(embed_dim, output_dim or embed_dim, device=device)
-
-    def forward(self, x, mask):
-        """
-        Args:
-            x (torch.Tensor): (B, L, D) tensor of input tokens.
-            mask (torch.Tensor): (B, L) boolean tensor indicating which tokens are not padding.
-
-        NOTE: We assume x does not require gradients.
-
-        Returns:
-            x (torch.Tensor): (B, D) tensor of pooled tokens.
-        """
-        D = x.size(2)
-
-        # Construct attention mask, shape: (B, 1, num_queries=1, num_keys=1+L).
-        attn_mask = mask[:, None, None, :].bool()  # (B, 1, 1, L).
-        attn_mask = F.pad(attn_mask, (1, 0), value=True)  # (B, 1, 1, 1+L).
-
-        # Average non-padding token features. These will be used as the query.
-        x_pool = pool_tokens(x, mask, keepdim=True)  # (B, 1, D)
-
-        # Concat pooled features to input sequence.
-        x = torch.cat([x_pool, x], dim=1)  # (B, L+1, D)
-
-        # Compute queries, keys, values. Only the mean token is used to create a query.
-        kv = self.to_kv(x)  # (B, L+1, 2 * D)
-        q = self.to_q(x[:, 0])  # (B, D)
-
-        # Extract heads.
-        head_dim = D // self.num_heads
-        kv = kv.unflatten(2, (2, self.num_heads, head_dim))  # (B, 1+L, 2, H, head_dim)
-        kv = kv.transpose(1, 3)  # (B, H, 2, 1+L, head_dim)
-        k, v = kv.unbind(2)  # (B, H, 1+L, head_dim)
-        q = q.unflatten(1, (self.num_heads, head_dim))  # (B, H, head_dim)
-        q = q.unsqueeze(2)  # (B, H, 1, head_dim)
-
-        # Compute attention.
-        x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0)  # (B, H, 1, head_dim)
-
-        # Concatenate heads and run output.
-        x = x.squeeze(2).flatten(1, 2)  # (B, D = H * head_dim)
-        x = self.to_out(x)
-        return x
-
-
-def pad_and_split_xy(xy, indices, B, N, L, dtype) -> Tuple[torch.Tensor, torch.Tensor]:
-    D = xy.size(1)
-
-    # Pad sequences to (B, N + L, dim).
-    assert indices.ndim == 1
-    indices = indices.unsqueeze(1).expand(-1, D)  # (total,) -> (total, num_heads * head_dim)
-    output = torch.zeros(B * (N + L), D, device=xy.device, dtype=dtype)
-    output = torch.scatter(output, 0, indices, xy)
-    xy = output.view(B, N + L, D)
-
-    # Split visual and text tokens along the sequence length.
-    return torch.tensor_split(xy, (N,), dim=1)

fastvideo/models/mochi_genmo/mochi_preview/pipelines.py

@@ -1,682 +0,0 @@
-import json
-import os
-import random
-from abc import ABC, abstractmethod
-from contextlib import contextmanager
-from functools import partial
-from typing import Any, Dict, List, Literal, Optional, Union, cast
-
-import numpy as np
-import ray
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import repeat
-from safetensors import safe_open
-from safetensors.torch import load_file
-from torch import nn
-from torch.distributed.fsdp import (
-    BackwardPrefetch,
-    MixedPrecision,
-    ShardingStrategy,
-)
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp.wrap import (
-    lambda_auto_wrap_policy,
-    transformer_auto_wrap_policy,
-)
-from transformers import T5EncoderModel, T5Tokenizer
-from transformers.models.t5.modeling_t5 import T5Block
-
-import genmo.mochi_preview.dit.joint_model.context_parallel as cp
-from genmo.lib.progress import get_new_progress_bar, progress_bar
-from genmo.lib.utils import Timer
-from genmo.mochi_preview.vae.models import (
-    Decoder,
-    Encoder,
-    decode_latents,
-    decode_latents_tiled_full,
-    decode_latents_tiled_spatial,
-)
-from genmo.mochi_preview.vae.vae_stats import dit_latents_to_vae_latents
-
-
-def load_to_cpu(p, weights_only=True):
-    if p.endswith(".safetensors"):
-        return load_file(p)
-    else:
-        assert p.endswith(".pt")
-        return torch.load(p, map_location="cpu", weights_only=weights_only)
-
-
-def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
-    if linear_steps is None:
-        linear_steps = num_steps // 2
-    linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
-    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)
-    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
-    const = quadratic_coef * (linear_steps**2)
-    quadratic_sigma_schedule = [
-        quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps)
-    ]
-    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
-    sigma_schedule = [1.0 - x for x in sigma_schedule]
-    return sigma_schedule
-
-
-T5_MODEL = "google/t5-v1_1-xxl"
-MAX_T5_TOKEN_LENGTH = 256
-
-
-def setup_fsdp_sync(model, device_id, *, param_dtype, auto_wrap_policy) -> FSDP:
-    model = FSDP(
-        model,
-        sharding_strategy=ShardingStrategy.FULL_SHARD,
-        mixed_precision=MixedPrecision(
-            param_dtype=param_dtype,
-            reduce_dtype=torch.float32,
-            buffer_dtype=torch.float32,
-        ),
-        auto_wrap_policy=auto_wrap_policy,
-        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
-        limit_all_gathers=True,
-        device_id=device_id,
-        sync_module_states=True,
-        use_orig_params=True,
-    )
-    torch.cuda.synchronize()
-    return model
-
-
-class ModelFactory(ABC):
-    def __init__(self, **kwargs):
-        self.kwargs = kwargs
-
-    @abstractmethod
-    def get_model(self, *, local_rank: int, device_id: Union[int, Literal["cpu"]], world_size: int) -> Any:
-        assert isinstance(device_id, int) or device_id == "cpu", "device_id must be an integer or 'cpu'"
-        # FSDP does not work when the model is on the CPU
-        if device_id == "cpu":
-            assert world_size == 1, "CPU offload only supports single-GPU inference"
-
-
-class T5ModelFactory(ModelFactory):
-    def __init__(self, model_dir=None):
-        super().__init__()
-        self.model_dir = model_dir or T5_MODEL
-
-    def get_model(self, *, local_rank, device_id, world_size):
-        super().get_model(local_rank=local_rank, device_id=device_id, world_size=world_size)
-        model = T5EncoderModel.from_pretrained(self.model_dir)
-        if world_size > 1:
-            model = setup_fsdp_sync(
-                model,
-                device_id=device_id,
-                param_dtype=torch.float32,
-                auto_wrap_policy=partial(
-                    transformer_auto_wrap_policy,
-                    transformer_layer_cls={
-                        T5Block,
-                    },
-                ),
-            )
-        elif isinstance(device_id, int):
-            model = model.to(torch.device(f"cuda:{device_id}"))  # type: ignore
-        return model.eval()
-
-
-class DitModelFactory(ModelFactory):
-    def __init__(
-        self, *,
-        model_path: str,
-        model_dtype: str,
-        lora_path: Optional[str] = None,
-        attention_mode: Optional[str] = None
-    ):
-        # Infer attention mode if not specified
-        if attention_mode is None:
-            from genmo.lib.attn_imports import flash_varlen_attn  # type: ignore
-            attention_mode = "sdpa" if flash_varlen_attn is None else "flash"
-        print(f"Attention mode: {attention_mode}")
-
-        super().__init__(
-            model_path=model_path,
-            lora_path=lora_path,
-            model_dtype=model_dtype,
-            attention_mode=attention_mode
-        )
-
-    def get_model(
-        self,
-        *,
-        local_rank,
-        device_id,
-        world_size,
-        model_kwargs=None,
-        patch_model_fns=None,
-        strict_load=True,
-        load_checkpoint=True,
-        fast_init=True,
-    ):
-        from genmo.mochi_preview.dit.joint_model.asymm_models_joint import AsymmDiTJoint
-
-        if not model_kwargs:
-            model_kwargs = {}
-
-        lora_sd = None
-        lora_path = self.kwargs["lora_path"]
-        if lora_path is not None:
-            if lora_path.endswith(".safetensors"):
-                lora_sd = {}
-                with safe_open(lora_path, framework="pt") as f:
-                    for k in f.keys():
-                        lora_sd[k] = f.get_tensor(k)
-                    lora_kwargs = json.loads(f.metadata()["kwargs"])
-                    print(f"Loaded LoRA kwargs: {lora_kwargs}")
-            else:
-                lora = load_to_cpu(lora_path, weights_only=False)
-                lora_sd, lora_kwargs = lora["state_dict"], lora["kwargs"]
-
-            model_kwargs.update(cast(dict, lora_kwargs))
-
-        model_args = dict(
-            depth=48,
-            patch_size=2,
-            num_heads=24,
-            hidden_size_x=3072,
-            hidden_size_y=1536,
-            mlp_ratio_x=4.0,
-            mlp_ratio_y=4.0,
-            in_channels=12,
-            qk_norm=True,
-            qkv_bias=False,
-            out_bias=True,
-            patch_embed_bias=True,
-            timestep_mlp_bias=True,
-            timestep_scale=1000.0,
-            t5_feat_dim=4096,
-            t5_token_length=256,
-            rope_theta=10000.0,
-            attention_mode=self.kwargs["attention_mode"],
-            **model_kwargs,
-        )
-
-        if fast_init:
-            model: nn.Module = torch.nn.utils.skip_init(AsymmDiTJoint, **model_args)
-        else:
-            model: nn.Module = AsymmDiTJoint(**model_args)
-
-        for fn in patch_model_fns or []:
-            model = fn(model)
-
-        # FSDP syncs weights from rank 0 to all other ranks
-        if local_rank == 0 and load_checkpoint:
-            model_path = self.kwargs["model_path"]
-            sd = load_to_cpu(model_path)
-
-            # Load the state dictionary and capture the return value
-            load_result = model.load_state_dict(sd, strict=strict_load)
-            if not strict_load:
-                # Print mismatched keys
-                missing_keys = [k for k in load_result.missing_keys if ".lora_" not in k]
-                if missing_keys:
-                    print(f"Missing keys from {model_path}: {missing_keys}")
-                if load_result.unexpected_keys:
-                    print(f"Unexpected keys from {model_path}: {load_result.unexpected_keys}")
-
-            if lora_sd:
-                model.load_state_dict(lora_sd, strict=strict_load) # type: ignore
-
-        if world_size > 1:
-            assert self.kwargs["model_dtype"] == "bf16", "FP8 is not supported for multi-GPU inference"
-
-            model = setup_fsdp_sync(
-                model,
-                device_id=device_id,
-                param_dtype=torch.float32,
-                auto_wrap_policy=partial(
-                    lambda_auto_wrap_policy,
-                    lambda_fn=lambda m: m in model.blocks,
-                ),
-            )
-        elif isinstance(device_id, int):
-            model = model.to(torch.device(f"cuda:{device_id}"))
-        return model.eval()
-
-
-class DecoderModelFactory(ModelFactory):
-    def __init__(self, *, model_path: str):
-        super().__init__(model_path=model_path)
-
-    def get_model(self, *, local_rank=0, device_id=0, world_size=1):
-        # TODO(ved): Set flag for torch.compile
-        # TODO(ved): Use skip_init
-
-        decoder = Decoder(
-            out_channels=3,
-            base_channels=128,
-            channel_multipliers=[1, 2, 4, 6],
-            temporal_expansions=[1, 2, 3],
-            spatial_expansions=[2, 2, 2],
-            num_res_blocks=[3, 3, 4, 6, 3],
-            latent_dim=12,
-            has_attention=[False, False, False, False, False],
-            output_norm=False,
-            nonlinearity="silu",
-            output_nonlinearity="silu",
-            causal=True,
-        )
-        # VAE is not FSDP-wrapped
-        state_dict = load_file(self.kwargs["model_path"])
-        decoder.load_state_dict(state_dict, strict=True)
-        device = torch.device(f"cuda:{device_id}") if isinstance(device_id, int) else "cpu"
-        decoder.eval().to(device)
-        return decoder
-
-
-class EncoderModelFactory(ModelFactory):
-    def __init__(self, *, model_path: str):
-        super().__init__(model_path=model_path)
-
-    def get_model(self, *, local_rank=0, device_id=0, world_size=1):
-        # TODO(ved): Set flag for torch.compile
-        # TODO(ved): Use skip_init
-
-        # We don't FSDP the encoder b/c it is small
-        encoder = Encoder(
-            in_channels=15,
-            base_channels=64,
-            channel_multipliers=[1, 2, 4, 6],
-            num_res_blocks=[3, 3, 4, 6, 3],
-            latent_dim=12,
-            temporal_reductions=[1, 2, 3],
-            spatial_reductions=[2, 2, 2],
-            prune_bottlenecks=[False, False, False, False, False],
-            has_attentions=[False, True, True, True, True],
-            affine=True,
-            bias=True,
-            input_is_conv_1x1=True,
-            padding_mode="replicate",
-        )
-        state_dict = load_file(self.kwargs["model_path"])
-        encoder.load_state_dict(state_dict, strict=True)
-        device = torch.device(f"cuda:{device_id}") if isinstance(device_id, int) else "cpu"
-        encoder.eval().to(device)
-        return encoder
-
-
-def get_conditioning(
-    tokenizer: T5Tokenizer,
-    encoder: Encoder,
-    device: torch.device,
-    batch_inputs: bool,
-    *,
-    prompt: str,
-    negative_prompt: str,
-):
-    if batch_inputs:
-        return dict(
-            batched=get_conditioning_for_prompts(
-                tokenizer, encoder, device, [prompt, negative_prompt]
-            )
-        )
-    else:
-        cond_input = get_conditioning_for_prompts(tokenizer, encoder, device, [prompt])
-        null_input = get_conditioning_for_prompts(tokenizer, encoder, device, [negative_prompt])
-        return dict(cond=cond_input, null=null_input)
-
-
-def get_conditioning_for_prompts(tokenizer, encoder, device, prompts: List[str]):
-    assert len(prompts) in [1, 2]  # [neg] or [pos] or [pos, neg]
-    B = len(prompts)
-    t5_toks = tokenizer(
-        prompts,
-        padding="max_length",
-        truncation=True,
-        max_length=MAX_T5_TOKEN_LENGTH,
-        return_tensors="pt",
-        return_attention_mask=True,
-    )
-    caption_input_ids_t5 = t5_toks["input_ids"]
-    caption_attention_mask_t5 = t5_toks["attention_mask"].bool()
-    del t5_toks
-
-    assert caption_input_ids_t5.shape == (B, MAX_T5_TOKEN_LENGTH)
-    assert caption_attention_mask_t5.shape == (B, MAX_T5_TOKEN_LENGTH)
-
-    # Special-case empty negative prompt by zero-ing it
-    if prompts[-1] == "":
-        caption_input_ids_t5[-1] = 0
-        caption_attention_mask_t5[-1] = False
-
-    caption_input_ids_t5 = caption_input_ids_t5.to(device, non_blocking=True)
-    caption_attention_mask_t5 = caption_attention_mask_t5.to(device, non_blocking=True)
-
-    y_mask = [caption_attention_mask_t5]
-    y_feat = [encoder(caption_input_ids_t5, caption_attention_mask_t5).last_hidden_state.detach()]
-    # Sometimes returns a tensor, othertimes a tuple, not sure why
-    # See: https://huggingface.co/genmo/mochi-1-preview/discussions/3
-    assert tuple(y_feat[-1].shape) == (B, MAX_T5_TOKEN_LENGTH, 4096)
-    assert y_feat[-1].dtype == torch.float32
-
-    return dict(y_mask=y_mask, y_feat=y_feat)
-
-
-def compute_packed_indices(
-    device: torch.device, text_mask: torch.Tensor, num_latents: int
-) -> Dict[str, Union[torch.Tensor, int]]:
-    """
-    Based on https://github.com/Dao-AILab/flash-attention/blob/765741c1eeb86c96ee71a3291ad6968cfbf4e4a1/flash_attn/bert_padding.py#L60-L80
-
-    Args:
-        num_latents: Number of latent tokens
-        text_mask: (B, L) List of boolean tensor indicating which text tokens are not padding.
-
-    Returns:
-        packed_indices: Dict with keys for Flash Attention:
-            - valid_token_indices_kv: up to (B * (N + L),) tensor of valid token indices (non-padding)
-                                   in the packed sequence.
-            - cu_seqlens_kv: (B + 1,) tensor of cumulative sequence lengths in the packed sequence.
-            - max_seqlen_in_batch_kv: int of the maximum sequence length in the batch.
-    """
-    # Create an expanded token mask saying which tokens are valid across both visual and text tokens.
-    PATCH_SIZE = 2
-    num_visual_tokens = num_latents // (PATCH_SIZE**2)
-    assert num_visual_tokens > 0
-
-    mask = F.pad(text_mask, (num_visual_tokens, 0), value=True)  # (B, N + L)
-    seqlens_in_batch = mask.sum(dim=-1, dtype=torch.int32)  # (B,)
-    valid_token_indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten()  # up to (B * (N + L),)
-    assert valid_token_indices.size(0) >= text_mask.size(0) * num_visual_tokens  # At least (B * N,)
-    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
-    max_seqlen_in_batch = seqlens_in_batch.max().item()
-
-    return {
-        "cu_seqlens_kv": cu_seqlens.to(device, non_blocking=True),
-        "max_seqlen_in_batch_kv": cast(int, max_seqlen_in_batch),
-        "valid_token_indices_kv": valid_token_indices.to(device, non_blocking=True),
-    }
-
-
-def assert_eq(x, y, msg=None):
-    assert x == y, f"{msg or 'Assertion failed'}: {x} != {y}"
-
-
-def sample_model(device, dit, conditioning, **args):
-    random.seed(args["seed"])
-    np.random.seed(args["seed"])
-    torch.manual_seed(args["seed"])
-
-    generator = torch.Generator(device=device)
-    generator.manual_seed(args["seed"])
-
-    w, h, t = args["width"], args["height"], args["num_frames"]
-    sample_steps = args["num_inference_steps"]
-    cfg_schedule = args["cfg_schedule"]
-    sigma_schedule = args["sigma_schedule"]
-
-    assert_eq(len(cfg_schedule), sample_steps, "cfg_schedule must have length sample_steps")
-    assert_eq((t - 1) % 6, 0, "t - 1 must be divisible by 6")
-    assert_eq(
-        len(sigma_schedule),
-        sample_steps + 1,
-        "sigma_schedule must have length sample_steps + 1",
-    )
-
-    B = 1
-    SPATIAL_DOWNSAMPLE = 8
-    TEMPORAL_DOWNSAMPLE = 6
-    IN_CHANNELS = 12
-    latent_t = ((t - 1) // TEMPORAL_DOWNSAMPLE) + 1
-    latent_w, latent_h = w // SPATIAL_DOWNSAMPLE, h // SPATIAL_DOWNSAMPLE
-
-    z = torch.randn(
-        (B, IN_CHANNELS, latent_t, latent_h, latent_w),
-        device=device,
-        dtype=torch.float32,
-    )
-
-    num_latents = latent_t * latent_h * latent_w
-    cond_batched = cond_text = cond_null = None
-    if "cond" in conditioning:
-        cond_text = conditioning["cond"]
-        cond_null = conditioning["null"]
-        cond_text["packed_indices"] = compute_packed_indices(device, cond_text["y_mask"][0], num_latents)
-        cond_null["packed_indices"] = compute_packed_indices(device, cond_null["y_mask"][0], num_latents)
-    else:
-        cond_batched = conditioning["batched"]
-        cond_batched["packed_indices"] = compute_packed_indices(device, cond_batched["y_mask"][0], num_latents)
-        z = repeat(z, "b ... -> (repeat b) ...", repeat=2)
-
-    def model_fn(*, z, sigma, cfg_scale):
-        if cond_batched:
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                out = dit(z, sigma, **cond_batched)
-            out_cond, out_uncond = torch.chunk(out, chunks=2, dim=0)
-        else:
-            nonlocal cond_text, cond_null
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                out_cond = dit(z, sigma, **cond_text)
-                out_uncond = dit(z, sigma, **cond_null)
-        assert out_cond.shape == out_uncond.shape
-        out_uncond = out_uncond.to(z)
-        out_cond = out_cond.to(z)
-        return out_uncond + cfg_scale * (out_cond - out_uncond)
-
-    # Euler sampler w/ customizable sigma schedule & cfg scale
-    for i in get_new_progress_bar(range(0, sample_steps), desc="Sampling"):
-        sigma = sigma_schedule[i]
-        dsigma = sigma - sigma_schedule[i + 1]
-
-        # `pred` estimates `z_0 - eps`.
-        pred = model_fn(
-            z=z,
-            sigma=torch.full([B] if cond_text else [B * 2], sigma, device=z.device),
-            cfg_scale=cfg_schedule[i],
-        )
-        assert pred.dtype == torch.float32
-        z = z + dsigma * pred
-
-    z = z[:B] if cond_batched else z
-    return dit_latents_to_vae_latents(z)
-
-
-@contextmanager
-def move_to_device(model: nn.Module, target_device, *, enabled=True):
-    if not enabled:
-        yield
-        return
-
-    og_device = next(model.parameters()).device
-    if og_device == target_device:
-        print(f"move_to_device is a no-op model is already on {target_device}")
-    else:
-        print(f"moving model from {og_device} -> {target_device}")
-
-    model.to(target_device)
-    yield
-    if og_device != target_device:
-        print(f"moving model from {target_device} -> {og_device}")
-    model.to(og_device)
-
-
-def t5_tokenizer(model_dir=None):
-    return T5Tokenizer.from_pretrained(model_dir or T5_MODEL, legacy=False)
-
-
-class MochiSingleGPUPipeline:
-    def __init__(
-        self,
-        *,
-        text_encoder_factory: ModelFactory,
-        dit_factory: ModelFactory,
-        decoder_factory: ModelFactory,
-        cpu_offload: Optional[bool] = False,
-        decode_type: str = "full",
-        decode_args: Optional[Dict[str, Any]] = None,
-        fast_init=True,
-        strict_load=True
-    ):
-        self.device = torch.device("cuda:0")
-        self.tokenizer = t5_tokenizer(text_encoder_factory.model_dir)
-        t = Timer()
-        self.cpu_offload = cpu_offload
-        self.decode_args = decode_args or {}
-        self.decode_type = decode_type
-        init_id = "cpu" if cpu_offload else 0
-        with t("load_text_encoder"):
-            self.text_encoder = text_encoder_factory.get_model(
-                local_rank=0,
-                device_id=init_id,
-                world_size=1,
-            )
-        with t("load_dit"):
-            self.dit = dit_factory.get_model(local_rank=0, device_id=init_id, world_size=1, fast_init=fast_init, strict_load=strict_load) # type: ignore
-        with t("load_vae"):
-            self.decoder = decoder_factory.get_model(local_rank=0, device_id=init_id, world_size=1)
-        t.print_stats()
-
-    def __call__(self, batch_cfg, prompt, negative_prompt, **kwargs):
-        with torch.inference_mode():
-            print_max_memory = lambda: print(
-                f"Max memory reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB"
-            )
-            print_max_memory()
-
-            with move_to_device(self.text_encoder, self.device):
-                conditioning = get_conditioning(
-                    tokenizer=self.tokenizer,
-                    encoder=self.text_encoder,
-                    device=self.device,
-                    batch_inputs=batch_cfg,
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                )
-            print_max_memory()
-
-            with move_to_device(self.dit, self.device):
-                latents = sample_model(self.device, self.dit, conditioning, **kwargs)
-            print_max_memory()
-
-            with move_to_device(self.decoder, self.device):
-                if self.decode_type == "tiled_full":
-                    frames = decode_latents_tiled_full(
-                        self.decoder, latents, **self.decode_args)
-                elif self.decode_type == "tiled_spatial":
-                    frames = decode_latents_tiled_spatial(
-                        self.decoder, latents, **self.decode_args,
-                        num_tiles_w=4, num_tiles_h=2)
-                else:
-                    frames = decode_latents(self.decoder, latents)
-            print_max_memory()
-            return frames.cpu().numpy()
-
-
-def cast_dit(model, dtype):
-    for name, module in model.named_modules():
-        if isinstance(module, nn.Linear):
-            assert any(
-                n in name for n in ["mlp", "t5", "mod_", "attn.qkv_", "attn.proj_", "final_layer"]
-            ), f"Unexpected linear layer: {name}"
-            module.to(dtype=dtype)
-        elif isinstance(module, nn.Conv2d):
-            assert "x_embedder.proj" in name, f"Unexpected conv2d layer: {name}"
-            module.to(dtype=dtype)
-    return model
-
-
-### ALL CODE BELOW HERE IS FOR MULTI-GPU MODE ###
-
-
-# In multi-gpu mode, all models must belong to a device which has a predefined context parallel group
-# So it doesn't make sense to work with models individually
-class MultiGPUContext:
-    def __init__(
-        self,
-        *,
-        text_encoder_factory,
-        dit_factory,
-        decoder_factory,
-        device_id,
-        local_rank,
-        world_size,
-    ):
-        t = Timer()
-        self.device = torch.device(f"cuda:{device_id}")
-        print(f"Initializing rank {local_rank+1}/{world_size}")
-        assert world_size > 1, f"Multi-GPU mode requires world_size > 1, got {world_size}"
-        os.environ["MASTER_ADDR"] = "127.0.0.1"
-        os.environ["MASTER_PORT"] = "29500"
-        with t("init_process_group"):
-            dist.init_process_group(
-                "nccl",
-                rank=local_rank,
-                world_size=world_size,
-                device_id=self.device,  # force non-lazy init
-            )
-        pg = dist.group.WORLD
-        cp.set_cp_group(pg, list(range(world_size)), local_rank)
-        distributed_kwargs = dict(local_rank=local_rank, device_id=device_id, world_size=world_size)
-        self.world_size = world_size
-        self.tokenizer = t5_tokenizer(text_encoder_factory.model_dir)
-        with t("load_text_encoder"):
-            self.text_encoder = text_encoder_factory.get_model(**distributed_kwargs)
-        with t("load_dit"):
-            self.dit = dit_factory.get_model(**distributed_kwargs)
-        with t("load_vae"):
-            self.decoder = decoder_factory.get_model(**distributed_kwargs)
-        self.local_rank = local_rank
-        t.print_stats()
-
-    def run(self, *, fn, **kwargs):
-        return fn(self, **kwargs)
-
-
-class MochiMultiGPUPipeline:
-    def __init__(
-        self,
-        *,
-        text_encoder_factory: ModelFactory,
-        dit_factory: ModelFactory,
-        decoder_factory: ModelFactory,
-        world_size: int,
-    ):
-        ray.init()
-        RemoteClass = ray.remote(MultiGPUContext)
-        self.ctxs = [
-            RemoteClass.options(num_gpus=1).remote(
-                text_encoder_factory=text_encoder_factory,
-                dit_factory=dit_factory,
-                decoder_factory=decoder_factory,
-                world_size=world_size,
-                device_id=0,
-                local_rank=i,
-            )
-            for i in range(world_size)
-        ]
-        for ctx in self.ctxs:
-            ray.get(ctx.__ray_ready__.remote())
-
-    def __call__(self, **kwargs):
-        def sample(ctx, *, batch_cfg, prompt, negative_prompt, **kwargs):
-            with progress_bar(type="ray_tqdm", enabled=ctx.local_rank == 0), torch.inference_mode():
-                conditioning = get_conditioning(
-                    ctx.tokenizer,
-                    ctx.text_encoder,
-                    ctx.device,
-                    batch_cfg,
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                )
-                latents = sample_model(ctx.device, ctx.dit, conditioning=conditioning, **kwargs)
-                if ctx.local_rank == 0:
-                    torch.save(latents, "latents.pt")
-                frames = decode_latents(ctx.decoder, latents)
-            return frames.cpu().numpy()
-
-        return ray.get([ctx.run.remote(fn=sample, **kwargs, show_progress=i == 0) for i, ctx in enumerate(self.ctxs)])[
-            0
-        ]

fastvideo/models/mochi_genmo/mochi_preview/vae/cp_conv.py

@@ -1,155 +0,0 @@
-from typing import Tuple, Union
-
-import torch
-import torch.distributed as dist
-import torch.nn.functional as F
-
-import genmo.mochi_preview.dit.joint_model.context_parallel as cp
-
-
-def cast_tuple(t, length=1):
-    return t if isinstance(t, tuple) else ((t,) * length)
-
-
-def cp_pass_frames(x: torch.Tensor, frames_to_send: int) -> torch.Tensor:
-    """
-    Forward pass that handles communication between ranks for inference.
-    Args:
-        x: Tensor of shape (B, C, T, H, W)
-        frames_to_send: int, number of frames to communicate between ranks
-    Returns:
-        output: Tensor of shape (B, C, T', H, W)
-    """
-    cp_rank, cp_world_size = cp.get_cp_rank_size()
-    if frames_to_send == 0 or cp_world_size == 1:
-        return x
-
-    group = cp.get_cp_group()
-    global_rank = dist.get_rank()
-
-    # Send to next rank
-    if cp_rank < cp_world_size - 1:
-        assert x.size(2) >= frames_to_send
-        tail = x[:, :, -frames_to_send:].contiguous()
-        dist.send(tail, global_rank + 1, group=group)
-
-    # Receive from previous rank
-    if cp_rank > 0:
-        B, C, _, H, W = x.shape
-        recv_buffer = torch.empty(
-            (B, C, frames_to_send, H, W),
-            dtype=x.dtype,
-            device=x.device,
-        )
-        dist.recv(recv_buffer, global_rank - 1, group=group)
-        x = torch.cat([recv_buffer, x], dim=2)
-
-    return x
-
-
-def _pad_to_max(x: torch.Tensor, max_T: int) -> torch.Tensor:
-    if max_T > x.size(2):
-        pad_T = max_T - x.size(2)
-        pad_dims = (0, 0, 0, 0, 0, pad_T)
-        return F.pad(x, pad_dims)
-    return x
-
-
-def gather_all_frames(x: torch.Tensor) -> torch.Tensor:
-    """
-    Gathers all frames from all processes for inference.
-    Args:
-        x: Tensor of shape (B, C, T, H, W)
-    Returns:
-        output: Tensor of shape (B, C, T_total, H, W)
-    """
-    cp_rank, cp_size = cp.get_cp_rank_size()
-    if cp_size == 1:
-        return x
-
-    cp_group = cp.get_cp_group()
-
-    # Ensure the tensor is contiguous for collective operations
-    x = x.contiguous()
-
-    # Get the local time dimension size
-    local_T = x.size(2)
-    local_T_tensor = torch.tensor([local_T], device=x.device, dtype=torch.int64)
-
-    # Gather all T sizes from all processes
-    all_T = [torch.zeros(1, dtype=torch.int64, device=x.device) for _ in range(cp_size)]
-    dist.all_gather(all_T, local_T_tensor, group=cp_group)
-    all_T = [t.item() for t in all_T]
-
-    # Pad the tensor at the end of the time dimension to match max_T
-    max_T = max(all_T)
-    x = _pad_to_max(x, max_T).contiguous()
-
-    # Prepare a list to hold the gathered tensors
-    gathered_x = [torch.zeros_like(x).contiguous() for _ in range(cp_size)]
-
-    # Perform the all_gather operation
-    dist.all_gather(gathered_x, x, group=cp_group)
-
-    # Slice each gathered tensor back to its original T size
-    for idx, t_size in enumerate(all_T):
-        gathered_x[idx] = gathered_x[idx][:, :, :t_size]
-
-    return torch.cat(gathered_x, dim=2)
-
-
-def excessive_memory_usage(input: torch.Tensor, max_gb: float = 2.0) -> bool:
-    """Estimate memory usage based on input tensor size and data type."""
-    element_size = input.element_size()  # Size in bytes of each element
-    memory_bytes = input.numel() * element_size
-    memory_gb = memory_bytes / 1024**3
-    return memory_gb > max_gb
-
-
-class ContextParallelCausalConv3d(torch.nn.Conv3d):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size: Union[int, Tuple[int, int, int]],
-        stride: Union[int, Tuple[int, int, int]],
-        **kwargs,
-    ):
-        kernel_size = cast_tuple(kernel_size, 3)
-        stride = cast_tuple(stride, 3)
-        height_pad = (kernel_size[1] - 1) // 2
-        width_pad = (kernel_size[2] - 1) // 2
-
-        super().__init__(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            stride=stride,
-            dilation=(1, 1, 1),
-            padding=(0, height_pad, width_pad),
-            **kwargs,
-        )
-
-    def forward(self, x: torch.Tensor):
-        cp_rank, cp_world_size = cp.get_cp_rank_size()
-
-        context_size = self.kernel_size[0] - 1
-        if cp_rank == 0:
-            mode = "constant" if self.padding_mode == "zeros" else self.padding_mode
-            x = F.pad(x, (0, 0, 0, 0, context_size, 0), mode=mode)
-
-        if cp_world_size == 1:
-            return super().forward(x)
-
-        if all(s == 1 for s in self.stride):
-            # Receive some frames from previous rank.
-            x = cp_pass_frames(x, context_size)
-            return super().forward(x)
-
-        # Less efficient implementation for strided convs.
-        # All gather x, infer and chunk.
-        x = gather_all_frames(x)  # [B, C, k - 1 + global_T, H, W]
-        x = super().forward(x)
-        x_chunks = x.tensor_split(cp_world_size, dim=2)
-        assert len(x_chunks) == cp_world_size
-        return x_chunks[cp_rank]

fastvideo/models/mochi_genmo/mochi_preview/vae/latent_dist.py

@@ -1,35 +0,0 @@
-"""Container for latent space posterior."""
-
-import torch
-
-
-class LatentDistribution:
-    def __init__(self, mean: torch.Tensor, logvar: torch.Tensor):
-        """Initialize latent distribution.
-
-        Args:
-            mean: Mean of the distribution. Shape: [B, C, T, H, W].
-            logvar: Logarithm of variance of the distribution. Shape: [B, C, T, H, W].
-        """
-        assert mean.shape == logvar.shape
-        self.mean = mean
-        self.logvar = logvar
-
-    def sample(self, temperature=1.0, generator: torch.Generator = None, noise=None):
-        if temperature == 0.0:
-            return self.mean
-
-        if noise is None:
-            noise = torch.randn(self.mean.shape, device=self.mean.device, dtype=self.mean.dtype, generator=generator)
-        else:
-            assert noise.device == self.mean.device
-            noise = noise.to(self.mean.dtype)
-
-        if temperature != 1.0:
-            raise NotImplementedError(f"Temperature {temperature} is not supported.")
-
-        # Just Gaussian sample with no scaling of variance.
-        return noise * torch.exp(self.logvar * 0.5) + self.mean
-
-    def mode(self):
-        return self.mean