fix test

.github/workflows/build-image-template.yml

@@ -0,0 +1,106 @@
+name: Build Image Template
+
+on:
+  workflow_call:
+    inputs:
+      python_version:
+        required: true
+        type: string
+      dockerfile_path:
+        required: true
+        type: string
+      tag_suffix:
+        required: true
+        type: string
+
+jobs:
+  build-and-push:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      
+      - name: Free up disk space
+        run: |
+          # Display initial space
+          echo "Initial disk space:"
+          df -h
+          
+          # Remove large directories directly
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          sudo rm -rf /usr/share/swift
+          sudo rm -rf /usr/local/lib/node_modules
+          sudo rm -rf /usr/local/share/powershell
+          sudo rm -rf /usr/share/rust
+          sudo rm -rf /usr/local/.ghcup
+          
+          # Remove cached files
+          sudo rm -rf /var/lib/apt/lists/*
+          sudo rm -rf /var/cache/apt/archives/*
+          
+          # Clean Docker
+          docker system prune -af --volumes
+          
+          # Display available space after cleanup
+          echo "Disk space after cleanup:"
+          df -h
+      
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      
+      - name: Prepare tags
+        id: prepare-tags
+        run: |
+          SHORT_SHA=$(echo ${{ github.sha }} | cut -c1-7)
+          
+          TAGS="type=raw,value=${{ inputs.tag_suffix }}-latest"
+          TAGS="${TAGS}\ntype=raw,value=${{ inputs.tag_suffix }}-sha-${SHORT_SHA}"
+          
+          # Set Python 3.10 as the default image
+          if [[ "${{ inputs.python_version }}" == "3.10" ]]; then
+            TAGS="${TAGS}\ntype=raw,value=latest"
+          fi
+          
+          {
+            echo "tags<<EOF"
+            echo -e "$TAGS"
+            echo "EOF" 
+          } >> $GITHUB_OUTPUT
+      
+      - name: Extract metadata for Docker
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ghcr.io/${{ github.repository }}/fastvideo-dev
+          tags: ${{ steps.prepare-tags.outputs.tags }}
+      
+      - name: Build and push Docker image
+        id: build-push
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: ${{ inputs.dockerfile_path }}
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+      
+      - name: Success message
+        run: |
+          echo "✅ Python ${{ inputs.python_version }} image successfully built and pushed to ghcr.io/${{ github.repository }}/fastvideo-dev:${{ inputs.tag_suffix }}-latest"
+          echo "To run tests with this image, manually trigger the 'Run Tests' workflow." 

.github/workflows/build-image.yml

@@ -1,78 +1,52 @@
-name: Build and Push Docker Image
+name: Build and Push Docker Images
 
 on:
-  workflow_dispatch:  # Only manual triggers
+  workflow_dispatch:
+    inputs:
+      python_3_10:
+        description: 'Build Python 3.10 image'
+        required: false
+        default: false
+        type: boolean
+      python_3_11:
+        description: 'Build Python 3.11 image'
+        required: false
+        default: false
+        type: boolean
+      python_3_12:
+        description: 'Build Python 3.12 image'
+        required: false
+        default: false
+        type: boolean
+
+permissions:
+  contents: read
+  packages: write
 
 jobs:
-  build-and-push:
-    runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      packages: write
-    
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-      
-      - name: Free up disk space
-        run: |
-          # Display initial space
-          echo "Initial disk space:"
-          df -h
-          
-          # Remove large directories directly
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /usr/local/lib/android
-          sudo rm -rf /opt/ghc
-          sudo rm -rf /usr/local/share/boost
-          sudo rm -rf /usr/share/swift
-          sudo rm -rf /usr/local/lib/node_modules
-          sudo rm -rf /usr/local/share/powershell
-          sudo rm -rf /usr/share/rust
-          sudo rm -rf /usr/local/.ghcup
-          
-          # Remove cached files
-          sudo rm -rf /var/lib/apt/lists/*
-          sudo rm -rf /var/cache/apt/archives/*
-          
-          # Clean Docker
-          docker system prune -af --volumes
-          
-          # Display available space after cleanup
-          echo "Disk space after cleanup:"
-          df -h
-      
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-      
-      - name: Login to GitHub Container Registry
-        uses: docker/login-action@v3
-        with:
-          registry: ghcr.io
-          username: ${{ github.repository_owner }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-      
-      - name: Extract metadata for Docker
-        id: meta
-        uses: docker/metadata-action@v5
-        with:
-          images: ghcr.io/${{ github.repository }}/fastvideo-dev
-          tags: |
-            type=raw,value=latest
-            type=sha,format=short
-      
-      - name: Build and push Docker image
-        id: build-push
-        uses: docker/build-push-action@v6
-        with:
-          context: .
-          push: true
-          tags: ${{ steps.meta.outputs.tags }}
-          labels: ${{ steps.meta.outputs.labels }}
-          cache-from: type=gha
-          cache-to: type=gha,mode=max
-      
-      - name: Success message
-        run: |
-          echo "✅ Image successfully built and pushed to ghcr.io/${{ github.repository }}/fastvideo-dev:latest"
-          echo "To run tests with this image, manually trigger the 'Run Tests' workflow." 
+  build-python-3-10:
+    if: ${{ github.event.inputs.python_3_10 == 'true' }}
+    uses: ./.github/workflows/build-image-template.yml
+    with:
+      python_version: '3.10'
+      dockerfile_path: docker/Dockerfile.python3.10
+      tag_suffix: py3.10
+    secrets: inherit
+
+  build-python-3-11:
+    if: ${{ github.event.inputs.python_3_11 == 'true' }}
+    uses: ./.github/workflows/build-image-template.yml
+    with:
+      python_version: '3.11'
+      dockerfile_path: docker/Dockerfile.python3.11
+      tag_suffix: py3.11
+    secrets: inherit
+
+  build-python-3-12:
+    if: ${{ github.event.inputs.python_3_12 == 'true' }}
+    uses: ./.github/workflows/build-image-template.yml
+    with:
+      python_version: '3.12'
+      dockerfile_path: docker/Dockerfile.python3.12
+      tag_suffix: py3.12
+    secrets: inherit

.pre-commit-config.yaml

@@ -19,9 +19,11 @@ exclude: |
     fastvideo/sample/.*|
     fastvideo/train\.py|
     fastvideo/utils/.*|
-    fastvideo/v1/examples/.*|
+    examples/.*|
+    fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py|
     .github/workflows/fastvideo-publish.yml|
-    .github/workflows/sta-publish.yml
+    .github/workflows/sta-publish.yml|
+    .github/workflows/build-image-template.yml
   )
 repos:
 - repo: https://github.com/google/yapf

docker/Dockerfile.python3.10

@@ -29,7 +29,7 @@ RUN echo "# Placeholder" > README.md
 
 RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
     conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
-    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.0.post2 --no-build-isolation && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
     conda clean -afy
 
 COPY . .

docker/Dockerfile.python3.11

@@ -0,0 +1,48 @@
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+RUN conda create --name fastvideo-dev python=3.11.11 -y
+
+SHELL ["/bin/bash", "-c"]
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject.toml ./
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
+    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    conda clean -afy
+
+COPY . .
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+
+# Remove authentication headers
+RUN git config --unset-all http.https://github.com/.extraheader || true
+
+# Set up automatic conda environment activation for all shells
+RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    # Ensure .bashrc is sourced for SSH login shells
+    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+EXPOSE 22

docker/Dockerfile.python3.12

@@ -0,0 +1,48 @@
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+RUN conda create --name fastvideo-dev python=3.12.9 -y
+
+SHELL ["/bin/bash", "-c"]
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject.toml ./
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
+    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    conda clean -afy
+
+COPY . .
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+
+# Remove authentication headers
+RUN git config --unset-all http.https://github.com/.extraheader || true
+
+# Set up automatic conda environment activation for all shells
+RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    # Ensure .bashrc is sourced for SSH login shells
+    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+EXPOSE 22

docs/requirements-docs.txt

@@ -1,25 +1,15 @@
-sphinx==6.2.1
-sphinx-argparse==0.4.0
-sphinx-book-theme==1.0.1
+sphinx==7.4.7
+sphinx-argparse==0.5.2
+sphinx-autodoc2==0.5.0
+sphinx-book-theme==1.1.4
 sphinx-copybutton==0.5.2
 sphinx-design==0.6.1
 sphinx-togglebutton==0.3.2
 myst-parser==3.0.1
 msgspec
-cloudpickle
+commonmark # Required by sphinx-argparse when using :markdownhelp:
 
 # packages to install to build the documentation
 cachetools
-pydantic >= 2.8
 -f https://download.pytorch.org/whl/cpu
-torch
-py-cpuinfo
-transformers
-mistral_common >= 1.5.4
-aiohttp
-starlette
-openai # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
-fastapi # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
-partial-json-parser # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
-requests
-zmq
+torch

docs/source/api/summary.md

@@ -0,0 +1,19 @@
+# Summary
+
+## Video Generator
+
+```{autodoc2-summary}
+    fastvideo.VideoGenerator
+```
+
+## Initialization Configuration
+
+```{autodoc2-summary}
+    fastvideo.v1.configs.pipelines.PipelineConfig
+```
+
+## Sampling Configuration
+
+```{autodoc2-summary}
+    fastvideo.v1.configs.sample.SamplingParam
+```

docs/source/autodoc2_docstring_parser.py

@@ -0,0 +1,22 @@
+# type: ignore
+# SPDX-License-Identifier: Apache-2.0
+from docutils import nodes
+from myst_parser.parsers.sphinx_ import MystParser
+from sphinx.ext.napoleon import docstring
+
+
+class NapoleonParser(MystParser):
+
+    def parse(self, input_string: str, document: nodes.document) -> None:
+        # Get the Sphinx configuration
+        config = document.settings.env.config
+
+        parsed_content = str(
+            docstring.GoogleDocstring(
+                str(docstring.NumpyDocstring(input_string, config)),
+                config,
+            ))
+        return super().parse(parsed_content, document)
+
+
+Parser = NapoleonParser

docs/source/conf.py

@@ -13,17 +13,19 @@
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 
 import datetime
-import inspect
 import logging
 import os
+import re
 import sys
+from pathlib import Path
 from typing import Optional
 
 import requests
-from sphinx.ext import autodoc
 
 logger = logging.getLogger(__name__)
-sys.path.append(os.path.abspath("../.."))
+REPO_ROOT = Path(__file__).resolve().parent.parent.parent
+print(os.path.abspath(REPO_ROOT))
+sys.path.append(os.path.abspath(REPO_ROOT))
 
 # -- Project information -----------------------------------------------------
 
@@ -41,8 +43,7 @@ extensions = [
     "sphinx.ext.linkcode",
     "sphinx.ext.intersphinx",
     "sphinx_copybutton",
-    "sphinx.ext.autodoc",
-    "sphinx.ext.autosummary",
+    "autodoc2",
     "myst_parser",
     "sphinxarg.ext",
     "sphinx_design",
@@ -50,6 +51,31 @@ extensions = [
 ]
 myst_enable_extensions = [
     "colon_fence",
+    "fieldlist",
+]
+autodoc2_packages = [
+    {
+        "path": "../../fastvideo",
+        "exclude_dirs": ["__pycache__", "third_party"],
+    },
+]
+autodoc2_output_dir = "api"
+autodoc2_render_plugin = "myst"
+autodoc2_hidden_objects = ["dunder", "private", "inherited"]
+autodoc2_docstring_parser_regexes = [
+    (".*", "docs.source.autodoc2_docstring_parser"),
+]
+autodoc2_sort_names = True
+autodoc2_index_template = None
+autodoc2_skip_module_regexes = [
+    "fastvideo.dataset",
+    "fastvideo.distill",
+    "fastvideo.data_preprocess",
+    "fastvideo.models",
+    "fastvideo.sample",
+    "fastvideo.utils",
+    "fastvideo.distill_adv",
+    "fastvideo.train",
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -78,6 +104,11 @@ html_theme_options = {
     'repository_url': 'https://github.com/hao-ai-lab/FastVideo/',
     'use_repository_button': True,
     'use_edit_page_button': True,
+    # Prevents the full API being added to the left sidebar of every page.
+    # Reduces build time by 2.5x and reduces build size from ~225MB to ~95MB.
+    'collapse_navbar': True,
+    # Makes API visible in the right sidebar on API reference pages.
+    'show_toc_level': 3,
 }
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
@@ -160,38 +191,38 @@ def linkcode_resolve(domain, info):
         return None
     if not info['module']:
         return None
-    module = info['module']
 
-    # try to determine the correct file and line number to link to
-    obj = sys.modules[module]
+    # Get path from module name
+    file = Path(f"{info['module'].replace('.', '/')}.py")
+    path = REPO_ROOT / file
+    if not path.exists():
+        path = REPO_ROOT / file.with_suffix("") / "__init__.py"
+    if not path.exists():
+        return None
 
-    # get as specific as we can
-    lineno: int = 0
-    filename: str = ""
-    try:
-        for part in info['fullname'].split('.'):
-            obj = getattr(obj, part)
+    # Get the line number of the object
+    with open(path) as f:
+        lines = f.readlines()
+    name = info['fullname'].split(".")[-1]
+    pattern = fr"^( {{4}})*((def|class) )?{name}\b.*"
+    for lineno, line in enumerate(lines, 1):
+        if not line or line.startswith("#"):
+            continue
+        if re.match(pattern, line):
+            break
 
-            if not (inspect.isclass(obj) or inspect.isfunction(obj)
-                    or inspect.ismethod(obj)):
-                obj = obj.__class__  # type: ignore[assignment]
+    # If the line number is not found, return None
+    if lineno == len(lines):
+        return None
 
-            lineno = inspect.getsourcelines(obj)[1]
-            filename = (inspect.getsourcefile(obj)
-                        or f"{filename}.py").split("FastVideo/", 1)[1]
-    except Exception:
-        # For some things, like a class member, won't work, so
-        # we'll use the line number of the parent (the class)
-        pass
-
-    if filename.startswith("checkouts/"):
+    # If the line number is found, create the URL
+    filename = path.relative_to(REPO_ROOT)
+    if "checkouts" in path.parts:
         # a PR build on readthedocs
-        pr_number = filename.split("/")[1]
-        filename = filename.split("/", 2)[2]
+        pr_number = REPO_ROOT.name
         base, branch = get_repo_base_and_branch(pr_number)
         if base and branch:
             return f"https://github.com/{base}/blob/{branch}/{filename}#L{lineno}"
-
     # Otherwise, link to the source file on the main branch
     return f"https://github.com/hao-ai-lab/FastVideo/blob/main/{filename}#L{lineno}"
 
@@ -203,6 +234,8 @@ autodoc_mock_imports = [
     "cpuinfo",
     "cv2",
     "torch",
+    "huggingface_hub",
+    "torchvision",
     "transformers",
     "psutil",
     "prometheus_client",
@@ -231,18 +264,6 @@ for mock_target in autodoc_mock_imports:
             "been loaded into sys.modules when the sphinx build starts.",
             mock_target)
 
-
-class MockedClassDocumenter(autodoc.ClassDocumenter):
-    """Remove note about base class when a class is derived from object."""
-
-    def add_line(self, line: str, source: str, *lineno: int) -> None:
-        if line == "   Bases: :py:class:`object`":
-            return
-        super().add_line(line, source, *lineno)
-
-
-autodoc.ClassDocumenter = MockedClassDocumenter
-
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
     "typing_extensions":
@@ -254,7 +275,4 @@ intersphinx_mapping = {
     "psutil": ("https://psutil.readthedocs.io/en/stable", None),
 }
 
-autodoc_preserve_defaults = True
-autodoc_warningiserror = True
-
 navigation_with_keys = False

docs/source/contributing/developer_env/docker.md

@@ -1,3 +1,4 @@
+(docker)=
 # 🐳 Using the FastVideo Docker Image
 
 If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:

docs/source/contributing/overview.md

@@ -39,7 +39,7 @@ Now you can install FastVideo and setup git hooks for running linting. By using
 pip install -e .[dev]
 
 # Can also install flash-attn (optional)
-pip install flash-attn==2.7.0.post2 --no-build-isolation 
+pip install flash-attn==2.7.4.post1 --no-build-isolation 
 
 # Linting, formatting and static type checking
 pre-commit install --hook-type pre-commit --hook-type commit-msg

docs/source/generate_examples.py

@@ -9,7 +9,7 @@ from typing import Optional
 
 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
-EXAMPLE_DIR = ROOT_DIR / "fastvideo/v1/examples"
+EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
 
 

docs/source/getting_started/installation.md

@@ -4,32 +4,33 @@
 
 FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
 
-FastVideo has been tested on the following GPUs, but it should work on any GPUs that supports CUDA 12.4+, please create an issue if you discover any issues:
-- RTX 4090
-- A40
-- L40S
-- A100
-- H100
-
 ## Requirements
 
-- OS: Linux
-- Python: 3.10-3.12
-- CUDA 12.4+
+- **OS: Linux**
+- **Python: 3.10-3.12**
+- **CUDA 12.4**
+- **At least 1 NVIDIA GPU**
 
-## Installation Options
-
-### Option 1: Quick Install
+## Quick start
+We recommend using an environment manager such as [Conda](#miniconda-recommended)
 
 ```bash
 pip install fastvideo
 ```
 
-### Option 2: Installation from Source
+Also optionally install flash-attn:
 
-We recommend using a Python environment such as Conda.
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
 
-#### 1. [Optional] Install Miniconda (if not already installed)
+[Click here to install from source](#installation-from-source)
+
+## Environment Setup
+
+We highly recommend using a `conda` or `virtualenv` environment:
+### Miniconda (Recommended)
+#### 1. Install Miniconda (if not already installed)
 
 ```bash
 wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
@@ -37,20 +38,22 @@ bash Miniconda3-latest-Linux-x86_64.sh
 source ~/.bashrc
 ```
 
-#### 2. [Optional] Create and activate a Conda environment for FastVideo
+#### 2. Create and activate a Conda environment for FastVideo
 
 ```bash
 conda create -n fastvideo python=3.10 -y
 conda activate fastvideo
 ```
 
-#### 3. Clone the FastVideo repository
+## Installation from Source
+
+### 1. Clone the FastVideo repository
 
 ```bash
 git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 ```
 
-#### 4. Install FastVideo
+#### 2. Install FastVideo
 
 Basic installation:
 
@@ -63,13 +66,16 @@ pip install -e .
 ### Flash Attention
 
 ```bash
-pip install flash-attn==2.7.0.post2 --no-build-isolation
+pip install flash-attn==2.7.4.post1 --no-build-isolation
 ```
 
-### Sliding Tile Attention (STA) (Requires CUDA 12.4+ and H100)
+### Sliding Tile Attention (STA) (Requires CUDA 12.4 and H100)
 
 To try Sliding Tile Attention (optional), please follow the instructions in [csrc/sliding_tile_attention/README.md](#sta-installation) to install STA.
 
+## Docker Images
+We also have prebuilt docker images with FastVideo dependencies pre-installed:
+[Docker Images](#docker)
 ## Development Environment Setup
 
 If you're planning to contribute to FastVideo please see the following page:
@@ -78,8 +84,7 @@ If you're planning to contribute to FastVideo please see the following page:
 ## Hardware Requirements
 
 ### For Basic Inference
-- NVIDIA GPU with CUDA support
-- Minimum 20GB VRAM for quantized models (e.g., single RTX 4090)
+- NVIDIA GPU with CUDA 12.4support
 
 ### For Lora Finetuning
 - 40GB GPU memory each for 2 GPUs with lora

docs/source/getting_started/quick_start.md

@@ -0,0 +1,143 @@
+# Quick Start
+
+This page contains instructions and code to get you quickly started with video generation using FastVideo.
+
+## Requirements
+- **OS: Linux**
+- **Python: 3.10-3.12**
+- **CUDA 12.4**
+- **At least 1 NVIDIA GPU**
+
+## Installation
+
+We recommend using a environment manager such as `Conda`.
+
+```bash
+pip install fastvideo
+```
+
+Also see the [Installation Guide](installation.md).
+
+## Generating Your First Video
+Here's a minimal example to generate a video using the default settings. All of the following code snippets can be directly copied into a Python file and executed with
+
+```bash
+python example.py
+```
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(prompt)
+if __name__ == '__main__':
+    main()
+```
+
+The generated video will be saved in the current directory under `outputs/` by default.
+
+## Customizing Generation
+
+You can customize various parameters when generating videos:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam
+
+def main():
+    # Create the generator
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,
+    )
+
+    # Create and customize sampling parameters
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    # How many frames to generate
+    sampling_param.num_frames = 45
+    # How many steps we denoise the video
+    sampling_param.num_inference_steps = 30
+    # How strongly does the video to conform to the prompt
+    sampling_param.guidance_scale = 7.5
+
+    # Optional: provide an initial image for image-to-video generation
+    sampling_param.image_path = "path/to/your/image.jpg"  # Optional
+
+    # Generate video with custom parameters
+    prompt = "A beautiful sunset over a calm ocean, with gentle waves."
+    video = generator.generate_video(
+        prompt, 
+        sampling_param=sampling_param, 
+        output_path="my_videos/", # controls where videos are saved
+        return_frames=True # also return frames from this call (defaults to False)
+    )
+
+    # `video` now contains frames
+if __name__ == '__main__':
+    main()
+```
+
+## Available Models
+
+FastVideo supports various models for text-to-video generation:
+
+- `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` - A popular text-to-video model
+- `FastVideo/FastHunyuan-Diffusers` - A high-performance model for video generation
+
+## Advanced Configuration
+
+You can use PipelineConfig for more advanced customization:
+
+```python
+from fastvideo import VideoGenerator, PipelineConfig
+
+# Load the default configuration for a model
+config = PipelineConfig.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+
+# Modify configuration settings
+config.vae_config.scale_factor = 16
+config.vae_config.precision = "fp16"
+
+# Create generator with custom config
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=4,
+    pipeline_config=config,
+)
+
+# Generate video
+prompt = "A futuristic cityscape with flying cars and neon signs."
+video = generator.generate_video(prompt)
+```
+
+## Multi-GPU Setup
+
+FastVideo automatically handles distributed setup when multiple GPUs are specified:
+
+```python
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=4,  # Will use 4 GPUs in parallel
+)
+```
+
+## Hardware Requirements
+
+- **Basic Inference**: Minimum 20GB VRAM for quantized models (e.g., single RTX 4090)
+- **LoRA Finetuning**: 40GB GPU memory each for 2 GPUs with LoRA
+- **Full Finetuning/Distillation**: Multiple high-memory GPUs recommended (e.g., H100)
+
+## Next Steps
+
+- Explore the [API Reference](../api/index.md) for more details
+- Learn about [Advanced Inference Options](../inference/overview_back.md)
+- See [Examples](../examples/index.md) for more usage scenarios

docs/source/getting_started/v1_api.md

@@ -0,0 +1,18 @@
+# V1 API
+
+FastVideo's V1 API provides a streamlined interface for video generation tasks with powerful customization options. This page documents the primary components of the API.
+
+## Video Generator
+
+This class will be the primary Python API for generating videos and images.
+
+```{autodoc2-summary}
+    fastvideo.VideoGenerator
+```
+
+VideoGenerator.from_pretrained() should be the primary way of creating a new video generator.
+
+````{py:method} from_pretrained(model_path: str, device: typing.optional[str] = none, torch_dtype: typing.optional[torch.dtype] = none, pipeline_config: typing.optional[typing.union[str | fastvideo.v1.configs.pipelines.pipelineconfig]] = none, **kwargs) -> v1.entrypoints.video_generator.videogenerator
+:canonical: v1.entrypoints.video_generator.videogenerator.from_pretrained
+:classmethod:
+```

docs/source/index.md

@@ -51,14 +51,17 @@ Dev in progress and highly experimental.
 :maxdepth: 1
 
 getting_started/installation
-<!-- getting_started/examples/examples_index -->
+getting_started/quick_start
+getting_started/v1_api
 :::
 
 :::{toctree}
 :caption: Inference
 :maxdepth: 1
 
-inference/examples/examples_inference_index
+inference/examples/basic
+inference/examples/gradio
+inference/examples/optimizations
 inference/v0_inference
 :::
 
@@ -96,6 +99,14 @@ contributing/developer_env/index
 contributing/add_pipeline
 :::
 
+:::{toctree}
+:caption: API Reference
+:maxdepth: 2
+
+api/summary
+api/fastvideo/fastvideo
+:::
+
 ## Indices and tables
 
 - {ref}`genindex`

examples/inference/basic/README.md

@@ -1,3 +1,41 @@
-# Basic
+# Basic Video Generation Tutorial
+The `VideoGenerator` class provides the primary Python interface for doing offline video generation, which is interacting with a diffusion pipeline without using a separate inference api server.
 
-The class provides the main python interface for using FastVideo's inference pipeline.
+## Requirements
+- At least a single NVIDIA GPU with CUDA 12.4.
+- Python 3.10-3.12
+
+## Installation
+If you have not installed FastVideo, please following these [instructions](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) first.
+
+## Usage
+The first script in this example shows the most basic usage of FastVideo. If you are new to Python and FastVideo, you should start here.
+
+```bash
+# if you have not cloned the directory:
+git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+
+python fastvideo/v1/examples/inference/basic/basic.py
+```
+
+## Basic Walkthrough
+
+All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines!
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,
+    )
+
+    prompt = ("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.")
+    video = generator.generate_video(prompt)
+
+if __name__ == "__main__":
+    main()
+```

examples/inference/basic/basic.py

@@ -1 +1,41 @@
-print('Hello, world!')
+from fastvideo import VideoGenerator
+
+# from fastvideo.v1.configs.sample import SamplingParam
+
+
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        # if num_gpus > 1, FastVideo will automatically handle distributed setup
+        num_gpus=1,
+    )
+
+    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    # sampling_param.num_frames = 45
+    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    # Generate videos with the same simple API, regardless of GPU count
+    prompt = (
+        "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."
+    )
+    video = generator.generate_video(prompt)
+    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+    prompt2 = (
+        "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.")
+    video2 = generator.generate_video(prompt2)
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/default_args.py

@@ -0,0 +1,62 @@
+from fastvideo import VideoGenerator
+
+
+def main():
+
+    # This is the config class for the model initialization
+    config = PipelineConfig.from_pretrained("FastVideo/FastHunyuan-Diffusers")
+    # can be used to dump the config to a yaml file
+    config.dump_to_yaml("config.yaml")
+    print(config)
+    # {
+    #    'vae_config': {
+    #        'scale_factor': 8,
+    #        'sp': True,
+    #        'tiling': True,
+    #        'precision': 'fp16'
+    #    },
+    #    'text_encoder_config': {
+    #        'precision': 'fp16'
+    #    },
+    #    'dit_config': {
+    #        'precision': 'fp16'
+    #    },
+    #    'inference_args': {
+    #        'guidance_scale': 7.5,
+    #        'num_inference_steps': 5,
+    #        'seed': 1024,
+    #        'guidance_rescale': 0.0,
+    #        'flow_shift': 17,
+    #        'num_inference_steps': 5,
+    #    }
+    # }
+
+    config.vae_config.scale_factor = 16
+
+    # FastVideo will automatically used the optimal default arguments for the model
+    # If a local path is provided, FastVideo will make a best effort attempt to
+    # identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/FastHunyuan-Diffusers",
+        num_gpus=4,
+        config=config,
+        # or
+        config_path="config.yaml",
+    )
+
+    sampling_param = SamplingParam.from_pretrained(
+        "FastVideo/FastHunyuan-Diffusers")
+    sampling_param.num_inference_steps = 5
+
+    # Generate videos with the same simple API, regardless of GPU count
+    prompt = "A beautiful woman in a red dress walking down a street"
+    video = generator.generate_video(prompt,
+                                     sampling_param=sampling_param,
+                                     num_inference_steps=6)
+
+    video2 = generator.generate_video(prompt2)
+    prompt2 = "A beautiful woman in a blue dress walking down a street"
+
+
+if __name__ == "__main__":
+    main()

examples/inference/gradio/README.md

@@ -1,4 +1,4 @@
-# FastVideo VideoGenerator Gradio Demo
+# FastVideo Gradio Demo
 
 This is a Gradio-based web interface for generating videos using the FastVideo framework. The demo allows users to create videos from text prompts with various customization options.
 
@@ -16,7 +16,8 @@ The demo uses the FastVideo framework to generate videos based on text prompts.
 ## Requirements
 
 - Linux-based OS
-- Python 3.10
+- Python 3.10-3.12
+- NVIDIA GPU
 - Cuda 12.4
 - FastVideo
 
@@ -40,15 +41,6 @@ This will start a web server at `http://0.0.0.0:7860` where you can access the i
 
 ## Model Initialization
 
-```python
-args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
-
-generator = VideoGenerator.from_pretrained(
-    model_path=args.model_path,
-    num_gpus=args.num_gpus
-)
-```
-
 This demo initializes a `VideoGenerator` with the minimum required arguments for inference. Users can seamlessly adjust inference options between generations, including prompts, resolution, video length, or even the number of inference steps, *without ever needing to reload the model*.
 
 ## Video Generation
@@ -78,4 +70,4 @@ The interface is built with several components:
 - **Guidance Scale**: Control how closely the generation follows the prompt
 - **Inference Steps**: More steps can improve quality but take longer
 - **Negative Prompt**: Specify what you don't want to see in the video
-- **Seed**: Control randomness for reproducible results
+- **Seed**: Control randomness for reproducible results

examples/inference/gradio/gradio_demo.py

@@ -0,0 +1,169 @@
+import argparse
+import os
+from copy import deepcopy
+
+import gradio as gr
+import torch
+
+from fastvideo import VideoGenerator
+from fastvideo.v1.configs.sample.base import SamplingParam
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="FastVideo Gradio Demo")
+    parser.add_argument("--model_path",
+                        type=str,
+                        default="FastVideo/FastHunyuan-diffusers",
+                        help="Path to the model")
+    parser.add_argument("--num_gpus",
+                        type=int,
+                        default=1,
+                        help="Number of GPUs to use")
+    parser.add_argument("--output_path",
+                        type=str,
+                        default="outputs",
+                        help="Path to save generated videos")
+    parsed_args = parser.parse_args()
+
+    # args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
+
+    generator = VideoGenerator.from_pretrained(
+        model_path=parsed_args.model_path, num_gpus=parsed_args.num_gpus)
+
+    default_params = SamplingParam.from_pretrained(parsed_args.model_path)
+
+    def generate_video(
+        prompt,
+        negative_prompt,
+        use_negative_prompt,
+        seed,
+        guidance_scale,
+        num_frames,
+        height,
+        width,
+        num_inference_steps,
+        randomize_seed=False,
+    ):
+        params = deepcopy(default_params)
+        params.prompt = prompt
+        params.negative_prompt = negative_prompt
+        params.seed = seed
+        params.guidance_scale = guidance_scale
+        params.num_frames = num_frames
+        params.height = height
+        params.width = width
+        params.num_inference_steps = num_inference_steps
+
+        if randomize_seed:
+            params.seed = torch.randint(0, 1000000, (1, )).item()
+
+        if not use_negative_prompt:
+            params.negative_prompt = None
+
+        generator.generate_video(prompt=prompt, sampling_param=params)
+
+        output_path = os.path.join(parsed_args.output_path,
+                                   f"{params.prompt[:100]}.mp4")
+
+        return output_path, params.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.",
+    ]
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# FastVideo Inference Demo")
+
+        with gr.Group():
+            with gr.Row():
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your prompt",
+                    container=False,
+                )
+                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=default_params.height,
+                    )
+                    width = gr.Slider(label="Width",
+                                      minimum=256,
+                                      maximum=1024,
+                                      step=32,
+                                      value=default_params.width)
+
+                with gr.Row():
+                    num_frames = gr.Slider(
+                        label="Number of Frames",
+                        minimum=21,
+                        maximum=163,
+                        value=default_params.num_frames,
+                    )
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=1,
+                        maximum=12,
+                        value=default_params.guidance_scale,
+                    )
+                    num_inference_steps = gr.Slider(
+                        label="Inference Steps",
+                        minimum=4,
+                        maximum=100,
+                        value=default_params.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,
+                    )
+
+                seed = gr.Slider(label="Seed",
+                                 minimum=0,
+                                 maximum=1000000,
+                                 step=1,
+                                 value=default_params.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,
+            outputs=default_params.negative_prompt,
+        )
+
+        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],
+        )
+
+    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)

examples/inference/optimizations/README.md

@@ -0,0 +1,104 @@
+# Speeding Up Generation
+
+This page describes the various options for speeding up generation times.
+
+## Table of Contents
+- Optimized Attention Backends
+    - [Flash Attention](#optimizations-flash)
+    - [Sliding Tile Attention](#optimizations-sta)
+    - [Sage Attention](#optimizations-sage)
+
+- Caching Techniques
+    - [TeaCache](#optimizations-teacache)
+
+
+(optimizations-backends)=
+## Attention Backends
+
+`attention_example.py` shows how to set `FASTVIDEO_ATTENTION_BACKEND` env var to change attention backends. To run this example:
+```bash
+python examples/inference/optimizations/attention_example.py
+```
+
+In python, set the `FASTVIDEO_ATTENTION_BACKEND` before instantiating `VideoGenerator` like this:
+
+```python
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
+```
+
+You can also set the env var when running any of the other example like this:
+```bash
+FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
+```
+
+(optimizations-flash)=
+### Flash Attention
+
+(optimizations-sta)=
+### Sliding Tile Attention
+
+(optimizations-sage)=
+### Sage Attention
+
+
+(optimizations-teacache)=
+## Teacache 
+TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo.
+
+### What is TeaCache?
+
+See the official [TeaCache](https://github.com/ali-vilab/TeaCache) repo and their paper for more details.
+
+
+### How to Enable TeaCache
+
+Enabling TeaCache is straightforward - simply add the `enable_teacache=True` parameter to your `generate_video()` call:
+
+```python
+# ... previous code
+generator.generate_video(
+    prompt="Your prompt here",
+    sampling_param=params,
+    enable_teacache=True
+)
+# more code ...
+```
+
+### Complete Example
+
+At the bottom is a complete example of using TeaCache for faster video generation. You can run it using the following command:
+
+```bash
+python examples/inference/optimizations/teacache_example.py
+```
+
+### Advanced Configuration
+
+While TeaCache works well with default settings, you can fine-tune its behavior by adjusting the threshold value:
+
+1. Lower threshold values (e.g., 0.1) will result in more skipped calculations and faster generation with slightly more potential for quality degradation
+2. Higher threshold values (e.g., 0.15-0.23) will skip fewer calculations but maintain quality closer to the original 
+
+Note that the optimal threshold depends on your specific model and content.
+
+## Benchmarking different optimizations
+
+To benchmark the performance improvement, try generating the same video with and without TeaCache enabled and compare the generation times:
+
+```python
+# Without TeaCache
+start_time = time.time()
+generator.generate_video(prompt="Your prompt", enable_teacache=False)
+standard_time = time.time() - start_time
+
+# With TeaCache
+start_time = time.time()
+generator.generate_video(prompt="Your prompt", enable_teacache=True)
+teacache_time = time.time() - start_time
+
+print(f"Standard generation: {standard_time:.2f} seconds")
+print(f"TeaCache generation: {teacache_time:.2f} seconds")
+print(f"Speedup: {standard_time/teacache_time:.2f}x")
+```
+
+Note: If you want to benchmark different attention backends, you'll need to reinstantiate `VideoGenerator`.

examples/inference/optimizations/attention_example.py

@@ -0,0 +1,33 @@
+import os
+import time
+
+from fastvideo import VideoGenerator
+
+def main():
+    # set the attention backend 
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
+
+    start_time = time.time()
+    gen = VideoGenerator.from_pretrained(
+        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,
+    )
+    load_time = time.time() - start_time
+    print(f"Model loading time: {load_time:.2f} seconds")
+
+    gen_start_time = time.time()
+
+    gen.generate_video(
+        prompt=
+        "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.",
+        seed=1024,
+        output_path="example_outputs/")
+    
+    generation_time = time.time() - gen_start_time
+    print(f"Video generation time: {generation_time:.2f} seconds")
+
+    total_time = time.time() - start_time
+    print(f"Total execution time: {total_time:.2f} seconds")
+
+if __name__ == "__main__":
+    main()

examples/inference/optimizations/teacache_example.py

@@ -0,0 +1,44 @@
+import time
+
+from fastvideo import VideoGenerator, SamplingParam
+
+
+def main():
+    start_time = time.time()
+
+    gen = VideoGenerator.from_pretrained(
+        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,
+        use_cpu_offload=False,
+    )
+    load_time = time.time() - start_time
+    print(f"Model loading time: {load_time:.2f} seconds")
+
+    gen_start_time = time.time()
+
+    params = SamplingParam.from_pretrained(
+        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    )
+    # this controls the threshold for the tea cache
+    params.teacache_params.teacache_thresh = 0.08
+    gen.generate_video(
+        prompt=
+        "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.",
+        sampling_param=params,
+        height=480,
+        width=832,
+        num_frames=61,  # 85 ,77 
+        num_inference_steps=50,
+        enable_teacache=True,
+        seed=1024,
+        output_path="example_outputs/")
+    
+    generation_time = time.time() - gen_start_time
+    print(f"Video generation time: {generation_time:.2f} seconds")
+
+    total_time = time.time() - start_time
+    print(f"Total execution time: {total_time:.2f} seconds")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/v1/attention/backends/abstract.py

@@ -56,20 +56,6 @@ class AttentionMetadata:
     # Current step of diffusion process
     current_timestep: int
 
-    # @property
-    # @abstractmethod
-    # def inference_metadata(self) -> Optional["AttentionMetadata"]:
-    #     """Return the attention metadata that's required to run prefill
-    #     attention."""
-    #     pass
-
-    # @property
-    # @abstractmethod
-    # def training_metadata(self) -> Optional["AttentionMetadata"]:
-    #     """Return the attention metadata that's required to run decode
-    #     attention."""
-    #     pass
-
     def asdict_zerocopy(self,
                         skip_fields: Optional[Set[str]] = None
                         ) -> Dict[str, Any]:
@@ -86,55 +72,6 @@ class AttentionMetadata:
 
 T = TypeVar("T", bound=AttentionMetadata)
 
-# class AttentionState(ABC, Generic[T]):
-#     """Holds attention backend-specific objects reused during the
-#     lifetime of the model runner."""
-
-#     @abstractmethod
-#     def __init__(self, runner: "ModelRunnerBase"):
-#         ...
-
-#     @abstractmethod
-#     @contextmanager
-#     def graph_capture(self, max_batch_size: int):
-#         """Context manager used when capturing CUDA graphs."""
-#         yield
-
-#     @abstractmethod
-#     def graph_clone(self, batch_size: int) -> "AttentionState[T]":
-#         """Clone attention state to save in CUDA graph metadata."""
-#         ...
-
-#     @abstractmethod
-#     def graph_capture_get_metadata_for_batch(
-#             self,
-#             batch_size: int,
-#             is_encoder_decoder_model: bool = False) -> T:
-#         """Get attention metadata for CUDA graph capture of batch_size."""
-#         ...
-
-#     @abstractmethod
-#     def get_graph_input_buffers(
-#             self,
-#             attn_metadata: T,
-#             is_encoder_decoder_model: bool = False) -> Dict[str, Any]:
-#         """Get attention-specific input buffers for CUDA graph capture."""
-#         ...
-
-#     @abstractmethod
-#     def prepare_graph_input_buffers(
-#             self,
-#             input_buffers: Dict[str, Any],
-#             attn_metadata: T,
-#             is_encoder_decoder_model: bool = False) -> None:
-#         """In-place modify input buffers dict for CUDA graph replay."""
-#         ...
-
-#     @abstractmethod
-#     def begin_forward(self, model_input: "ModelRunnerInputBase") -> None:
-#         """Prepare state for forward pass."""
-#         ...
-
 
 class AttentionMetadataBuilder(ABC, Generic[T]):
     """Abstract class for attention metadata builders."""

fastvideo/v1/configs/fasthunyuan_t2v.json

@@ -0,0 +1,48 @@
+{
+  "embedded_cfg_scale": 6,
+  "flow_shift": 17,
+  "use_cpu_offload": false,
+  "disable_autocast": false,
+  "precision": "bf16",
+  "vae_precision": "fp16",
+  "vae_tiling": true,
+  "vae_sp": true,
+  "vae_config": {
+    "load_encoder": false,
+    "load_decoder": true,
+    "tile_sample_min_height": 256,
+    "tile_sample_min_width": 256,
+    "tile_sample_min_num_frames": 16,
+    "tile_sample_stride_height": 192,
+    "tile_sample_stride_width": 192,
+    "tile_sample_stride_num_frames": 12,
+    "blend_num_frames": 4,
+    "use_tiling": true,
+    "use_temporal_tiling": true,
+    "use_parallel_tiling": true
+  },
+  "dit_config": {
+    "prefix": "Hunyuan",
+    "quant_config": null
+  },
+  "text_encoder_precisions": [
+    "fp16",
+    "fp16"
+  ],
+  "text_encoder_configs": [
+    {
+      "prefix": "llama",
+      "quant_config": null,
+      "lora_config": null
+    },
+    {
+      "prefix": "clip",
+      "quant_config": null,
+      "lora_config": null,
+      "num_hidden_layers_override": null,
+      "require_post_norm": null
+    }
+  ],
+  "mask_strategy_file_path": null,
+  "enable_torch_compile": false
+}

fastvideo/v1/configs/models/base.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, fields
+from dataclasses import dataclass, field, fields
 from typing import Any, Dict
 
 from fastvideo.v1.logger import init_logger
@@ -18,7 +18,7 @@ class ArchConfig:
 class ModelConfig:
     # Every model config parameter can be categorized into either ArchConfig or everything else
     # Diffuser/Transformer parameters
-    arch_config: ArchConfig = ArchConfig()
+    arch_config: ArchConfig = field(default_factory=ArchConfig)
 
     # FastVideo-specific parameters here
     # i.e. STA, quantization, teacache
@@ -30,8 +30,26 @@ class ModelConfig:
         raise AttributeError(
             f"'{type(self).__name__}' object has no attribute '{name}'")
 
+    def __getstate__(self):
+        # Return a dictionary of attributes to pickle
+        # Convert to dict and exclude any problematic attributes
+        state = self.__dict__.copy()
+        return state
+
+    def __setstate__(self, state):
+        # Restore instance attributes from the unpickled state
+        self.__dict__.update(state)
+
     # This should be used only when loading from transformers/diffusers
     def update_model_arch(self, source_model_dict: Dict[str, Any]) -> None:
+        # Remove all keys that start with "_"
+        keys_to_remove = [
+            key for key in list(source_model_dict.keys())
+            if str(key).startswith("_")
+        ]
+        for key in keys_to_remove:
+            source_model_dict.pop(key)
+
         arch_config = self.arch_config
         valid_fields = {f.name for f in fields(arch_config)}
 

fastvideo/v1/configs/models/dits/__init__.py

@@ -1,4 +1,5 @@
+from fastvideo.v1.configs.models.dits.flux import FluxImageConfig
 from fastvideo.v1.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.v1.configs.models.dits.wanvideo import WanVideoConfig
 
-__all__ = ["HunyuanVideoConfig", "WanVideoConfig"]
+__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "FluxImageConfig"]

fastvideo/v1/configs/models/dits/base.py

@@ -9,6 +9,7 @@ from fastvideo.v1.platforms import _Backend
 @dataclass
 class DiTArchConfig(ArchConfig):
     _fsdp_shard_conditions: list = field(default_factory=list)
+    _compile_conditions: list = field(default_factory=list)
     _param_names_mapping: dict = field(default_factory=dict)
     _supported_attention_backends: Tuple[_Backend,
                                          ...] = (_Backend.SLIDING_TILE_ATTN,
@@ -20,10 +21,14 @@ class DiTArchConfig(ArchConfig):
     num_attention_heads: int = 0
     num_channels_latents: int = 0
 
+    def __post_init__(self) -> None:
+        if not self._compile_conditions:
+            self._compile_conditions = self._fsdp_shard_conditions.copy()
+
 
 @dataclass
 class DiTConfig(ModelConfig):
-    arch_config: DiTArchConfig = DiTArchConfig()
+    arch_config: DiTArchConfig = field(default_factory=DiTArchConfig)
 
     # FastVideoDiT-specific parameters
     prefix: str = ""

fastvideo/v1/configs/models/dits/flux.py

@@ -0,0 +1,130 @@
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+import torch
+
+from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class FluxImageArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    _param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # 1. context_embedder to txt_in mapping:
+            r"^context_embedder\.(.*)$":
+            r"txt_in.\1",
+
+            # 2. x_embedder to img_in mapping:
+            r"^x_embedder\.(.*)$":
+            r"img_in.\1",
+
+            # 3. Top-level time_text_embed mappings:
+            r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
+            r"time_in.mlp.fc_in.\1",
+            r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
+            r"time_in.mlp.fc_out.\1",
+            r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$":
+            r"guidance_in.mlp.fc_in.\1",
+            r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$":
+            r"guidance_in.mlp.fc_out.\1",
+            r"^time_text_embed\.text_embedder\.linear_1\.(.*)$":
+            r"txt2_in.fc_in.\1",
+            r"^time_text_embed\.text_embedder\.linear_2\.(.*)$":
+            r"txt2_in.fc_out.\1",
+
+            # 4. transformer_blocks mapping:
+            r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$":
+            r"double_blocks.\1.img_mod.linear.\2",
+            r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$":
+            r"double_blocks.\1.txt_mod.linear.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
+            r"double_blocks.\1.img_attn_q_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
+            r"double_blocks.\1.img_attn_k_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
+            r"double_blocks.\1.img_attn_proj.\2",
+            # Corrected: merge attn.to_add_out into the main projection.
+            r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$":
+            r"double_blocks.\1.txt_attn_proj.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$":
+            r"double_blocks.\1.txt_attn_q_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$":
+            r"double_blocks.\1.txt_attn_k_norm.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.img_mlp.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.img_mlp.fc_out.\2",
+            r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.txt_mlp.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.txt_mlp.fc_out.\2",
+
+            # 5. single_transformer_blocks mapping:
+            r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
+            r"single_blocks.\1.q_norm.\2",
+            r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
+            r"single_blocks.\1.k_norm.\2",
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+            (r"single_blocks.\1.linear1.\2", 0, 4),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+            (r"single_blocks.\1.linear1.\2", 1, 4),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+            (r"single_blocks.\1.linear1.\2", 2, 4),
+            r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$":
+            (r"single_blocks.\1.linear1.\2", 3, 4),
+            # Corrected: map proj_out to modulation.linear rather than a separate proj_out branch.
+            r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$":
+            r"single_blocks.\1.linear2.\2",
+            r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$":
+            r"single_blocks.\1.modulation.linear.\2",
+
+            # 6. Final layers mapping:
+            r"^norm_out\.linear\.(.*)$":
+            r"final_layer.adaLN_modulation.linear.\1",
+            r"^proj_out\.(.*)$":
+            r"final_layer.linear.\1",
+        })
+
+    patch_size: int = 1
+    in_channels: int = 64
+    out_channels: Optional[int] = None
+    num_layers: int = 19
+    num_single_layers: int = 38
+    attention_head_dim: int = 128
+    num_attention_heads: int = 24
+    joint_attention_dim: int = 4096
+    pooled_projection_dim: int = 768
+    guidance_embeds: bool = False
+    axes_dims_rope: Tuple[int, ...] = (16, 56, 56)
+    rope_theta: int = 10000
+    dtype: Optional[torch.dtype] = torch.bfloat16
+
+    def __post_init__(self):
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.in_channels // 4
+
+
+@dataclass
+class FluxImageConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=FluxImageArchConfig)
+
+    prefix: str = "Flux"

fastvideo/v1/configs/models/dits/hunyuanvideo.py

@@ -18,12 +18,19 @@ def is_refiner_block(n: str, m) -> bool:
     return "refiner" in n and str.isdigit(n.split(".")[-1])
 
 
+def is_txt_in(n: str, m) -> bool:
+    return n.split(".")[-1] == "txt_in"
+
+
 @dataclass
 class HunyuanVideoArchConfig(DiTArchConfig):
     _fsdp_shard_conditions: list = field(
         default_factory=lambda:
         [is_double_block, is_single_block, is_refiner_block])
 
+    _compile_conditions: list = field(
+        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
+
     _param_names_mapping: dict = field(
         default_factory=lambda: {
             # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
@@ -158,12 +165,13 @@ class HunyuanVideoArchConfig(DiTArchConfig):
     qk_norm: str = "rms_norm"
 
     def __post_init__(self):
+        super().__post_init__()
         self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
         self.num_channels_latents: int = self.in_channels
 
 
 @dataclass
 class HunyuanVideoConfig(DiTConfig):
-    arch_config: DiTArchConfig = HunyuanVideoArchConfig()
+    arch_config: DiTArchConfig = field(default_factory=HunyuanVideoArchConfig)
 
     prefix: str = "Hunyuan"

fastvideo/v1/configs/models/dits/wanvideo.py

@@ -70,6 +70,7 @@ class WanVideoArchConfig(DiTArchConfig):
     rope_max_seq_len: int = 1024
 
     def __post_init__(self):
+        super().__post_init__()
         self.out_channels = self.out_channels or self.in_channels
         self.hidden_size = self.num_attention_heads * self.attention_head_dim
         self.num_channels_latents = self.in_channels if self.added_kv_proj_dim is None else self.out_channels
@@ -77,6 +78,6 @@ class WanVideoArchConfig(DiTArchConfig):
 
 @dataclass
 class WanVideoConfig(DiTConfig):
-    arch_config: DiTArchConfig = WanVideoArchConfig()
+    arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig)
 
     prefix: str = "Wan"

fastvideo/v1/configs/models/encoders/__init__.py

@@ -1,4 +1,5 @@
-from fastvideo.v1.configs.models.encoders.base import (EncoderConfig,
+from fastvideo.v1.configs.models.encoders.base import (BaseEncoderOutput,
+                                                       EncoderConfig,
                                                        ImageEncoderConfig,
                                                        TextEncoderConfig)
 from fastvideo.v1.configs.models.encoders.clip import (CLIPTextConfig,
@@ -8,5 +9,6 @@ from fastvideo.v1.configs.models.encoders.t5 import T5Config
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
-    "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig", "T5Config"
+    "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig",
+    "T5Config"
 ]

fastvideo/v1/configs/models/encoders/base.py

@@ -1,5 +1,7 @@
 from dataclasses import dataclass, field
-from typing import Any, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
 from fastvideo.v1.configs.quantization import QuantizationConfig
@@ -30,15 +32,33 @@ class TextEncoderArchConfig(EncoderArchConfig):
     scalable_attention: bool = True
     tie_word_embeddings: bool = False
 
+    tokenizer_kwargs: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        self.tokenizer_kwargs = {
+            "truncation": True,
+            "max_length": self.text_len,
+            "return_tensors": "pt",
+        }
+
 
 @dataclass
 class ImageEncoderArchConfig(EncoderArchConfig):
     pass
 
 
+@dataclass
+class BaseEncoderOutput:
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    pooler_output: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attention_mask: Optional[torch.Tensor] = None
+
+
 @dataclass
 class EncoderConfig(ModelConfig):
-    arch_config: ArchConfig = EncoderArchConfig()
+    arch_config: ArchConfig = field(default_factory=EncoderArchConfig)
 
     prefix: str = ""
     quant_config: Optional[QuantizationConfig] = None
@@ -47,9 +67,9 @@ class EncoderConfig(ModelConfig):
 
 @dataclass
 class TextEncoderConfig(EncoderConfig):
-    arch_config: ArchConfig = TextEncoderArchConfig()
+    arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
 
 
 @dataclass
 class ImageEncoderConfig(EncoderConfig):
-    arch_config: ArchConfig = ImageEncoderArchConfig()
+    arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig)

fastvideo/v1/configs/models/encoders/clip.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional
 
 from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
@@ -48,7 +48,8 @@ class CLIPVisionArchConfig(ImageEncoderArchConfig):
 
 @dataclass
 class CLIPTextConfig(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = CLIPTextArchConfig()
+    arch_config: TextEncoderArchConfig = field(
+        default_factory=CLIPTextArchConfig)
 
     num_hidden_layers_override: Optional[int] = None
     require_post_norm: Optional[bool] = None
@@ -57,7 +58,8 @@ class CLIPTextConfig(TextEncoderConfig):
 
 @dataclass
 class CLIPVisionConfig(ImageEncoderConfig):
-    arch_config: ImageEncoderArchConfig = CLIPVisionArchConfig()
+    arch_config: ImageEncoderArchConfig = field(
+        default_factory=CLIPVisionArchConfig)
 
     num_hidden_layers_override: Optional[int] = None
     require_post_norm: Optional[bool] = None

fastvideo/v1/configs/models/encoders/llama.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional
 
 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
@@ -35,6 +35,6 @@ class LlamaArchConfig(TextEncoderArchConfig):
 
 @dataclass
 class LlamaConfig(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = LlamaArchConfig()
+    arch_config: TextEncoderArchConfig = field(default_factory=LlamaArchConfig)
 
     prefix: str = "llama"

fastvideo/v1/configs/models/encoders/t5.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional
 
 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
@@ -31,15 +31,25 @@ class T5ArchConfig(TextEncoderArchConfig):
 
     # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
     def __post_init__(self):
+        super().__post_init__()
         act_info = self.feed_forward_proj.split("-")
         self.dense_act_fn: str = act_info[-1]
         self.is_gated_act: bool = act_info[0] == "gated"
         if self.feed_forward_proj == "gated-gelu":
             self.dense_act_fn = "gelu_new"
 
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": self.text_len,
+            "add_special_tokens": True,
+            "return_attention_mask": True,
+            "return_tensors": "pt",
+        }
+
 
 @dataclass
 class T5Config(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = T5ArchConfig()
+    arch_config: TextEncoderArchConfig = field(default_factory=T5ArchConfig)
 
     prefix: str = "t5"

fastvideo/v1/configs/models/vaes/__init__.py

@@ -1,7 +1,9 @@
 from fastvideo.v1.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
+from fastvideo.v1.configs.models.vaes.image_vae import ImageVAEConfig
 from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
 
 __all__ = [
     "HunyuanVAEConfig",
     "WanVAEConfig",
+    "ImageVAEConfig",
 ]

fastvideo/v1/configs/models/vaes/base.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Union
 
 import torch
@@ -16,7 +16,7 @@ class VAEArchConfig(ArchConfig):
 
 @dataclass
 class VAEConfig(ModelConfig):
-    arch_config: VAEArchConfig = VAEArchConfig()
+    arch_config: VAEArchConfig = field(default_factory=VAEArchConfig)
 
     # FastVideoVAE-specific parameters
     load_encoder: bool = True

fastvideo/v1/configs/models/vaes/hunyuanvae.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Tuple
 
 from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
@@ -37,4 +37,4 @@ class HunyuanVAEArchConfig(VAEArchConfig):
 
 @dataclass
 class HunyuanVAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = HunyuanVAEArchConfig()
+    arch_config: VAEArchConfig = field(default_factory=HunyuanVAEArchConfig)

fastvideo/v1/configs/models/vaes/image_vae.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class ImageVAEArchConfig(VAEArchConfig):
+    in_channels: int = 3
+    out_channels: int = 3
+    down_block_types: Tuple[str] = ("DownEncoderBlock2D", )
+    up_block_types: Tuple[str] = ("UpDecoderBlock2D", )
+    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
+    scaling_factor: float = 0.18215
+    shift_factor: Optional[float] = None
+    latents_mean: Optional[Tuple[float]] = None
+    latents_std: Optional[Tuple[float]] = None
+    force_upcast: float = True
+    use_quant_conv: bool = True
+    use_post_quant_conv: bool = True
+    mid_block_add_attention: bool = True
+
+    def __post_init__(self):
+        self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) -
+                                                  1)
+        self.temporal_compression_ratio = 1
+
+
+@dataclass
+class ImageVAEConfig(VAEConfig):
+    arch_config: VAEArchConfig = field(default_factory=ImageVAEArchConfig)
+
+    # overrides VAEConfig
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False

fastvideo/v1/configs/models/vaes/wanvae.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Tuple
 
 import torch
@@ -63,7 +63,7 @@ class WanVAEArchConfig(VAEArchConfig):
 
 @dataclass
 class WanVAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = WanVAEArchConfig()
+    arch_config: VAEArchConfig = field(default_factory=WanVAEArchConfig)
     use_feature_cache: bool = True
 
     use_tiling: bool = False

fastvideo/v1/configs/pipelines/__init__.py

@@ -5,10 +5,12 @@ from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
 from fastvideo.v1.configs.pipelines.registry import (
     get_pipeline_config_cls_for_name)
 from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
-                                                WanT2V480PConfig)
+                                                WanI2V720PConfig,
+                                                WanT2V480PConfig,
+                                                WanT2V720PConfig)
 
 __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
     "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
-    "get_pipeline_config_cls_for_name"
+    "WanT2V720PConfig", "WanI2V720PConfig", "get_pipeline_config_cls_for_name"
 ]

fastvideo/v1/configs/pipelines/base.py

@@ -1,15 +1,26 @@
 import json
-from dataclasses import asdict, dataclass, fields
-from typing import Any, Dict, Optional
+from dataclasses import asdict, dataclass, field, fields
+from typing import Any, Callable, Dict, Optional, Tuple, cast
+
+import torch
 
 from fastvideo.v1.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
                                          VAEConfig)
+from fastvideo.v1.configs.models.encoders import BaseEncoderOutput
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import shallow_asdict
 
 logger = init_logger(__name__)
 
 
+def preprocess_text(prompt: str) -> str:
+    return prompt
+
+
+def postprocess_text(output: BaseEncoderOutput) -> torch.tensor:
+    raise NotImplementedError
+
+
 @dataclass
 class PipelineConfig:
     """Base configuration for all pipeline architectures."""
@@ -26,17 +37,26 @@ class PipelineConfig:
     vae_precision: str = "fp16"
     vae_tiling: bool = True
     vae_sp: bool = True
-    vae_config: VAEConfig = VAEConfig()
+    vae_config: VAEConfig = field(default_factory=VAEConfig)
 
     # DiT configuration
-    dit_config: DiTConfig = DiTConfig()
+    dit_config: DiTConfig = field(default_factory=DiTConfig)
 
     # Text encoder configuration
-    text_encoder_precision: str = "fp16"
-    text_encoder_config: EncoderConfig = EncoderConfig()
+    text_encoder_precisions: Tuple[str, ...] = field(
+        default_factory=lambda: ("fp16", ))
+    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (EncoderConfig(), ))
+    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (preprocess_text, ))
+    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (postprocess_text, ))
 
     # STA (Spatial-Temporal Attention) parameters
     mask_strategy_file_path: Optional[str] = None
+
+    # Compilation
     enable_torch_compile: bool = False
 
     @classmethod
@@ -52,16 +72,32 @@ class PipelineConfig:
                 model_path)
             pipeline_config = cls()
 
-        return pipeline_config
+        return cast(PipelineConfig, pipeline_config)
 
     def dump_to_json(self, file_path: str):
         output_dict = shallow_asdict(self)
+        del_keys = []
         for key, value in output_dict.items():
             if isinstance(value, ModelConfig):
                 model_dict = asdict(value)
                 # Model Arch Config should be hidden away from the users
                 model_dict.pop("arch_config")
                 output_dict[key] = model_dict
+            elif isinstance(value, tuple) and all(
+                    isinstance(v, ModelConfig) for v in value):
+                model_dicts = []
+                for v in value:
+                    model_dict = asdict(v)
+                    # Model Arch Config should be hidden away from the users
+                    model_dict.pop("arch_config")
+                    model_dicts.append(model_dict)
+                output_dict[key] = model_dicts
+            elif isinstance(value, tuple) and all(callable(f) for f in value):
+                # Skip dumping functions
+                del_keys.append(key)
+
+        for key in del_keys:
+            output_dict.pop(key, None)
 
         with open(file_path, "w") as f:
             json.dump(output_dict, f, indent=2)
@@ -82,6 +118,14 @@ class PipelineConfig:
                 # If it's a nested ModelConfig, update it recursively
                 if isinstance(current_value, ModelConfig):
                     current_value.update_model_config(new_value)
+                elif isinstance(current_value, tuple) and all(
+                        isinstance(v, ModelConfig) for v in current_value):
+                    assert len(current_value) == len(
+                        new_value
+                    ), "Users shouldn't delete or add text encoder config objects in your json"
+                    for target_config, source_config in zip(
+                            current_value, new_value):
+                        target_config.update_model_config(source_config)
                 else:
                     setattr(self, key, new_value)
 

fastvideo/v1/configs/pipelines/flux.py

@@ -0,0 +1,68 @@
+from dataclasses import dataclass, field
+from typing import Callable, Tuple
+
+import torch
+
+from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.v1.configs.models.dits import FluxImageConfig
+from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
+                                                  CLIPTextConfig, T5Config)
+from fastvideo.v1.configs.models.vaes import ImageVAEConfig
+from fastvideo.v1.configs.pipelines.base import PipelineConfig
+
+
+def t5_preprocess_text(prompt: str) -> str:
+    return prompt
+
+
+def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
+    hidden_state: torch.tensor = outputs.last_hidden_state
+    assert torch.isnan(hidden_state).sum() == 0
+    prompt_embeds_tensor: torch.tensor = torch.stack([
+        torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
+        for u in hidden_state
+    ],
+                                                     dim=0)
+    return prompt_embeds_tensor
+
+
+def clip_preprocess_text(prompt: str) -> str:
+    return prompt
+
+
+def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
+    pooler_output: torch.tensor = outputs.pooler_output
+    return pooler_output
+
+
+@dataclass
+class FluxConfig(PipelineConfig):
+    """Base configuration for Flux pipeline architecture."""
+
+    # FluxConfig-specific parameters with defaults
+    # DiT
+    dit_config: DiTConfig = field(default_factory=FluxImageConfig)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=ImageVAEConfig)
+    # Denoising stage
+    embedded_cfg_scale: float = 3.5
+
+    # Text encoding stage
+    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (CLIPTextConfig(), T5Config()))
+    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (clip_preprocess_text, t5_preprocess_text))
+    postprocess_text_funcs: Tuple[
+        Callable[[BaseEncoderOutput], torch.tensor],
+        ...] = field(default_factory=lambda:
+                     (clip_postprocess_text, t5_postprocess_text))
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp16"
+    text_encoder_precisions: Tuple[str, ...] = field(
+        default_factory=lambda: ("bf16", "bf16"))
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True

fastvideo/v1/configs/pipelines/hunyuan.py

@@ -1,11 +1,59 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Callable, Tuple, TypedDict
+
+import torch
 
 from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.v1.configs.models.dits import HunyuanVideoConfig
-from fastvideo.v1.configs.models.encoders import CLIPTextConfig, LlamaConfig
+from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
+                                                  CLIPTextConfig, LlamaConfig)
 from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 
+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|>")
+
+
+class PromptTemplate(TypedDict):
+    template: str
+    crop_start: int
+
+
+prompt_template_video: PromptTemplate = {
+    "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
+    "crop_start": 95,
+}
+
+
+def llama_preprocess_text(prompt: str) -> str:
+    return prompt_template_video["template"].format(prompt)
+
+
+def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
+    hidden_state_skip_layer = 2
+    assert outputs.hidden_states is not None
+    hidden_states: Tuple[torch.Tensor, ...] = outputs.hidden_states
+    last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer +
+                                                      1)]
+    crop_start = prompt_template_video.get("crop_start", -1)
+    last_hidden_state = last_hidden_state[:, crop_start:]
+    return last_hidden_state
+
+
+def clip_preprocess_text(prompt: str) -> str:
+    return prompt
+
+
+def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
+    pooler_output: torch.tensor = outputs.pooler_output
+    return pooler_output
+
 
 @dataclass
 class HunyuanConfig(PipelineConfig):
@@ -13,25 +61,28 @@ class HunyuanConfig(PipelineConfig):
 
     # HunyuanConfig-specific parameters with defaults
     # DiT
-    dit_config: DiTConfig = HunyuanVideoConfig()
+    dit_config: DiTConfig = field(default_factory=HunyuanVideoConfig)
     # VAE
-    vae_config: VAEConfig = HunyuanVAEConfig()
+    vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig)
     # Denoising stage
     embedded_cfg_scale: int = 6
     flow_shift: int = 7
 
     # Text encoding stage
-    text_encoder_config: EncoderConfig = LlamaConfig()
+    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
+    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (llama_preprocess_text, clip_preprocess_text))
+    postprocess_text_funcs: Tuple[
+        Callable[[BaseEncoderOutput], torch.tensor],
+        ...] = field(default_factory=lambda:
+                     (llama_postprocess_text, clip_postprocess_text))
 
     # Precision for each component
     precision: str = "bf16"
     vae_precision: str = "fp16"
-    text_encoder_precision: str = "fp16"
-
-    # HunyuanConfig-specific added parameters
-    # Secondary text encoder
-    text_encoder_config_2: EncoderConfig = CLIPTextConfig()
-    text_encoder_precision_2: str = "fp16"
+    text_encoder_precisions: Tuple[str, ...] = field(
+        default_factory=lambda: ("fp16", "fp16"))
 
     def __post_init__(self):
         self.vae_config.load_encoder = False

fastvideo/v1/configs/pipelines/registry.py

@@ -4,10 +4,13 @@ import os
 from typing import Callable, Dict, Optional, Type
 
 from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.v1.configs.pipelines.flux import FluxConfig
 from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
                                                     HunyuanConfig)
 from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
-                                                WanT2V480PConfig)
+                                                WanI2V720PConfig,
+                                                WanT2V480PConfig,
+                                                WanT2V720PConfig)
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import (maybe_download_model_index,
                                 verify_model_config_and_directory)
@@ -19,7 +22,11 @@ WEIGHT_CONFIG_REGISTRY: Dict[str, Type[PipelineConfig]] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
     "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
     "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
-    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
+    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
+    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
+    "black-forest-labs/FLUX.1-dev": FluxConfig,
+    "black-forest-labs/FLUX.1-schnell": FluxConfig,
     # Add other specific weight variants
 }
 
@@ -28,6 +35,7 @@ PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
     "hunyuan": lambda id: "hunyuan" in id.lower(),
     "wanpipeline": lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
+    "flux": lambda id: "flux" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -38,6 +46,7 @@ PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {
     "wanpipeline":
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
+    "flux": FluxConfig,
     # Other fallbacks by architecture
 }
 

fastvideo/v1/configs/pipelines/wan.py

@@ -1,21 +1,39 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Callable, Tuple
+
+import torch
 
 from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.v1.configs.models.dits import WanVideoConfig
-from fastvideo.v1.configs.models.encoders import CLIPVisionConfig, T5Config
+from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
+                                                  CLIPVisionConfig, T5Config)
 from fastvideo.v1.configs.models.vaes import WanVAEConfig
 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 
 
+def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
+    mask: torch.tensor = outputs.attention_mask
+    hidden_state: torch.tensor = outputs.last_hidden_state
+    seq_lens = mask.gt(0).sum(dim=1).long()
+    assert torch.isnan(hidden_state).sum() == 0
+    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens)]
+    prompt_embeds_tensor: torch.tensor = torch.stack([
+        torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
+        for u in prompt_embeds
+    ],
+                                                     dim=0)
+    return prompt_embeds_tensor
+
+
 @dataclass
 class WanT2V480PConfig(PipelineConfig):
     """Base configuration for Wan T2V 1.3B pipeline architecture."""
 
     # WanConfig-specific parameters with defaults
     # DiT
-    dit_config: DiTConfig = WanVideoConfig()
+    dit_config: DiTConfig = field(default_factory=WanVideoConfig)
     # VAE
-    vae_config: VAEConfig = WanVAEConfig()
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
     vae_tiling: bool = False
     vae_sp: bool = False
 
@@ -26,12 +44,17 @@ class WanT2V480PConfig(PipelineConfig):
     flow_shift: int = 3
 
     # Text encoding stage
-    text_encoder_config: EncoderConfig = T5Config()
+    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (T5Config(), ))
+    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (t5_postprocess_text, ))
 
     # Precision for each component
     precision: str = "bf16"
     vae_precision: str = "fp16"
-    text_encoder_precision: str = "fp32"
+    text_encoder_precisions: Tuple[str, ...] = field(
+        default_factory=lambda: ("fp32", ))
 
     # WanConfig-specific added parameters
 
@@ -40,6 +63,16 @@ class WanT2V480PConfig(PipelineConfig):
         self.vae_config.load_decoder = True
 
 
+@dataclass
+class WanT2V720PConfig(WanT2V480PConfig):
+    """Base configuration for Wan T2V 14B 720P pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+
+    # Denoising stage
+    flow_shift: int = 5
+
+
 @dataclass
 class WanI2V480PConfig(WanT2V480PConfig):
     """Base configuration for Wan I2V 14B 480P pipeline architecture."""
@@ -47,9 +80,20 @@ class WanI2V480PConfig(WanT2V480PConfig):
     # WanConfig-specific parameters with defaults
 
     # Precision for each component
-    image_encoder_config: EncoderConfig = CLIPVisionConfig()
+    image_encoder_config: EncoderConfig = field(
+        default_factory=CLIPVisionConfig)
     image_encoder_precision: str = "fp32"
 
     def __post_init__(self):
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
+
+
+@dataclass
+class WanI2V720PConfig(WanI2V480PConfig):
+    """Base configuration for Wan I2V 14B 720P pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+
+    # Denoising stage
+    flow_shift: int = 5

fastvideo/v1/configs/sample/base.py

@@ -8,6 +8,9 @@ logger = init_logger(__name__)
 
 @dataclass
 class SamplingParam:
+    """
+    Sampling parameters for video generation.
+    """
     # All fields below are copied from ForwardBatch
     data_type: str = "video"
 
@@ -26,6 +29,7 @@ class SamplingParam:
 
     # Original dimensions (before VAE scaling)
     num_frames: int = 125
+    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
     height: int = 720
     width: int = 1280
     fps: int = 24
@@ -35,6 +39,9 @@ class SamplingParam:
     guidance_scale: float = 1.0
     guidance_rescale: float = 0.0
 
+    # TeaCache parameters
+    enable_teacache: bool = False
+
     # Misc
     save_video: bool = True
     return_frames: bool = False
@@ -70,3 +77,8 @@ class SamplingParam:
             sampling_param = cls()
 
         return sampling_param
+
+
+@dataclass
+class CacheParams:
+    cache_type: str = "none"

fastvideo/v1/configs/sample/hunyuan.py

@@ -1,6 +1,7 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.v1.configs.sample.teacache import TeaCacheParams
 
 
 @dataclass
@@ -14,6 +15,14 @@ class HunyuanSamplingParam(SamplingParam):
 
     guidance_scale: float = 1.0
 
+    teacache_params: TeaCacheParams = field(
+        default_factory=lambda: TeaCacheParams(
+            teacache_thresh=0.15,
+            coefficients=[
+                7.33226126e+02, -4.01131952e+02, 6.75869174e+01,
+                -3.14987800e+00, 9.61237896e-02
+            ]))
+
 
 @dataclass
 class FastHunyuanSamplingParam(HunyuanSamplingParam):

fastvideo/v1/configs/sample/registry.py

@@ -3,8 +3,10 @@ from typing import Any, Callable, Dict, Optional
 
 from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                                  HunyuanSamplingParam)
-from fastvideo.v1.configs.sample.wan import (WanI2V480PSamplingParam,
-                                             WanT2V480PSamplingParam)
+from fastvideo.v1.configs.sample.wan import (WanI2V_14B_480P_SamplingParam,
+                                             WanI2V_14B_720P_SamplingParam,
+                                             WanT2V_1_3B_SamplingParam,
+                                             WanT2V_14B_SamplingParam)
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import (maybe_download_model_index,
                                 verify_model_config_and_directory)
@@ -14,8 +16,10 @@ logger = init_logger(__name__)
 SAMPLING_PARAM_REGISTRY: Dict[str, Any] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
     "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PSamplingParam,
-    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PSamplingParam
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam,
+    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
+    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
     # Add other specific weight variants
 }
 
@@ -32,8 +36,8 @@ SAMPLING_FALLBACK_PARAM: Dict[str, Any] = {
     "hunyuan":
     HunyuanSamplingParam,  # Base Hunyuan config as fallback for any Hunyuan variant
     "wanpipeline":
-    WanT2V480PSamplingParam,  # Base Wan config as fallback for any Wan variant
-    "wanimagetovideo": WanI2V480PSamplingParam,
+    WanT2V_1_3B_SamplingParam,  # Base Wan config as fallback for any Wan variant
+    "wanimagetovideo": WanI2V_14B_480P_SamplingParam,
     # Other fallbacks by architecture
 }
 

fastvideo/v1/configs/sample/teacache.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.sample.base import CacheParams
+
+
+@dataclass
+class TeaCacheParams(CacheParams):
+    cache_type: str = "teacache"
+    teacache_thresh: float = 0.0
+    coefficients: list[float] = field(default_factory=list)
+
+
+@dataclass
+class WanTeaCacheParams(CacheParams):
+    # Unfortunately, TeaCache is very different for Wan than other models
+    cache_type: str = "teacache"
+    teacache_thresh: float = 0.0
+    use_ret_steps: bool = True
+    ret_steps_coeffs: list[float] = field(default_factory=list)
+    non_ret_steps_coeffs: list[float] = field(default_factory=list)
+
+    @property
+    def coefficients(self) -> list[float]:
+        if self.use_ret_steps:
+            return self.ret_steps_coeffs
+        else:
+            return self.non_ret_steps_coeffs
+
+    @property
+    def ret_steps(self) -> int:
+        if self.use_ret_steps:
+            return 5 * 2
+        else:
+            return 1 * 2
+
+    def get_cutoff_steps(self, num_inference_steps: int) -> int:
+        if self.use_ret_steps:
+            return num_inference_steps * 2
+        else:
+            return num_inference_steps * 2 - 2

fastvideo/v1/configs/sample/wan.py

@@ -1,10 +1,11 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.v1.configs.sample.teacache import WanTeaCacheParams
 
 
 @dataclass
-class WanT2V480PSamplingParam(SamplingParam):
+class WanT2V_1_3B_SamplingParam(SamplingParam):
     # Video parameters
     height: int = 480
     width: int = 832
@@ -16,9 +17,79 @@ class WanT2V480PSamplingParam(SamplingParam):
     negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
     num_inference_steps: int = 50
 
+    teacache_params: WanTeaCacheParams = field(
+        default_factory=lambda: WanTeaCacheParams(
+            teacache_thresh=0.08,
+            ret_steps_coeffs=[
+                -5.21862437e+04, 9.23041404e+03, -5.28275948e+02,
+                1.36987616e+01, -4.99875664e-02
+            ],
+            non_ret_steps_coeffs=[
+                2.39676752e+03, -1.31110545e+03, 2.01331979e+02,
+                -8.29855975e+00, 1.37887774e-01
+            ]))
+
 
 @dataclass
-class WanI2V480PSamplingParam(WanT2V480PSamplingParam):
+class WanT2V_14B_SamplingParam(SamplingParam):
+    # Video parameters
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 81
+    fps: int = 16
+
+    # Denoising stage
+    guidance_scale: float = 5.0
+    negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+    num_inference_steps: int = 50
+
+    teacache_params: WanTeaCacheParams = field(
+        default_factory=lambda: WanTeaCacheParams(
+            teacache_thresh=0.20,
+            use_ret_steps=False,
+            ret_steps_coeffs=[
+                -3.03318725e+05, 4.90537029e+04, -2.65530556e+03,
+                5.87365115e+01, -3.15583525e-01
+            ],
+            non_ret_steps_coeffs=[
+                -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
+                -13.02252404
+            ]))
+
+
+@dataclass
+class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParam):
     # Denoising stage
     guidance_scale: float = 5.0
     num_inference_steps: int = 40
+
+    teacache_params: WanTeaCacheParams = field(
+        default_factory=lambda: WanTeaCacheParams(
+            teacache_thresh=0.26,
+            ret_steps_coeffs=[
+                -3.03318725e+05, 4.90537029e+04, -2.65530556e+03,
+                5.87365115e+01, -3.15583525e-01
+            ],
+            non_ret_steps_coeffs=[
+                -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
+                -13.02252404
+            ]))
+
+
+@dataclass
+class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam):
+    # Denoising stage
+    guidance_scale: float = 5.0
+    num_inference_steps: int = 40
+
+    teacache_params: WanTeaCacheParams = field(
+        default_factory=lambda: WanTeaCacheParams(
+            teacache_thresh=0.3,
+            ret_steps_coeffs=[
+                -3.03318725e+05, 4.90537029e+04, -2.65530556e+03,
+                5.87365115e+01, -3.15583525e-01
+            ],
+            non_ret_steps_coeffs=[
+                -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
+                -13.02252404
+            ]))

fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json

@@ -26,12 +26,16 @@
     "prefix": "Wan",
     "quant_config": null
   },
-  "text_encoder_precision": "fp32",
-  "text_encoder_config": {
-    "prefix": "t5",
-    "quant_config": null,
-    "lora_config": null
-  },
+  "text_encoder_precisions": [
+    "fp32"
+  ],
+  "text_encoder_configs": [
+    {
+      "prefix": "t5",
+      "quant_config": null,
+      "lora_config": null
+    }
+  ],
   "mask_strategy_file_path": null,
   "enable_torch_compile": false
 }

fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json

@@ -26,12 +26,16 @@
     "prefix": "Wan",
     "quant_config": null
   },
-  "text_encoder_precision": "fp32",
-  "text_encoder_config": {
-    "prefix": "t5",
-    "quant_config": null,
-    "lora_config": null
-  },
+  "text_encoder_precisions": [
+    "fp32"
+  ],
+  "text_encoder_configs": [
+    {
+      "prefix": "t5",
+      "quant_config": null,
+      "lora_config": null
+    }
+  ],
   "mask_strategy_file_path": null,
   "enable_torch_compile": false,
   "image_encoder_config": {

fastvideo/v1/entrypoints/video_generator.py

@@ -6,9 +6,9 @@ This module provides a consolidated interface for generating videos using
 diffusion models.
 """
 
+import gc
 import os
 import time
-from dataclasses import asdict
 from typing import Any, Dict, List, Optional, Union
 
 import imageio
@@ -210,12 +210,15 @@ class VideoGenerator:
               guidance_scale: {sampling_param.guidance_scale}
                     n_tokens: {n_tokens}
                   flow_shift: {fastvideo_args.flow_shift}
-     embedded_guidance_scale: {fastvideo_args.embedded_cfg_scale}"""
+     embedded_guidance_scale: {fastvideo_args.embedded_cfg_scale}
+                  save_video: {sampling_param.save_video}
+                  output_path: {sampling_param.output_path}
+        """ # type: ignore[attr-defined]
         logger.info(debug_str)
 
         # Prepare batch
         batch = ForwardBatch(
-            **asdict(sampling_param),
+            **shallow_asdict(sampling_param),
             eta=0.0,
             n_tokens=n_tokens,
             extra={},
@@ -243,7 +246,7 @@ class VideoGenerator:
             if save_path:
                 os.makedirs(os.path.dirname(save_path), exist_ok=True)
                 video_path = os.path.join(save_path, f"{prompt[:100]}.mp4")
-                imageio.mimsave(video_path, frames, fps=batch.fps)
+                imageio.mimsave(video_path, frames, fps=batch.fps, format="mp4")
                 logger.info("Saved video to %s", video_path)
             else:
                 logger.warning("No output path provided, video not saved")
@@ -257,3 +260,12 @@ class VideoGenerator:
                 "size": (target_height, target_width, batch.num_frames),
                 "generation_time": gen_time
             }
+
+    def shutdown(self):
+        """
+        Shutdown the video generator.
+        """
+        self.executor.shutdown()
+        del self.executor
+        gc.collect()
+        torch.cuda.empty_cache()

fastvideo/v1/examples/inference/basic/README.md

@@ -1,28 +0,0 @@
-# Basic Video Generation Tutorial
-The `VideoGenerator` class provides the primary Python interface for doing offline video generation, which is interacting with a diffusion pipeline without using a separate inference api server.
-
-
-## Usage
-The first script in this example shows the most basic usage of FastVideo. If you are new to Python and FastVideo, you should start here.
-
-```bash
-python fastvideo/v1/examples/inference/basic/basic.py
-```
-
-## Basic Walkthrough
-
-All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines! 
-
-```python
-from fastvideo import VideoGenerator
-
-generator = VideoGenerator.from_pretrained(
-    "FastVideo/FastHunyuan-Diffusers",
-    num_gpus=2,
-)
-
-prompt = "A beautiful woman in a red dress walking down a street"
-video = generator.generate_video(prompt)
-```
-
-More to come! These examples and APIs are still under construction!

fastvideo/v1/examples/inference/basic/basic.py

@@ -7,7 +7,7 @@ def main():
     # If a local path is provided, FastVideo will make a best effort
     # attempt to identify the optimal arguments.
     generator = VideoGenerator.from_pretrained(
-        "FastVideo/FastHunyuan-Diffusers",
+        "FastVideo/FastHunyuan-diffusers",
         # if num_gpus > 1, FastVideo will automatically handle distributed setup
         num_gpus=4,
     )
@@ -27,4 +27,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

fastvideo/v1/examples/inference/gradio/gradio_demo.py

@@ -1,141 +0,0 @@
-import os
-import gradio as gr
-import torch
-
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo import VideoGenerator
-
-
-if __name__ == "__main__":
-    args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
-
-    generator = VideoGenerator.from_pretrained(
-        model_path=args.model_path,
-        num_gpus=args.num_gpus
-    )
-
-    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()
-
-        if not use_negative_prompt:
-            negative_prompt = None
-
-        generator.generate_video(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=num_inference_steps,
-            num_frames=num_frames,
-            height=height,
-            width=width,
-            guidance_scale=guidance_scale,
-            seed=seed
-        )
-
-        output_path = os.path.join(args.output_path, f"{prompt[:100]}.mp4")
-
-        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.",
-    ]
-
-    with gr.Blocks() as demo:
-        gr.Markdown("# FastVideo Inference Demo")
-
-        with gr.Group():
-            with gr.Row():
-                prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    max_lines=1,
-                    placeholder="Enter your prompt",
-                    container=False,
-                )
-                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,
-                    )
-                    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=21,
-                        maximum=163,
-                        value=45,
-                    )
-                    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=6,
-                    )
-
-                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,
-                )
-
-                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,
-            outputs=negative_prompt,
-        )
-
-        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],
-        )
-
-    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)

fastvideo/v1/fastvideo_args.py

@@ -5,7 +5,8 @@
 import argparse
 import dataclasses
 from contextlib import contextmanager
-from typing import List, Optional
+from dataclasses import field
+from typing import Any, Callable, List, Optional, Tuple
 
 from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.v1.logger import init_logger
@@ -14,11 +15,22 @@ from fastvideo.v1.utils import FlexibleArgumentParser
 logger = init_logger(__name__)
 
 
+def preprocess_text(prompt: str) -> str:
+    return prompt
+
+
+def postprocess_text(output: Any) -> Any:
+    raise NotImplementedError
+
+
 @dataclasses.dataclass
 class FastVideoArgs:
     # Model and path configuration
     model_path: str
 
+    # Cache strategy
+    cache_strategy: str = "none"
+
     # Distributed executor backend
     distributed_executor_backend: str = "mp"
 
@@ -41,7 +53,7 @@ class FastVideoArgs:
     output_type: str = "pil"
 
     # DiT configuration
-    dit_config: DiTConfig = DiTConfig()
+    dit_config: DiTConfig = field(default_factory=DiTConfig)
     precision: str = "bf16"
 
     # VAE configuration
@@ -49,19 +61,21 @@ class FastVideoArgs:
     vae_tiling: bool = True  # Might change in between forward passes
     vae_sp: bool = False  # Might change in between forward passes
     # vae_scale_factor: Optional[int] = None # Deprecated
-    vae_config: VAEConfig = VAEConfig()
+    vae_config: VAEConfig = field(default_factory=VAEConfig)
 
     # Image encoder configuration
     image_encoder_precision: str = "fp32"
-    image_encoder_config: EncoderConfig = EncoderConfig()
+    image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)
 
     # Text encoder configuration
-    text_encoder_precision: str = "fp16"
-    text_encoder_config: EncoderConfig = EncoderConfig()
-
-    # Secondary text encoder
-    text_encoder_config_2: EncoderConfig = EncoderConfig()
-    text_encoder_precision_2: str = "fp16"
+    text_encoder_precisions: Tuple[str, ...] = field(
+        default_factory=lambda: ("fp16", ))
+    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (EncoderConfig(), ))
+    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (preprocess_text, ))
+    postprocess_text_funcs: Tuple[Callable[[Any], Any], ...] = field(
+        default_factory=lambda: (postprocess_text, ))
 
     # STA (Spatial-Temporal Attention) parameters
     mask_strategy_file_path: Optional[str] = None
@@ -204,10 +218,11 @@ class FastVideoArgs:
 
         parser.add_argument(
             "--text-encoder-precision",
+            nargs="+",
             type=str,
-            default=FastVideoArgs.text_encoder_precision,
+            default=FastVideoArgs.text_encoder_precisions,
             choices=["fp32", "fp16", "bf16"],
-            help="Precision for text encoder",
+            help="Precision for each text encoder",
         )
 
         # Image encoder config
@@ -219,16 +234,6 @@ class FastVideoArgs:
             help="Precision for image encoder",
         )
 
-        # Secondary text encoder
-
-        parser.add_argument(
-            "--text-encoder-precision-2",
-            type=str,
-            default=FastVideoArgs.text_encoder_precision_2,
-            choices=["fp32", "fp16", "bf16"],
-            help="Precision for secondary text encoder",
-        )
-
         # STA (Spatial-Temporal Attention) parameters
         parser.add_argument(
             "--mask-strategy-file-path",
@@ -311,6 +316,27 @@ class FastVideoArgs:
                 "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
             )
 
+        if len(self.text_encoder_configs) != len(self.text_encoder_precisions):
+            raise ValueError(
+                f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text encoder precisions ({len(self.text_encoder_precisions)})"
+            )
+
+        if len(self.text_encoder_configs) != len(self.preprocess_text_funcs):
+            raise ValueError(
+                f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
+            )
+
+        if len(self.preprocess_text_funcs) != len(self.postprocess_text_funcs):
+            raise ValueError(
+                f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
+            )
+
+        if self.enable_torch_compile and self.num_gpus > 1:
+            logger.warning(
+                "Currently torch compile does not work with multi-gpu. Setting enable_torch_compile to False"
+            )
+            self.enable_torch_compile = False
+
 
 _current_fastvideo_args = None
 

fastvideo/v1/forward_context.py

@@ -11,6 +11,7 @@ import torch
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 
 if TYPE_CHECKING:
     from fastvideo.v1.attention import AttentionMetadata
@@ -30,11 +31,13 @@ batchsize_forward_time: defaultdict = defaultdict(list)
 #
 @dataclass
 class ForwardContext:
+    current_timestep: int
     # TODO(will): check this arg
     # copy from vllm_config.compilation_config.static_forward_context
     # attn_layers: Dict[str, Any]
     # TODO: extend to support per-layer dynamic forward context
     attn_metadata: "AttentionMetadata"  # set dynamically for each forward pass
+    forward_batch: Optional[ForwardBatch] = None
 
 
 _forward_context: Optional[ForwardContext] = None
@@ -52,6 +55,7 @@ def get_forward_context() -> ForwardContext:
 @contextmanager
 def set_forward_context(current_timestep,
                         attn_metadata,
+                        forward_batch: Optional[ForwardBatch] = None,
                         fastvideo_args: Optional[FastVideoArgs] = None):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
@@ -63,7 +67,9 @@ def set_forward_context(current_timestep,
         forward_start_time = time.perf_counter()
     global _forward_context
     prev_context = _forward_context
-    _forward_context = ForwardContext(attn_metadata=attn_metadata)
+    _forward_context = ForwardContext(current_timestep=current_timestep,
+                                      attn_metadata=attn_metadata,
+                                      forward_batch=forward_batch)
     try:
         yield
     finally:
@@ -76,10 +82,6 @@ def set_forward_context(current_timestep,
             else:
                 # for v1 attention backends
                 batchsize = attn_metadata.num_input_tokens
-            # we use synchronous scheduling right now,
-            # adding a sync point here should not affect
-            # scheduling of the next batch
-            torch.cuda.synchronize()
             now = time.perf_counter()
             # time measurement is in milliseconds
             batchsize_forward_time[batchsize].append(

fastvideo/v1/models/dits/base.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from abc import ABC, abstractmethod
-from typing import List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 from torch import nn
@@ -12,6 +12,7 @@ from fastvideo.v1.platforms import _Backend
 # TODO
 class BaseDiT(nn.Module, ABC):
     _fsdp_shard_conditions: list = []
+    _compile_conditions: list = []
     _param_names_mapping: dict
     hidden_size: int
     num_attention_heads: int
@@ -22,7 +23,8 @@ class BaseDiT(nn.Module, ABC):
 
     def __init_subclass__(cls) -> None:
         required_class_attrs = [
-            "_fsdp_shard_conditions", "_param_names_mapping"
+            "_fsdp_shard_conditions", "_param_names_mapping",
+            "_compile_conditions"
         ]
         super().__init_subclass__()
         for attr in required_class_attrs:
@@ -63,3 +65,58 @@ class BaseDiT(nn.Module, ABC):
     @property
     def supported_attention_backends(self) -> Tuple[_Backend, ...]:
         return self._supported_attention_backends
+
+
+class CachableDiT(BaseDiT):
+    """
+    An intermediate base class that adds TeaCache optimization functionality to DiT models.
+    TeaCache accelerates inference by selectively skipping redundant computation when consecutive
+    diffusion steps are similar enough.
+    """
+    # These are required class attributes that should be overridden by concrete implementations
+    _fsdp_shard_conditions = []
+    _param_names_mapping = {}
+    # Ensure these instance attributes are properly defined in subclasses
+    hidden_size: int
+    num_attention_heads: int
+    num_channels_latents: int
+    # always supports torch_sdpa
+    _supported_attention_backends: Tuple[
+        _Backend, ...] = DiTConfig()._supported_attention_backends
+
+    def __init__(self, config: DiTConfig, **kwargs) -> None:
+        super().__init__(config, **kwargs)
+
+        self.cnt = 0
+        self.teacache_thresh = 0
+        self.coefficients: list[float] = []
+
+        # NOTE(will): Only wan2.1 needs these, so we are hardcoding it here
+        if self.config.prefix == "wan":
+            self.use_ret_steps = self.config.cache_config.use_ret_steps
+            self.is_even = False
+            self.previous_e0_even: torch.Tensor | None = None
+            self.previous_e0_odd: torch.Tensor | None = None
+            self.previous_residual_even: torch.Tensor | None = None
+            self.previous_residual_odd: torch.Tensor | None = None
+            self.accumulated_rel_l1_distance_even = 0
+            self.accumulated_rel_l1_distance_odd = 0
+            self.should_calc_even = True
+            self.should_calc_odd = True
+        else:
+            self.accumulated_rel_l1_distance = 0
+            self.previous_modulated_input = None
+            self.previous_residual = None
+
+    def maybe_cache_states(self, hidden_states: torch.Tensor,
+                           original_hidden_states: torch.Tensor) -> None:
+        pass
+
+    def should_skip_forward_for_cached_states(self,
+                                              **kwargs: dict[str, Any]) -> bool:
+        return False
+
+    def retrieve_cached_states(self,
+                               hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError(
+            "maybe_retrieve_cached_states is not implemented")

fastvideo/v1/models/dits/flux.py

@@ -0,0 +1,560 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from diffusers.models.normalization import RMSNorm
+
+from fastvideo.v1.attention import DistributedAttention
+from fastvideo.v1.configs.models.dits import FluxImageConfig
+from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
+                                           ScaleResidualLayerNormScaleShift)
+from fastvideo.v1.layers.linear import ReplicatedLinear
+from fastvideo.v1.layers.mlp import MLP
+from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
+                                                  get_rotary_pos_embed)
+from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
+                                                  TimestepEmbedder)
+from fastvideo.v1.models.dits.base import BaseDiT
+from fastvideo.v1.platforms import _Backend
+
+
+class MMDoubleStreamBlock(nn.Module):
+    """
+    A multimodal DiT block with separate modulation for text and image/video,
+    using distributed attention and linear layers.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        mlp_ratio: float = 4.0,
+        dtype: Optional[torch.dtype] = None,
+        supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.deterministic = False
+        self.num_attention_heads = num_attention_heads
+        head_dim = hidden_size // num_attention_heads
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        # Image modulation components
+        self.img_mod = ModulateProjection(
+            hidden_size,
+            factor=6,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.img_mod",
+        )
+
+        # Fused operations for image stream
+        self.img_attn_norm = LayerNormScaleShift(hidden_size,
+                                                 norm_type="layer",
+                                                 elementwise_affine=False,
+                                                 dtype=dtype)
+        self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            elementwise_affine=False,
+            dtype=dtype)
+        self.img_mlp_residual = ScaleResidual()
+
+        # Image attention components
+        self.img_attn_qkv = ReplicatedLinear(hidden_size,
+                                             hidden_size * 3,
+                                             bias=True,
+                                             params_dtype=dtype,
+                                             prefix=f"{prefix}.img_attn_qkv")
+
+        self.img_attn_q_norm = RMSNorm(head_dim, eps=1e-6)
+        self.img_attn_k_norm = RMSNorm(head_dim, eps=1e-6)
+
+        self.img_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype,
+                                              prefix=f"{prefix}.img_attn_proj")
+
+        self.img_mlp = MLP(hidden_size,
+                           mlp_hidden_dim,
+                           bias=True,
+                           dtype=dtype,
+                           prefix=f"{prefix}.img_mlp")
+
+        # Text modulation components
+        self.txt_mod = ModulateProjection(
+            hidden_size,
+            factor=6,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.txt_mod",
+        )
+
+        # Fused operations for text stream
+        self.txt_attn_norm = LayerNormScaleShift(hidden_size,
+                                                 norm_type="layer",
+                                                 elementwise_affine=False,
+                                                 dtype=dtype)
+        self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            elementwise_affine=False,
+            dtype=dtype)
+        self.txt_mlp_residual = ScaleResidual()
+
+        # Text attention components
+        self.txt_attn_qkv = ReplicatedLinear(hidden_size,
+                                             hidden_size * 3,
+                                             bias=True,
+                                             params_dtype=dtype)
+
+        # QK norm layers for text
+        self.txt_attn_q_norm = RMSNorm(head_dim, eps=1e-6)
+        self.txt_attn_k_norm = RMSNorm(head_dim, eps=1e-6)
+
+        self.txt_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype)
+
+        self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)
+
+        # Distributed attention
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn")
+
+    def forward(
+        self,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        freqs_cis_img: Tuple[torch.Tensor, torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Process modulation vectors
+        img_mod_outputs = self.img_mod(vec)
+        (
+            img_attn_shift,
+            img_attn_scale,
+            img_attn_gate,
+            img_mlp_shift,
+            img_mlp_scale,
+            img_mlp_gate,
+        ) = torch.chunk(img_mod_outputs, 6, dim=-1)
+
+        txt_mod_outputs = self.txt_mod(vec)
+        (
+            txt_attn_shift,
+            txt_attn_scale,
+            txt_attn_gate,
+            txt_mlp_shift,
+            txt_mlp_scale,
+            txt_mlp_gate,
+        ) = torch.chunk(txt_mod_outputs, 6, dim=-1)
+
+        # Prepare image for attention using fused operation
+        img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale)
+        # Get QKV for image
+        img_qkv, _ = self.img_attn_qkv(img_attn_input)
+        batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1]
+
+        # Split QKV
+        img_qkv = img_qkv.view(batch_size, image_seq_len, 3,
+                               self.num_attention_heads, -1)
+        img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :,
+                                                                          2]
+
+        # 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 rotary embeddings for image
+        cos, sin = freqs_cis_img
+        img_q, img_k = _apply_rotary_emb(
+            img_q, cos, sin,
+            is_neox_style=False), _apply_rotary_emb(img_k,
+                                                    cos,
+                                                    sin,
+                                                    is_neox_style=False)
+        # Prepare text for attention using fused operation
+        txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale)
+
+        # Get QKV for text
+        txt_qkv, _ = self.txt_attn_qkv(txt_attn_input)
+        batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1]
+
+        # Split QKV
+        txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3,
+                               self.num_attention_heads, -1)
+        txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :,
+                                                                          2]
+
+        # Apply QK-Norm if needed
+        txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
+        txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
+
+        # Run distributed attention
+        img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v)
+        img_attn_out, _ = self.img_attn_proj(
+            img_attn.view(batch_size, image_seq_len, -1))
+        # Use fused operation for residual connection, normalization, and modulation
+        img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(
+            img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale)
+
+        # Process image MLP
+        img_mlp_out = self.img_mlp(img_mlp_input)
+        img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)
+
+        # Process text attention output
+        txt_attn_out, _ = self.txt_attn_proj(
+            txt_attn.reshape(batch_size, text_seq_len, -1))
+
+        # Use fused operation for residual connection, normalization, and modulation
+        txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(
+            txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale)
+
+        # Process text MLP
+        txt_mlp_out = self.txt_mlp(txt_mlp_input)
+        txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)
+
+        return img, txt
+
+
+class MMSingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers using distributed attention
+    and tensor parallelism.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        mlp_ratio: float = 4.0,
+        dtype: Optional[torch.dtype] = None,
+        supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.deterministic = False
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        head_dim = hidden_size // num_attention_heads
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp_hidden_dim = mlp_hidden_dim
+
+        # Combined QKV and MLP input projection
+        self.linear1 = ReplicatedLinear(hidden_size,
+                                        hidden_size * 3 + mlp_hidden_dim,
+                                        bias=True,
+                                        params_dtype=dtype,
+                                        prefix=f"{prefix}.linear1")
+
+        # Combined projection and MLP output
+        self.linear2 = ReplicatedLinear(hidden_size + mlp_hidden_dim,
+                                        hidden_size,
+                                        bias=True,
+                                        params_dtype=dtype,
+                                        prefix=f"{prefix}.linear2")
+
+        # QK norm layers
+        self.q_norm = RMSNorm(head_dim, eps=1e-6)
+        self.k_norm = RMSNorm(head_dim, eps=1e-6)
+
+        # Fused operations with better naming
+        self.input_norm_scale_shift = LayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            eps=1e-6,
+            elementwise_affine=False,
+            dtype=dtype)
+        self.output_residual = ScaleResidual()
+
+        # Activation function
+        self.mlp_act = nn.GELU(approximate="tanh")
+
+        # Modulation
+        self.modulation = ModulateProjection(hidden_size,
+                                             factor=3,
+                                             act_layer="silu",
+                                             dtype=dtype,
+                                             prefix=f"{prefix}.modulation")
+
+        # Distributed attention
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn")
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        vec: torch.Tensor,
+        txt_len: int,
+        freqs_cis_img: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        # Process modulation
+        mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
+
+        # Apply pre-norm and modulation using fused operation
+        x_mod = self.input_norm_scale_shift(x, mod_shift, mod_scale)
+
+        # Get combined projections
+        linear1_out, _ = self.linear1(x_mod)
+
+        # Split into QKV and MLP parts
+        qkv, mlp = torch.split(linear1_out,
+                               [3 * self.hidden_size, self.mlp_hidden_dim],
+                               dim=-1)
+
+        # Process QKV
+        batch_size, seq_len = qkv.shape[0], qkv.shape[1]
+        qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+
+        # Apply QK-Norm
+        q = self.q_norm(q).to(v.dtype)
+        k = self.k_norm(k).to(v.dtype)
+
+        # Split into image and text parts
+        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:]
+        # Apply rotary embeddings to image parts
+        cos, sin = freqs_cis_img
+        img_q, img_k = _apply_rotary_emb(
+            img_q, cos, sin,
+            is_neox_style=False), _apply_rotary_emb(img_k,
+                                                    cos,
+                                                    sin,
+                                                    is_neox_style=False)
+
+        # Run distributed attention
+        img_attn_output, txt_attn_output = self.attn(img_q, img_k, img_v, txt_q,
+                                                     txt_k, txt_v)
+        attn_output = torch.cat((img_attn_output, txt_attn_output),
+                                dim=1).view(batch_size, seq_len, -1)
+        # Process MLP activation
+        mlp_output = self.mlp_act(mlp)
+
+        # Combine attention and MLP outputs
+        combined = torch.cat((attn_output, mlp_output), dim=-1)
+
+        # Final projection
+        output, _ = self.linear2(combined)
+
+        # Apply residual connection with gating using fused operation
+        return self.output_residual(x, output, mod_gate)
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT that projects features to pixel space.
+    """
+
+    def __init__(self,
+                 hidden_size,
+                 patch_size,
+                 out_channels,
+                 dtype=None,
+                 prefix: str = "") -> None:
+        super().__init__()
+
+        # Normalization
+        self.norm_final = nn.LayerNorm(hidden_size,
+                                       eps=1e-6,
+                                       elementwise_affine=False,
+                                       dtype=dtype)
+
+        output_dim = patch_size**3 * out_channels
+
+        self.linear = ReplicatedLinear(hidden_size,
+                                       output_dim,
+                                       bias=True,
+                                       params_dtype=dtype,
+                                       prefix=f"{prefix}.linear")
+
+        # Modulation
+        self.adaLN_modulation = ModulateProjection(
+            hidden_size,
+            factor=2,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.adaLN_modulation")
+
+    def forward(self, img, vec):
+        scale, shift = self.adaLN_modulation(vec).chunk(2, dim=-1)
+        img = self.norm_final(img) * (1.0 +
+                                      scale.unsqueeze(1)) + shift.unsqueeze(1)
+        img, _ = self.linear(img)
+        return img
+
+
+class FluxTransformer2DModel(BaseDiT):
+    _fsdp_shard_conditions = FluxImageConfig()._fsdp_shard_conditions
+    _supported_attention_backends = FluxImageConfig(
+    )._supported_attention_backends
+    _param_names_mapping = FluxImageConfig()._param_names_mapping
+
+    def __init__(self, config: FluxImageConfig) -> None:
+        super().__init__(config=config)
+
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.num_channels_latents
+        self.text_states_dim = config.joint_attention_dim
+        self.text_states_dim_2 = config.pooled_projection_dim
+        self.rope_dim_list = list(config.axes_dims_rope)
+        self.rope_theta = config.rope_theta
+        self.out_channels = config.out_channels
+        self.patch_size = config.patch_size
+
+        self.img_in = ReplicatedLinear(config.in_channels,
+                                       self.hidden_size,
+                                       params_dtype=config.dtype,
+                                       prefix=f"{config.prefix}.img_in")
+        self.txt_in = ReplicatedLinear(self.text_states_dim,
+                                       self.hidden_size,
+                                       params_dtype=config.dtype,
+                                       prefix=f"{config.prefix}.txt_in")
+        self.time_in = TimestepEmbedder(self.hidden_size,
+                                        act_layer="silu",
+                                        dtype=config.dtype,
+                                        prefix=f"{config.prefix}.time_in")
+        self.txt2_in = MLP(self.text_states_dim_2,
+                           self.hidden_size,
+                           self.hidden_size,
+                           act_type="silu",
+                           dtype=config.dtype,
+                           prefix=f"{config.prefix}.txt2_in")
+        self.guidance_in = (TimestepEmbedder(
+            self.hidden_size,
+            act_layer="silu",
+            dtype=config.dtype,
+            prefix=f"{config.prefix}.guidance_in")
+                            if config.guidance_embeds else None)
+
+        # Double blocks
+        self.double_blocks = nn.ModuleList([
+            MMDoubleStreamBlock(
+                self.hidden_size,
+                self.num_attention_heads,
+                dtype=config.dtype,
+                supported_attention_backends=self._supported_attention_backends,
+                prefix=f"{config.prefix}.double_blocks.{i}")
+            for i in range(config.num_layers)
+        ])
+
+        # Single blocks
+        self.single_blocks = nn.ModuleList([
+            MMSingleStreamBlock(
+                self.hidden_size,
+                self.num_attention_heads,
+                dtype=config.dtype,
+                supported_attention_backends=self._supported_attention_backends,
+                prefix=f"{config.prefix}.single_blocks.{i+config.num_layers}")
+            for i in range(config.num_single_layers)
+        ])
+
+        self.final_layer = FinalLayer(config.hidden_size,
+                                      self.patch_size,
+                                      self.out_channels,
+                                      dtype=config.dtype,
+                                      prefix=f"{config.prefix}.final_layer")
+
+        self.__post_init__()
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+                timestep: torch.LongTensor,
+                encoder_hidden_states_image: Optional[Union[
+                    torch.Tensor, List[torch.Tensor]]] = None,
+                guidance=None,
+                **kwargs):
+        """
+        Forward pass of the FluxTransformer2DModel.
+        
+        Args:
+            hidden_states: Input image latents [B, N, C]
+            encoder_hidden_states: Text embeddings [B, L, D]
+            timestep: Diffusion timestep
+            guidance: Guidance scale for CFG
+            
+        Returns:
+            Tuple of (output)
+        """
+        h = kwargs.pop("height_latents") or None
+        w = kwargs.pop("width_latents") or None
+        assert h is not None and w is not None
+
+        img = x = hidden_states
+
+        # Match diffusers implementation by multiplying timestep by 1000
+        t = timestep.to(img.dtype)
+
+        # Split text embeddings - first token is global, rest are per-token
+        txt = encoder_hidden_states[1]
+        text_states_2 = encoder_hidden_states[0]
+
+        # Get spatial dimensions
+        # _, _, oh, ow = img.shape
+        th, tw = (h // self.patch_size // 2, w // self.patch_size // 2)
+
+        # Get rotary embeddings
+        freqs_cos_img, freqs_sin_img = get_rotary_pos_embed(
+            (1, th, tw),
+            self.hidden_size,
+            self.num_attention_heads,
+            self.rope_dim_list,
+            self.rope_theta,
+            shard_dim=1,
+        )
+        freqs_cos_img = freqs_cos_img.to(img.device)
+        freqs_sin_img = freqs_sin_img.to(img.device)
+        freqs_cis_img = (freqs_cos_img, freqs_sin_img)
+
+        # Prepare modulation vectors
+        vec = self.time_in(t)
+
+        # Add text modulation
+        vec = vec + self.txt2_in(text_states_2)
+
+        # Add guidance modulation
+        if self.guidance_in is not None and guidance is not None:
+            vec = vec + self.guidance_in(guidance)
+
+        # embed text and image
+        img, _ = self.img_in(img)
+        txt, _ = self.txt_in(txt)
+        img_seq_len = img.shape[1]
+        txt_seq_len = txt.shape[1]
+
+        # Process through double stream blocks
+        for index, block in enumerate(self.double_blocks):
+            double_block_args = [img, txt, vec, freqs_cis_img]
+            img, txt = block(*double_block_args)
+
+        # Merge txt and img to pass through single stream blocks
+        x = torch.cat((img, txt), 1)
+
+        # Process through single stream blocks
+        for index, block in enumerate(self.single_blocks):
+            single_block_args = [x, vec, txt_seq_len, freqs_cis_img]
+            x = block(*single_block_args)
+
+        # Extract image features
+        img = x[:, :img_seq_len, ...]
+
+        # Final layer
+        img = self.final_layer(img, vec)
+
+        return img

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -2,13 +2,16 @@
 
 from typing import List, Optional, Tuple, Union
 
+import numpy as np
 import torch
 import torch.nn as nn
 
 from fastvideo.v1.attention import DistributedAttention, LocalAttention
 from fastvideo.v1.configs.models.dits import HunyuanVideoConfig
+from fastvideo.v1.configs.sample.teacache import TeaCacheParams
 from fastvideo.v1.distributed.parallel_state import (
     get_sequence_model_parallel_world_size)
+from fastvideo.v1.forward_context import get_forward_context
 from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
                                            ScaleResidualLayerNormScaleShift)
 from fastvideo.v1.layers.linear import ReplicatedLinear
@@ -19,7 +22,8 @@ from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
 from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
                                                   PatchEmbed, TimestepEmbedder,
                                                   unpatchify)
-from fastvideo.v1.models.dits.base import BaseDiT
+from fastvideo.v1.models.dits.base import CachableDiT
+from fastvideo.v1.models.utils import modulate
 from fastvideo.v1.platforms import _Backend
 
 
@@ -418,7 +422,7 @@ class MMSingleStreamBlock(nn.Module):
         return self.output_residual(x, output, mod_gate)
 
 
-class HunyuanVideoTransformer3DModel(BaseDiT):
+class HunyuanVideoTransformer3DModel(CachableDiT):
     """
     HunyuanVideo Transformer backbone adapted for distributed training.
     
@@ -433,6 +437,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
 
     # shard single stream, double stream blocks, and refiner_blocks
     _fsdp_shard_conditions = HunyuanVideoConfig()._fsdp_shard_conditions
+    _compile_conditions = HunyuanVideoConfig()._compile_conditions
     _supported_attention_backends = HunyuanVideoConfig(
     )._supported_attention_backends
     _param_names_mapping = HunyuanVideoConfig()._param_names_mapping
@@ -555,6 +560,11 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
         Returns:
             Tuple of (output)
         """
+        forward_context = get_forward_context()
+        forward_batch = forward_context.forward_batch
+        assert forward_batch is not None
+        enable_teacache = forward_batch.enable_teacache
+
         if guidance is None:
             guidance = torch.tensor([6016.0],
                                     device=hidden_states.device,
@@ -602,26 +612,40 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
         img_seq_len = img.shape[1]
 
         freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
-        # Process through double stream blocks
-        for index, block in enumerate(self.double_blocks):
-            double_block_args = [img, txt, vec, freqs_cis]
-            img, txt = block(*double_block_args)
-        # Merge txt and img to pass through single stream blocks
-        x = torch.cat((img, txt), 1)
 
-        # Process through single stream blocks
-        if len(self.single_blocks) > 0:
-            for index, block in enumerate(self.single_blocks):
-                single_block_args = [
-                    x,
-                    vec,
-                    txt_seq_len,
-                    freqs_cis,
-                ]
-                x = block(*single_block_args)
+        should_skip_forward = self.should_skip_forward_for_cached_states(
+            img=img, vec=vec)
+
+        if should_skip_forward:
+            img = self.retrieve_cached_states(img)
+        else:
+            if enable_teacache:
+                original_img = img.clone()
+
+            # Process through double stream blocks
+            for index, block in enumerate(self.double_blocks):
+                double_block_args = [img, txt, vec, freqs_cis]
+                img, txt = block(*double_block_args)
+            # Merge txt and img to pass through single stream blocks
+            x = torch.cat((img, txt), 1)
+
+            # Process through single stream blocks
+            if len(self.single_blocks) > 0:
+                for index, block in enumerate(self.single_blocks):
+                    single_block_args = [
+                        x,
+                        vec,
+                        txt_seq_len,
+                        freqs_cis,
+                    ]
+                    x = block(*single_block_args)
+
+            # Extract image features
+            img = x[:, :img_seq_len, ...]
+
+            if enable_teacache:
+                self.maybe_cache_states(img, original_img)
 
-        # Extract image features
-        img = x[:, :img_seq_len, ...]
         # Final layer processing
         img = self.final_layer(img, vec)
         # Unpatchify to get original shape
@@ -629,6 +653,99 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
 
         return img
 
+    def maybe_cache_states(self, hidden_states: torch.Tensor,
+                           original_hidden_states: torch.Tensor) -> None:
+        self.previous_residual = hidden_states - original_hidden_states
+
+    def should_skip_forward_for_cached_states(self, **kwargs) -> bool:
+
+        forward_context = get_forward_context()
+        forward_batch = forward_context.forward_batch
+        assert forward_batch is not None
+        current_timestep = forward_context.current_timestep
+        enable_teacache = forward_batch.enable_teacache
+
+        if not enable_teacache:
+            return False
+        raise NotImplementedError(
+            "teacache is not supported yet for HunyuanVideo")
+
+        teacache_params = forward_batch.teacache_params
+        assert teacache_params is not None, "teacache_params is not initialized"
+        assert isinstance(
+            teacache_params,
+            TeaCacheParams), "teacache_params is not a TeaCacheParams"
+        num_inference_steps = forward_batch.num_inference_steps
+        teache_thresh = teacache_params.teacache_thresh
+
+        coefficients = teacache_params.coefficients
+
+        if current_timestep == 0:
+            self.cnt = 0
+
+        inp = kwargs["img"].clone()
+        vec_ = kwargs["vec"].clone()
+        # convert to DTensor
+        vec_ = torch.distributed.tensor.DTensor.from_local(
+            vec_,
+            torch.distributed.DeviceMesh(
+                "cuda",
+                list(range(get_sequence_model_parallel_world_size())),
+                mesh_dim_names=("dp", )),
+            [torch.distributed.tensor.Replicate()])
+
+        inp = torch.distributed.tensor.DTensor.from_local(
+            inp,
+            torch.distributed.DeviceMesh(
+                "cuda",
+                list(range(get_sequence_model_parallel_world_size())),
+                mesh_dim_names=("dp", )),
+            [torch.distributed.tensor.Replicate()])
+
+        # txt_ = kwargs["txt"].clone()
+
+        # inp = img.clone()
+        # vec_ = vec.clone()
+        # txt_ = txt.clone()
+        (
+            img_mod1_shift,
+            img_mod1_scale,
+            img_mod1_gate,
+            img_mod2_shift,
+            img_mod2_scale,
+            img_mod2_gate,
+        ) = self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1)
+        normed_inp = self.double_blocks[0].img_attn_norm.norm(inp)
+        modulated_inp = modulate(normed_inp,
+                                 shift=img_mod1_shift,
+                                 scale=img_mod1_scale)
+        if self.cnt == 0 or self.cnt == num_inference_steps - 1:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0
+        else:
+            coefficients = [
+                7.33226126e+02, -4.01131952e+02, 6.75869174e+01,
+                -3.14987800e+00, 9.61237896e-02
+            ]
+            rescale_func = np.poly1d(coefficients)
+            assert self.previous_modulated_input is not None, "previous_modulated_input is not initialized"
+            self.accumulated_rel_l1_distance += rescale_func(
+                ((modulated_inp - self.previous_modulated_input).abs().mean() /
+                 self.previous_modulated_input.abs().mean()).cpu().item())
+            if self.accumulated_rel_l1_distance < teache_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = modulated_inp
+        self.cnt += 1
+
+        return not should_calc
+
+    def retrieve_cached_states(self,
+                               hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states + self.previous_residual
+
 
 class SingleTokenRefiner(nn.Module):
     """

fastvideo/v1/models/dits/wanvideo.py

@@ -3,13 +3,16 @@
 import math
 from typing import List, Optional, Tuple, Union
 
+import numpy as np
 import torch
 import torch.nn as nn
 
 from fastvideo.v1.attention import DistributedAttention, LocalAttention
 from fastvideo.v1.configs.models.dits import WanVideoConfig
+from fastvideo.v1.configs.sample.wan import WanTeaCacheParams
 from fastvideo.v1.distributed.parallel_state import (
     get_sequence_model_parallel_world_size)
+from fastvideo.v1.forward_context import get_forward_context
 from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, RMSNorm,
                                            ScaleResidual,
                                            ScaleResidualLayerNormScaleShift)
@@ -21,7 +24,7 @@ from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
                                                   get_rotary_pos_embed)
 from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
                                                   PatchEmbed, TimestepEmbedder)
-from fastvideo.v1.models.dits.base import BaseDiT
+from fastvideo.v1.models.dits.base import CachableDiT
 from fastvideo.v1.platforms import _Backend
 
 
@@ -350,8 +353,9 @@ class WanTransformerBlock(nn.Module):
         return hidden_states
 
 
-class WanTransformer3DModel(BaseDiT):
+class WanTransformer3DModel(CachableDiT):
     _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
+    _compile_conditions = WanVideoConfig()._compile_conditions
     _supported_attention_backends = WanVideoConfig(
     )._supported_attention_backends
     _param_names_mapping = WanVideoConfig()._param_names_mapping
@@ -419,6 +423,10 @@ class WanTransformer3DModel(BaseDiT):
                     torch.Tensor, List[torch.Tensor]]] = None,
                 guidance=None,
                 **kwargs) -> torch.Tensor:
+        forward_batch = get_forward_context().forward_batch
+        assert forward_batch is not None
+        enable_teacache = forward_batch.enable_teacache
+
         orig_dtype = hidden_states.dtype
         if not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
@@ -462,16 +470,32 @@ class WanTransformer3DModel(BaseDiT):
                 [encoder_hidden_states_image, encoder_hidden_states], dim=1)
 
         assert encoder_hidden_states.dtype == orig_dtype
+
         # 4. Transformer blocks
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-            for block in self.blocks:
-                hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, encoder_hidden_states, timestep_proj,
-                    freqs_cis)
+        # if caching is enabled, we might be able to skip the forward pass
+        should_skip_forward = self.should_skip_forward_for_cached_states(
+            timestep_proj=timestep_proj, temb=temb)
+
+        if should_skip_forward:
+            hidden_states = self.retrieve_cached_states(hidden_states)
         else:
-            for block in self.blocks:
-                hidden_states = block(hidden_states, encoder_hidden_states,
-                                      timestep_proj, freqs_cis)
+            # if teacache is enabled, we need to cache the original hidden states
+            if enable_teacache:
+                original_hidden_states = hidden_states.clone()
+
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                for block in self.blocks:
+                    hidden_states = self._gradient_checkpointing_func(
+                        block, hidden_states, encoder_hidden_states,
+                        timestep_proj, freqs_cis)
+            else:
+                for block in self.blocks:
+                    hidden_states = block(hidden_states, encoder_hidden_states,
+                                          timestep_proj, freqs_cis)
+
+            # if teacache is enabled, we need to cache the original hidden states
+            if enable_teacache:
+                self.maybe_cache_states(hidden_states, original_hidden_states)
 
         # 5. Output norm, projection & unpatchify
         shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
@@ -487,3 +511,96 @@ class WanTransformer3DModel(BaseDiT):
         output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
 
         return output
+
+    def maybe_cache_states(self, hidden_states: torch.Tensor,
+                           original_hidden_states: torch.Tensor) -> None:
+        if self.is_even:
+            self.previous_residual_even = hidden_states.squeeze(
+                0) - original_hidden_states
+        else:
+            self.previous_residual_odd = hidden_states.squeeze(
+                0) - original_hidden_states
+
+    def should_skip_forward_for_cached_states(self, **kwargs) -> bool:
+
+        forward_context = get_forward_context()
+        forward_batch = forward_context.forward_batch
+        assert forward_batch is not None
+        if not forward_batch.enable_teacache:
+            return False
+        teacache_params = forward_batch.teacache_params
+        assert teacache_params is not None, "teacache_params is not initialized"
+        assert isinstance(
+            teacache_params,
+            WanTeaCacheParams), "teacache_params is not a WanTeaCacheParams"
+        current_timestep = forward_context.current_timestep
+        num_inference_steps = forward_batch.num_inference_steps
+
+        # initialize the coefficients, cutoff_steps, and ret_steps
+        coefficients = teacache_params.coefficients
+        use_ret_steps = teacache_params.use_ret_steps
+        cutoff_steps = teacache_params.get_cutoff_steps(num_inference_steps)
+        ret_steps = teacache_params.ret_steps
+        teacache_thresh = teacache_params.teacache_thresh
+
+        if current_timestep == 0:
+            self.cnt = 0
+
+        timestep_proj = kwargs["timestep_proj"]
+        temb = kwargs["temb"]
+        modulated_inp = timestep_proj if use_ret_steps else temb
+
+        if self.cnt % 2 == 0:  # even -> condition
+            self.is_even = True
+            if self.cnt < ret_steps or self.cnt >= cutoff_steps:
+                self.should_calc_even = True
+                self.accumulated_rel_l1_distance_even = 0
+            else:
+                assert self.previous_e0_even is not None, "previous_e0_even is not initialized"
+                assert self.accumulated_rel_l1_distance_even is not None, "accumulated_rel_l1_distance_even is not initialized"
+                rescale_func = np.poly1d(coefficients)
+                self.accumulated_rel_l1_distance_even += rescale_func(
+                    ((modulated_inp - self.previous_e0_even).abs().mean() /
+                     self.previous_e0_even.abs().mean()).cpu().item())
+                if self.accumulated_rel_l1_distance_even < teacache_thresh:
+                    self.should_calc_even = False
+                else:
+                    self.should_calc_even = True
+                    self.accumulated_rel_l1_distance_even = 0
+            self.previous_e0_even = modulated_inp.clone()
+
+        else:  # odd -> unconditon
+            self.is_even = False
+            if self.cnt < ret_steps or self.cnt >= cutoff_steps:
+                self.should_calc_odd = True
+                self.accumulated_rel_l1_distance_odd = 0
+            else:
+                assert self.previous_e0_odd is not None, "previous_e0_odd is not initialized"
+                assert self.accumulated_rel_l1_distance_odd is not None, "accumulated_rel_l1_distance_odd is not initialized"
+                rescale_func = np.poly1d(coefficients)
+                self.accumulated_rel_l1_distance_odd += rescale_func(
+                    ((modulated_inp - self.previous_e0_odd).abs().mean() /
+                     self.previous_e0_odd.abs().mean()).cpu().item())
+                if self.accumulated_rel_l1_distance_odd < teacache_thresh:
+                    self.should_calc_odd = False
+                else:
+                    self.should_calc_odd = True
+                    self.accumulated_rel_l1_distance_odd = 0
+            self.previous_e0_odd = modulated_inp.clone()
+        self.cnt += 1
+        should_skip_forward = False
+        if self.is_even:
+            if not self.should_calc_even:
+                should_skip_forward = True
+        else:
+            if not self.should_calc_odd:
+                should_skip_forward = True
+
+        return should_skip_forward
+
+    def retrieve_cached_states(self,
+                               hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.is_even:
+            return hidden_states + self.previous_residual_even
+        else:
+            return hidden_states + self.previous_residual_odd

fastvideo/v1/models/encoders/base.py

@@ -1,16 +1,20 @@
-from typing import Tuple
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple
 
+import torch
 from torch import nn
 
-from fastvideo.v1.configs.models.encoders import EncoderConfig
+from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
+                                                  ImageEncoderConfig,
+                                                  TextEncoderConfig)
 from fastvideo.v1.platforms import _Backend
 
 
-class BaseEncoder(nn.Module):
+class TextEncoder(nn.Module, ABC):
     _supported_attention_backends: Tuple[
-        _Backend, ...] = EncoderConfig()._supported_attention_backends
+        _Backend, ...] = TextEncoderConfig()._supported_attention_backends
 
-    def __init__(self, config: EncoderConfig) -> None:
+    def __init__(self, config: TextEncoderConfig) -> None:
         super().__init__()
         self.config = config
         if not self.supported_attention_backends:
@@ -18,7 +22,36 @@ class BaseEncoder(nn.Module):
                 f"Subclass {self.__class__.__name__} must define _supported_attention_backends"
             )
 
-    def forward(self, *args, **kwargs):
+    @abstractmethod
+    def forward(self,
+                input_ids: Optional[torch.Tensor],
+                position_ids: Optional[torch.Tensor] = None,
+                attention_mask: Optional[torch.Tensor] = None,
+                inputs_embeds: Optional[torch.Tensor] = None,
+                output_hidden_states: Optional[bool] = None,
+                **kwargs) -> BaseEncoderOutput:
+        pass
+
+    @property
+    def supported_attention_backends(self) -> Tuple[_Backend, ...]:
+        return self._supported_attention_backends
+
+
+class ImageEncoder(nn.Module, ABC):
+    _supported_attention_backends: Tuple[
+        _Backend, ...] = ImageEncoderConfig()._supported_attention_backends
+
+    def __init__(self, config: ImageEncoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        if not self.supported_attention_backends:
+            raise ValueError(
+                f"Subclass {self.__class__.__name__} must define _supported_attention_backends"
+            )
+
+    @abstractmethod
+    def forward(self, pixel_values: torch.Tensor,
+                **kwargs) -> BaseEncoderOutput:
         pass
 
     @property

fastvideo/v1/models/encoders/clip.py

@@ -7,11 +7,11 @@ from typing import Iterable, Optional, Set, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers.modeling_outputs import BaseModelOutputWithPooling
 
 # from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask
 from fastvideo.v1.attention import LocalAttention
-from fastvideo.v1.configs.models.encoders import (CLIPTextConfig,
+from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
+                                                  CLIPTextConfig,
                                                   CLIPVisionConfig)
 from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.distributed import (divide,
@@ -20,7 +20,7 @@ from fastvideo.v1.layers.activation import get_act_fn
 from fastvideo.v1.layers.linear import (ColumnParallelLinear, QKVParallelLinear,
                                         RowParallelLinear)
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.encoders.base import BaseEncoder
+from fastvideo.v1.models.encoders.base import ImageEncoder, TextEncoder
 from fastvideo.v1.models.encoders.vision import resolve_visual_encoder_outputs
 # TODO: support quantization
 # from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -348,13 +348,12 @@ class CLIPTextTransformer(nn.Module):
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.Tensor],
         position_ids: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
         output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+    ) -> BaseEncoderOutput:
         r"""
         Returns:
 
@@ -362,7 +361,6 @@ class CLIPTextTransformer(nn.Module):
         output_hidden_states = (output_hidden_states
                                 if output_hidden_states is not None else
                                 self.config.output_hidden_states)
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         if input_ids is None:
             raise ValueError("You have to specify input_ids")
@@ -421,11 +419,7 @@ class CLIPTextTransformer(nn.Module):
                               ) == self.eos_token_id).int().argmax(dim=-1),
             ]
 
-        if not return_dict:
-            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
-
-        # return last_hidden_state
-        return BaseModelOutputWithPooling(
+        return BaseEncoderOutput(
             last_hidden_state=last_hidden_state,
             pooler_output=pooled_output,
             hidden_states=encoder_outputs,
@@ -433,7 +427,7 @@ class CLIPTextTransformer(nn.Module):
         )
 
 
-class CLIPTextModel(BaseEncoder):
+class CLIPTextModel(TextEncoder):
 
     def __init__(
         self,
@@ -446,21 +440,21 @@ class CLIPTextModel(BaseEncoder):
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.Tensor],
         position_ids: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
         output_hidden_states: Optional[bool] = None,
-    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        **kwargs,
+    ) -> BaseEncoderOutput:
 
-        return self.text_model(
+        outputs: BaseEncoderOutput = self.text_model(
             input_ids=input_ids,
             attention_mask=attention_mask,
             position_ids=position_ids,
-            output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            return_dict=None,
         )
+        return outputs
 
     def load_weights(self, weights: Iterable[Tuple[str,
                                                    torch.Tensor]]) -> Set[str]:
@@ -571,7 +565,7 @@ class CLIPVisionTransformer(nn.Module):
         return encoder_outputs
 
 
-class CLIPVisionModel(BaseEncoder):
+class CLIPVisionModel(ImageEncoder):
     config_class = CLIPVisionConfig
     main_input_name = "pixel_values"
     packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
@@ -589,8 +583,11 @@ class CLIPVisionModel(BaseEncoder):
         self,
         pixel_values: torch.Tensor,
         feature_sample_layers: Optional[list[int]] = None,
-    ) -> torch.Tensor:
-        return self.vision_model(pixel_values, feature_sample_layers)
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        last_hidden_state = self.vision_model(pixel_values,
+                                              feature_sample_layers)
+        return BaseEncoderOutput(last_hidden_state=last_hidden_state)
 
     @property
     def device(self):

fastvideo/v1/models/encoders/llama.py

@@ -27,12 +27,11 @@ from typing import Any, Dict, Iterable, Optional, Set, Tuple
 
 import torch
 from torch import nn
-from transformers.modeling_outputs import BaseModelOutputWithPast
 
 # from vllm.model_executor.layers.quantization import QuantizationConfig
 from fastvideo.v1.attention import LocalAttention
 # from ..utils import (extract_layer_index)
-from fastvideo.v1.configs.models.encoders import LlamaConfig
+from fastvideo.v1.configs.models.encoders import BaseEncoderOutput, LlamaConfig
 from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.distributed import get_tensor_model_parallel_world_size
 from fastvideo.v1.layers.activation import SiluAndMul
@@ -41,7 +40,7 @@ from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
                                         QKVParallelLinear, RowParallelLinear)
 from fastvideo.v1.layers.rotary_embedding import get_rope
 from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from fastvideo.v1.models.encoders.base import BaseEncoder
+from fastvideo.v1.models.encoders.base import TextEncoder
 from fastvideo.v1.models.loader.weight_utils import (default_weight_loader,
                                                      maybe_remap_kv_scale_name)
 
@@ -275,7 +274,7 @@ class LlamaDecoderLayer(nn.Module):
         return hidden_states, residual
 
 
-class LlamaModel(BaseEncoder):
+class LlamaModel(TextEncoder):
 
     def __init__(
         self,
@@ -320,11 +319,12 @@ class LlamaModel(BaseEncoder):
     def forward(
         self,
         input_ids: Optional[torch.Tensor],
-        positions: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
         output_hidden_states: Optional[bool] = None,
-    ) -> torch.Tensor:
+        **kwargs,
+    ) -> BaseEncoderOutput:
         output_hidden_states = (output_hidden_states
                                 if output_hidden_states is not None else
                                 self.config.output_hidden_states)
@@ -334,10 +334,10 @@ class LlamaModel(BaseEncoder):
             hidden_states = self.get_input_embeddings(input_ids)
         residual = None
 
-        if positions is None:
-            positions = torch.arange(0,
-                                     hidden_states.shape[1],
-                                     device=hidden_states.device).unsqueeze(0)
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, hidden_states.shape[1],
+                device=hidden_states.device).unsqueeze(0)
 
         all_hidden_states: Optional[Tuple[Any, ...]] = (
         ) if output_hidden_states else None
@@ -347,7 +347,8 @@ class LlamaModel(BaseEncoder):
                 all_hidden_states += (
                     hidden_states, ) if residual is None else (hidden_states +
                                                                residual, )
-            hidden_states, residual = layer(positions, hidden_states, residual)
+            hidden_states, residual = layer(position_ids, hidden_states,
+                                            residual)
 
         hidden_states, _ = self.norm(hidden_states, residual)
 
@@ -357,7 +358,7 @@ class LlamaModel(BaseEncoder):
 
         # TODO(will): maybe unify the output format with other models and use
         # our own class
-        output = BaseModelOutputWithPast(
+        output = BaseEncoderOutput(
             last_hidden_state=hidden_states,
             # past_key_values=past_key_values if use_cache else None,
             hidden_states=all_hidden_states,

fastvideo/v1/models/encoders/t5.py

@@ -27,7 +27,7 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 
-from fastvideo.v1.configs.models.encoders import T5Config
+from fastvideo.v1.configs.models.encoders import BaseEncoderOutput, T5Config
 from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.distributed import (get_tensor_model_parallel_rank,
                                       get_tensor_model_parallel_world_size)
@@ -36,7 +36,7 @@ from fastvideo.v1.layers.layernorm import RMSNorm
 from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
                                         QKVParallelLinear, RowParallelLinear)
 from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from fastvideo.v1.models.encoders.base import BaseEncoder
+from fastvideo.v1.models.encoders.base import TextEncoder
 from fastvideo.v1.models.loader.weight_utils import default_weight_loader
 
 
@@ -499,7 +499,7 @@ class T5Stack(nn.Module):
         return hidden_states
 
 
-class T5EncoderModel(BaseEncoder):
+class T5EncoderModel(TextEncoder):
 
     def __init__(self, config: T5Config, prefix: str = ""):
         super().__init__(config)
@@ -524,22 +524,21 @@ class T5EncoderModel(BaseEncoder):
 
     def forward(
         self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        output_attentions: Optional[bool] = None,
+        input_ids: Optional[torch.Tensor],
+        position_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
         output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> torch.Tensor:
+        **kwargs,
+    ) -> BaseEncoderOutput:
         attn_metadata = AttentionMetadata(None)
-        encoder_outputs = self.encoder(
+        hidden_states = self.encoder(
             input_ids=input_ids,
             attention_mask=attention_mask,
             attn_metadata=attn_metadata,
         )
 
-        return encoder_outputs
+        return BaseEncoderOutput(last_hidden_state=hidden_states)
 
     def load_weights(self, weights: Iterable[Tuple[str,
                                                    torch.Tensor]]) -> Set[str]:
@@ -587,7 +586,7 @@ class T5EncoderModel(BaseEncoder):
         return loaded_params
 
 
-class UMT5EncoderModel(BaseEncoder):
+class UMT5EncoderModel(TextEncoder):
 
     def __init__(self, config: T5Config, prefix: str = ""):
         super().__init__(config)
@@ -612,22 +611,24 @@ class UMT5EncoderModel(BaseEncoder):
 
     def forward(
         self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        output_attentions: Optional[bool] = None,
+        input_ids: Optional[torch.Tensor],
+        position_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
         output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> torch.Tensor:
+        **kwargs,
+    ) -> BaseEncoderOutput:
         attn_metadata = AttentionMetadata(None)
-        encoder_outputs = self.encoder(
+        hidden_states = self.encoder(
             input_ids=input_ids,
             attention_mask=attention_mask,
             attn_metadata=attn_metadata,
         )
 
-        return encoder_outputs
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            attention_mask=attention_mask,
+        )
 
     def load_weights(self, weights: Iterable[Tuple[str,
                                                    torch.Tensor]]) -> Set[str]:

fastvideo/v1/models/loader/component_loader.py

@@ -216,14 +216,15 @@ class TextEncoderLoader(ComponentLoader):
         model_config.pop("torch_dtype", None)
         logger.info("HF Model config: %s", model_config)
 
+        # @TODO(Wei): Better way to handle this?
         try:
-            encoder_config = fastvideo_args.text_encoder_config
+            encoder_config = fastvideo_args.text_encoder_configs[0]
             encoder_config.update_model_arch(model_config)
-            encoder_precision = fastvideo_args.text_encoder_precision
+            encoder_precision = fastvideo_args.text_encoder_precisions[0]
         except Exception:
-            encoder_config = fastvideo_args.text_encoder_config_2
+            encoder_config = fastvideo_args.text_encoder_configs[1]
             encoder_config.update_model_arch(model_config)
-            encoder_precision = fastvideo_args.text_encoder_precision_2
+            encoder_precision = fastvideo_args.text_encoder_precisions[1]
 
         target_device = torch.device(fastvideo_args.device_str)
         # TODO(will): add support for other dtypes
@@ -370,7 +371,6 @@ class TransformerLoader(ComponentLoader):
             raise ValueError(
                 "Model config does not contain a _class_name attribute. "
                 "Only diffusers format is supported.")
-        config.pop("_diffusers_version")
 
         # Config from Diffusers supersedes fastvideo's model config
         dit_config = fastvideo_args.dit_config
@@ -398,6 +398,25 @@ class TransformerLoader(ComponentLoader):
                                 device=fastvideo_args.device,
                                 cpu_offload=fastvideo_args.use_cpu_offload,
                                 default_dtype=default_dtype)
+        if fastvideo_args.enable_torch_compile:
+            logger.info("Torch Compile enabled for DiT")
+            for n, m in reversed(list(model.named_modules())):
+                if any([
+                        compile_condition(n, m)
+                        for compile_condition in model._compile_conditions
+                ]):
+                    parts = n.split(".")
+                    parent = model
+                    attr = parts[-1]
+                    for part in parts[:-1]:
+                        if part.isdigit():
+                            parent = parent[int(part)]
+                        else:
+                            parent = getattr(parent, part)
+                    if attr.isdigit():
+                        parent[int(attr)] = torch.compile(m)
+                    else:
+                        setattr(parent, attr, torch.compile(m))
 
         total_params = sum(p.numel() for p in model.parameters())
         logger.info("Loaded model with %.2fB parameters", total_params / 1e9)

fastvideo/v1/models/registry.py

@@ -23,6 +23,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "HunyuanVideoTransformer3DModel":
     ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+    "FluxTransformer2DModel": ("dits", "flux", "FluxTransformer2DModel"),
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -34,6 +35,7 @@ _TEXT_ENCODER_MODELS = {
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "LlamaModel": ("encoders", "llama", "LlamaModel"),
     "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"),
+    "T5EncoderModel": ("encoders", "t5", "T5EncoderModel"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -45,12 +47,13 @@ _VAE_MODELS = {
     "AutoencoderKLHunyuanVideo":
     ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
     "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),
+    "AutoencoderKL": ("vaes", "image_vae", "AutoencoderKL"),
 }
 
 _SCHEDULERS = {
     "FlowMatchEulerDiscreteScheduler":
     ("schedulers", "scheduling_flow_match_euler_discrete",
-     "FlowMatchDiscreteScheduler"),
+     "FlowMatchEulerDiscreteScheduler"),
     "UniPCMultistepScheduler":
     ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"),
 }

fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -1,3 +1,4 @@
+# type: ignore
 # SPDX-License-Identifier: Apache-2.0
 
 # Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
@@ -19,13 +20,16 @@
 #
 # ==============================================================================
 
+import math
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
+import numpy as np
+import scipy
 import torch
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
-from diffusers.utils import BaseOutput, logging
+from diffusers.utils import BaseOutput, is_scipy_available, logging
 
 from fastvideo.v1.models.schedulers.base import BaseScheduler
 
@@ -33,7 +37,7 @@ logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
 @dataclass
-class FlowMatchDiscreteSchedulerOutput(BaseOutput):
+class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
     """
     Output class for the scheduler's `step` function output.
 
@@ -46,7 +50,8 @@ class FlowMatchDiscreteSchedulerOutput(BaseOutput):
     prev_sample: torch.FloatTensor
 
 
-class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
+class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
+                                      BaseScheduler):
     """
     Euler scheduler.
 
@@ -56,16 +61,37 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     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.
+        use_dynamic_shifting (`bool`, defaults to False):
+            Whether to apply timestep shifting on-the-fly based on the image resolution.
+        base_shift (`float`, defaults to 0.5):
+            Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent
+            with desired output.
+        max_shift (`float`, defaults to 1.15):
+            Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be
+            more exaggerated or stylized.
+        base_image_seq_len (`int`, defaults to 256):
+            The base image sequence length.
+        max_image_seq_len (`int`, defaults to 4096):
+            The maximum image sequence length.
+        invert_sigmas (`bool`, defaults to False):
+            Whether to invert the sigmas.
+        shift_terminal (`float`, defaults to None):
+            The end value of the shifted timestep schedule.
+        use_karras_sigmas (`bool`, defaults to False):
+            Whether to use Karras sigmas for step sizes in the noise schedule during sampling.
+        use_exponential_sigmas (`bool`, defaults to False):
+            Whether to use exponential sigmas for step sizes in the noise schedule during sampling.
+        use_beta_sigmas (`bool`, defaults to False):
+            Whether to use beta sigmas for step sizes in the noise schedule during sampling.
+        time_shift_type (`str`, defaults to "exponential"):
+            The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
+        stochastic_sampling (`bool`, defaults to False):
+            Whether to use stochastic sampling.
     """
 
-    _compatibles: list[Any] = []
+    _compatibles = []
     order = 1
 
     @register_to_config
@@ -73,31 +99,62 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self,
         num_train_timesteps: int = 1000,
         shift: float = 1.0,
-        reverse: bool = True,
-        solver: str = "euler",
-        n_tokens: Optional[int] = None,
-        **kwargs,
+        use_dynamic_shifting: bool = False,
+        base_shift: Optional[float] = 0.5,
+        max_shift: Optional[float] = 1.15,
+        base_image_seq_len: Optional[int] = 256,
+        max_image_seq_len: Optional[int] = 4096,
+        invert_sigmas: bool = False,
+        shift_terminal: Optional[float] = None,
+        use_karras_sigmas: Optional[bool] = False,
+        use_exponential_sigmas: Optional[bool] = False,
+        use_beta_sigmas: Optional[bool] = False,
+        time_shift_type: str = "exponential",
+        stochastic_sampling: bool = False,
     ):
-        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: int | None = None
-        self._begin_index = 0
-
-        self.supported_solver = ["euler"]
-        if solver not in self.supported_solver:
+        if self.config.use_beta_sigmas and not is_scipy_available():
+            raise ImportError(
+                "Make sure to install scipy if you want to use beta sigmas.")
+        if sum([
+                self.config.use_beta_sigmas, self.config.use_exponential_sigmas,
+                self.config.use_karras_sigmas
+        ]) > 1:
             raise ValueError(
-                f"Solver {solver} not supported. Supported solvers: {self.supported_solver}"
+                "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
             )
+        if time_shift_type not in {"exponential", "linear"}:
+            raise ValueError(
+                "`time_shift_type` must either be 'exponential' or 'linear'.")
 
-        BaseScheduler.__init__(self)
+        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
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = sigmas * num_train_timesteps
+
+        self._step_index = None
+        self._begin_index = None
+
+        self._shift = shift
+
+        self.sigmas = sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def shift(self):
+        """
+        The value used for shifting.
+        """
+        return self._shift
 
     @property
     def step_index(self):
@@ -124,42 +181,207 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         """
         self._begin_index = begin_index
 
+    def set_shift(self, shift: float):
+        self._shift = shift
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        noise: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        Forward process in flow-matching
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(sample.device,
+                                                   dtype=torch.float32)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            timestep = timestep.to(sample.device)
+
+        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps) for t in timestep
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        sample = sigma * noise + (1.0 - sigma) * sample
+
+        return sample
+
     def _sigma_to_t(self, sigma):
         return sigma * self.config.num_train_timesteps
 
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        if self.config.time_shift_type == "exponential":
+            return self._time_shift_exponential(mu, sigma, t)
+        elif self.config.time_shift_type == "linear":
+            return self._time_shift_linear(mu, sigma, t)
+
+    def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
+        r"""
+        Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config
+        value.
+
+        Reference:
+        https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51
+
+        Args:
+            t (`torch.Tensor`):
+                A tensor of timesteps to be stretched and shifted.
+
+        Returns:
+            `torch.Tensor`:
+                A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.
+        """
+        one_minus_z = 1 - t
+        scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
+        stretched_t = 1 - (one_minus_z / scale_factor)
+        return stretched_t
+
     def set_timesteps(
         self,
-        num_inference_steps: int,
+        num_inference_steps: Optional[int] = None,
         device: Union[str, torch.device] = None,
-        n_tokens: int = 0,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[float] = None,
+        timesteps: Optional[List[float]] = None,
+        **kwargs,
     ):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
         Args:
-            num_inference_steps (`int`):
+            num_inference_steps (`int`, *optional*):
                 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.
+            sigmas (`List[float]`, *optional*):
+                Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
+                automatically.
+            mu (`float`, *optional*):
+                Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
+                shifting.
+            timesteps (`List[float]`, *optional*):
+                Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
+                automatically.
         """
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                "`mu` must be passed when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is not None and timesteps is not None and len(sigmas) != len(
+                timesteps):
+            raise ValueError(
+                "`sigmas` and `timesteps` should have the same length")
+
+        if num_inference_steps is not None:
+            if (sigmas is not None and len(sigmas) != num_inference_steps) or (
+                    timesteps is not None
+                    and len(timesteps) != num_inference_steps):
+                raise ValueError(
+                    "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided"
+                )
+        else:
+            num_inference_steps = len(sigmas) if sigmas is not None else len(
+                timesteps)
+
         self.num_inference_steps = num_inference_steps
 
-        sigmas = torch.linspace(1, 0, num_inference_steps + 1)
-        sigmas = self.sd3_time_shift(sigmas)
+        # 1. Prepare default sigmas
+        is_timesteps_provided = timesteps is not None
 
-        if not self.config.reverse:
-            sigmas = 1 - sigmas
+        if is_timesteps_provided:
+            timesteps = np.array(timesteps).astype(np.float32)
 
+        if sigmas is None:
+            if timesteps is None:
+                timesteps = np.linspace(self._sigma_to_t(self.sigma_max),
+                                        self._sigma_to_t(self.sigma_min),
+                                        num_inference_steps)
+            sigmas = timesteps / self.config.num_train_timesteps
+        else:
+            sigmas = np.array(sigmas).astype(np.float32)
+            num_inference_steps = len(sigmas)
+
+        # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
+        #    "exponential" or "linear" type is applied
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)
+        else:
+            sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)
+
+        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
+        if self.config.shift_terminal:
+            sigmas = self.stretch_shift_to_terminal(sigmas)
+
+        # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
+        if self.config.use_karras_sigmas:
+            sigmas = self._convert_to_karras(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+        elif self.config.use_exponential_sigmas:
+            sigmas = self._convert_to_exponential(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+        elif self.config.use_beta_sigmas:
+            sigmas = self._convert_to_beta(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
+        # 5. Convert sigmas and timesteps to tensors and move to specified device
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+        if not is_timesteps_provided:
+            timesteps = sigmas * self.config.num_train_timesteps
+        else:
+            timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32,
+                                                       device=device)
+
+        # 6. Append the terminal sigma value.
+        #    If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
+        #    `invert_sigmas` flag can be set to `True`. This case is only required in Mochi
+        if self.config.invert_sigmas:
+            sigmas = 1.0 - sigmas
+            timesteps = sigmas * self.config.num_train_timesteps
+            sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)])
+        else:
+            sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+
+        self.timesteps = timesteps
         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
+        self._begin_index = None
 
-    def index_for_timestep(self, timestep, schedule_timesteps=None) -> int:
+    def scale_model_input(self,
+                          sample: torch.Tensor,
+                          timestep: Optional[int] = None) -> torch.Tensor:
+        return sample
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
         if schedule_timesteps is None:
             schedule_timesteps = self.timesteps
 
@@ -171,14 +393,9 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         # case we start in the middle of the denoising schedule (e.g. for image-to-image)
         pos = 1 if len(indices) > 1 else 0
 
-        idx: int = indices[pos].item()
+        return indices[pos].item()
 
-        return idx
-
-    def set_shift(self, shift: float) -> None:
-        self.config.shift = shift
-
-    def _init_step_index(self, timestep) -> None:
+    def _init_step_index(self, timestep):
         if self.begin_index is None:
             if isinstance(timestep, torch.Tensor):
                 timestep = timestep.to(self.timesteps.device)
@@ -186,22 +403,19 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         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,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: Optional[torch.Generator] = None,
+        per_token_timesteps: Optional[torch.Tensor] = None,
         return_dict: bool = True,
-        **kwargs,
-    ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
+    ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, 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).
@@ -213,24 +427,30 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
                 The current discrete timestep in the diffusion chain.
             sample (`torch.FloatTensor`):
                 A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
             generator (`torch.Generator`, *optional*):
                 A random number generator.
-            n_tokens (`int`, *optional*):
-                Number of tokens in the input sequence.
+            per_token_timesteps (`torch.Tensor`, *optional*):
+                The timesteps for each token in the sample.
             return_dict (`bool`):
-                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
-                tuple.
+                Whether or not to return a
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] 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.
+            [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`,
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned,
+                otherwise a tuple is returned where the first element is the sample tensor.
         """
 
         if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
             raise ValueError((
                 "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-                " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
                 " one of the `scheduler.timesteps` as a timestep."), )
 
         if self.step_index is None:
@@ -239,24 +459,132 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         # Upcast to avoid precision issues when computing prev_sample
         sample = sample.to(torch.float32)
 
-        assert self.step_index is not None
-        dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
+        if per_token_timesteps is not None:
+            per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps
 
-        if self.config.solver == "euler":
-            prev_sample = sample + model_output.to(torch.float32) * dt
+            sigmas = self.sigmas[:, None, None]
+            lower_mask = sigmas < per_token_sigmas[None] - 1e-6
+            lower_sigmas = lower_mask * sigmas
+            lower_sigmas, _ = lower_sigmas.max(dim=0)
+
+            current_sigma = per_token_sigmas[..., None]
+            next_sigma = lower_sigmas[..., None]
+            dt = current_sigma - next_sigma
         else:
-            raise ValueError(
-                f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}"
-            )
+            sigma_idx = self.step_index
+            sigma = self.sigmas[sigma_idx]
+            sigma_next = self.sigmas[sigma_idx + 1]
+
+            current_sigma = sigma
+            next_sigma = sigma_next
+            dt = sigma_next - sigma
+
+        if self.config.stochastic_sampling:
+            x0 = sample - current_sigma * model_output
+            noise = torch.randn_like(sample)
+            prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
+        else:
+            prev_sample = sample + dt * model_output
 
         # upon completion increase step index by one
-        assert self._step_index is not None
         self._step_index += 1
+        if per_token_timesteps is None:
+            # Cast sample back to model compatible dtype
+            prev_sample = prev_sample.to(model_output.dtype)
 
         if not return_dict:
             return (prev_sample, )
 
-        return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(self, in_sigmas: torch.Tensor,
+                           num_inference_steps) -> torch.Tensor:
+        """Constructs the noise schedule of Karras et al. (2022)."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min**(1 / rho)
+        max_inv_rho = sigma_max**(1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
+    def _convert_to_exponential(self, in_sigmas: torch.Tensor,
+                                num_inference_steps: int) -> torch.Tensor:
+        """Constructs an exponential noise schedule."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.exp(
+            np.linspace(math.log(sigma_max), math.log(sigma_min),
+                        num_inference_steps))
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
+    def _convert_to_beta(self,
+                         in_sigmas: torch.Tensor,
+                         num_inference_steps: int,
+                         alpha: float = 0.6,
+                         beta: float = 0.6) -> torch.Tensor:
+        """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)"""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.array([
+            sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [
+                scipy.stats.beta.ppf(timestep, alpha, beta)
+                for timestep in 1 - np.linspace(0, 1, num_inference_steps)
+            ]
+        ])
+        return sigmas
+
+    def _time_shift_exponential(self, mu, sigma, t):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+    def _time_shift_linear(self, mu, sigma, t):
+        return mu / (mu + (1 / t - 1)**sigma)
 
     def __len__(self):
         return self.config.num_train_timesteps

fastvideo/v1/models/utils.py

@@ -118,3 +118,27 @@ def extract_layer_index(layer_name: str) -> int:
     assert len(int_vals) == 1, (f"layer name {layer_name} should"
                                 " only contain one integer")
     return int_vals[0]
+
+
+def modulate(x: torch.Tensor,
+             shift: Optional[torch.Tensor] = None,
+             scale: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """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))  # type: ignore[union-attr]
+    elif scale is None:
+        return x + shift.unsqueeze(1)  # type: ignore[union-attr]
+    else:
+        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(
+            1)  # type: ignore[union-attr]

fastvideo/v1/models/vaes/image_vae.py

@@ -0,0 +1,556 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from diffusers
+
+# 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.
+
+# 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.
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention_processor import (ADDED_KV_ATTENTION_PROCESSORS,
+                                                  CROSS_ATTENTION_PROCESSORS,
+                                                  Attention, AttentionProcessor,
+                                                  AttnAddedKVProcessor,
+                                                  AttnProcessor,
+                                                  FusedAttnProcessor2_0)
+from diffusers.models.autoencoders.vae import Decoder, Encoder
+
+from fastvideo.v1.configs.models.vaes import ImageVAEConfig
+from fastvideo.v1.models.vaes.common import DiagonalGaussianDistribution
+
+
+class AutoencoderKL(nn.Module):
+    r"""
+    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
+        out_channels (int,  *optional*, defaults to 3): Number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+            Tuple of downsample block types.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+            Tuple of upsample block types.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
+            Tuple of block output channels.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
+        sample_size (`int`, *optional*, defaults to `32`): Sample input size.
+        scaling_factor (`float`, *optional*, defaults to 0.18215):
+            The component-wise standard deviation of the trained latent space computed using the first batch of the
+            training set. This is used to scale the latent space to have unit variance when training the diffusion
+            model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
+            diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
+            / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
+            Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
+        force_upcast (`bool`, *optional*, default to `True`):
+            If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
+            can be fine-tuned / trained to a lower range without losing too much precision in which case
+            `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+        mid_block_add_attention (`bool`, *optional*, default to `True`):
+            If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
+            mid_block will only have resnet blocks
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
+
+    def __init__(
+        self,
+        config: ImageVAEConfig,
+    ):
+        nn.Module.__init__(self)
+        self.shift_factor = config.shift_factor
+        self.scaling_factor = config.scaling_factor
+
+        if config.load_encoder:
+            # pass init params to Encoder
+            self.encoder = Encoder(
+                in_channels=config.in_channels,
+                out_channels=config.latent_channels,
+                down_block_types=config.down_block_types,
+                block_out_channels=config.block_out_channels,
+                layers_per_block=config.layers_per_block,
+                act_fn=config.act_fn,
+                norm_num_groups=config.norm_num_groups,
+                double_z=True,
+                mid_block_add_attention=config.mid_block_add_attention,
+            )
+            self.quant_conv = nn.Conv2d(2 * config.latent_channels, 2 *
+                                        config.latent_channels,
+                                        1) if config.use_quant_conv else None
+
+        if config.load_decoder:
+            # pass init params to Decoder
+            self.decoder = Decoder(
+                in_channels=config.latent_channels,
+                out_channels=config.out_channels,
+                up_block_types=config.up_block_types,
+                block_out_channels=config.block_out_channels,
+                layers_per_block=config.layers_per_block,
+                norm_num_groups=config.norm_num_groups,
+                act_fn=config.act_fn,
+                mid_block_add_attention=config.mid_block_add_attention,
+            )
+            self.post_quant_conv = nn.Conv2d(
+                config.latent_channels, config.latent_channels,
+                1) if config.use_post_quant_conv else None
+
+        self.use_slicing = False
+        self.use_tiling = False
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_size = config.sample_size
+        sample_size = (config.sample_size[0] if isinstance(
+            config.sample_size, (list, tuple)) else config.sample_size)
+        self.tile_latent_min_size = int(
+            sample_size / (2**(len(config.block_out_channels) - 1)))
+        self.tile_overlap_factor = 0.25
+
+    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 images.
+        """
+        self.use_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.enable_tiling(False)
+
+    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.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: dict[str, AttentionProcessor] = {}
+
+        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)
+
+    # Copied from diffusers.models.unets.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)
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, height, width = x.shape
+
+        if self.use_tiling and (width > self.tile_sample_min_size
+                                or height > self.tile_sample_min_size):
+            return self._tiled_encode(x)
+
+        enc = self.encoder(x)
+        if self.quant_conv is not None:
+            enc = self.quant_conv(enc)
+
+        return enc
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            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. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            enc = torch.cat(encoded_slices)
+        else:
+            enc = self._encode(x)
+
+        enc = DiagonalGaussianDistribution(enc)
+
+        return enc
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size
+                                or z.shape[-2] > self.tile_latent_min_size):
+            return self.tiled_decode(z)
+
+        if self.post_quant_conv is not None:
+            z = self.post_quant_conv(z)
+
+        dec = self.decoder(z)
+
+        return dec
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        Decode a batch of images.
+
+        Args:
+            z (`torch.Tensor`): 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)
+
+        return 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[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 _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images 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 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.Tensor`): Input batch of images.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+
+        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 the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[:, :, i:i + self.tile_sample_min_size,
+                         j:j + self.tile_sample_min_size]
+                tile = self.encoder(tile)
+                if self.quant_conv:
+                    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=3))
+
+        enc = torch.cat(result_rows, dim=2)
+        return enc
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images 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 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.Tensor`): Input batch of images.
+            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 the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[:, :, i:i + self.tile_sample_min_size,
+                         j:j + self.tile_sample_min_size]
+                tile = self.encoder(tile)
+                if self.quant_conv:
+                    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=3))
+
+        moments = torch.cat(result_rows, dim=2)
+        enc = DiagonalGaussianDistribution(moments)
+
+        return enc
+
+    def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): 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 64x64 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[3], overlap_size):
+                tile = z[:, :, i:i + self.tile_latent_min_size,
+                         j:j + self.tile_latent_min_size]
+                if self.post_quant_conv:
+                    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=3))
+
+        dec = torch.cat(result_rows, dim=2)
+
+        return dec
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+        """
+        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
+
+        return dec
+
+    # Copied from diffusers.models.unets.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)
+
+        self.set_attn_processor(FusedAttnProcessor2_0())
+
+    # Copied from diffusers.models.unets.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/v1/pipelines/README.md

@@ -1,3 +1,3 @@
 # Adding a New Custom Pipeline
 
-Please see
+Please see documentation [here](https://hao-ai-lab.github.io/FastVideo/contributing/add_pipeline.html)

fastvideo/v1/pipelines/composed_pipeline_base.py

@@ -128,8 +128,12 @@ class ComposedPipelineBase(ABC):
         modules_config = deepcopy(self.config)
 
         # remove keys that are not pipeline modules
-        modules_config.pop("_class_name")
-        modules_config.pop("_diffusers_version")
+        keys_to_remove = [
+            key for key in list(modules_config.keys())
+            if str(key).startswith("_")
+        ]
+        for key in keys_to_remove:
+            modules_config.pop(key)
 
         # some sanity checks
         assert len(

fastvideo/v1/pipelines/flux/custom_stages.py

@@ -0,0 +1,359 @@
+import importlib
+
+import numpy as np
+import torch
+from einops import rearrange
+
+from fastvideo.v1.attention import get_attn_backend
+from fastvideo.v1.distributed import (get_sequence_model_parallel_rank,
+                                      get_sequence_model_parallel_world_size)
+from fastvideo.v1.distributed.communication_op import (
+    sequence_model_parallel_all_gather)
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines.stages.base import PipelineStage
+from fastvideo.v1.pipelines.stages.denoising import DenoisingStage
+from fastvideo.v1.platforms import _Backend
+from fastvideo.v1.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class TimestepsPreparationPreStage(PipelineStage):
+
+    def __init__(self, scheduler) -> None:
+        super().__init__()
+        self.scheduler = scheduler
+
+    def calculate_shift(self,
+                        image_seq_len,
+                        base_seq_len: int = 256,
+                        max_seq_len: int = 4096,
+                        base_shift: float = 0.5,
+                        max_shift: float = 1.15):
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        mu = image_seq_len * m + b
+        return mu
+
+    def forward(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        assert batch.height is not None
+        assert batch.width is not None
+
+        batch.sigmas = np.linspace(
+            1.0, 1 / batch.num_inference_steps,
+            batch.num_inference_steps) if batch.sigmas is None else batch.sigmas
+        spatial_compression_ratio = fastvideo_args.vae_config.arch_config.spatial_compression_ratio
+        batch.extra_set_timesteps_kwargs["mu"] = (
+            batch.extra_set_timesteps_kwargs.get("mu", None)
+            or self.calculate_shift(
+                (batch.height // spatial_compression_ratio // 2) *
+                (batch.width // spatial_compression_ratio // 2),
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 4096),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.15),
+            ))
+        return batch
+
+
+class DenoisingPreprocessingStage(PipelineStage):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        # [B, in_channels // 4, 1, H, W] -> [B, H // 2, W // 2, in_channels]
+        assert batch.latents is not None
+        b, c, _, h, w = batch.latents.shape
+        latents = batch.latents.view(b, c, h // 2, 2, w // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(b, (h // 2) * (w // 2), c * 4)
+        batch.latents = latents
+        return batch
+
+
+class DenoisingPostprocessingStage(PipelineStage):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        assert batch.latents is not None
+        assert batch.height_latents is not None
+        assert batch.width_latents is not None
+
+        latents = batch.latents
+        # Skip decoding if output type is latent
+        if fastvideo_args.output_type == "latent":
+            latents = latents
+        else:
+            # [B, (H // 2) * (W // 2), in_channels] -> [B, in_channels // 4, 1, H, W]
+            b, _, c = latents.shape
+            h, w = batch.height_latents, batch.width_latents
+            latents = latents.view(b, h // 2, w // 2, c // 4, 2, 2)
+            latents = latents.permute(0, 3, 1, 4, 2, 5)
+            latents = latents.reshape(b, c // 4, h, w)
+            # latents = latents.squeeze(2)
+
+        batch.latents = latents
+        return batch
+
+
+st_attn_available = False
+spec = importlib.util.find_spec("st_attn")
+if spec is not None:
+    st_attn_available = True
+
+    from fastvideo.v1.attention.backends.sliding_tile_attn import (
+        SlidingTileAttentionBackend)
+
+
+class FluxDenoisingStage(DenoisingStage):
+    """
+    Stage for running the denoising loop in diffusion pipelines.
+    
+    This stage handles the iterative denoising process that transforms
+    the initial noise into the final output.
+    """
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Run the denoising loop.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with denoised latents.
+        """
+        # Prepare extra step kwargs for scheduler
+        extra_step_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.step,
+            {
+                "generator": batch.generator,
+                "eta": batch.eta
+            },
+        )
+
+        # Setup precision and autocast settings
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        # Handle sequence parallelism if enabled
+        world_size, rank = get_sequence_model_parallel_world_size(
+        ), get_sequence_model_parallel_rank()
+        sp_group = world_size > 1
+        if sp_group:
+            latents = rearrange(batch.latents,
+                                "b (n s) c -> b n s c",
+                                n=world_size).contiguous()
+            latents = latents[:, rank, :, :]
+            batch.latents = latents
+            if batch.image_latent is not None:
+                image_latent = rearrange(batch.image_latent,
+                                         "b (n s) c -> b n s c",
+                                         n=world_size).contiguous()
+                image_latent = image_latent[:, rank, :, :]
+                batch.image_latent = image_latent
+
+        # Get timesteps and calculate warmup steps
+        timesteps = batch.timesteps
+        # TODO(will): remove this once we add input/output validation for stages
+        if timesteps is None:
+            raise ValueError("Timesteps must be provided")
+        num_inference_steps = batch.num_inference_steps
+        num_warmup_steps = len(
+            timesteps) - num_inference_steps * self.scheduler.order
+
+        # Create 3D list for mask strategy
+        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, layer, h = map(int, key.split('_'))
+                result[t][layer][h] = value
+            return result
+
+        # Prepare image latents and embeddings for I2V generation
+        image_embeds = batch.image_embeds
+        if len(image_embeds) > 0:
+            assert torch.isnan(image_embeds[0]).sum() == 0
+            image_embeds = [
+                image_embed.to(target_dtype) for image_embed in image_embeds
+            ]
+
+        image_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_image": image_embeds,
+                "height_latents": batch.height_latents,
+                "width_latents": batch.width_latents,
+            },
+        )
+
+        # Get latents and embeddings
+        latents = batch.latents
+        prompt_embeds = batch.prompt_embeds
+        assert torch.isnan(prompt_embeds[0]).sum() == 0
+        if batch.do_classifier_free_guidance:
+            neg_prompt_embeds = batch.negative_prompt_embeds
+            assert neg_prompt_embeds is not None
+            assert torch.isnan(neg_prompt_embeds[0]).sum() == 0
+
+        # Run denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # Skip if interrupted
+                if hasattr(self, 'interrupt') and self.interrupt:
+                    continue
+
+                # Expand latents for I2V
+                latent_model_input = latents.to(target_dtype)
+                if batch.image_latent is not None:
+                    latent_model_input = torch.cat(
+                        [latent_model_input, batch.image_latent],
+                        dim=1).to(target_dtype)
+                assert torch.isnan(latent_model_input).sum() == 0
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t)
+
+                # Prepare inputs for transformer
+                t_expand = t.repeat(latent_model_input.shape[0])
+                guidance_expand = (torch.tensor(
+                    [fastvideo_args.embedded_cfg_scale] *
+                    latent_model_input.shape[0],
+                    dtype=torch.float32,
+                    device=fastvideo_args.device,
+                ).to(target_dtype) * 1000.0 if fastvideo_args.embedded_cfg_scale
+                                   is not None else None)
+
+                # Predict noise residual
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+
+                    # TODO(will-refactor): all of this should be in the stage's init
+                    attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
+                    self.attn_backend = get_attn_backend(
+                        head_size=attn_head_size,
+                        dtype=torch.float16,  # TODO(will): hack
+                        supported_attention_backends=(
+                            _Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN,
+                            _Backend.TORCH_SDPA)  # hack
+                    )
+                    if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
+                        )
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
+                            )
+                            # TODO(will): clean this up
+                            attn_metadata = self.attn_metadata_builder.build(
+                                current_timestep=i,
+                                forward_batch=batch,
+                                fastvideo_args=fastvideo_args,
+                            )
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    else:
+                        attn_metadata = None
+                    # TODO(will): finalize the interface. vLLM uses this to
+                    # support torch dynamo compilation. They pass in
+                    # attn_metadata, vllm_config, and num_tokens. We can pass in
+                    # fastvideo_args or training_args, and attn_metadata.
+                    with set_forward_context(
+                            current_timestep=i,
+                            attn_metadata=attn_metadata,
+                            # fastvideo_args=fastvideo_args
+                    ):
+                        # Run transformer
+                        noise_pred = self.transformer(
+                            latent_model_input,
+                            prompt_embeds,
+                            t_expand,
+                            guidance=guidance_expand,
+                            **image_kwargs,
+                        )
+
+                    # Apply guidance
+                    if batch.do_classifier_free_guidance:
+                        with set_forward_context(
+                                current_timestep=i,
+                                attn_metadata=attn_metadata,
+                                # fastvideo_args=fastvideo_args
+                        ):
+                            # Run transformer
+                            noise_pred_uncond = self.transformer(
+                                latent_model_input,
+                                neg_prompt_embeds,
+                                t_expand,
+                                guidance=guidance_expand,
+                                **image_kwargs,
+                            )
+                        noise_pred_text = noise_pred
+                        noise_pred = noise_pred_uncond + batch.guidance_scale * (
+                            noise_pred_text - noise_pred_uncond)
+
+                        # Apply guidance rescale if needed
+                        if batch.guidance_rescale > 0.0:
+                            # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                            noise_pred = self.rescale_noise_cfg(
+                                noise_pred,
+                                noise_pred_text,
+                                guidance_rescale=batch.guidance_rescale,
+                            )
+
+                    # Compute the previous noisy sample
+                    latents = self.scheduler.step(noise_pred,
+                                                  t,
+                                                  latents,
+                                                  **extra_step_kwargs,
+                                                  return_dict=False)[0]
+
+                # Update progress bar
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and
+                    (i + 1) % self.scheduler.order == 0
+                        and progress_bar is not None):
+                    progress_bar.update()
+
+        # Gather results if using sequence parallelism
+        if sp_group:
+            latents = sequence_model_parallel_all_gather(latents, dim=1)
+
+        # Update batch with final latents
+        batch.latents = latents
+
+        if fastvideo_args.use_cpu_offload:
+            self.transformer.to('cpu')
+            torch.cuda.empty_cache()
+
+        return batch
+
+
+class ImageOutputStage(PipelineStage):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        output = batch.output
+        output = output.unsqueeze(2)
+        batch.output = output
+        return batch

fastvideo/v1/pipelines/flux/flux_pipeline.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Flux image diffusion pipeline implementation.
+
+This module contains an implementation of the Flux image diffusion pipeline
+using the modular pipeline architecture.
+"""
+
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.v1.pipelines.flux.custom_stages import (
+    DenoisingPostprocessingStage, DenoisingPreprocessingStage,
+    FluxDenoisingStage, ImageOutputStage, TimestepsPreparationPreStage)
+from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage,
+                                           InputValidationStage,
+                                           LatentPreparationStage,
+                                           TextEncodingStage,
+                                           TimestepPreparationStage)
+
+logger = init_logger(__name__)
+
+
+class FluxPipeline(ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae",
+        "transformer", "scheduler"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2"),
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2"),
+                           ],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timesteps_preparation_pre_stage",
+                       stage=TimestepsPreparationPreStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="modulation",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="denoising_preprocessing_stage",
+                       stage=DenoisingPreprocessingStage())
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=FluxDenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="denoising_postprocessing_stage",
+                       stage=DenoisingPostprocessingStage())
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="output_stage", stage=ImageOutputStage())
+
+
+EntryClass = FluxPipeline

fastvideo/v1/pipelines/hunyuan/hunyuan_pipeline.py

@@ -9,11 +9,10 @@ using the modular pipeline architecture.
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.v1.pipelines.stages import (CLIPTextEncodingStage,
-                                           ConditioningStage, DecodingStage,
+from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage,
                                            DenoisingStage, InputValidationStage,
                                            LatentPreparationStage,
-                                           LlamaEncodingStage,
+                                           TextEncodingStage,
                                            TimestepPreparationStage)
 
 # TODO(will): move PRECISION_TO_TYPE to better place
@@ -35,15 +34,15 @@ class HunyuanVideoPipeline(ComposedPipelineBase):
                        stage=InputValidationStage())
 
         self.add_stage(stage_name="prompt_encoding_stage_primary",
-                       stage=LlamaEncodingStage(
-                           text_encoder=self.get_module("text_encoder"),
-                           tokenizer=self.get_module("tokenizer"),
-                       ))
-
-        self.add_stage(stage_name="prompt_encoding_stage_secondary",
-                       stage=CLIPTextEncodingStage(
-                           text_encoder=self.get_module("text_encoder_2"),
-                           tokenizer=self.get_module("tokenizer_2"),
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2")
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2")
+                           ],
                        ))
 
         self.add_stage(stage_name="conditioning_stage",

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -12,6 +12,9 @@ from typing import Any, Dict, List, Optional, Union
 
 import torch
 
+from fastvideo.v1.configs.sample.teacache import (TeaCacheParams,
+                                                  WanTeaCacheParams)
+
 
 @dataclass
 class ForwardBatch:
@@ -66,6 +69,7 @@ class ForwardBatch:
     height_latents: Optional[int] = None
     width_latents: Optional[int] = None
     num_frames: int = 1  # Default for image models
+    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
 
     # Original dimensions (before VAE scaling)
     height: Optional[int] = None
@@ -75,6 +79,7 @@ class ForwardBatch:
     # Timesteps
     timesteps: Optional[torch.Tensor] = None
     timestep: Optional[Union[torch.Tensor, float, int]] = None
+    extra_set_timesteps_kwargs: Dict[str, Any] = field(default_factory=dict)
     step_index: Optional[int] = None
 
     # Scheduler parameters
@@ -102,6 +107,10 @@ class ForwardBatch:
     save_video: bool = True
     return_frames: bool = False
 
+    # TeaCache parameters
+    enable_teacache: bool = False
+    teacache_params: Optional[TeaCacheParams | WanTeaCacheParams] = None
+
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
 

fastvideo/v1/pipelines/stages/__init__.py

@@ -7,19 +7,15 @@ complete diffusion pipelines.
 """
 
 from fastvideo.v1.pipelines.stages.base import PipelineStage
-from fastvideo.v1.pipelines.stages.clip_image_encoding import (
-    CLIPImageEncodingStage)
-from fastvideo.v1.pipelines.stages.clip_text_encoding import (
-    CLIPTextEncodingStage)
 from fastvideo.v1.pipelines.stages.conditioning import ConditioningStage
 from fastvideo.v1.pipelines.stages.decoding import DecodingStage
 from fastvideo.v1.pipelines.stages.denoising import DenoisingStage
 from fastvideo.v1.pipelines.stages.encoding import EncodingStage
+from fastvideo.v1.pipelines.stages.image_encoding import ImageEncodingStage
 from fastvideo.v1.pipelines.stages.input_validation import InputValidationStage
 from fastvideo.v1.pipelines.stages.latent_preparation import (
     LatentPreparationStage)
-from fastvideo.v1.pipelines.stages.llama_encoding import LlamaEncodingStage
-from fastvideo.v1.pipelines.stages.t5_encoding import T5EncodingStage
+from fastvideo.v1.pipelines.stages.text_encoding import TextEncodingStage
 from fastvideo.v1.pipelines.stages.timestep_preparation import (
     TimestepPreparationStage)
 
@@ -32,8 +28,6 @@ __all__ = [
     "DenoisingStage",
     "EncodingStage",
     "DecodingStage",
-    "LlamaEncodingStage",
-    "T5EncodingStage",
-    "CLIPTextEncodingStage",
-    "CLIPImageEncodingStage",
+    "ImageEncodingStage",
+    "TextEncodingStage",
 ]

fastvideo/v1/pipelines/stages/clip_text_encoding.py

@@ -1,92 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Prompt encoding stages for diffusion pipelines.
-
-This module contains implementations of prompt encoding stages for diffusion pipelines.
-"""
-
-import torch
-
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.forward_context import set_forward_context
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.pipelines.stages.base import PipelineStage
-
-logger = init_logger(__name__)
-
-
-class CLIPTextEncodingStage(PipelineStage):
-    """
-    Stage for encoding text prompts into embeddings for diffusion models.
-    
-    This stage handles the encoding of text prompts into the embedding space
-    expected by the diffusion model.
-    """
-
-    def __init__(self, text_encoder, tokenizer) -> None:
-        """
-        Initialize the prompt encoding stage.
-        
-        Args:
-            enable_logging: Whether to enable logging for this stage.
-            is_secondary: Whether this is a secondary text encoder.
-        """
-        super().__init__()
-        self.tokenizer = tokenizer
-        self.text_encoder = text_encoder
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        """
-        Encode the prompt into text encoder hidden states.
-        
-        Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
-            
-        Returns:
-            The batch with encoded prompt embeddings.
-        """
-        if fastvideo_args.use_cpu_offload:
-            self.text_encoder = self.text_encoder.to(fastvideo_args.device)
-
-        text_inputs = self.tokenizer(
-            batch.prompt,
-            truncation=True,
-            # better way to handle this?
-            max_length=77,
-            return_tensors="pt",
-        )
-        with set_forward_context(current_timestep=0, attn_metadata=None):
-            outputs = self.text_encoder(input_ids=text_inputs["input_ids"].to(
-                fastvideo_args.device), )
-        prompt_embeds = outputs["pooler_output"]
-
-        batch.prompt_embeds.append(prompt_embeds)
-
-        if batch.do_classifier_free_guidance:
-            negative_text_inputs = self.tokenizer(
-                batch.negative_prompt,
-                truncation=True,
-                # better way to handle this?
-                max_length=77,
-                return_tensors="pt",
-            )
-            with set_forward_context(current_timestep=0, attn_metadata=None):
-                negative_outputs = self.text_encoder(
-                    input_ids=negative_text_inputs["input_ids"].to(
-                        fastvideo_args.device), )
-            negative_prompt_embeds = negative_outputs["pooler_output"]
-
-            assert batch.negative_prompt_embeds is not None
-            batch.negative_prompt_embeds.append(negative_prompt_embeds)
-
-        if fastvideo_args.use_cpu_offload:
-            self.text_encoder.to('cpu')
-            torch.cuda.empty_cache()
-
-        return batch

fastvideo/v1/pipelines/stages/decoding.py

@@ -7,7 +7,6 @@ import torch
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.utils import PRECISION_TO_TYPE
@@ -23,8 +22,8 @@ class DecodingStage(PipelineStage):
     output format (e.g., pixel values).
     """
 
-    def __init__(self, vae: ParallelTiledVAE) -> None:
-        self.vae: ParallelTiledVAE = vae
+    def __init__(self, vae) -> None:
+        self.vae = vae
 
     def forward(
         self,
@@ -41,6 +40,8 @@ class DecodingStage(PipelineStage):
         Returns:
             The batch with decoded outputs.
         """
+        self.vae = self.vae.to(fastvideo_args.device)
+
         latents = batch.latents
         # TODO(will): remove this once we add input/output validation for stages
         if latents is None:
@@ -95,4 +96,7 @@ class DecodingStage(PipelineStage):
         if hasattr(self, 'maybe_free_model_hooks'):
             self.maybe_free_model_hooks()
 
+        self.vae.to("cpu")
+        torch.cuda.empty_cache()
+
         return batch

fastvideo/v1/pipelines/stages/denoising.py

@@ -64,9 +64,7 @@ class DenoisingStage(PipelineStage):
         Returns:
             The batch with denoised latents.
         """
-        # If use cpu offload, need to load the model back into gpu again
-        if fastvideo_args.use_cpu_offload:
-            self.transformer = self.transformer.to(fastvideo_args.device)
+        assert batch.latents is not None
         # Prepare extra step kwargs for scheduler
         extra_step_kwargs = self.prepare_extra_func_kwargs(
             self.scheduler.step,
@@ -86,16 +84,28 @@ class DenoisingStage(PipelineStage):
         ), get_sequence_model_parallel_rank()
         sp_group = world_size > 1
         if sp_group:
-            latents = rearrange(batch.latents,
-                                "b t (n s) h w -> b t n s h w",
-                                n=world_size).contiguous()
-            latents = latents[:, :, rank, :, :, :]
+            if batch.latents.shape[2] == 1:
+                latents = rearrange(batch.latents,
+                                    "b t f (n s) w -> b t f n s w",
+                                    n=world_size).contiguous()
+                latents = latents[:, :, :, rank, :, :]
+            else:
+                latents = rearrange(batch.latents,
+                                    "b t (n s) h w -> b t n s h w",
+                                    n=world_size).contiguous()
+                latents = latents[:, :, rank, :, :, :]
             batch.latents = latents
             if batch.image_latent is not None:
-                image_latent = rearrange(batch.image_latent,
-                                         "b t (n s) h w -> b t n s h w",
-                                         n=world_size).contiguous()
-                image_latent = image_latent[:, :, rank, :, :, :]
+                if batch.image_latent.shape[2] == 1:
+                    image_latent = rearrange(batch.image_latent,
+                                             "b t f (n s) w -> b t f n s w",
+                                             n=world_size).contiguous()
+                    image_latent = image_latent[:, :, :, rank, :, :]
+                else:
+                    image_latent = rearrange(batch.image_latent,
+                                             "b t (n s) h w -> b t n s h w",
+                                             n=world_size).contiguous()
+                    image_latent = image_latent[:, :, rank, :, :, :]
                 batch.image_latent = image_latent
 
         # Get timesteps and calculate warmup steps
@@ -207,6 +217,7 @@ class DenoisingStage(PipelineStage):
                     with set_forward_context(
                             current_timestep=i,
                             attn_metadata=attn_metadata,
+                            forward_batch=batch,
                             # fastvideo_args=fastvideo_args
                     ):
                         # Run transformer
@@ -223,6 +234,7 @@ class DenoisingStage(PipelineStage):
                         with set_forward_context(
                                 current_timestep=i,
                                 attn_metadata=attn_metadata,
+                                forward_batch=batch,
                                 # fastvideo_args=fastvideo_args
                         ):
                             # Run transformer
@@ -262,7 +274,11 @@ class DenoisingStage(PipelineStage):
 
         # Gather results if using sequence parallelism
         if sp_group:
-            latents = sequence_model_parallel_all_gather(latents, dim=2)
+            if latents.shape[2] == 1:
+                # image latents
+                latents = sequence_model_parallel_all_gather(latents, dim=3)
+            else:
+                latents = sequence_model_parallel_all_gather(latents, dim=2)
 
         # Update batch with final latents
         batch.latents = latents
@@ -286,7 +302,9 @@ class DenoisingStage(PipelineStage):
         """
         extra_step_kwargs = {}
         for k, v in kwargs.items():
-            accepts = k in set(inspect.signature(func).parameters.keys())
+            accepts = (k in set(inspect.signature(func).parameters.keys())
+                       or "kwargs" in set(
+                           inspect.signature(func).parameters.keys()))
             if accepts:
                 extra_step_kwargs[k] = v
         return extra_step_kwargs

fastvideo/v1/pipelines/stages/image_encoding.py

@@ -17,7 +17,7 @@ from fastvideo.v1.pipelines.stages.base import PipelineStage
 logger = init_logger(__name__)
 
 
-class CLIPImageEncodingStage(PipelineStage):
+class ImageEncodingStage(PipelineStage):
     """
     Stage for encoding image prompts into embeddings for diffusion models.
     
@@ -60,7 +60,8 @@ class CLIPImageEncodingStage(PipelineStage):
         image_inputs = self.image_processor(
             images=image, return_tensors="pt").to(fastvideo_args.device)
         with set_forward_context(current_timestep=0, attn_metadata=None):
-            image_embeds = self.image_encoder(**image_inputs)
+            outputs = self.image_encoder(**image_inputs)
+            image_embeds = outputs.last_hidden_state
 
         batch.image_embeds.append(image_embeds)
 

fastvideo/v1/pipelines/stages/input_validation.py

@@ -3,6 +3,8 @@
 Input validation stage for diffusion pipelines.
 """
 
+import math
+
 import torch
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -86,4 +88,24 @@ class InputValidationStage(PipelineStage):
                 f"Guidance scale must be positive, but got {batch.guidance_scale}"
             )
 
+        # Adjust number of frames based on number of GPUs
+        temporal_scale_factor = fastvideo_args.vae_config.arch_config.temporal_compression_ratio
+        orig_latent_num_frames = (batch.num_frames -
+                                  1) // temporal_scale_factor + 1
+        if orig_latent_num_frames % fastvideo_args.num_gpus != 0:
+            if batch.num_frames_round_down:
+                new_latent_num_frames = math.floor(
+                    orig_latent_num_frames /
+                    fastvideo_args.num_gpus) * fastvideo_args.num_gpus
+            else:
+                new_latent_num_frames = math.ceil(
+                    orig_latent_num_frames /
+                    fastvideo_args.num_gpus) * fastvideo_args.num_gpus
+            new_num_frames = (new_latent_num_frames -
+                              1) * temporal_scale_factor + 1
+            logger.info(
+                "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
+                batch.num_frames, new_num_frames, fastvideo_args.num_gpus)
+            batch.num_frames = new_num_frames
+
         return batch

fastvideo/v1/pipelines/stages/latent_preparation.py

@@ -102,7 +102,8 @@ class LatentPreparationStage(PipelineStage):
 
         # Update batch with prepared latents
         batch.latents = latents
-
+        batch.height_latents = latents.shape[-2]
+        batch.width_latents = latents.shape[-1]
         return batch
 
     def adjust_video_length(self, batch: ForwardBatch,

fastvideo/v1/pipelines/stages/llama_encoding.py

@@ -1,130 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Prompt encoding stages for diffusion pipelines.
-
-This module contains implementations of prompt encoding stages for diffusion pipelines.
-"""
-
-from typing import TypedDict
-
-import torch
-
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.forward_context import set_forward_context
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.pipelines.stages.base import PipelineStage
-
-logger = init_logger(__name__)
-# TODO: this this hunyuan specific. Have an config for each model's special default hyperparameters
-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|>")
-
-
-class PromptTemplate(TypedDict):
-    template: str
-    crop_start: int
-
-
-prompt_template_video: PromptTemplate = {
-    "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
-    "crop_start": 95,
-}
-
-
-class LlamaEncodingStage(PipelineStage):
-    """
-    Stage for encoding text prompts into embeddings for diffusion models.
-    
-    This stage handles the encoding of text prompts into the embedding space
-    expected by the diffusion model.
-    """
-
-    def __init__(self, text_encoder, tokenizer) -> None:
-        """
-        Initialize the prompt encoding stage.
-        
-        Args:
-            enable_logging: Whether to enable logging for this stage.
-            is_secondary: Whether this is a secondary text encoder.
-        """
-        super().__init__()
-        self.text_encoder = text_encoder
-        self.tokenizer = tokenizer
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        """
-        Encode the prompt into text encoder hidden states.
-        
-        Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
-            
-        Returns:
-            The batch with encoded prompt embeddings.
-        """
-        if fastvideo_args.use_cpu_offload:
-            self.text_encoder = self.text_encoder.to(fastvideo_args.device)
-
-        text = prompt_template_video["template"].format(batch.prompt)
-        text_inputs = self.tokenizer(
-            text,
-            truncation=True,
-            # better way to handle this?
-            max_length=256,
-            return_tensors="pt",
-        )
-        hidden_state_skip_layer = 2
-        with set_forward_context(current_timestep=0, attn_metadata=None):
-            outputs = self.text_encoder(
-                input_ids=text_inputs["input_ids"].to(fastvideo_args.device),
-                output_hidden_states=hidden_state_skip_layer is not None,
-            )
-
-        last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer +
-                                                    1)]
-        crop_start = prompt_template_video.get("crop_start", -1)
-        last_hidden_state = last_hidden_state[:, crop_start:]
-        batch.prompt_embeds.append(last_hidden_state)
-
-        if batch.do_classifier_free_guidance:
-            negative_text = prompt_template_video["template"].format(
-                batch.negative_prompt)
-            negative_text_inputs = self.tokenizer(
-                negative_text,
-                truncation=True,
-                # better way to handle this?
-                max_length=256,
-                return_tensors="pt",
-            )
-            hidden_state_skip_layer = 2
-            with set_forward_context(current_timestep=0, attn_metadata=None):
-                negative_outputs = self.text_encoder(
-                    input_ids=negative_text_inputs["input_ids"].to(
-                        fastvideo_args.device),
-                    output_hidden_states=hidden_state_skip_layer is not None,
-                )
-
-            negative_last_hidden_state = negative_outputs.hidden_states[-(
-                hidden_state_skip_layer + 1)]
-            crop_start = prompt_template_video.get("crop_start", -1)
-            negative_last_hidden_state = negative_last_hidden_state[:,
-                                                                    crop_start:]
-            assert batch.negative_prompt_embeds is not None
-            batch.negative_prompt_embeds.append(negative_last_hidden_state)
-
-        if fastvideo_args.use_cpu_offload:
-            self.text_encoder.to('cpu')
-            torch.cuda.empty_cache()
-
-        return batch

fastvideo/v1/pipelines/stages/t5_encoding.py

@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Prompt encoding stages for diffusion pipelines.
-
-This module contains implementations of prompt encoding stages for diffusion pipelines.
-"""
-import torch
-
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.forward_context import set_forward_context
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.pipelines.stages.base import PipelineStage
-
-logger = init_logger(__name__)
-
-
-class T5EncodingStage(PipelineStage):
-    """
-    Stage for encoding text prompts into embeddings for diffusion models.
-    
-    This stage handles the encoding of text prompts into the embedding space
-    expected by the diffusion model.
-    """
-
-    def __init__(self, text_encoder, tokenizer) -> None:
-        """
-        Initialize the prompt encoding stage.
-        
-        Args:
-            enable_logging: Whether to enable logging for this stage.
-            is_secondary: Whether this is a secondary text encoder.
-        """
-        super().__init__()
-        self.text_encoder = text_encoder
-        self.tokenizer = tokenizer
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        """
-        Encode the prompt into text encoder hidden states.
-        
-        Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
-            
-        Returns:
-            The batch with encoded prompt embeddings.
-        """
-        if fastvideo_args.use_cpu_offload:
-            self.text_encoder = self.text_encoder.to(fastvideo_args.device)
-
-        text = batch.prompt
-        text_inputs = self.tokenizer(
-            text,
-            padding="max_length",
-            truncation=True,
-            max_length=512,
-            add_special_tokens=True,
-            return_attention_mask=True,
-            return_tensors="pt",
-        ).to(fastvideo_args.device)
-        text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
-        seq_lens = mask.gt(0).sum(dim=1).long()
-        with set_forward_context(current_timestep=0, attn_metadata=None):
-            outputs = self.text_encoder(
-                input_ids=text_input_ids,
-                attention_mask=mask,
-            )
-        assert torch.isnan(outputs).sum() == 0
-        prompt_embeds = [u[:v] for u, v in zip(outputs, seq_lens)]
-        prompt_embeds = torch.stack([
-            torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
-            for u in prompt_embeds
-        ],
-                                    dim=0)
-        batch.prompt_embeds.append(prompt_embeds)
-
-        if batch.do_classifier_free_guidance:
-            negative_text = batch.negative_prompt
-            negative_text_inputs = self.tokenizer(
-                negative_text,
-                padding="max_length",
-                truncation=True,
-                max_length=512,
-                add_special_tokens=True,
-                return_attention_mask=True,
-                return_tensors="pt",
-            ).to(fastvideo_args.device)
-            text_input_ids, mask = negative_text_inputs.input_ids, negative_text_inputs.attention_mask
-            seq_lens = mask.gt(0).sum(dim=1).long()
-            with set_forward_context(current_timestep=0, attn_metadata=None):
-                negative_outputs = self.text_encoder(
-                    input_ids=text_input_ids,
-                    attention_mask=mask,
-                )
-            assert torch.isnan(negative_outputs).sum() == 0
-            neg_prompt_embeds = [
-                u[:v] for u, v in zip(negative_outputs, seq_lens)
-            ]
-            neg_prompt_embeds = torch.stack([
-                torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
-                for u in neg_prompt_embeds
-            ],
-                                            dim=0)
-            assert batch.negative_prompt_embeds is not None
-            batch.negative_prompt_embeds.append(neg_prompt_embeds)
-
-        if fastvideo_args.use_cpu_offload:
-            self.text_encoder.to('cpu')
-            torch.cuda.empty_cache()
-
-        return batch

fastvideo/v1/pipelines/stages/text_encoding.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Prompt encoding stages for diffusion pipelines.
+
+This module contains implementations of prompt encoding stages for diffusion pipelines.
+"""
+
+import torch
+
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines.stages.base import PipelineStage
+
+logger = init_logger(__name__)
+
+
+class TextEncodingStage(PipelineStage):
+    """
+    Stage for encoding text prompts into embeddings for diffusion models.
+    
+    This stage handles the encoding of text prompts into the embedding space
+    expected by the diffusion model.
+    """
+
+    def __init__(self, text_encoders, tokenizers) -> None:
+        """
+        Initialize the prompt encoding stage.
+        
+        Args:
+            enable_logging: Whether to enable logging for this stage.
+            is_secondary: Whether this is a secondary text encoder.
+        """
+        super().__init__()
+        self.tokenizers = tokenizers
+        self.text_encoders = text_encoders
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Encode the prompt into text encoder hidden states.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with encoded prompt embeddings.
+        """
+        assert len(self.tokenizers) == len(self.text_encoders)
+        assert len(self.text_encoders) == len(
+            fastvideo_args.text_encoder_configs)
+
+        for tokenizer, text_encoder, encoder_config, preprocess_func, postprocess_func in zip(
+                self.tokenizers, self.text_encoders,
+                fastvideo_args.text_encoder_configs,
+                fastvideo_args.preprocess_text_funcs,
+                fastvideo_args.postprocess_text_funcs):
+            if fastvideo_args.use_cpu_offload:
+                text_encoder = text_encoder.to(fastvideo_args.device)
+
+            assert isinstance(batch.prompt, str)
+            text = preprocess_func(batch.prompt)
+            text_inputs = tokenizer(text, **encoder_config.tokenizer_kwargs).to(
+                fastvideo_args.device)
+            input_ids = text_inputs["input_ids"]
+            attention_mask = text_inputs["attention_mask"]
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                outputs = text_encoder(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                )
+            prompt_embeds = postprocess_func(outputs)
+
+            batch.prompt_embeds.append(prompt_embeds)
+
+            if batch.do_classifier_free_guidance:
+                assert isinstance(batch.negative_prompt, str)
+                negative_text = preprocess_func(batch.negative_prompt)
+                negative_text_inputs = tokenizer(
+                    negative_text,
+                    **encoder_config.tokenizer_kwargs).to(fastvideo_args.device)
+                negative_input_ids = negative_text_inputs["input_ids"]
+                negative_attention_mask = negative_text_inputs["attention_mask"]
+                with set_forward_context(current_timestep=0,
+                                         attn_metadata=None):
+                    negative_outputs = text_encoder(
+                        input_ids=negative_input_ids,
+                        attention_mask=negative_attention_mask,
+                        output_hidden_states=True,
+                    )
+                negative_prompt_embeds = postprocess_func(negative_outputs)
+
+                assert batch.negative_prompt_embeds is not None
+                batch.negative_prompt_embeds.append(negative_prompt_embeds)
+
+            if fastvideo_args.use_cpu_offload:
+                text_encoder.to('cpu')
+                torch.cuda.empty_cache()
+
+        return batch

fastvideo/v1/pipelines/stages/timestep_preparation.py

@@ -47,9 +47,9 @@ class TimestepPreparationStage(PipelineStage):
         timesteps = batch.timesteps
         sigmas = batch.sigmas
         n_tokens = batch.n_tokens
+        extra_set_timesteps_kwargs = batch.extra_set_timesteps_kwargs
 
         # Prepare extra kwargs for set_timesteps
-        extra_set_timesteps_kwargs = {}
         if n_tokens is not None and "n_tokens" in inspect.signature(
                 scheduler.set_timesteps).parameters:
             extra_set_timesteps_kwargs["n_tokens"] = n_tokens

fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py

@@ -12,9 +12,9 @@ from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
 
 # isort: off
 from fastvideo.v1.pipelines.stages import (
-    CLIPImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage,
+    ImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage,
     EncodingStage, InputValidationStage, LatentPreparationStage,
-    T5EncodingStage, TimestepPreparationStage)
+    TextEncodingStage, TimestepPreparationStage)
 # isort: on
 
 logger = init_logger(__name__)
@@ -34,13 +34,13 @@ class WanImageToVideoPipeline(ComposedPipelineBase):
                        stage=InputValidationStage())
 
         self.add_stage(stage_name="prompt_encoding_stage",
-                       stage=T5EncodingStage(
-                           text_encoder=self.get_module("text_encoder"),
-                           tokenizer=self.get_module("tokenizer"),
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
                        ))
 
         self.add_stage(stage_name="image_encoding_stage",
-                       stage=CLIPImageEncodingStage(
+                       stage=ImageEncodingStage(
                            image_encoder=self.get_module("image_encoder"),
                            image_processor=self.get_module("image_processor"),
                        ))

fastvideo/v1/pipelines/wan/wan_pipeline.py

@@ -12,7 +12,7 @@ from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage,
                                            DenoisingStage, InputValidationStage,
                                            LatentPreparationStage,
-                                           T5EncodingStage,
+                                           TextEncodingStage,
                                            TimestepPreparationStage)
 
 # TODO(will): move PRECISION_TO_TYPE to better place
@@ -33,9 +33,9 @@ class WanPipeline(ComposedPipelineBase):
                        stage=InputValidationStage())
 
         self.add_stage(stage_name="prompt_encoding_stage",
-                       stage=T5EncodingStage(
-                           text_encoder=self.get_module("text_encoder"),
-                           tokenizer=self.get_module("tokenizer"),
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
                        ))
 
         self.add_stage(stage_name="conditioning_stage",

fastvideo/v1/tests/encoders/test_clip_encoder.py

@@ -40,8 +40,8 @@ def test_clip_encoder():
     - Produce nearly identical outputs for the same input prompts
     """
     args = FastVideoArgs(model_path="openai/clip-vit-large-patch14",
-                         text_encoder_precision_2="fp16",
-                         text_encoder_config_2=CLIPTextConfig())
+                         text_encoder_precisions=("fp16",),
+                         text_encoder_configs=(CLIPTextConfig(),))
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
     logger.info("Loading models from %s", args.model_path)
@@ -66,7 +66,7 @@ def test_clip_encoder():
 
     # Load the HuggingFace implementation directly
     # model2 = CLIPTextModel(hf_config)
-    model2 = model2.to(torch.float16)
+    # model2 = model2.to(torch.float16)
     model2 = model2.to(device)
     model2.eval()
 

fastvideo/v1/tests/encoders/test_llama_encoder.py

@@ -40,8 +40,8 @@ def test_llama_encoder():
     - Produce nearly identical outputs for the same input prompts
     """
     args = FastVideoArgs(model_path="meta-llama/Llama-2-7b-hf",
-                         precision="float16",
-                         text_encoder_config=LlamaConfig())
+                         text_encoder_precisions=("fp16",),
+                         text_encoder_configs=(LlamaConfig(),))
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -62,7 +62,7 @@ def test_llama_encoder():
     model2 = loader.load(TEXT_ENCODER_PATH, "", args)
 
     # Convert to float16 and move to device
-    model2 = model2.to(torch.float16)
+    # model2 = model2.to(torch.float16)
     model2 = model2.to(device)
     model2.eval()
 

fastvideo/v1/tests/encoders/test_t5_encoder.py

@@ -9,7 +9,7 @@ from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.loader.component_loader import TextEncoderLoader
-from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.utils import maybe_download_model, PRECISION_TO_TYPE
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.configs.models.encoders import T5Config
 
@@ -33,17 +33,18 @@ def test_t5_encoder():
     print(hf_config)
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.float16
+    precision_str = "fp32"
+    precision = PRECISION_TO_TYPE[precision_str]
     model1 = UMT5EncoderModel.from_pretrained(TEXT_ENCODER_PATH).to(
         precision).to(device).eval()
     tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
 
-    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH, text_encoder_config=T5Config(), device_str="cuda")
+    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH, text_encoder_configs=(T5Config(),), text_encoder_precisions=(precision_str,), device_str="cuda")
     loader = TextEncoderLoader()
     model2 = loader.load(TEXT_ENCODER_PATH, "", args)
 
     # Convert to float16 and move to device
-    model2 = model2.to(precision)
+    # model2 = model2.to(precision)
     model2 = model2.to(device)
     model2.eval()
 
@@ -113,7 +114,7 @@ def test_t5_encoder():
                 outputs2 = model2(
                     input_ids=tokens.input_ids,
                     attention_mask=tokens.attention_mask,
-                )
+                ).last_hidden_state
 
             # Compare last hidden states
             last_hidden_state1 = outputs1[tokens.attention_mask == 1]

fastvideo/v1/tests/ssim/test_inference_similarity.py

@@ -4,9 +4,10 @@ import os
 
 import pytest
 
-from fastvideo.v1.entrypoints.cli.utils import launch_distributed
+from fastvideo import VideoGenerator
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.tests.ssim.compute_ssim import compute_video_ssim_torchvision
+from fastvideo.v1.worker.multiproc_executor import MultiprocExecutor
 
 logger = init_logger(__name__)
 
@@ -44,7 +45,7 @@ WAN_T2V_PARAMS = {
     "vae_sp": True,
     "fps": 24,
     "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
-    "text-encoder-precision": "fp32",
+    "text-encoder-precision": ("fp32",)
 }
 
 WAN_I2V_PARAMS = {
@@ -63,7 +64,7 @@ WAN_I2V_PARAMS = {
     "vae_sp": True,
     "fps": 24,
     "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
-    "text-encoder-precision": "fp32",
+    "text-encoder-precision": ("fp32",)
 }
 
 MODEL_TO_PARAMS = {
@@ -150,49 +151,39 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     BASE_PARAMS = I2V_MODEL_TO_PARAMS[model_id]
     num_inference_steps = BASE_PARAMS["num_inference_steps"]
     image_path = I2V_IMAGE_PATHS[I2V_TEST_PROMPTS.index(prompt)]
-    launch_args = [
-        "--num-inference-steps",
-        str(num_inference_steps),
-        "--prompt",
-        prompt,
-        "--output-path",
-        output_dir,
-        "--image_path",
-        image_path,
-        "--model-path",
-        BASE_PARAMS["model_path"],
-        "--height",
-        str(BASE_PARAMS["height"]),
-        "--width",
-        str(BASE_PARAMS["width"]),
-        "--num-frames",
-        str(BASE_PARAMS["num_frames"]),
-        "--guidance-scale",
-        str(BASE_PARAMS["guidance_scale"]),
-        "--embedded-cfg-scale",
-        str(BASE_PARAMS["embedded_cfg_scale"]),
-        "--flow-shift",
-        str(BASE_PARAMS["flow_shift"]),
-        "--seed",
-        str(BASE_PARAMS["seed"]),
-        "--sp-size",
-        str(BASE_PARAMS["sp_size"]),
-        "--tp-size",
-        str(BASE_PARAMS["tp_size"]),
-        "--fps",
-        str(BASE_PARAMS["fps"]),
-    ]
 
-    if BASE_PARAMS["vae_sp"]:
-        launch_args.append("--vae-sp")
-    if "neg_prompt" in BASE_PARAMS.keys():
-        launch_args.append("--neg_prompt")
-        launch_args.append(BASE_PARAMS["neg_prompt"])
-    if "text-encoder-precision" in BASE_PARAMS.keys():
-        launch_args.append("--text-encoder-precision")
-        launch_args.append(BASE_PARAMS["text-encoder-precision"])
+    init_kwargs = {
+        "num_gpus": BASE_PARAMS["num_gpus"],
+        "flow_shift": BASE_PARAMS["flow_shift"],
+        "sp_size": BASE_PARAMS["sp_size"],
+        "tp_size": BASE_PARAMS["tp_size"],
+    }
+    if BASE_PARAMS.get("vae_sp"):
+        init_kwargs["vae_sp"] = True
+        init_kwargs["vae_tiling"] = True
+    if "text-encoder-precision" in BASE_PARAMS:
+        init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"]
 
-    launch_distributed(num_gpus=BASE_PARAMS["num_gpus"], args=launch_args)
+    generation_kwargs = {
+        "num_inference_steps": num_inference_steps,
+        "output_path": output_dir,
+        "image_path": image_path,
+        "height": BASE_PARAMS["height"],
+        "width": BASE_PARAMS["width"],
+        "num_frames": BASE_PARAMS["num_frames"],
+        "guidance_scale": BASE_PARAMS["guidance_scale"],
+        "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"],
+        "seed": BASE_PARAMS["seed"],
+        "fps": BASE_PARAMS["fps"],
+    }
+    if "neg_prompt" in BASE_PARAMS:
+        generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]
+
+    generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs)
+    generator.generate_video(prompt, **generation_kwargs)
+    
+    if isinstance(generator.executor, MultiprocExecutor):
+        generator.executor.shutdown()
 
     assert os.path.exists(
         output_dir), f"Output video was not generated at {output_dir}"
@@ -237,7 +228,7 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     if not success:
         logger.error("Failed to write SSIM results to file")
 
-    min_acceptable_ssim = 1
+    min_acceptable_ssim = 0.97
     assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"
 
 @pytest.mark.parametrize("prompt", TEST_PROMPTS)
@@ -260,47 +251,38 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
 
     BASE_PARAMS = MODEL_TO_PARAMS[model_id]
     num_inference_steps = BASE_PARAMS["num_inference_steps"]
-    launch_args = [
-        "--num-inference-steps",
-        str(num_inference_steps),
-        "--prompt",
-        prompt,
-        "--output-path",
-        output_dir,
-        "--model-path",
-        BASE_PARAMS["model_path"],
-        "--height",
-        str(BASE_PARAMS["height"]),
-        "--width",
-        str(BASE_PARAMS["width"]),
-        "--num-frames",
-        str(BASE_PARAMS["num_frames"]),
-        "--guidance-scale",
-        str(BASE_PARAMS["guidance_scale"]),
-        "--embedded-cfg-scale",
-        str(BASE_PARAMS["embedded_cfg_scale"]),
-        "--flow-shift",
-        str(BASE_PARAMS["flow_shift"]),
-        "--seed",
-        str(BASE_PARAMS["seed"]),
-        "--sp-size",
-        str(BASE_PARAMS["sp_size"]),
-        "--tp-size",
-        str(BASE_PARAMS["tp_size"]),
-        "--fps",
-        str(BASE_PARAMS["fps"]),
-    ]
 
-    if BASE_PARAMS["vae_sp"]:
-        launch_args.append("--vae-sp")
-    if "neg_prompt" in BASE_PARAMS.keys():
-        launch_args.append("--neg_prompt")
-        launch_args.append(BASE_PARAMS["neg_prompt"])
-    if "text-encoder-precision" in BASE_PARAMS.keys():
-        launch_args.append("--text-encoder-precision")
-        launch_args.append(BASE_PARAMS["text-encoder-precision"])
+    init_kwargs = {
+        "num_gpus": BASE_PARAMS["num_gpus"],
+        "flow_shift": BASE_PARAMS["flow_shift"],
+        "sp_size": BASE_PARAMS["sp_size"],
+        "tp_size": BASE_PARAMS["tp_size"],
+    }
+    if BASE_PARAMS.get("vae_sp"):
+        init_kwargs["vae_sp"] = True
+        init_kwargs["vae_tiling"] = True
+    if "text-encoder-precision" in BASE_PARAMS:
+        init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"]
 
-    launch_distributed(num_gpus=BASE_PARAMS["num_gpus"], args=launch_args)
+    generation_kwargs = {
+        "num_inference_steps": num_inference_steps,
+        "output_path": output_dir,
+        "height": BASE_PARAMS["height"],
+        "width": BASE_PARAMS["width"],
+        "num_frames": BASE_PARAMS["num_frames"],
+        "guidance_scale": BASE_PARAMS["guidance_scale"],
+        "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"],
+        "seed": BASE_PARAMS["seed"],
+        "fps": BASE_PARAMS["fps"],
+    }
+    if "neg_prompt" in BASE_PARAMS:
+        generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]
+
+    generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs)
+    generator.generate_video(prompt, **generation_kwargs)
+
+    if isinstance(generator.executor, MultiprocExecutor):
+        generator.executor.shutdown()
 
     assert os.path.exists(
         output_dir), f"Output video was not generated at {output_dir}"
@@ -345,5 +327,5 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     if not success:
         logger.error("Failed to write SSIM results to file")
 
-    min_acceptable_ssim = 1
+    min_acceptable_ssim = 0.97
     assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"

fastvideo/v1/tests/transformers/test_flux.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import numpy as np
+import pytest
+import torch
+from diffusers import FluxTransformer2DModel
+
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import TransformerLoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.configs.models.dits import FluxImageConfig
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "black-forest-labs/FLUX.1-dev"
+# BASE_MODEL_PATH = "/home/test/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev/snapshots/0ef5fff789c832c5c7f4e127f94c8b54bbcced44/"
+
+
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_flux_transformer():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
+                         use_cpu_offload=False,
+                         precision=precision_str)
+    args.device = device
+    args.dit_config = FluxImageConfig()
+
+    loader = TransformerLoader()
+    model2 = loader.load(TRANSFORMER_PATH, "", args).to(device, dtype=precision)
+
+    model1 = FluxTransformer2DModel.from_pretrained(
+        TRANSFORMER_PATH, device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+
+    total_params = sum(p.numel() for p in model1.parameters())
+    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
+    weight_sum_model1 = sum(
+        p.to(torch.float64).sum().item() for p in model1.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model1 = weight_sum_model1 / total_params
+    logger.info("Model 1 weight sum: %s", weight_sum_model1)
+    logger.info("Model 1 weight mean: %s", weight_mean_model1)
+
+    # Calculate weight sum for model2 (converting to float64 to avoid overflow)
+    total_params_model2 = sum(p.numel() for p in model2.parameters())
+    weight_sum_model2 = sum(
+        p.to(torch.float64).sum().item() for p in model2.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model2 = weight_sum_model2 / total_params_model2
+    logger.info("Model 2 weight sum: %s", weight_sum_model2)
+    logger.info("Model 2 weight mean: %s", weight_mean_model2)
+
+    weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2)
+    logger.info("Weight sum difference: %s", weight_sum_diff)
+    weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2)
+    logger.info("Weight mean difference: %s", weight_mean_diff)
+
+    # Set both models to eval mode
+    model1 = model1.eval()
+    model2 = model2.eval()
+
+    # Create identical inputs for both models
+    batch_size = 1
+    seq_len = 30
+
+    # Video latents [B, N, C]
+    hidden_states = torch.randn(batch_size,
+                                64 * 64,
+                                64,
+                                device=device, dtype=precision)
+
+    # Text embeddings [B, L, D] (including global token)
+    encoder_hidden_states = torch.randn(batch_size,
+                                        seq_len,
+                                        4096,
+                                        device=device,
+                                        dtype=precision)
+    pooled_text_embeds = torch.randn(batch_size,
+                                    768,
+                                    device=device,
+                                    dtype=precision)
+
+    # Timestep
+    timestep = torch.tensor([0.5], device=device, dtype=precision)
+    guidance_1 = torch.tensor([3.5], device=device, dtype=torch.float32).expand(batch_size)
+    guidance_2 = torch.tensor([3.5], device=device, dtype=torch.float32).expand(batch_size).to(precision) * 1000
+
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    with torch.amp.autocast('cuda', dtype=precision):
+        print("Running model1 inference...")
+        output1 = model1(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            pooled_projections=pooled_text_embeds,
+            timestep=timestep,
+            txt_ids=torch.zeros(encoder_hidden_states.shape[1], 3).to(device=device, dtype=precision),
+            img_ids=_prepare_latent_image_ids(batch_size, 64, 64, device, precision),
+            guidance=guidance_1,
+            return_dict=False,
+        )[0]
+        print("Running model2 inference...")
+        with set_forward_context(
+                current_timestep=0,
+                attn_metadata=None,
+        ):
+            output2 = model2(hidden_states=hidden_states,
+                             encoder_hidden_states=[pooled_text_embeds,
+                                                   encoder_hidden_states],
+                             guidance=guidance_2,
+                             timestep=timestep * 1000,
+                             height_latents=128,
+                             width_latents=128)
+
+
+    # Check if outputs have the same shape
+    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
+
+    # Check if outputs are similar (allowing for small numerical differences)
+    max_diff = torch.max(torch.abs(output1 - output2))
+    mean_diff = torch.mean(torch.abs(output1 - output2))
+    logger.info("Max Diff: %s", max_diff.item())
+    logger.info("Mean Diff: %s", mean_diff.item())
+    assert max_diff < 4, f"Maximum difference between outputs: {max_diff.item()}"
+    assert mean_diff < 5e-1, f"Mean difference between outputs: {mean_diff.item()}"