mp training

Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com>
Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com>
Co-authored-by: Will Lin <wlsaidhi@gmail.com>

.github/workflows/pr-test.yml

@@ -83,193 +83,81 @@ jobs:
     if: >-
       (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.encoder-test == 'true') || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_encoder_test == 'true')
-    runs-on: ubuntu-latest
-    environment: runpod-runners
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Set up SSH key
-        run: |
-          mkdir -p ~/.ssh
-          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
-          chmod 600 ~/.ssh/id_rsa
-          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Run tests on RunPod
-        env:
-          JOB_ID: "encoder-test"
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        timeout-minutes: 30
-        run: >-
-          python .github/scripts/runpod_api.py 
-          --gpu-type "NVIDIA A40"
-          --gpu-count 1
-          --volume-size 100
-          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
-
-      - name: Terminate RunPod Instances
-        if: ${{ always() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: "encoder-test"
-        run: python .github/scripts/runpod_cleanup.py
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "encoder-test"
+      gpu_type: "NVIDIA A40"
+      gpu_count: 1
+      volume_size: 100
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
+      timeout_minutes: 30
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
 
   vae-test:
     needs: change-filter
     if: >-
       (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.vae-test == 'true') || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_vae_test == 'true')
-    runs-on: ubuntu-latest
-    environment: runpod-runners
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Set up SSH key
-        run: |
-          mkdir -p ~/.ssh
-          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
-          chmod 600 ~/.ssh/id_rsa
-          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Run tests on RunPod
-        env:
-          JOB_ID: "vae-test"
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        timeout-minutes: 30
-        run: >-
-          python .github/scripts/runpod_api.py 
-          --gpu-type "NVIDIA A40"
-          --gpu-count 1
-          --volume-size 100
-          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
-
-      - name: Terminate RunPod Instances
-        if: ${{ always() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: "vae-test"
-        run: python .github/scripts/runpod_cleanup.py
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "vae-test"
+      gpu_type: "NVIDIA A40"
+      gpu_count: 1
+      volume_size: 100
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+      timeout_minutes: 30
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
 
   transformer-test:
     needs: change-filter
     if: >-
       (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true')
-    runs-on: ubuntu-latest
-    environment: runpod-runners
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Set up SSH key
-        run: |
-          mkdir -p ~/.ssh
-          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
-          chmod 600 ~/.ssh/id_rsa
-          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Run tests on RunPod
-        env:
-          JOB_ID: "transformer-test"
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        timeout-minutes: 30
-        run: >-
-          python .github/scripts/runpod_api.py 
-          --gpu-type "NVIDIA L40S"
-          --gpu-count 1
-          --volume-size 100
-          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
-
-      - name: Terminate RunPod Instances
-        if: ${{ always() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: "transformer-test"
-        run: python .github/scripts/runpod_cleanup.py
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "transformer-test"
+      gpu_type: "NVIDIA L40S"
+      gpu_count: 1
+      volume_size: 100
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+      timeout_minutes: 30
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
       
   ssim-test:
     needs: change-filter
     if: >-
       (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
-    runs-on: ubuntu-latest
-    environment: runpod-runners
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Set up SSH key
-        run: |
-          mkdir -p ~/.ssh
-          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
-          chmod 600 ~/.ssh/id_rsa
-          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Run tests on RunPod
-        env:
-          JOB_ID: "ssim-test"
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        timeout-minutes: 45
-        run: >-
-          python .github/scripts/runpod_api.py 
-          --gpu-type "NVIDIA A40"
-          --gpu-count 2
-          --disk-size 200
-          --volume-size 200
-          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
-
-      - name: Terminate RunPod Instances
-        if: ${{ always() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: "ssim-test"
-        run: python .github/scripts/runpod_cleanup.py
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [
+          {version: "3.10", tag: "latest"},
+          {version: "3.11", tag: "py3.11-latest"},
+          {version: "3.12", tag: "py3.12-latest"}
+        ]
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "ssim-test-py${{ matrix.python-version.version }}"
+      gpu_type: "NVIDIA A40"
+      gpu_count: 2
+      volume_size: 200
+      disk_size: 200
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
+      timeout_minutes: 60
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
 
   runpod-cleanup:
     needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
@@ -289,7 +177,7 @@ jobs:
 
       - name: Cleanup all RunPod instances
         env:
-          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test"]' # JSON array of job IDs
+          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]'
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
         run: python .github/scripts/runpod_cleanup.py

.github/workflows/runpod-test.yml

@@ -0,0 +1,91 @@
+name: RunPod Test
+
+on:
+  workflow_call:
+    inputs:
+      job_id:
+        required: true
+        type: string
+        description: "Unique identifier for this test job"
+      gpu_type:
+        required: true
+        type: string
+        description: "GPU type to use (e.g. NVIDIA A40, NVIDIA L40S)"
+      gpu_count:
+        required: true
+        type: number
+        description: "Number of GPUs to use"
+      volume_size:
+        required: false
+        type: number
+        default: 20
+        description: "Volume size in GB"
+      disk_size:
+        required: false
+        type: number
+        default: 20
+        description: "Disk size in GB"
+      image:
+        required: true
+        type: string
+        description: "Docker image to use"
+      test_command:
+        required: true
+        type: string
+        description: "Command to run tests"
+      timeout_minutes:
+        required: false
+        type: number
+        default: 30
+        description: "Timeout in minutes"
+    secrets:
+      RUNPOD_API_KEY:
+        required: true
+      RUNPOD_PRIVATE_KEY:
+        required: true
+
+jobs:
+  run-test:
+    runs-on: ubuntu-latest
+    environment: runpod-runners
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Set up SSH key
+        run: |
+          mkdir -p ~/.ssh
+          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
+          chmod 600 ~/.ssh/id_rsa
+          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
+
+      - name: Install dependencies
+        run: pip install requests
+
+      - name: Run tests on RunPod
+        env:
+          JOB_ID: ${{ inputs.job_id }}
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: ${{ inputs.timeout_minutes }}
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "${{ inputs.gpu_type }}"
+          --gpu-count ${{ inputs.gpu_count }}
+          --volume-size ${{ inputs.volume_size }}
+          --disk-size ${{ inputs.disk_size }}
+          --image "${{ inputs.image }}"
+          --test-command "${{ inputs.test_command }}"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: ${{ inputs.job_id }}
+        run: python .github/scripts/runpod_cleanup.py 

.pre-commit-config.yaml

@@ -22,7 +22,8 @@ exclude: |
     examples/.*|
     .github/workflows/fastvideo-publish.yml|
     .github/workflows/sta-publish.yml|
-    .github/workflows/build-image-template.yml
+    .github/workflows/build-image-template.yml|
+    docs/source/inference/support_matrix.md
   )
 repos:
 - repo: https://github.com/google/yapf

README.md

@@ -2,62 +2,109 @@
 <img src=assets/logo.jpg width="30%"/>
 </div>
 
-FastVideo is a lightweight framework for accelerating large video diffusion models.
+**FastVideo is a unified framework for accelerated video generation.**
+
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
 </p>
 
-https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
+<div align="center">
+<img src=assets/perf.png width="90%"/>
+</div>
 
-FastVideo currently offers: (with more to come)
+## Key Features
 
-- [NEW!] V1 inference API available. Full announcement coming soon!
-- [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
-- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
 
-Dev in progress and highly experimental.
-
-## Change Log
-- ```2025/02/20```: FastVideo now supports STA on [StepVideo](https://github.com/stepfun-ai/Step-Video-T2V) with 3.4X speedup!
-- ```2025/02/18```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- ```2025/01/13```: Support Lora finetuning for HunyuanVideo.
-- ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
-- ```2024/12/17```: `FastVideo` v0.0.1 is released.
-
 ## Getting Started
+We recommend using an environment manager such as `Conda` to create a clean environment:
+
+```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
+```
+
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
+
+## Inference
+### Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```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,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
+
+```bash
+python example.py
+```
+
+For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html).
+
+### Other docs:
 
-- [Install FastVideo](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
 - [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview.html)
 - [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
 
-### Inference
-- [Quick Start](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html)
-- V1 Inference API Guide (Coming soon!)
-
-### Distillation and Finetuning
+## Distillation and Finetuning
 - [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
 - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetuning.html)
 
-### Deprecated APIs
-- [V0 Inference (Deprecated)](https://hao-ai-lab.github.io/FastVideo/inference/v0_inference.html)
-
 ## 📑 Development Plan
 
 <!-- - More distillation methods -->
   <!-- - [ ] Add Distribution Matching Distillation -->
 - More models support
   <!-- - [ ] Add CogvideoX model -->
-  - [ ] Add StepVideo to V1
+  - [x] Add StepVideo to V1
 - Optimization features
-  - [ ] Teacache in V1
-  - [ ] SageAttention in V1
+  - [x] Teacache in V1
+  - [x] SageAttention in V1
 - Code updates
-  - [ ] V1 Configuration API
+  - [x] V1 Configuration API
   - [ ] Support Training in V1
   <!-- - [ ] fp8 support -->
   <!-- - [ ] faster load model and save model support -->

collect_env.py

@@ -0,0 +1,772 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# type: ignore
+# ruff: noqa
+# code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py
+# and vllm: https://github.com/vllm-project/vllm/blob/main/vllm/collect_env.py
+
+import datetime
+import locale
+import os
+import re
+import subprocess
+import sys
+# Unlike the rest of the PyTorch this file must be python2 compliant.
+# This script outputs relevant system environment info
+# Run it with `python collect_env.py` or `python -m torch.utils.collect_env`
+from collections import namedtuple
+
+from fastvideo.v1.envs import environment_variables
+
+try:
+    import torch
+    TORCH_AVAILABLE = True
+except (ImportError, NameError, AttributeError, OSError):
+    TORCH_AVAILABLE = False
+
+# System Environment Information
+SystemEnv = namedtuple(
+    'SystemEnv',
+    [
+        'torch_version',
+        'is_debug_build',
+        'cuda_compiled_version',
+        'gcc_version',
+        'clang_version',
+        'cmake_version',
+        'os',
+        'libc_version',
+        'python_version',
+        'python_platform',
+        'is_cuda_available',
+        'cuda_runtime_version',
+        'cuda_module_loading',
+        'nvidia_driver_version',
+        'nvidia_gpu_models',
+        'cudnn_version',
+        'pip_version',  # 'pip' or 'pip3'
+        'pip_packages',
+        'conda_packages',
+        'hip_compiled_version',
+        'hip_runtime_version',
+        'miopen_runtime_version',
+        'caching_allocator_config',
+        'is_xnnpack_available',
+        'cpu_info',
+        'fastvideo_version',
+        'fastvideo_build_flags',
+        'gpu_topo',
+        'env_vars',
+    ])
+
+DEFAULT_CONDA_PATTERNS = {
+    "torch",
+    "numpy",
+    "cudatoolkit",
+    "soumith",
+    "mkl",
+    "magma",
+    "triton",
+    "optree",
+    "nccl",
+    "transformers",
+    "zmq",
+    "nvidia",
+    "pynvml",
+}
+
+DEFAULT_PIP_PATTERNS = {
+    "torch",
+    "numpy",
+    "mypy",
+    "flake8",
+    "triton",
+    "optree",
+    "onnx",
+    "nccl",
+    "transformers",
+    "zmq",
+    "nvidia",
+    "pynvml",
+}
+
+
+def run(command):
+    """Return (return-code, stdout, stderr)."""
+    shell = True if type(command) is str else False
+    p = subprocess.Popen(command,
+                         stdout=subprocess.PIPE,
+                         stderr=subprocess.PIPE,
+                         shell=shell)
+    raw_output, raw_err = p.communicate()
+    rc = p.returncode
+    if get_platform() == 'win32':
+        enc = 'oem'
+    else:
+        enc = locale.getpreferredencoding()
+    output = raw_output.decode(enc)
+    if command == 'nvidia-smi topo -m':
+        # don't remove the leading whitespace of `nvidia-smi topo -m`
+        #   because they are meaningful
+        output = output.rstrip()
+    else:
+        output = output.strip()
+    err = raw_err.decode(enc)
+    return rc, output, err.strip()
+
+
+def run_and_read_all(run_lambda, command):
+    """Run command using run_lambda; reads and returns entire output if rc is 0."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    return out
+
+
+def run_and_parse_first_match(run_lambda, command, regex):
+    """Run command using run_lambda, returns the first regex match if it exists."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    match = re.search(regex, out)
+    if match is None:
+        return None
+    return match.group(1)
+
+
+def run_and_return_first_line(run_lambda, command):
+    """Run command using run_lambda and returns first line if output is not empty."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    return out.split('\n')[0]
+
+
+def get_conda_packages(run_lambda, patterns=None):
+    if patterns is None:
+        patterns = DEFAULT_CONDA_PATTERNS
+    conda = os.environ.get('CONDA_EXE', 'conda')
+    out = run_and_read_all(run_lambda, "{} list".format(conda))
+    if out is None:
+        return out
+
+    return "\n".join(line for line in out.splitlines()
+                     if not line.startswith("#") and any(name in line
+                                                         for name in patterns))
+
+
+def get_gcc_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'gcc --version', r'gcc (.*)')
+
+
+def get_clang_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'clang --version',
+                                     r'clang version (.*)')
+
+
+def get_cmake_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'cmake --version',
+                                     r'cmake (.*)')
+
+
+def get_nvidia_driver_version(run_lambda):
+    if get_platform() == 'darwin':
+        cmd = 'kextstat | grep -i cuda'
+        return run_and_parse_first_match(run_lambda, cmd,
+                                         r'com[.]nvidia[.]CUDA [(](.*?)[)]')
+    smi = get_nvidia_smi()
+    return run_and_parse_first_match(run_lambda, smi, r'Driver Version: (.*?) ')
+
+
+def get_gpu_info(run_lambda):
+    if get_platform() == 'darwin' or (TORCH_AVAILABLE and hasattr(
+            torch.version, 'hip') and torch.version.hip is not None):
+        if TORCH_AVAILABLE and torch.cuda.is_available():
+            if torch.version.hip is not None:
+                prop = torch.cuda.get_device_properties(0)
+                if hasattr(prop, "gcnArchName"):
+                    gcnArch = " ({})".format(prop.gcnArchName)
+                else:
+                    gcnArch = "NoGCNArchNameOnOldPyTorch"
+            else:
+                gcnArch = ""
+            return torch.cuda.get_device_name(None) + gcnArch
+        return None
+    smi = get_nvidia_smi()
+    uuid_regex = re.compile(r' \(UUID: .+?\)')
+    rc, out, _ = run_lambda(smi + ' -L')
+    if rc != 0:
+        return None
+    # Anonymize GPUs by removing their UUID
+    return re.sub(uuid_regex, '', out)
+
+
+def get_running_cuda_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'nvcc --version',
+                                     r'release .+ V(.*)')
+
+
+def get_cudnn_version(run_lambda):
+    """Return a list of libcudnn.so; it's hard to tell which one is being used."""
+    if get_platform() == 'win32':
+        system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
+        cuda_path = os.environ.get('CUDA_PATH', "%CUDA_PATH%")
+        where_cmd = os.path.join(system_root, 'System32', 'where')
+        cudnn_cmd = '{} /R "{}\\bin" cudnn*.dll'.format(where_cmd, cuda_path)
+    elif get_platform() == 'darwin':
+        # CUDA libraries and drivers can be found in /usr/local/cuda/. See
+        # https://docs.nvidia.com/cuda/cuda-installation-guide-mac-os-x/index.html#install
+        # https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#installmac
+        # Use CUDNN_LIBRARY when cudnn library is installed elsewhere.
+        cudnn_cmd = 'ls /usr/local/cuda/lib/libcudnn*'
+    else:
+        cudnn_cmd = 'ldconfig -p | grep libcudnn | rev | cut -d" " -f1 | rev'
+    rc, out, _ = run_lambda(cudnn_cmd)
+    # find will return 1 if there are permission errors or if not found
+    if len(out) == 0 or (rc != 1 and rc != 0):
+        l = os.environ.get('CUDNN_LIBRARY')
+        if l is not None and os.path.isfile(l):
+            return os.path.realpath(l)
+        return None
+    files_set = set()
+    for fn in out.split('\n'):
+        fn = os.path.realpath(fn)  # eliminate symbolic links
+        if os.path.isfile(fn):
+            files_set.add(fn)
+    if not files_set:
+        return None
+    # Alphabetize the result because the order is non-deterministic otherwise
+    files = sorted(files_set)
+    if len(files) == 1:
+        return files[0]
+    result = '\n'.join(files)
+    return 'Probably one of the following:\n{}'.format(result)
+
+
+def get_nvidia_smi():
+    # Note: nvidia-smi is currently available only on Windows and Linux
+    smi = 'nvidia-smi'
+    if get_platform() == 'win32':
+        system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
+        program_files_root = os.environ.get('PROGRAMFILES', 'C:\\Program Files')
+        legacy_path = os.path.join(program_files_root, 'NVIDIA Corporation',
+                                   'NVSMI', smi)
+        new_path = os.path.join(system_root, 'System32', smi)
+        smis = [new_path, legacy_path]
+        for candidate_smi in smis:
+            if os.path.exists(candidate_smi):
+                smi = '"{}"'.format(candidate_smi)
+                break
+    return smi
+
+
+def get_fastvideo_version():
+    return ""
+    from fastvideo import __version__, __version_tuple__
+
+    if __version__ == "dev":
+        return "N/A (dev)"
+    version_str = __version_tuple__[-1]
+    if isinstance(version_str, str) and version_str.startswith('g'):
+        # it's a dev build
+        if '.' in version_str:
+            # it's a dev build containing local changes
+            git_sha = version_str.split('.')[0][1:]
+            date = version_str.split('.')[-1][1:]
+            return f"{__version__} (git sha: {git_sha}, date: {date})"
+        else:
+            # it's a dev build without local changes
+            git_sha = version_str[1:]  # type: ignore
+            return f"{__version__} (git sha: {git_sha})"
+    return __version__
+
+
+def summarize_fastvideo_build_flags():
+    # This could be a static method if the flags are constant, or dynamic if you need to check environment variables, etc.
+    return 'CUDA Archs: {}; ROCm: {}; Neuron: {}'.format(
+        os.environ.get('TORCH_CUDA_ARCH_LIST', 'Not Set'),
+        'Enabled' if os.environ.get('ROCM_HOME') else 'Disabled',
+        'Enabled' if os.environ.get('NEURON_CORES') else 'Disabled',
+    )
+
+
+def get_gpu_topo(run_lambda):
+    output = None
+
+    if get_platform() == 'linux':
+        output = run_and_read_all(run_lambda, 'nvidia-smi topo -m')
+        if output is None:
+            output = run_and_read_all(run_lambda, 'rocm-smi --showtopo')
+
+    return output
+
+
+# example outputs of CPU infos
+#  * linux
+#    Architecture:            x86_64
+#      CPU op-mode(s):        32-bit, 64-bit
+#      Address sizes:         46 bits physical, 48 bits virtual
+#      Byte Order:            Little Endian
+#    CPU(s):                  128
+#      On-line CPU(s) list:   0-127
+#    Vendor ID:               GenuineIntel
+#      Model name:            Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
+#        CPU family:          6
+#        Model:               106
+#        Thread(s) per core:  2
+#        Core(s) per socket:  32
+#        Socket(s):           2
+#        Stepping:            6
+#        BogoMIPS:            5799.78
+#        Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr
+#                             sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl
+#                             xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16
+#                             pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand
+#                             hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced
+#                             fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap
+#                             avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1
+#                             xsaves wbnoinvd ida arat avx512vbmi pku ospke avx512_vbmi2 gfni vaes vpclmulqdq
+#                             avx512_vnni avx512_bitalg tme avx512_vpopcntdq rdpid md_clear flush_l1d arch_capabilities
+#    Virtualization features:
+#      Hypervisor vendor:     KVM
+#      Virtualization type:   full
+#    Caches (sum of all):
+#      L1d:                   3 MiB (64 instances)
+#      L1i:                   2 MiB (64 instances)
+#      L2:                    80 MiB (64 instances)
+#      L3:                    108 MiB (2 instances)
+#    NUMA:
+#      NUMA node(s):          2
+#      NUMA node0 CPU(s):     0-31,64-95
+#      NUMA node1 CPU(s):     32-63,96-127
+#    Vulnerabilities:
+#      Itlb multihit:         Not affected
+#      L1tf:                  Not affected
+#      Mds:                   Not affected
+#      Meltdown:              Not affected
+#      Mmio stale data:       Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
+#      Retbleed:              Not affected
+#      Spec store bypass:     Mitigation; Speculative Store Bypass disabled via prctl and seccomp
+#      Spectre v1:            Mitigation; usercopy/swapgs barriers and __user pointer sanitization
+#      Spectre v2:            Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
+#      Srbds:                 Not affected
+#      Tsx async abort:       Not affected
+#  * win32
+#    Architecture=9
+#    CurrentClockSpeed=2900
+#    DeviceID=CPU0
+#    Family=179
+#    L2CacheSize=40960
+#    L2CacheSpeed=
+#    Manufacturer=GenuineIntel
+#    MaxClockSpeed=2900
+#    Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
+#    ProcessorType=3
+#    Revision=27142
+#
+#    Architecture=9
+#    CurrentClockSpeed=2900
+#    DeviceID=CPU1
+#    Family=179
+#    L2CacheSize=40960
+#    L2CacheSpeed=
+#    Manufacturer=GenuineIntel
+#    MaxClockSpeed=2900
+#    Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
+#    ProcessorType=3
+#    Revision=27142
+
+
+def get_cpu_info(run_lambda):
+    rc, out, err = 0, '', ''
+    if get_platform() == 'linux':
+        rc, out, err = run_lambda('lscpu')
+    elif get_platform() == 'win32':
+        rc, out, err = run_lambda(
+            'wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID, \
+        CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE'
+        )
+    elif get_platform() == 'darwin':
+        rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string")
+    cpu_info = 'None'
+    if rc == 0:
+        cpu_info = out
+    else:
+        cpu_info = err
+    return cpu_info
+
+
+def get_platform():
+    if sys.platform.startswith('linux'):
+        return 'linux'
+    elif sys.platform.startswith('win32'):
+        return 'win32'
+    elif sys.platform.startswith('cygwin'):
+        return 'cygwin'
+    elif sys.platform.startswith('darwin'):
+        return 'darwin'
+    else:
+        return sys.platform
+
+
+def get_mac_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'sw_vers -productVersion',
+                                     r'(.*)')
+
+
+def get_windows_version(run_lambda):
+    system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
+    wmic_cmd = os.path.join(system_root, 'System32', 'Wbem', 'wmic')
+    findstr_cmd = os.path.join(system_root, 'System32', 'findstr')
+    return run_and_read_all(
+        run_lambda,
+        '{} os get Caption | {} /v Caption'.format(wmic_cmd, findstr_cmd))
+
+
+def get_lsb_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'lsb_release -a',
+                                     r'Description:\t(.*)')
+
+
+def check_release_file(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'cat /etc/*-release',
+                                     r'PRETTY_NAME="(.*)"')
+
+
+def get_os(run_lambda):
+    from platform import machine
+    platform = get_platform()
+
+    if platform == 'win32' or platform == 'cygwin':
+        return get_windows_version(run_lambda)
+
+    if platform == 'darwin':
+        version = get_mac_version(run_lambda)
+        if version is None:
+            return None
+        return 'macOS {} ({})'.format(version, machine())
+
+    if platform == 'linux':
+        # Ubuntu/Debian based
+        desc = get_lsb_version(run_lambda)
+        if desc is not None:
+            return '{} ({})'.format(desc, machine())
+
+        # Try reading /etc/*-release
+        desc = check_release_file(run_lambda)
+        if desc is not None:
+            return '{} ({})'.format(desc, machine())
+
+        return '{} ({})'.format(platform, machine())
+
+    # Unknown platform
+    return platform
+
+
+def get_python_platform():
+    import platform
+    return platform.platform()
+
+
+def get_libc_version():
+    import platform
+    if get_platform() != 'linux':
+        return 'N/A'
+    return '-'.join(platform.libc_ver())
+
+
+def get_pip_packages(run_lambda, patterns=None):
+    """Return `pip list` output. Note: will also find conda-installed pytorch and numpy packages."""
+    if patterns is None:
+        patterns = DEFAULT_PIP_PATTERNS
+
+    def run_with_pip():
+        try:
+            import importlib.util
+            pip_spec = importlib.util.find_spec('pip')
+            pip_available = pip_spec is not None
+        except ImportError:
+            pip_available = False
+
+        if pip_available:
+            cmd = [sys.executable, '-mpip', 'list', '--format=freeze']
+        elif os.environ.get("UV") is not None:
+            print("uv is set")
+            cmd = ["uv", "pip", "list", "--format=freeze"]
+        else:
+            raise RuntimeError(
+                "Could not collect pip list output (pip or uv module not available)"
+            )
+
+        out = run_and_read_all(run_lambda, cmd)
+        return "\n".join(line for line in out.splitlines()
+                         if any(name in line for name in patterns))
+
+    pip_version = 'pip3' if sys.version[0] == '3' else 'pip'
+    out = run_with_pip()
+    return pip_version, out
+
+
+def get_cachingallocator_config():
+    ca_config = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '')
+    return ca_config
+
+
+def get_cuda_module_loading_config():
+    if TORCH_AVAILABLE and torch.cuda.is_available():
+        torch.cuda.init()
+        config = os.environ.get('CUDA_MODULE_LOADING', '')
+        return config
+    else:
+        return "N/A"
+
+
+def is_xnnpack_available():
+    if TORCH_AVAILABLE:
+        import torch.backends.xnnpack
+        return str(torch.backends.xnnpack.enabled)  # type: ignore[attr-defined]
+    else:
+        return "N/A"
+
+
+def get_env_vars():
+    env_vars = ''
+    secret_terms = ('secret', 'token', 'api', 'access', 'password')
+    report_prefix = ("TORCH", "NCCL", "PYTORCH", "CUDA", "CUBLAS", "CUDNN",
+                     "OMP_", "MKL_", "NVIDIA")
+    for k, v in os.environ.items():
+        if any(term in k.lower() for term in secret_terms):
+            continue
+        if k in environment_variables:
+            env_vars = env_vars + "{}={}".format(k, v) + "\n"
+        if k.startswith(report_prefix):
+            env_vars = env_vars + "{}={}".format(k, v) + "\n"
+
+    return env_vars
+
+
+def get_env_info():
+    run_lambda = run
+    pip_version, pip_list_output = get_pip_packages(run_lambda)
+
+    if TORCH_AVAILABLE:
+        version_str = torch.__version__
+        debug_mode_str = str(torch.version.debug)
+        cuda_available_str = str(torch.cuda.is_available())
+        cuda_version_str = torch.version.cuda
+        if not hasattr(torch.version,
+                       'hip') or torch.version.hip is None:  # cuda version
+            hip_compiled_version = hip_runtime_version = miopen_runtime_version = 'N/A'
+        else:  # HIP version
+
+            def get_version_or_na(cfg, prefix):
+                _lst = [s.rsplit(None, 1)[-1] for s in cfg if prefix in s]
+                return _lst[0] if _lst else 'N/A'
+
+            cfg = torch._C._show_config().split('\n')
+            hip_runtime_version = get_version_or_na(cfg, 'HIP Runtime')
+            miopen_runtime_version = get_version_or_na(cfg, 'MIOpen')
+            cuda_version_str = 'N/A'
+            hip_compiled_version = torch.version.hip
+    else:
+        version_str = debug_mode_str = cuda_available_str = cuda_version_str = 'N/A'
+        hip_compiled_version = hip_runtime_version = miopen_runtime_version = 'N/A'
+
+    sys_version = sys.version.replace("\n", " ")
+
+    conda_packages = get_conda_packages(run_lambda)
+
+    fastvideo_version = get_fastvideo_version()
+    fastvideo_build_flags = summarize_fastvideo_build_flags()
+    gpu_topo = get_gpu_topo(run_lambda)
+
+    return SystemEnv(
+        torch_version=version_str,
+        is_debug_build=debug_mode_str,
+        python_version='{} ({}-bit runtime)'.format(
+            sys_version,
+            sys.maxsize.bit_length() + 1),
+        python_platform=get_python_platform(),
+        is_cuda_available=cuda_available_str,
+        cuda_compiled_version=cuda_version_str,
+        cuda_runtime_version=get_running_cuda_version(run_lambda),
+        cuda_module_loading=get_cuda_module_loading_config(),
+        nvidia_gpu_models=get_gpu_info(run_lambda),
+        nvidia_driver_version=get_nvidia_driver_version(run_lambda),
+        cudnn_version=get_cudnn_version(run_lambda),
+        hip_compiled_version=hip_compiled_version,
+        hip_runtime_version=hip_runtime_version,
+        miopen_runtime_version=miopen_runtime_version,
+        pip_version=pip_version,
+        pip_packages=pip_list_output,
+        conda_packages=conda_packages,
+        os=get_os(run_lambda),
+        libc_version=get_libc_version(),
+        gcc_version=get_gcc_version(run_lambda),
+        clang_version=get_clang_version(run_lambda),
+        cmake_version=get_cmake_version(run_lambda),
+        caching_allocator_config=get_cachingallocator_config(),
+        is_xnnpack_available=is_xnnpack_available(),
+        cpu_info=get_cpu_info(run_lambda),
+        fastvideo_version=fastvideo_version,
+        fastvideo_build_flags=fastvideo_build_flags,
+        gpu_topo=gpu_topo,
+        env_vars=get_env_vars(),
+    )
+
+
+env_info_fmt = """
+PyTorch version: {torch_version}
+Is debug build: {is_debug_build}
+CUDA used to build PyTorch: {cuda_compiled_version}
+ROCM used to build PyTorch: {hip_compiled_version}
+
+OS: {os}
+GCC version: {gcc_version}
+Clang version: {clang_version}
+CMake version: {cmake_version}
+Libc version: {libc_version}
+
+Python version: {python_version}
+Python platform: {python_platform}
+Is CUDA available: {is_cuda_available}
+CUDA runtime version: {cuda_runtime_version}
+CUDA_MODULE_LOADING set to: {cuda_module_loading}
+GPU models and configuration: {nvidia_gpu_models}
+Nvidia driver version: {nvidia_driver_version}
+cuDNN version: {cudnn_version}
+HIP runtime version: {hip_runtime_version}
+MIOpen runtime version: {miopen_runtime_version}
+Is XNNPACK available: {is_xnnpack_available}
+
+CPU:
+{cpu_info}
+
+Versions of relevant libraries:
+{pip_packages}
+{conda_packages}
+""".strip()
+
+# both the above code and the following code use `strip()` to
+# remove leading/trailing whitespaces, so we need to add a newline
+# in between to separate the two sections
+env_info_fmt += "\n"
+
+env_info_fmt += """
+FastVideo Version: {fastvideo_version}
+FastVideo Build Flags:
+{fastvideo_build_flags}
+GPU Topology:
+{gpu_topo}
+
+{env_vars}
+""".strip()
+
+
+def pretty_str(envinfo):
+
+    def replace_nones(dct, replacement='Could not collect'):
+        for key in dct.keys():
+            if dct[key] is not None:
+                continue
+            dct[key] = replacement
+        return dct
+
+    def replace_bools(dct, true='Yes', false='No'):
+        for key in dct.keys():
+            if dct[key] is True:
+                dct[key] = true
+            elif dct[key] is False:
+                dct[key] = false
+        return dct
+
+    def prepend(text, tag='[prepend]'):
+        lines = text.split('\n')
+        updated_lines = [tag + line for line in lines]
+        return '\n'.join(updated_lines)
+
+    def replace_if_empty(text, replacement='No relevant packages'):
+        if text is not None and len(text) == 0:
+            return replacement
+        return text
+
+    def maybe_start_on_next_line(string):
+        # If `string` is multiline, prepend a \n to it.
+        if string is not None and len(string.split('\n')) > 1:
+            return '\n{}\n'.format(string)
+        return string
+
+    mutable_dict = envinfo._asdict()
+
+    # If nvidia_gpu_models is multiline, start on the next line
+    mutable_dict['nvidia_gpu_models'] = \
+        maybe_start_on_next_line(envinfo.nvidia_gpu_models)
+
+    # If the machine doesn't have CUDA, report some fields as 'No CUDA'
+    dynamic_cuda_fields = [
+        'cuda_runtime_version',
+        'nvidia_gpu_models',
+        'nvidia_driver_version',
+    ]
+    all_cuda_fields = dynamic_cuda_fields + ['cudnn_version']
+    all_dynamic_cuda_fields_missing = all(mutable_dict[field] is None
+                                          for field in dynamic_cuda_fields)
+    if TORCH_AVAILABLE and not torch.cuda.is_available(
+    ) and all_dynamic_cuda_fields_missing:
+        for field in all_cuda_fields:
+            mutable_dict[field] = 'No CUDA'
+        if envinfo.cuda_compiled_version is None:
+            mutable_dict['cuda_compiled_version'] = 'None'
+
+    # Replace True with Yes, False with No
+    mutable_dict = replace_bools(mutable_dict)
+
+    # Replace all None objects with 'Could not collect'
+    mutable_dict = replace_nones(mutable_dict)
+
+    # If either of these are '', replace with 'No relevant packages'
+    mutable_dict['pip_packages'] = replace_if_empty(
+        mutable_dict['pip_packages'])
+    mutable_dict['conda_packages'] = replace_if_empty(
+        mutable_dict['conda_packages'])
+
+    # Tag conda and pip packages with a prefix
+    # If they were previously None, they'll show up as ie '[conda] Could not collect'
+    if mutable_dict['pip_packages']:
+        mutable_dict['pip_packages'] = prepend(
+            mutable_dict['pip_packages'], '[{}] '.format(envinfo.pip_version))
+    if mutable_dict['conda_packages']:
+        mutable_dict['conda_packages'] = prepend(mutable_dict['conda_packages'],
+                                                 '[conda] ')
+    mutable_dict['cpu_info'] = envinfo.cpu_info
+    return env_info_fmt.format(**mutable_dict)
+
+
+def get_pretty_env_info():
+    return pretty_str(get_env_info())
+
+
+def main():
+    print("Collecting environment information...")
+    output = get_pretty_env_info()
+    print(output)
+
+    if TORCH_AVAILABLE and hasattr(torch, 'utils') and hasattr(
+            torch.utils, '_crash_handler'):
+        minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR
+        if sys.platform == "linux" and os.path.exists(minidump_dir):
+            dumps = [
+                os.path.join(minidump_dir, dump)
+                for dump in os.listdir(minidump_dir)
+            ]
+            latest = max(dumps, key=os.path.getctime)
+            ctime = os.path.getctime(latest)
+            creation_time = datetime.datetime.fromtimestamp(ctime).strftime(
+                '%Y-%m-%d %H:%M:%S')
+            msg = "\n*** Detected a minidump at {} created on {}, ".format(latest, creation_time) + \
+                  "if this is related to your bug please include it when you file a report ***"
+            print(msg, file=sys.stderr)
+
+
+if __name__ == '__main__':
+    main()

distill_v1.sh

@@ -0,0 +1,46 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+DATA_DIR=./data
+# IP=[MASTER NODE IP]
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+    # --gradient_checkpointing\
+    # --pretrained_model_name_or_path hunyuanvideo-community/HunyuanVideo \
+    # --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+torchrun --nnodes 1 --nproc_per_node 4\
+    fastvideo/v1/pipelines/training_pipeline.py\
+    --inference_mode False\
+    --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --cache_dir "/home/ray/.cache"\
+    --data_json_path "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/videos2caption.json"\
+    --validation_prompt_dir "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/validation"\
+    --train_batch_size=1\
+    --num_latent_t 1 \
+    --sp_size 4 \
+    --tp_size 4 \
+    --train_sp_batch_size 1\
+    --dataloader_num_workers 4\
+    --gradient_accumulation_steps=1\
+    --max_train_steps=320\
+    --learning_rate=1e-6\
+    --mixed_precision="bf16"\
+    --checkpointing_steps=64\
+    --validation_steps 64\
+    --validation_sampling_steps "2,4,8" \
+    --checkpoints_total_limit 3\
+    --allow_tf32\
+    --ema_start_step 0\
+    --cfg 0.0\
+    --log_validation\
+    --output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_16_HD"\
+    --tracker_project_name Hunyuan_Distill \
+    --num_height 720 \
+    --num_width 1280 \
+    --num_frames  125 \
+    --shift 17 \
+    --validation_guidance_scale "1.0" \
+    --num_euler_timesteps 50 \
+    --multi_phased_distill_schedule "4000-1" \
+    --not_apply_cfg_solver \
+    --master_weight_type "bf16"

docs/source/getting_started/installation.md

@@ -11,11 +11,52 @@ FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
 - **CUDA 12.4**
 - **At least 1 NVIDIA GPU**
 
-## Quick start
-We recommend using an environment manager such as [Conda](#miniconda-recommended)
+## Set up using Python
+### Create a new Python environment
+
+#### Conda
+You can create a new python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html)
+##### 1. Install Miniconda (if not already installed)
+
+```bash
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+bash Miniconda3-latest-Linux-x86_64.sh
+source ~/.bashrc
+```
+
+##### 2. Create and activate a Conda environment for FastVideo
+
+```bash
+# (Recommended) Create a new conda environment.
+conda create -n fastvideo python=3.12 -y
+conda activate fastvideo
+```
+
+:::{note}
+[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
+:::
+
+#### uv
+
+:::{tip}
+We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
+:::
+
+Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:
+
+```console
+# (Recommended) Create a new uv environment. Use `--seed` to install `pip` and `setuptools` in the environment.
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+### Installation
 
 ```bash
 pip install fastvideo
+
+# or if you are using uv
+uv pip install fastvideo
 ```
 
 Also optionally install flash-attn:
@@ -24,30 +65,9 @@ Also optionally install flash-attn:
 pip install flash-attn==2.7.4.post1 --no-build-isolation
 ```
 
-[Click here to install from source](#installation-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
-bash Miniconda3-latest-Linux-x86_64.sh
-source ~/.bashrc
-```
-
-#### 2. Create and activate a Conda environment for FastVideo
-
-```bash
-conda create -n fastvideo python=3.10 -y
-conda activate fastvideo
-```
-
-## Installation from Source
-
-### 1. Clone the FastVideo repository
+#### 1. Clone the FastVideo repository
 
 ```bash
 git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
@@ -59,23 +79,23 @@ Basic installation:
 
 ```bash
 pip install -e .
+
+# or if you are using uv
+uv pip install -e .
 ```
 
-## Optional Dependencies
+### Optional Dependencies
 
-### Flash Attention
+#### Flash Attention
 
 ```bash
 pip install flash-attn==2.7.4.post1 --no-build-isolation
 ```
 
-### 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
+## Set up using Docker
 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:
@@ -84,7 +104,7 @@ If you're planning to contribute to FastVideo please see the following page:
 ## Hardware Requirements
 
 ### For Basic Inference
-- NVIDIA GPU with CUDA 12.4support
+- NVIDIA GPU with CUDA 12.4 support
 
 ### For Lora Finetuning
 - 40GB GPU memory each for 2 GPUs with lora

docs/source/getting_started/quick_start.md

@@ -1,143 +0,0 @@
-# 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

@@ -10,9 +10,74 @@ This class will be the primary Python API for generating videos and images.
     fastvideo.VideoGenerator
 ```
 
-VideoGenerator.from_pretrained() should be the primary way of creating a new video generator.
+`````{py:class} VideoGenerator(fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.v1.worker.executor.Executor], log_stats: bool)
+:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator
 
-````{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
+```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+`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) -> fastvideo.v1.entrypoints.video_generator.VideoGenerator
+:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
 :classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+## Configuring FastVideo
+
+The follow two classes `PipelineConfig` and `SamplingParam` are used to configure initialization and sampling parameters, respectively.
+
+### PipelineConfig
+```{autodoc2-summary}
+    fastvideo.PipelineConfig
+```
+
+`````{py:class} PipelineConfig
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.pipelines.base.PipelineConfig
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+:classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+````{py:method} dump_to_json(file_path: str)
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+### SamplingParam
+
+```{autodoc2-summary}
+    fastvideo.SamplingParam
+```
+
+`````{py:class} SamplingParam
+:canonical: fastvideo.v1.configs.sample.base.SamplingParam
+
+```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.sample.base.SamplingParam
+:canonical: fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+:classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
 ```

docs/source/index.md

@@ -9,7 +9,7 @@
 
 :::{raw} html
 <p style="text-align:center">
-<strong>FastVideo is a lightweight framework for accelerating large video diffusion models.
+<strong>FastVideo is a unified framework for accelerated video generation.
 </strong>
 </p>
 
@@ -21,27 +21,31 @@
 </p>
 :::
 
-FastVideo is a lightweight framework for accelerating large video diffusion models developed by the [Hao AI Lab](https://hao-ai-lab.github.io/).
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <div style="text-align: center;">
-  <video controls width="800">
-    <source src="https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1" type="video/mp4">
-    Your browser does not support the video tag.
-  </video>
+  <img src=_static/images/perf.png width="100%"/>
 </div>
 
-FastVideo currently offers: (with more to come)
+## Key Features
 
-- [NEW!] V1 inference API available. Full announcement coming soon!
-- [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
-- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
 
-Dev in progress and highly experimental.
-
 ## Documentation
 
 % How to start using FastVideo?
@@ -51,17 +55,20 @@ Dev in progress and highly experimental.
 :maxdepth: 1
 
 getting_started/installation
-getting_started/quick_start
-getting_started/v1_api
+<!-- getting_started/v1_api -->
 :::
 
 :::{toctree}
 :caption: Inference
 :maxdepth: 1
 
-inference/examples/basic
-inference/examples/gradio
-inference/examples/optimizations
+inference/inference_quick_start
+inference/configuration
+inference/optimizations
+inference/support_matrix
+inference/examples/examples_inference_index
+inference/cli
+inference/add_pipeline
 inference/v0_inference
 :::
 
@@ -96,14 +103,13 @@ design/overview
 
 contributing/overview
 contributing/developer_env/index
-contributing/add_pipeline
 :::
 
 :::{toctree}
 :caption: API Reference
 :maxdepth: 2
 
-api/summary
+<!-- api/summary -->
 api/fastvideo/fastvideo
 :::
 

docs/source/inference/add_pipeline.md

@@ -1,6 +1,6 @@
 (add-pipeline)=
 
-# 🏗️ Adding a New Diffusion Pipeline
+# 🏗️ Adding a New Pipeline
 
 This guide explains how to implement a custom diffusion pipeline in FastVideo, leveraging the framework's modular architecture for high-performance video generation.
 

docs/source/inference/cli.md

@@ -0,0 +1,151 @@
+# FastVideo CLI Inference
+
+The FastVideo CLI provides a quick way to access the FastVideo inference pipeline for video generation. For more advanced usage,
+see the Python interface [here](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html).
+
+## Basic Usage
+
+The basic command to generate a video is:
+
+```bash
+fastvideo generate --model-path {MODEL_PATH} --prompt {PROMPT}
+```
+
+### Required Parameters
+
+- `--model-path {MODEL_PATH}`: Path to the model or model ID
+- `--prompt {PROMPT}`: Text description for the video you want to generate
+
+## Common Arguments
+
+To see all the options, you can use the `--help` flag:
+
+```bash
+fastvideo generate --help
+```
+
+### Hardware Configuration
+
+- `--num-gpus {NUM_GPUS}`: Number of GPUs to use
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (Typically should match the number of GPUs)
+- `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)
+
+#### Video Configuration
+
+- `--height {HEIGHT}`: Height of the generated video
+- `--width {WIDTH}`: Width of the generated video
+- `--num-frames {NUM_FRAMES}`: Number of frames to generate
+- `--fps {FPS}`: Frames per second for the saved video
+
+#### Generation Parameters
+
+- `--num-inference-steps {STEPS}`: Number of denoising steps
+- `--negative-prompt {PROMPT}`: Negative prompt to guide generation away from certain concepts
+- `--seed {SEED}`: Random seed for reproducible generation
+
+#### Output Options
+
+- `--output-path {PATH}`: Directory to save the generated video
+- `--save-video`: Whether to save the video to disk
+- `--return-frames`: Whether to return the raw frames
+
+## Using Configuration Files
+
+Instead of specifying all parameters on the command line, you can use a configuration file:
+
+```bash
+fastvideo generate --config {CONFIG_FILE_PATH}
+```
+
+The config file should be in JSON or YAML format with the same parameter names as the CLI options. Command-line arguments will take precedence over settings in the configuration file, allowing you to override specific values while keeping the rest from the config file.
+
+Example configuration file (config.json):
+
+```json
+{
+    "model_path": "FastVideo/FastHunyuan-diffusers",
+    "prompt": "A beautiful woman in a red dress walking down a street",
+    "output_path": "outputs/",
+    "num_gpus": 2,
+    "sp_size": 2,
+    "tp_size": 2,
+    "num_frames": 45,
+    "height": 720,
+    "width": 1280,
+    "num_inference_steps": 6,
+    "seed": 1024,
+    "fps": 24,
+    "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
+    },
+    "text_encoder_precisions": [
+        "fp16",
+        "fp16"
+    ],
+    "mask_strategy_file_path": null,
+    "enable_torch_compile": false
+}
+```
+
+Or using YAML format (config.yaml):
+
+```yaml
+model_path: "FastVideo/FastHunyuan-diffusers"
+prompt: "A beautiful woman in a red dress walking down a street"
+output_path: "outputs/"
+num_gpus: 2
+sp_size: 2
+tp_size: 2
+num_frames: 45
+height: 720
+width: 1280
+num_inference_steps: 6
+seed: 1024
+fps: 24
+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
+text_encoder_precisions:
+  - "fp16"
+  - "fp16"
+mask_strategy_file_path: null
+enable_torch_compile: false
+```
+
+## Examples
+
+Generating a simple video:
+
+```bash
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A cat playing with a ball of yarn" --num-frames 45 --height 720 --width 1280 --num-inference-steps 6 --seed 1024 --output-path outputs/
+```
+
+Using a negative prompt to avoid certain elements:
+
+```bash
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A beautiful forest landscape" --negative-prompt "people, buildings, roads"
+```
+
+Combining command line arguments and a configuration file:
+
+```bash
+fastvideo generate --config config.json --prompt "A capybara lounging in a hammock"
+```
+
+## Troubleshooting
+
+- If you encounter CUDA out-of-memory errors, try reducing the video dimensions or number of frames, or the number of inference steps.
+- For reproducible results, set the same seed value between runs.

docs/source/inference/configuration.md

@@ -0,0 +1,77 @@
+(inference-configuration)=
+# Configuration
+
+## Multi-GPU Setup
+
+FastVideo automatically distributes the generation process when multiple GPUs are specified:
+
+```python
+# Will use 4 GPUs in parallel for faster generation
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=4,
+)
+```
+
+## Customizing Generation
+
+- `PipelineConfig`: Initialization time parameters
+- `SamplingParam`: Generation time parameters
+
+You can customize various parameters when generating videos using the `PipelineConfig` and `SamplingParam` class:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam, PipelineConfig
+
+def main():
+    model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    config = PipelineConfig.from_pretrained(model_name)
+    config.vae_precision = "fp16"
+    config.use_cpu_offload = True
+
+    # Create the generator
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        num_gpus=1,
+        pipeline_config=config
+    )
+
+    # 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
+    
+    # Video resolution (width, height)
+    sampling_param.width = 1024
+    sampling_param.height = 576
+    
+    # How many steps we denoise the video (higher = better quality, slower generation)
+    sampling_param.num_inference_steps = 30
+    
+    # How strongly the video conforms to the prompt (higher = more faithful to prompt)
+    sampling_param.guidance_scale = 7.5
+    
+    # Random seed for reproducibility
+    sampling_param.seed = 42  # Optional, leave unset for random results
+
+    # 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)
+        save_video=True
+    )
+
+    # If return_frames=True, video contains the generated frames as a NumPy array
+    print(f"Generated {len(video)} frames")
+
+if __name__ == '__main__':
+    main()
+```
+
+## Performance Optimization
+
+For configuring optimizations, please see our [optimizations guide](#inference-optimizations)

docs/source/inference/fasthunyuan.md

@@ -1,33 +0,0 @@
-(fasthunyuan)=
-
-# FastHunyuan
-## Inference FastHunyuan on single RTX4090
-We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan_hf_quantization.sh
-```
-
-For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
-
-| Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
-|--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
-| BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
-| INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
-| NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
-  
-For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
-
-## FastHunyuan
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan.sh
-```
-
-You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).

docs/source/inference/fastmochi.md

@@ -1,9 +0,0 @@
-(fastmochi)=
-
-# FastMochi
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffusers --local_dir=data/FastMochi-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_mochi_sp.sh

docs/source/inference/hunyuanvideo.md

@@ -1,18 +0,0 @@
-(hunyuanvideo)=
-
-# HunyuanVideo
-## Inference HunyuanVideo with Sliding Tile Attention
-First, download the model:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model 
-```
-
-We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
-
-```bash
-sh scripts/inference/inference_hunyuan_STA.sh
-```
-
-## Video Demos using STA + Teacache
-Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.

docs/source/inference/inference_quick_start.md

@@ -0,0 +1,124 @@
+# Inference Quick Start
+
+This page contains step-by-step instructions to get you quickly started with video generation using FastVideo.
+
+## Requirements
+- **OS**: Linux (Tested on Ubuntu 22.04+)
+- **Python**: 3.10-3.12
+- **CUDA**: 12.4
+- **GPU**: At least one NVIDIA GPU
+
+## Installation
+
+We recommend using an environment manager such as `Conda` to create a clean environment:
+
+```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
+```
+
+For advanced installation options, see the [Installation Guide](installation.md).
+
+## Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```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,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
+
+```bash
+python example.py
+```
+
+The generated video will be saved in the current directory under `my_videos/`.
+
+## Available Models
+
+Please see the [support matrix](#support-matrix) for the list of supported models and their available optimizations.
+
+## Image-to-Video Generation
+
+You can generate a video starting from an initial image:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam
+
+def main():
+    # Create the generator
+    model_name = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    generator = VideoGenerator.from_pretrained(model_name, num_gpus=1)
+
+    # Set up parameters with an initial image
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.num_frames = 107
+    sampling_param.image_strength = 0.8  # How much to preserve the original image (0-1)
+
+    # Generate video based on the image
+    prompt = "A photograph coming to life with gentle movement"
+    generator.generate_video(prompt, sampling_param=sampling_param,
+                             output_path="my_videos/",
+                             save_video=True)
+
+if __name__ == '__main__':
+    main()
+```
+
+## Troubleshooting
+
+Common issues and their solutions:
+
+### Out of Memory Errors
+If you encounter CUDA out of memory errors:
+- Reduce `num_frames` or video resolution
+- Enable memory optimization with `enable_model_cpu_offload`
+- Try a smaller model or use distilled versions
+- Use `num_gpus` > 1 if multiple GPUs are available
+
+### Slow Generation
+To speed up generation:
+- Reduce `num_inference_steps` (20-30 is usually sufficient)
+- Use half precision (`fp16`) for the VAE
+- Use multiple GPUs if available
+
+### Unexpected Results
+If the generated video doesn't match your prompt:
+- Try increasing `guidance_scale` (7.0-9.0 works well)
+- Make your prompt more detailed and specific
+- Experiment with different random seeds
+- Try a different model
+
+## Next Steps
+
+- Learn about [Advanced Inference Configurations](#inference-configuration)
+- Learn about using [Optimizations](#inference-optimizations)
+- See [Examples](../examples/examples_inference_index.md) for more usage scenarios
+- Join our [Community Discord](https://discord.gg/JA7cksDz86).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).

docs/source/inference/optimizations.md

@@ -0,0 +1,148 @@
+(inference-optimizations)=
+# Optimizations
+
+This page describes the various options for speeding up generation times in FastVideo.
+
+## 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
+
+### Available Backends
+- Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
+- Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
+- Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
+- Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
+
+### Configuring Backends
+
+There are two ways to configure the attention backend in FastVideo.
+
+#### 1. In Python
+In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this:
+
+```python
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
+```
+
+#### 2. In CLI
+You can also set the environment variable on the command line:
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
+```
+
+(optimizations-flash)=
+### Flash Attention
+
+**`FLASH_ATTN`**
+
+We recommend always installing [Flash Attention 2](https://github.com/Dao-AILab/flash-attention):
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#flashattention-3-beta-release) by compiling it from source (takes about 10 minutes for me):
+
+```bash
+git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention
+
+cd hopper
+pip install ninja 
+python setup.py install
+```
+
+:::{note}
+FastVideo will automatically detect and use `FA3` if it is installed when using `FLASH_ATTN` backend.
+:::
+
+(optimizations-sta)=
+### Sliding Tile Attention
+**`SLIDING_TILE_ATTN`**
+
+```bash
+pip install st_attn==0.0.4
+```
+
+Please see [this page](#sta-installation) for more installation instructions.
+
+(optimizations-sage)=
+### Sage Attention
+**`SAGE_ATTN`**
+
+To use [SageAttention](https://github.com/thu-ml/SageAttention) 2.1.1, please compile from source:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd sageattention 
+python setup.py install  # or pip install -e .
+```
+
+(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.perf_counter()
+generator.generate_video(prompt="Your prompt", enable_teacache=False)
+standard_time = time.perf_counter() - start_time
+
+# With TeaCache
+start_time = time.perf_counter()
+generator.generate_video(prompt="Your prompt", enable_teacache=True)
+teacache_time = time.perf_counter() - 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`.

docs/source/inference/stepvideo.md

@@ -1,16 +0,0 @@
-(stepvideo)=
-
-# StepVideo
-## Inference StepVideo with Sliding Tile Attention
-First, download the model:
-
-```
-python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model 
-```
-
-Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204×768×768 (currently, this is the only supported shape).
-
-```bash
-sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA
-sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
-```

docs/source/inference/support_matrix.md

@@ -0,0 +1,92 @@
+(support-matrix)=
+# Compatibility Matrix
+The table below shows every supported model and optimizations supported for them.
+
+The symbols used have the following meanings:
+
+- ✅ = Full compatibility
+- ❌ = No compatibility
+
+## Models x Optimization
+The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods and FastVideo will use the optimal default parameters when initializing and generating videos.
+
+:::{raw} html
+<style>
+  /* Make smaller to try to improve readability  */
+  td {
+    font-size: 0.9rem;
+    text-align: center;
+  }
+
+  th {
+    text-align: center;
+    font-size: 0.9rem;
+  }
+</style>
+:::
+
+:::{list-table}
+:header-rows: 1
+:stub-columns: 3
+:widths: auto
+:class: vertical-table-header
+
+- * Model Name
+  * HuggingFace Model ID
+  * Resolutions
+  * TeaCache
+  * Sliding Tile Attn
+  * Sage Attn
+- * HunyuanVideo
+  * `hunyuanvideo-community/HunyuanVideo`
+  * 720px1280p<br>544px960p
+  * ❌
+  * ✅
+  * ✅
+- * FastHunyuan
+  * `FastVideo/FastHunyuan-diffusers`
+  * 720px1280p<br>544px960p
+  * ❌
+  * ✅
+  * ✅
+- * Wan T2V 1.3B
+  * `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`
+  * 480P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan T2V 14B
+  * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
+  * 480P, 720P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan I2V 480P
+  * `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
+  * 480P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan I2V 720P
+  * `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers`
+  * 720P
+  * ✅
+  * ✅*
+  * ✅
+- * StepVideo T2V
+  * `FastVideo/stepvideo-t2v-diffusers`
+  * 768px768px204f<br>544px992px204f<br>544px992px136f
+  * ❌
+  * ❌
+  * ✅
+:::
+
+**Note**: there are some known quality issues with Wan2.1 + Sliding Tile Attn. We are working on fixing this issue.
+
+## Special requirements
+
+### StepVideo T2V
+- The self-attention in text-encoder (step_llm) only supports CUDA capabilities sm_80 sm_86 and sm_90
+
+### Sliding Tile Attention
+- Currently only Hopper GPUs (H100s) are supported.

docs/source/inference/wanvideo.md

@@ -1,44 +0,0 @@
-(wanvideo)=
-
-# WanVideo
-## Inference T2V with WanVideo
-First, download the model:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-T2V-1.3B-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model 
-```
-
-or
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-T2V-14B-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model 
-```
-
-Then run the inference using:
-
-```bash
-sh scripts/inference/v1_inference_wan.sh
-```
-
-Remember to set `MODEL_BASE` and `num_gpus` accordingly.
-
-## Inference I2V with WanVideo
-First, download the model:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model
-```
-
-or
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-I2V-14B-720P-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model
-```
-
-Then run the inference using:
-
-```bash
-sh scripts/inference/v1_inference_wan_i2v.sh
-```
-
-Remember to set `MODEL_BASE` and `num_gpus` accordingly.

examples/inference/basic/README.md

@@ -15,7 +15,7 @@ The first script in this example shows the most basic usage of FastVideo. If you
 # 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
+python examples/inference/basic/basic.py
 ```
 
 ## Basic Walkthrough

examples/inference/gradio/README.md

@@ -13,26 +13,12 @@ The demo uses the FastVideo framework to generate videos based on text prompts.
 
 ---
 
-## Requirements
-
-- Linux-based OS
-- Python 3.10-3.12
-- NVIDIA GPU
-- Cuda 12.4
-- FastVideo
-
-## Installation
-
-```bash
-pip install fastvideo
-```
-
 ## Usage
 
 Run the demo with:
 
 ```bash
-python fastvideo/v1/examples/inference/gradio/gradio_demo.py
+python examples/inference/gradio/gradio_demo.py
 ```
 
 This will start a web server at `http://0.0.0.0:7860` where you can access the interface.

examples/inference/optimizations/README.md

@@ -1,104 +1,9 @@
-# Speeding Up Generation
+# Optimization Examples
 
-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

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

examples/inference/optimizations/teacache_example.py

@@ -4,17 +4,17 @@ from fastvideo import VideoGenerator, SamplingParam
 
 
 def main():
-    start_time = time.time()
+    start_time = time.perf_counter()
 
     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
+    load_time = time.perf_counter() - start_time
     print(f"Model loading time: {load_time:.2f} seconds")
 
-    gen_start_time = time.time()
+    gen_start_time = time.perf_counter()
 
     params = SamplingParam.from_pretrained(
         model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
@@ -33,10 +33,10 @@ def main():
         seed=1024,
         output_path="example_outputs/")
     
-    generation_time = time.time() - gen_start_time
+    generation_time = time.perf_counter() - gen_start_time
     print(f"Video generation time: {generation_time:.2f} seconds")
 
-    total_time = time.time() - start_time
+    total_time = time.perf_counter() - start_time
     print(f"Total execution time: {total_time:.2f} seconds")
 
 

fastvideo/data_preprocess/preprocess.py

@@ -0,0 +1,122 @@
+import argparse
+import json
+import os
+
+import torch
+import torch.distributed as dist
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.utils import maybe_download_model, shallow_asdict
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+from fastvideo import PipelineConfig
+from fastvideo.v1.pipelines.preprocess_pipeline import PreprocessPipeline
+
+logger = init_logger(__name__)
+
+BASE_MODEL_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+
+def main(args):
+    # Assume using torchrun
+    local_rank = int(os.getenv("RANK", 0))
+    rank = int(os.environ.get("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    torch.cuda.set_device(local_rank)
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+
+    pipeline_config = PipelineConfig.from_pretrained(MODEL_PATH)
+    kwargs = {
+        "use_cpu_offload": False,
+        "vae_precision": "fp32",
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
+    }
+    pipeline_config_args = shallow_asdict(pipeline_config)
+    pipeline_config_args.update(kwargs)
+    fastvideo_args = FastVideoArgs(model_path=MODEL_PATH,
+                                   num_gpus=world_size,
+                                   device_str="cuda",
+                                   **pipeline_config_args,
+                                   )
+    fastvideo_args.check_fastvideo_args()
+    fastvideo_args.device = torch.device(f"cuda:{local_rank}")
+
+    pipeline = PreprocessPipeline(MODEL_PATH, fastvideo_args)
+    pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--model_type", type=str, default="mochi")
+    parser.add_argument("--data_merge_path", type=str, required=True)
+    parser.add_argument("--validation_prompt_txt", type=str)
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--preprocess_video_batch_size",
+        type=int,
+        default=2,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--preprocess_text_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--samples_per_file",
+        type=int,
+        default=64
+    )
+    parser.add_argument(
+        "--flush_frequency",
+        type=int,
+        default=256,
+        help="how often to save to parquet files"
+    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
+    parser.add_argument("--max_height", type=int, default=480)
+    parser.add_argument("--max_width", type=int, default=848)
+    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+    parser.add_argument("--dataset", default="t2v")
+    parser.add_argument("--train_fps", type=int, default=30)
+    parser.add_argument("--use_image_num", type=int, default=0)
+    parser.add_argument("--text_max_length", type=int, default=256)
+    parser.add_argument("--speed_factor", type=float, default=1.0)
+    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
+    # text encoder & vae & diffusion model
+    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    parser.add_argument("--cfg", type=float, default=0.0)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
+
+    args = parser.parse_args()
+    main(args)

fastvideo/data_preprocess/preprocess_text_embeddings_v1.py

@@ -0,0 +1,199 @@
+import argparse
+import json
+import os
+
+import torch
+import torch.distributed as dist
+from accelerate.logging import get_logger
+from diffusers.utils import export_to_video
+from diffusers.video_processor import VideoProcessor
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+
+# from fastvideo.utils.load import load_text_encoder, load_vae
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+from fastvideo.v1.configs.models.encoders.t5 import T5Config
+
+logger = get_logger(__name__)
+
+class T5dataset(Dataset):
+
+    def __init__(
+        self,
+        json_path,
+        vae_debug,
+    ):
+        self.json_path = json_path
+        self.vae_debug = vae_debug
+        with open(self.json_path, "r") as f:
+            train_dataset = json.load(f)
+            self.train_dataset = sorted(train_dataset, key=lambda x: x["latent_path"])
+
+    def __getitem__(self, idx):
+        caption = self.train_dataset[idx]["caption"]
+        filename = self.train_dataset[idx]["latent_path"].split(".")[0]
+        length = self.train_dataset[idx]["length"]
+        if self.vae_debug:
+            latents = torch.load(
+                os.path.join(args.output_dir, "latent", self.train_dataset[idx]["latent_path"]),
+                map_location="cpu",
+            )
+        else:
+            latents = []
+
+        return dict(caption=caption, latents=latents, filename=filename, length=length)
+
+    def __len__(self):
+        return len(self.train_dataset)
+
+
+def main(args):
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    rank = int(os.getenv("RANK", 0))
+    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    print("world_size", world_size, "local rank", local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    torch.cuda.set_device(device)
+    world_group = get_world_group()
+
+    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # torch.cuda.set_device(local_rank)
+    # if not dist.is_initialized():
+    #     dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+
+    videoprocessor = VideoProcessor(vae_scale_factor=8)
+    os.makedirs(args.output_dir, exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "video"), exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "prompt_embed"), exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True)
+
+    vae_precision = "fp16" 
+    text_encoder_precision = "fp32"
+    fastvideo_args = FastVideoArgs(model_path=args.model_path,
+                         use_cpu_offload=False,
+                         vae_precision=vae_precision,
+                         text_encoder_precisions=(text_encoder_precision,))
+    fastvideo_args.device = device
+    fastvideo_args.device_str = f"cuda:{local_rank}"
+
+    # fastvideo_args.dit_config = HunyuanVideoConfig()
+    fastvideo_args.vae_config = WanVAEConfig()
+    fastvideo_args.text_encoder_configs = (T5Config(),)
+
+    # vae_loader = VAELoader()
+    # vae = vae_loader.load_vae()
+    text_encoder_loader = TextEncoderLoader()
+    tokenizer_loader = TokenizerLoader()
+
+    model_path = args.model_path
+    path = maybe_download_model(model_path)
+    # PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    ENCODER_PATH = os.path.join(path, "text_encoder")
+    TOKENIZER_PATH = os.path.join(path, "tokenizer")
+    print(ENCODER_PATH)
+    text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args)
+    tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args)
+
+    latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
+    train_dataset = T5dataset(latents_json_path, args.vae_debug)
+    # text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
+    # vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
+    # vae.enable_tiling()
+    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+    )
+
+    json_data = []
+    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
+        with torch.inference_mode():
+            # with torch.autocast("cuda", dtype=torch.float32):
+            print(data["caption"])
+            text_inputs = tokenizer(data["caption"], **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to(
+                fastvideo_args.device)
+            input_ids = text_inputs["input_ids"]
+            attention_mask = text_inputs["attention_mask"]
+            outputs = text_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+            from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text
+            post_process_func = t5_postprocess_text
+            prompt_embeds = post_process_func(outputs)
+            prompt_attention_mask = attention_mask
+            if args.vae_debug:
+                latents = data["latents"]
+                video = vae.decode(latents.to(device), return_dict=False)[0]
+                video = videoprocessor.postprocess_video(video)
+            for idx, video_name in enumerate(data["filename"]):
+                prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt")
+                video_path = os.path.join(args.output_dir, "video", video_name + ".mp4")
+                prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask",
+                                                            video_name + ".pt")
+                # save latent
+                torch.save(prompt_embeds[idx], prompt_embed_path)
+                torch.save(prompt_attention_mask[idx], prompt_attention_mask_path)
+                print(f"sample {video_name} saved")
+                if args.vae_debug:
+                    export_to_video(video[idx], video_path, fps=16)
+                item = {}
+                item["length"] = int(data["length"][idx])
+                item["latent_path"] = video_name + ".pt"
+                item["prompt_embed_path"] = video_name + ".pt"
+                item["prompt_attention_mask"] = video_name + ".pt"
+                item["caption"] = data["caption"][idx]
+                json_data.append(item)
+    dist.barrier()
+    local_data = json_data
+    gathered_data = [None] * world_size
+    dist.all_gather_object(gathered_data, local_data)
+    if local_rank == 0:
+        # os.remove(latents_json_path)
+        all_json_data = [item for sublist in gathered_data for item in sublist]
+        with open(os.path.join(args.output_dir, "videos2caption.json"), "w") as f:
+            json.dump(all_json_data, f, indent=4)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    # parser.add_argument("--model_type", type=str, default="mochi")
+    # text encoder & vae & diffusion model
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        type=int,
+        default=1,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--vae_debug", action="store_true")
+    args = parser.parse_args()
+    main(args)

fastvideo/data_preprocess/preprocess_vae_latents_v1.py

@@ -0,0 +1,151 @@
+import argparse
+import json
+import os
+
+import torch
+# import torch.distributed as dist
+# from accelerate.logging import get_logger
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+
+from fastvideo.dataset import getdataset
+# from fastvideo.utils.load import load_vae
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import VAELoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+
+logger = init_logger(__name__)
+
+model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+path = maybe_download_model(model_path)
+# PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+VAE_PATH = os.path.join(path, "vae")
+print(VAE_PATH)
+
+
+def main(args):
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    rank = int(os.getenv("RANK", 0))
+    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    print("world_size", world_size, "local rank", local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    torch.cuda.set_device(device)
+    world_group = get_world_group()
+
+    vae_precision = "fp16" 
+    fastvideo_args = FastVideoArgs(model_path=VAE_PATH,
+                         use_cpu_offload=False,
+                         vae_precision=vae_precision)
+    fastvideo_args.device = device
+    # fastvideo_args.dit_config = HunyuanVideoConfig()
+    fastvideo_args.vae_config = WanVAEConfig()
+
+    train_dataset = getdataset(args)
+    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+    )
+
+
+    # encoder_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # torch.cuda.set_device(local_rank)
+    # if not dist.is_initialized():
+    #     dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+    vae_loader = VAELoader()
+    vae = vae_loader.load(VAE_PATH, "vae", fastvideo_args)
+    # vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
+    # vae.enable_tiling()
+    os.makedirs(args.output_dir, exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
+
+    json_data = []
+    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
+        with torch.inference_mode():
+            with torch.autocast("cuda", dtype=torch.float16):
+                latents = vae.encode(data["pixel_values"].to(device)).sample()
+            for idx, video_path in enumerate(data["path"]):
+                video_name = os.path.basename(video_path).split(".")[0]
+                latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt")
+                torch.save(latents[idx].to(torch.bfloat16), latent_path)
+                item = {}
+                item["length"] = latents[idx].shape[1]
+                item["latent_path"] = video_name + ".pt"
+                item["caption"] = data["text"][idx]
+                json_data.append(item)
+                print(f"{video_name} processed")
+    world_group.barrier()
+    local_data = json_data
+    gathered_data = [None] * world_size
+    for i in range(world_size):
+        if local_rank == i:
+            world_group.broadcast_object(local_data, src=i)
+        else:
+            gathered_data[i] = world_group.broadcast_object(None, src=i)
+    gathered_data[local_rank] = json_data
+    print(gathered_data)
+    if local_rank == 0:
+        all_json_data = [item for sublist in gathered_data for item in sublist]
+        with open(os.path.join(args.output_dir, "videos2caption_temp.json"), "w") as f:
+            json.dump(all_json_data, f, indent=4)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    # parser.add_argument("--model_type", type=str, default="mochi")
+    parser.add_argument("--data_merge_path", type=str, required=True)
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
+    parser.add_argument("--max_height", type=int, default=480)
+    parser.add_argument("--max_width", type=int, default=848)
+    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+    parser.add_argument("--dataset", default="t2v")
+    parser.add_argument("--train_fps", type=int, default=30)
+    parser.add_argument("--use_image_num", type=int, default=0)
+    parser.add_argument("--text_max_length", type=int, default=256)
+    parser.add_argument("--speed_factor", type=float, default=1.0)
+    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
+    # text encoder & vae & diffusion model
+    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    parser.add_argument("--cfg", type=float, default=0.0)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
+
+    args = parser.parse_args()
+    main(args)

fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py

@@ -0,0 +1,115 @@
+import argparse
+import os
+
+import torch
+# import torch.distributed as dist
+from accelerate.logging import get_logger
+
+# from fastvideo.utils.load import load_text_encoder
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+from fastvideo.v1.configs.models.encoders.t5 import T5Config
+
+logger = get_logger(__name__)
+
+
+def main(args):
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    rank = int(os.getenv("RANK", 0))
+    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    print("world_size", world_size, "local rank", local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    torch.cuda.set_device(device)
+    world_group = get_world_group()
+
+    vae_precision = "fp16" 
+    text_encoder_precision = "fp32"
+    fastvideo_args = FastVideoArgs(model_path=args.model_path,
+                         use_cpu_offload=False,
+                         vae_precision=vae_precision,
+                         text_encoder_precisions=(text_encoder_precision,))
+    fastvideo_args.device = device
+    fastvideo_args.device_str = f"cuda:{local_rank}"
+
+    # fastvideo_args.dit_config = HunyuanVideoConfig()
+    fastvideo_args.vae_config = WanVAEConfig()
+    fastvideo_args.text_encoder_configs = (T5Config(),)
+
+    # vae_loader = VAELoader()
+    # vae = vae_loader.load_vae()
+    text_encoder_loader = TextEncoderLoader()
+    tokenizer_loader = TokenizerLoader()
+
+    model_path = args.model_path
+    path = maybe_download_model(model_path)
+    # PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    ENCODER_PATH = os.path.join(path, "text_encoder")
+    TOKENIZER_PATH = os.path.join(path, "tokenizer")
+    print(ENCODER_PATH)
+    text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args)
+    tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args)
+
+    # text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
+    # autocast_type = torch.float16 if args.model_type == "hunyuan" else torch.bfloat16
+    # output_dir/validation/prompt_attention_mask
+    # output_dir/validation/prompt_embed
+    os.makedirs(os.path.join(args.output_dir, "validation"), exist_ok=True)
+    os.makedirs(
+        os.path.join(args.output_dir, "validation", "prompt_attention_mask"),
+        exist_ok=True,
+    )
+    os.makedirs(os.path.join(args.output_dir, "validation", "prompt_embed"), exist_ok=True)
+
+    with open(args.validation_prompt_txt, "r", encoding="utf-8") as file:
+        lines = file.readlines()
+    prompts = [line.strip() for line in lines]
+    for prompt in prompts:
+        with torch.inference_mode():
+            # with torch.autocast("cuda", dtype=autocast_type):
+            text_inputs = tokenizer(prompt, **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to(
+                fastvideo_args.device)
+            input_ids = text_inputs["input_ids"]
+            attention_mask = text_inputs["attention_mask"]
+            outputs = text_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+            from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text
+            post_process_func = t5_postprocess_text
+            prompt_embeds = post_process_func(outputs)
+            prompt_attention_mask = attention_mask
+
+            file_name = prompt.split(".")[0]
+            prompt_embed_path = os.path.join(args.output_dir, "validation", "prompt_embed", f"{file_name}.pt")
+            prompt_attention_mask_path = os.path.join(
+                args.output_dir,
+                "validation",
+                "prompt_attention_mask",
+                f"{file_name}.pt",
+            )
+            torch.save(prompt_embeds[0], prompt_embed_path)
+            torch.save(prompt_attention_mask[0], prompt_attention_mask_path)
+            print(f"sample {file_name} saved")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--validation_prompt_txt", type=str)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    args = parser.parse_args()
+    main(args)

fastvideo/distill.py

@@ -452,7 +452,7 @@ def main(args):
         return phase
 
     for step in range(init_steps + 1, args.max_train_steps + 1):
-        start_time = time.time()
+        start_time = time.perf_counter()
         assert args.multi_phased_distill_schedule is not None
         num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
 
@@ -482,7 +482,7 @@ def main(args):
             args.hunyuan_teacher_disable_cfg,
         )
 
-        step_time = time.time() - start_time
+        step_time = time.perf_counter() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 

fastvideo/distill/solver.py

@@ -7,7 +7,7 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from diffusers.utils import BaseOutput, logging
 
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
+# from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -38,6 +38,7 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
         linear_range=0.5,
     ):
         if linear_quadratic:
+            raise NotImplementedError("Linear quadratic schedule is not implemented")
             linear_steps = int(num_train_timesteps * linear_range)
             sigmas = linear_quadratic_schedule(num_train_timesteps, linear_quadratic_threshold, linear_steps)
             sigmas = torch.tensor(sigmas).to(dtype=torch.float32)

fastvideo/distill_adv.py

@@ -517,7 +517,7 @@ def main(args):
     for step in range(init_steps + 1, args.max_train_steps + 1):
         assert args.multi_phased_distill_schedule is not None
         num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
-        start_time = time.time()
+        start_time = time.perf_counter()
         (
             generator_loss,
             generator_grad_norm,
@@ -547,7 +547,7 @@ def main(args):
             args.discriminator_head_stride,
         )
 
-        step_time = time.time() - start_time
+        step_time = time.perf_counter() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 

fastvideo/distill_wan.py

@@ -0,0 +1,870 @@
+# !/bin/python3
+# isort: skip_file
+import argparse
+import math
+import os
+import time
+from collections import deque
+
+import torch
+import torch.distributed as dist
+import wandb
+from accelerate.utils import set_seed
+from diffusers.optimization import get_scheduler
+from diffusers.utils import check_min_version
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import ShardingStrategy
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from tqdm.auto import tqdm
+
+from fastvideo.dataset.latent_datasets import (LatentDataset,
+                                               latent_collate_function)
+from fastvideo.distill.solver import EulerSolver, extract_into_tensor
+from fastvideo.utils.checkpoint import (save_checkpoint, save_lora_checkpoint)
+from fastvideo.utils.communications import (broadcast,
+                                            sp_parallel_dataloader_wrapper)
+from fastvideo.utils.dataset_utils import LengthGroupedSampler
+from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group,
+                                             get_sequence_parallel_state)
+from fastvideo.utils.validation import log_validation
+
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.configs.models.dits import WanVideoConfig
+from fastvideo.v1.models.loader.component_loader import TransformerLoader, SchedulerLoader
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.31.0")
+
+logger = init_logger(__name__)
+
+BASE_MODEL_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+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")
+SCHEDULER_PATH = os.path.join(MODEL_PATH, "scheduler")
+
+
+def reshard_fsdp(model):
+    for m in FSDP.fsdp_modules(model):
+        if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
+            torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
+
+
+def get_norm(model_pred, norms, gradient_accumulation_steps):
+    fro_norm = (
+        torch.linalg.matrix_norm(model_pred, ord="fro") /  # codespell:ignore
+        gradient_accumulation_steps)
+    largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) /
+                              gradient_accumulation_steps)
+    absolute_mean = torch.mean(
+        torch.abs(model_pred)) / gradient_accumulation_steps
+    absolute_max = torch.max(
+        torch.abs(model_pred)) / gradient_accumulation_steps
+    dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG)
+    dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG)
+    dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG)
+    norms["fro"] += torch.mean(fro_norm).item()  # codespell:ignore
+    norms["largest singular value"] += torch.mean(largest_singular_value).item()
+    norms["absolute mean"] += absolute_mean.item()
+    norms["absolute max"] += absolute_max.item()
+
+
+def distill_one_step(
+    transformer,
+    model_type,
+    teacher_transformer,
+    ema_transformer,
+    optimizer,
+    lr_scheduler,
+    loader,
+    noise_scheduler,
+    solver,
+    noise_random_generator,
+    gradient_accumulation_steps,
+    sp_size,
+    max_grad_norm,
+    uncond_prompt_embed,
+    uncond_prompt_mask,
+    num_euler_timesteps,
+    multiphase,
+    not_apply_cfg_solver,
+    distill_cfg,
+    ema_decay,
+    pred_decay_weight,
+    pred_decay_type,
+    hunyuan_teacher_disable_cfg,
+):
+    total_loss = 0.0
+    optimizer.zero_grad()
+    model_pred_norm = {
+        "fro": 0.0,  # codespell:ignore
+        "largest singular value": 0.0,
+        "absolute mean": 0.0,
+        "absolute max": 0.0,
+    }
+    for _ in range(gradient_accumulation_steps):
+        (
+            latents,
+            encoder_hidden_states,
+            latents_attention_mask,
+            encoder_attention_mask,
+        ) = next(loader)
+        # model_input = normalize_dit_input(model_type, latents)
+        model_input = latents
+        noise = torch.randn_like(model_input)
+        bsz = model_input.shape[0]
+        index = torch.randint(0,
+                              num_euler_timesteps, (bsz, ),
+                              device=model_input.device).long()
+        if sp_size > 1:
+            broadcast(index)
+        # Add noise according to flow matching.
+        # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
+        sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
+        sigmas_prev = extract_into_tensor(solver.sigmas_prev, index,
+                                          model_input.shape)
+
+        timesteps = (sigmas *
+                     noise_scheduler.config.num_train_timesteps).view(-1)
+        # if squeeze to [], unsqueeze to [1]
+
+        timesteps_prev = (sigmas_prev *
+                          noise_scheduler.config.num_train_timesteps).view(-1)
+        noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
+        noisy_model_input = noisy_model_input.to(torch.bfloat16)
+
+        forward_batch = ForwardBatch(data_type="video", enable_teacache=False)
+        # Predict the noise residual
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            teacher_kwargs = {
+                "hidden_states": noisy_model_input,
+                "encoder_hidden_states": encoder_hidden_states,
+                "timestep": timesteps,
+                "encoder_attention_mask": encoder_attention_mask,  # B, L
+                "return_dict": False,
+            }
+            if hunyuan_teacher_disable_cfg:
+                teacher_kwargs["guidance"] = torch.tensor(
+                    [1000.0],
+                    device=noisy_model_input.device,
+                    dtype=torch.bfloat16)
+            with set_forward_context(current_timestep=0,
+                                     attn_metadata=None,
+                                     forward_batch=forward_batch):
+                with torch.autograd.graph.save_on_cpu(pin_memory=True):
+                    model_pred = transformer(**teacher_kwargs)
+
+        # if accelerator.is_main_process:
+        model_pred, end_index = solver.euler_style_multiphase_pred(
+            noisy_model_input, model_pred, index, multiphase)
+        with torch.no_grad():
+            w = distill_cfg
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                with set_forward_context(current_timestep=0,
+                                         attn_metadata=None,
+                                         forward_batch=forward_batch):
+                    cond_teacher_output = teacher_transformer(
+                        noisy_model_input,
+                        encoder_hidden_states,
+                        timesteps,
+                        encoder_attention_mask,  # B, L
+                        return_dict=False,
+                    ).float()
+            if not_apply_cfg_solver:
+                uncond_teacher_output = cond_teacher_output
+            else:
+                # Get teacher model prediction on noisy_latents and unconditional embedding
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    with set_forward_context(current_timestep=0,
+                                             attn_metadata=None,
+                                             forward_batch=forward_batch):
+                        uncond_teacher_output = teacher_transformer(
+                            noisy_model_input,
+                            uncond_prompt_embed.unsqueeze(0).expand(
+                                bsz, -1, -1),
+                            timesteps,
+                            uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
+                            return_dict=False,
+                        ).float()
+            teacher_output = uncond_teacher_output + w * (cond_teacher_output -
+                                                          uncond_teacher_output)
+            x_prev = solver.euler_step(noisy_model_input, teacher_output,
+                                       index).to(torch.bfloat16)
+
+        # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
+        with torch.no_grad():
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                if ema_transformer is not None:
+                    target_pred = ema_transformer(
+                        x_prev.float(),
+                        encoder_hidden_states,
+                        timesteps_prev,
+                        encoder_attention_mask,  # B, L
+                        return_dict=False,
+                    )[0]
+                else:
+                    with set_forward_context(current_timestep=0,
+                                             attn_metadata=None,
+                                             forward_batch=forward_batch):
+                        with torch.autograd.graph.save_on_cpu(pin_memory=True):
+                            target_pred = transformer(
+                                x_prev,
+                                encoder_hidden_states,
+                                timesteps_prev,
+                                encoder_attention_mask,  # B, L
+                                return_dict=False,
+                            )
+
+            target, end_index = solver.euler_style_multiphase_pred(
+                x_prev, target_pred, index, multiphase, True)
+
+        huber_c = 0.001
+        # loss = loss.mean()
+        loss = (torch.mean(
+            torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) -
+            huber_c) / gradient_accumulation_steps)
+        if pred_decay_weight > 0:
+            if pred_decay_type == "l1":
+                pred_decay_loss = (
+                    torch.mean(torch.sqrt(model_pred.float()**2)) *
+                    pred_decay_weight / gradient_accumulation_steps)
+                loss += pred_decay_loss
+            elif pred_decay_type == "l2":
+                # essnetially k2?
+                pred_decay_loss = (torch.mean(model_pred.float()**2) *
+                                   pred_decay_weight /
+                                   gradient_accumulation_steps)
+                loss += pred_decay_loss
+            else:
+                assert NotImplementedError("pred_decay_type is not implemented")
+
+        # calculate model_pred norm and mean
+        get_norm(model_pred.detach().float(), model_pred_norm,
+                 gradient_accumulation_steps)
+        loss.backward()
+
+        avg_loss = loss.detach().clone()
+        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
+        total_loss += avg_loss.item()
+
+    # update ema
+    if ema_transformer is not None:
+        reshard_fsdp(ema_transformer)
+        for p_averaged, p_model in zip(ema_transformer.parameters(),
+                                       transformer.parameters()):
+            with torch.no_grad():
+                p_averaged.copy_(
+                    torch.lerp(p_averaged.detach(), p_model.detach(),
+                               1 - ema_decay))
+
+    # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+    grad_norm = torch.nn.utils.clip_grad_norm_(transformer.parameters(),
+                                               max_norm=max_grad_norm)
+    optimizer.step()
+    lr_scheduler.step()
+
+    return total_loss, grad_norm.item(), model_pred_norm
+
+
+def main(args):
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    rank = int(os.environ.get("RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    torch.cuda.set_device(rank)
+    init_distributed_environment(world_size=world_size,
+                                 rank=rank,
+                                 local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=args.sp_size,
+                              sequence_model_parallel_size=args.sp_size)
+    fastvideo_args = FastVideoArgs(
+        model_path=MODEL_PATH,
+        num_gpus=world_size,
+        use_cpu_offload=False,
+        precision=args.master_weight_type,
+        dit_config=WanVideoConfig(),
+        device_str="cuda",
+    )
+    fastvideo_args.check_fastvideo_args()
+
+    device_str = f"cuda:{rank}"
+    device = torch.device(device_str)
+    fastvideo_args.device = device
+
+    # If passed along, set the training seed now. On GPU...
+    if args.seed is not None:
+        # TODO: t within the same seq parallel group should be the same. Noise should be different.
+        set_seed(args.seed + rank)
+    # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch.
+    noise_random_generator = None
+
+    # Handle the repository creation
+    if rank <= 0 and args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    # For mixed precision training we cast all non-trainable weights to half-precision
+    # as these weights are only used for inference, keeping weights in full precision is not required.
+
+    # Create model:
+
+    logger.info("--> loading model from %s", TRANSFORMER_PATH)
+
+    fastvideo_args.device = device
+    transformer_loader = TransformerLoader()
+    transformer = transformer_loader.load(TRANSFORMER_PATH, "", fastvideo_args)
+    transformer = transformer.train()
+    transformer.requires_grad_(True)
+
+    teacher_loader = TransformerLoader()
+    teacher_transformer = teacher_loader.load(TRANSFORMER_PATH, "",
+                                              fastvideo_args)
+    if args.use_ema:
+        ema_transformer = teacher_loader.load(TRANSFORMER_PATH, "",
+                                              fastvideo_args)
+    else:
+        ema_transformer = None
+
+    logger.info(
+        "  Total training parameters = %s M",
+        sum(p.numel()
+            for p in transformer.parameters() if p.requires_grad) / 1e6)
+    logger.info("--> model loaded")
+
+    teacher_transformer.requires_grad_(False)
+    if args.use_ema:
+        ema_transformer.requires_grad_(False)
+
+    # scheduler
+    noise_scheduler_loader = SchedulerLoader()
+    noise_scheduler = noise_scheduler_loader.load(SCHEDULER_PATH, "",
+                                                  fastvideo_args)
+    solver = EulerSolver(
+        noise_scheduler.sigmas.numpy()[::-1],
+        noise_scheduler.config.num_train_timesteps,
+        euler_timesteps=args.num_euler_timesteps,
+    )
+    solver.to(device)
+    params_to_optimize = transformer.parameters()
+    params_to_optimize = list(
+        filter(lambda p: p.requires_grad, params_to_optimize))
+
+    optimizer = torch.optim.AdamW(
+        params_to_optimize,
+        lr=args.learning_rate,
+        betas=(0.9, 0.999),
+        weight_decay=args.weight_decay,
+        eps=1e-8,
+    )
+
+    init_steps = 0
+    logger.info("optimizer: %s", optimizer)
+
+    # todo add lr scheduler
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * world_size,
+        num_training_steps=args.max_train_steps * world_size,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+        last_epoch=init_steps - 1,
+    )
+
+    train_dataset = LatentDataset(args.data_json_path, args.num_latent_t,
+                                  args.cfg)
+    uncond_prompt_embed = train_dataset.uncond_prompt_embed
+    uncond_prompt_mask = train_dataset.uncond_prompt_mask
+    sampler = (LengthGroupedSampler(
+        args.train_batch_size,
+        rank=rank,
+        world_size=world_size,
+        lengths=train_dataset.lengths,
+        group_frame=args.group_frame,
+        group_resolution=args.group_resolution,
+    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
+        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        collate_fn=latent_collate_function,
+        pin_memory=True,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+        drop_last=True,
+    )
+
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader) / args.gradient_accumulation_steps *
+        args.sp_size / args.train_sp_batch_size)
+    args.num_train_epochs = math.ceil(args.max_train_steps /
+                                      num_update_steps_per_epoch)
+
+    if rank <= 0:
+        project = args.tracker_project_name or "fastvideo"
+        wandb.init(project=project, config=args)
+
+    # Train!
+    total_batch_size = (world_size * args.gradient_accumulation_steps /
+                        args.sp_size * args.train_sp_batch_size)
+    logger.info("***** Running training *****")
+    logger.info("  Num examples = %s", len(train_dataset))
+    logger.info("  Dataloader size = %s", len(train_dataloader))
+    logger.info("  Num Epochs = %s", args.num_train_epochs)
+    logger.info("  Resume training from step %s", init_steps)
+    logger.info("  Instantaneous batch size per device = %s",
+                args.train_batch_size)
+    logger.info(
+        "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
+        total_batch_size)
+    logger.info("  Gradient Accumulation steps = %s",
+                args.gradient_accumulation_steps)
+    logger.info("  Total optimization steps = %s", args.max_train_steps)
+    logger.info(
+        "  Total training parameters per FSDP shard = %s B",
+        sum(p.numel()
+            for p in transformer.parameters() if p.requires_grad) / 1e9)
+    # print dtype
+    logger.info("  Master weight dtype: %s",
+                transformer.parameters().__next__().dtype)
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        assert NotImplementedError(
+            "resume_from_checkpoint is not supported now.")
+        # TODO
+
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=init_steps,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=local_rank > 0,
+    )
+
+    loader = sp_parallel_dataloader_wrapper(
+        train_dataloader,
+        device,
+        args.train_batch_size,
+        args.sp_size,
+        args.train_sp_batch_size,
+    )
+
+    step_times = deque(maxlen=100)
+
+    # todo future
+    for i in range(init_steps):
+        next(loader)
+
+    # log_validation(args, transformer, device,
+    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
+    def get_num_phases(multi_phased_distill_schedule, step):
+        # step-phase,step-phase
+        multi_phases = multi_phased_distill_schedule.split(",")
+        phase = multi_phases[-1].split("-")[-1]
+        for step_phases in multi_phases:
+            phase_step, phase = step_phases.split("-")
+            if step <= int(phase_step):
+                return int(phase)
+        return phase
+
+    for step in range(init_steps + 1, args.max_train_steps + 1):
+        start_time = time.time()
+        assert args.multi_phased_distill_schedule is not None
+        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
+
+        loss, grad_norm, pred_norm = distill_one_step(
+            transformer,
+            args.model_type,
+            teacher_transformer,
+            ema_transformer,
+            optimizer,
+            lr_scheduler,
+            loader,
+            noise_scheduler,
+            solver,
+            noise_random_generator,
+            args.gradient_accumulation_steps,
+            args.sp_size,
+            args.max_grad_norm,
+            uncond_prompt_embed,
+            uncond_prompt_mask,
+            args.num_euler_timesteps,
+            num_phases,
+            args.not_apply_cfg_solver,
+            args.distill_cfg,
+            args.ema_decay,
+            args.pred_decay_weight,
+            args.pred_decay_type,
+            args.hunyuan_teacher_disable_cfg,
+        )
+
+        step_time = time.time() - start_time
+        step_times.append(step_time)
+        avg_step_time = sum(step_times) / len(step_times)
+
+        progress_bar.set_postfix({
+            "loss": f"{loss:.4f}",
+            "step_time": f"{step_time:.2f}s",
+            "grad_norm": grad_norm,
+            "phases": num_phases,
+        })
+        progress_bar.update(1)
+        if rank <= 0:
+            wandb.log(
+                {
+                    "train_loss":
+                    loss,
+                    "learning_rate":
+                    lr_scheduler.get_last_lr()[0],
+                    "step_time":
+                    step_time,
+                    "avg_step_time":
+                    avg_step_time,
+                    "grad_norm":
+                    grad_norm,
+                    "pred_fro_norm":
+                    pred_norm["fro"],  # codespell:ignore
+                    "pred_largest_singular_value":
+                    pred_norm["largest singular value"],
+                    "pred_absolute_mean":
+                    pred_norm["absolute mean"],
+                    "pred_absolute_max":
+                    pred_norm["absolute max"],
+                },
+                step=step,
+            )
+        if step % args.checkpointing_steps == 0:
+            if args.use_lora:
+                # Save LoRA weights
+                save_lora_checkpoint(transformer, optimizer, rank,
+                                     args.output_dir, step)
+            else:
+                # Your existing checkpoint saving code
+                if args.use_ema:
+                    save_checkpoint(ema_transformer, rank, args.output_dir,
+                                    step)
+                else:
+                    save_checkpoint(transformer, rank, args.output_dir, step)
+            dist.barrier()
+        if args.log_validation and step % args.validation_steps == 0:
+            log_validation(
+                args,
+                transformer,
+                device,
+                torch.bfloat16,
+                step,
+                scheduler_type=args.scheduler_type,
+                shift=args.shift,
+                num_euler_timesteps=args.num_euler_timesteps,
+                linear_quadratic_threshold=args.linear_quadratic_threshold,
+                linear_range=args.linear_range,
+                ema=False,
+            )
+            if args.use_ema:
+                log_validation(
+                    args,
+                    ema_transformer,
+                    device,
+                    torch.bfloat16,
+                    step,
+                    scheduler_type=args.scheduler_type,
+                    shift=args.shift,
+                    num_euler_timesteps=args.num_euler_timesteps,
+                    linear_quadratic_threshold=args.linear_quadratic_threshold,
+                    linear_range=args.linear_range,
+                    ema=True,
+                )
+
+    if args.use_lora:
+        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir,
+                             args.max_train_steps)
+    else:
+        save_checkpoint(transformer, rank, args.output_dir,
+                        args.max_train_steps)
+
+    if get_sequence_parallel_state():
+        destroy_sequence_parallel_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model_type",
+                        type=str,
+                        default="mochi",
+                        help="The type of model to train.")
+
+    # dataset & dataloader
+    parser.add_argument("--data_json_path", type=str, required=True)
+    parser.add_argument("--num_height", type=int, default=480)
+    parser.add_argument("--num_width", type=int, default=848)
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=10,
+        help=
+        "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--num_latent_t",
+                        type=int,
+                        default=28,
+                        help="Number of latent timesteps.")
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+
+    # text encoder & vae & diffusion model
+    parser.add_argument("--pretrained_model_name_or_path", type=str)
+    parser.add_argument("--dit_model_name_or_path", type=str)
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+
+    # diffusion setting
+    parser.add_argument("--ema_decay", type=float, default=0.95)
+    parser.add_argument("--ema_start_step", type=int, default=0)
+    parser.add_argument("--cfg", type=float, default=0.1)
+
+    # validation & logs
+    parser.add_argument("--validation_prompt_dir", type=str)
+    parser.add_argument("--validation_sampling_steps", type=str, default="64")
+    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
+
+    parser.add_argument("--validation_steps", type=float, default=64)
+    parser.add_argument("--log_validation", action="store_true")
+    parser.add_argument("--tracker_project_name", type=str, default=None)
+    parser.add_argument("--seed",
+                        type=int,
+                        default=None,
+                        help="A seed for reproducible training.")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help=
+        "The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=("Max number of checkpoints to store."),
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=
+        ("Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+         " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+         " training using `--resume_from_checkpoint`."),
+    )
+    parser.add_argument("--shift", type=float, default=1.0)
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=
+        ("Whether training should be resumed from a previous checkpoint. Use a path saved by"
+         ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+         ),
+    )
+    parser.add_argument(
+        "--resume_from_lora_checkpoint",
+        type=str,
+        default=None,
+        help=
+        ("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
+         ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+         ),
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=
+        ("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+         " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
+
+    # optimizer & scheduler & Training
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help=
+        "Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help=
+        "Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-4,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help=
+        "Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_warmup_steps",
+        type=int,
+        default=10,
+        help="Number of steps for the warmup in the lr scheduler.",
+    )
+    parser.add_argument("--max_grad_norm",
+                        default=1.0,
+                        type=float,
+                        help="Max gradient norm.")
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help=
+        "Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument("--selective_checkpointing", type=float, default=1.0)
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=
+        ("Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+         " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+         ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=
+        ("Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+         " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+         " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+         ),
+    )
+    parser.add_argument(
+        "--use_cpu_offload",
+        action="store_true",
+        help=
+        "Whether to use CPU offload for param & gradient & optimizer states.",
+    )
+
+    parser.add_argument("--sp_size",
+                        type=int,
+                        default=1,
+                        help="For sequence parallel")
+    parser.add_argument(
+        "--train_sp_batch_size",
+        type=int,
+        default=1,
+        help="Batch size for sequence parallel training",
+    )
+
+    parser.add_argument(
+        "--use_lora",
+        action="store_true",
+        default=False,
+        help="Whether to use LoRA for finetuning.",
+    )
+    parser.add_argument("--lora_alpha",
+                        type=int,
+                        default=256,
+                        help="Alpha parameter for LoRA.")
+    parser.add_argument("--lora_rank",
+                        type=int,
+                        default=128,
+                        help="LoRA rank parameter. ")
+    parser.add_argument("--fsdp_sharding_startegy", default="full")
+
+    # lr_scheduler
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=
+        ('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+         ' "constant", "constant_with_warmup"]'),
+    )
+    parser.add_argument("--num_euler_timesteps", type=int, default=100)
+    parser.add_argument(
+        "--lr_num_cycles",
+        type=int,
+        default=1,
+        help="Number of cycles in the learning rate scheduler.",
+    )
+    parser.add_argument(
+        "--lr_power",
+        type=float,
+        default=1.0,
+        help="Power factor of the polynomial scheduler.",
+    )
+    parser.add_argument(
+        "--not_apply_cfg_solver",
+        action="store_true",
+        help="Whether to apply the cfg_solver.",
+    )
+    parser.add_argument("--distill_cfg",
+                        type=float,
+                        default=3.0,
+                        help="Distillation coefficient.")
+    # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
+    parser.add_argument("--scheduler_type",
+                        type=str,
+                        default="pcm",
+                        help="The scheduler type to use.")
+    parser.add_argument(
+        "--linear_quadratic_threshold",
+        type=float,
+        default=0.025,
+        help="Threshold for linear quadratic scheduler.",
+    )
+    parser.add_argument(
+        "--linear_range",
+        type=float,
+        default=0.5,
+        help="Range for linear quadratic scheduler.",
+    )
+    parser.add_argument("--weight_decay",
+                        type=float,
+                        default=0.001,
+                        help="Weight decay to apply.")
+    parser.add_argument("--use_ema",
+                        action="store_true",
+                        help="Whether to use EMA.")
+    parser.add_argument("--multi_phased_distill_schedule",
+                        type=str,
+                        default=None)
+    parser.add_argument("--pred_decay_weight", type=float, default=0.0)
+    parser.add_argument("--pred_decay_type", default="l1")
+    parser.add_argument("--hunyuan_teacher_disable_cfg", action="store_true")
+    parser.add_argument(
+        "--master_weight_type",
+        type=str,
+        default="fp32",
+        help="Weight type to use - fp32 or bf16.",
+    )
+    args = parser.parse_args()
+    main(args)

fastvideo/models/hunyuan/inference.py

@@ -452,7 +452,7 @@ class HunyuanVideoSampler(Inference):
         # ========================================================================
         # Pipeline inference
         # ========================================================================
-        start_time = time.time()
+        start_time = time.perf_counter()
         samples = self.pipeline(
             prompt=prompt,
             height=target_height,
@@ -476,7 +476,7 @@ class HunyuanVideoSampler(Inference):
         out_dict["samples"] = samples
         out_dict["prompts"] = prompt
 
-        gen_time = time.time() - start_time
+        gen_time = time.perf_counter() - start_time
         logger.info(f"Success, time: {gen_time}")
 
         return out_dict

fastvideo/models/mochi_hf/mochi_latents_utils.py

@@ -31,7 +31,7 @@ mochi_latents_std = torch.tensor([
 mochi_scaling_factor = 1.0
 
 
-def normalize_dit_input(model_type, latents):
+def normalize_dit_input(model_type, latents, args=None):
     if model_type == "mochi":
         latents_mean = mochi_latents_mean.to(latents.device, latents.dtype)
         latents_std = mochi_latents_std.to(latents.device, latents.dtype)
@@ -41,5 +41,16 @@ def normalize_dit_input(model_type, latents):
         return latents * 0.476986
     elif model_type == "hunyuan":
         return latents * 0.476986
+    elif model_type == "wan":
+        from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
+        vae_config = WanVAEConfig()
+        latents_mean = torch.tensor(vae_config.arch_config.latents_mean)
+        latents_std = 1.0 / torch.tensor(vae_config.arch_config.latents_std)
+
+
+        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(device=latents.device)
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device)
+        latents = ((latents.float() - latents_mean) * latents_std).to(latents)
+        return latents
     else:
         raise NotImplementedError(f"model_type {model_type} not supported")

fastvideo/train.py

@@ -364,7 +364,7 @@ def main(args):
     for i in range(init_steps):
         next(loader)
     for step in range(init_steps + 1, args.max_train_steps + 1):
-        start_time = time.time()
+        start_time = time.perf_counter()
         loss, grad_norm = train_one_step(
             transformer,
             args.model_type,
@@ -383,7 +383,7 @@ def main(args):
             args.mode_scale,
         )
 
-        step_time = time.time() - start_time
+        step_time = time.perf_counter() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 

fastvideo/utils/checkpoint.py

@@ -11,6 +11,7 @@ from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_
 from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import StateDictType
+import dataclasses
 
 from fastvideo.utils.logging_ import main_print
 
@@ -44,13 +45,50 @@ def save_checkpoint_optimizer(model, optimizer, rank, output_dir, step, discrimi
         optimizer_path = os.path.join(save_dir, "optimizer.pt")
         torch.save(optim_state, optimizer_path)
     else:
-        weight_path = os.path.join(save_dir, "discriminator_pytorch_model.safetensors")
+        weight_path = os.path.join(save_dstate_dictir, "discriminator_pytorch_model.safetensors")
         save_file(cpu_state, weight_path)
         optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt")
         torch.save(optim_state, optimizer_path)
     main_print(f"--> checkpoint saved at step {step}")
 
 
+def save_checkpoint_v1(transformer, rank, output_dir, step):
+    # from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+# from torch.distributed.fsdp import StateDictType, FullStateDictConfig
+
+    # Configure FSDP to save full state dict
+    FSDP.set_state_dict_type(
+        transformer,
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+    )
+
+    # Now get the state dict
+    cpu_state = transformer.state_dict()
+
+    # Save it (only on rank 0 since we used rank0_only=True)
+    # if torch.distributed.get_rank() == 0:
+    #     torch.save(state_dict, "model_checkpoint.pt")
+    if rank <= 0:
+        save_dir = os.path.join(output_dir, f"checkpoint-{step}")
+        os.makedirs(save_dir, exist_ok=True)
+        # save using safetensors
+        # weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
+        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.pt")
+        print(weight_path)
+        # save_file(cpu_state, weight_path)
+        torch.save(cpu_state, weight_path)
+        config_dict = transformer.hf_config
+        if "dtype" in config_dict:
+            del config_dict["dtype"]  # TODO
+        config_path = os.path.join(save_dir, "config.json")
+        # save dict as json
+        with open(config_path, "w") as f:
+            json.dump(config_dict, f, indent=4)
+    main_print(f"--> checkpoint saved at step {step}")
+
+
+
 def save_checkpoint(transformer, rank, output_dir, step):
     main_print(f"--> saving checkpoint at step {step}")
     with FSDP.state_dict_type(

fastvideo/v1/attention/layer.py

@@ -84,9 +84,6 @@ class DistributedAttention(nn.Module):
         # Check input shapes
         assert q.dim() == 4 and k.dim() == 4 and v.dim(
         ) == 4, "Expected 4D tensors"
-        # assert bs = 1
-        assert q.shape[
-            0] == 1, "Batch size must be 1, and there should be no padding tokens"
         batch_size, seq_len, num_heads, head_dim = q.shape
         local_rank = get_sequence_model_parallel_rank()
         world_size = get_sequence_model_parallel_world_size()

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

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

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

@@ -1,8 +1,8 @@
 from dataclasses import dataclass, field
-from typing import Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.configs.quantization import QuantizationConfig
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.platforms import _Backend
 
 
@@ -33,3 +33,24 @@ class DiTConfig(ModelConfig):
     # FastVideoDiT-specific parameters
     prefix: str = ""
     quant_config: Optional[QuantizationConfig] = None
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
+        """Add CLI arguments for DiTConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.prefix",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.prefix",
+            default=DiTConfig.prefix,
+            help="Prefix for the DiT model",
+        )
+
+        parser.add_argument(
+            f"--{prefix}.quant-config",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.quant_config",
+            default=None,
+            help="Quantization configuration for the DiT model",
+        )
+
+        return parser

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

@@ -0,0 +1,65 @@
+from dataclasses import dataclass, field
+from typing import List, Optional, Tuple, Union
+
+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 StepVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    _param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # transformer block
+            r"^transformer_blocks\.(\d+)\.norm1\.(weight|bias)$":
+            r"transformer_blocks.\1.norm1.norm.\2",
+            r"^transformer_blocks\.(\d+)\.norm2\.(weight|bias)$":
+            r"transformer_blocks.\1.norm2.norm.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.weight$":
+            r"transformer_blocks.\1.ff.fc_in.weight",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.weight$":
+            r"transformer_blocks.\1.ff.fc_out.weight",
+
+            # adanorm block
+            r"^adaln_single\.emb\.timestep_embedder\.linear_1\.(weight|bias)$":
+            r"adaln_single.emb.mlp.fc_in.\1",
+            r"^adaln_single\.emb\.timestep_embedder\.linear_2\.(weight|bias)$":
+            r"adaln_single.emb.mlp.fc_out.\1",
+
+            # caption projection
+            r"^caption_projection\.linear_1\.(weight|bias)$":
+            r"caption_projection.fc_in.\1",
+            r"^caption_projection\.linear_2\.(weight|bias)$":
+            r"caption_projection.fc_out.\1",
+        })
+
+    num_attention_heads: int = 48
+    attention_head_dim: int = 128
+    in_channels: int = 64
+    out_channels: Optional[int] = 64
+    num_layers: int = 48
+    dropout: float = 0.0
+    patch_size: int = 1
+    norm_type: str = "ada_norm_single"
+    norm_elementwise_affine: bool = False
+    norm_eps: float = 1e-6
+    caption_channels: Optional[Union[int, List[int], Tuple[int, ...]]] = field(
+        default_factory=lambda: [6144, 1024])
+    attention_type: Optional[str] = "torch"
+    use_additional_conditions: Optional[bool] = False
+
+    def __post_init__(self):
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.out_channels = self.in_channels if self.out_channels is None else self.out_channels
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class StepVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=StepVideoArchConfig)
+
+    prefix: str = "StepVideo"

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

@@ -4,7 +4,7 @@ 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
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.platforms import _Backend
 
 

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.stepvideovae import StepVideoVAEConfig
 from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
 
 __all__ = [
     "HunyuanVAEConfig",
     "WanVAEConfig",
+    "StepVideoVAEConfig",
 ]

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

@@ -1,9 +1,10 @@
 from dataclasses import dataclass, field
-from typing import Union
+from typing import Any, Union
 
 import torch
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.v1.utils import StoreBoolean
 
 
 @dataclass
@@ -33,6 +34,97 @@ class VAEConfig(ModelConfig):
     use_tiling: bool = True
     use_temporal_tiling: bool = True
     use_parallel_tiling: bool = True
+    use_temporal_scaling_frames: bool = True
 
     def __post_init__(self):
         self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "vae-config") -> Any:
+        """Add CLI arguments for VAEConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.load-encoder",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.load_encoder",
+            default=VAEConfig.load_encoder,
+            help="Whether to load the VAE encoder",
+        )
+        parser.add_argument(
+            f"--{prefix}.load-decoder",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.load_decoder",
+            default=VAEConfig.load_decoder,
+            help="Whether to load the VAE decoder",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-height",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_height",
+            default=VAEConfig.tile_sample_min_height,
+            help="Minimum height for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-width",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_width",
+            default=VAEConfig.tile_sample_min_width,
+            help="Minimum width for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_num_frames",
+            default=VAEConfig.tile_sample_min_num_frames,
+            help="Minimum number of frames for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-height",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_height",
+            default=VAEConfig.tile_sample_stride_height,
+            help="Stride height for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-width",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_width",
+            default=VAEConfig.tile_sample_stride_width,
+            help="Stride width for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_num_frames",
+            default=VAEConfig.tile_sample_stride_num_frames,
+            help="Stride number of frames for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.blend-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.blend_num_frames",
+            default=VAEConfig.blend_num_frames,
+            help="Number of frames to blend for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_tiling",
+            default=VAEConfig.use_tiling,
+            help="Whether to use tiling for VAE",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-temporal-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_temporal_tiling",
+            default=VAEConfig.use_temporal_tiling,
+            help="Whether to use temporal tiling for VAE",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-parallel-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_parallel_tiling",
+            default=VAEConfig.use_parallel_tiling,
+            help="Whether to use parallel tiling for VAE",
+        )
+
+        return parser

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

@@ -0,0 +1,28 @@
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class StepVideoVAEArchConfig(VAEArchConfig):
+    in_channels: int = 3
+    out_channels: int = 3
+    z_channels: int = 64
+    num_res_blocks: int = 2
+    version: int = 2
+    frame_len: int = 17
+    world_size: int = 1
+
+    spatial_compression_ratio: int = 16
+    temporal_compression_ratio: int = 8
+
+    scaling_factor: float = 1.0
+
+
+@dataclass
+class StepVideoVAEConfig(VAEConfig):
+    arch_config: VAEArchConfig = field(default_factory=StepVideoVAEArchConfig)
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+    use_temporal_scaling_frames: bool = False

fastvideo/v1/configs/pipelines/__init__.py

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

fastvideo/v1/configs/pipelines/hunyuan.py

@@ -68,6 +68,9 @@ class HunyuanConfig(PipelineConfig):
     embedded_cfg_scale: int = 6
     flow_shift: int = 7
 
+    # Video parameters
+    use_cpu_offload: bool = True
+
     # Text encoding stage
     text_encoder_configs: Tuple[EncoderConfig, ...] = field(
         default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))

fastvideo/v1/configs/pipelines/registry.py

@@ -6,6 +6,7 @@ from typing import Callable, Dict, Optional, Type
 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
                                                     HunyuanConfig)
+from fastvideo.v1.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
                                                 WanI2V720PConfig,
                                                 WanT2V480PConfig,
@@ -24,6 +25,7 @@ WEIGHT_CONFIG_REGISTRY: Dict[str, Type[PipelineConfig]] = {
     "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,
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
     # Add other specific weight variants
 }
 
@@ -32,6 +34,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(),
+    "stepvideo": lambda id: "stepvideo" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -42,6 +45,7 @@ PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {
     "wanpipeline":
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
+    "stepvideo": StepVideoT2VConfig
     # Other fallbacks by architecture
 }
 

fastvideo/v1/configs/pipelines/stepvideo.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.models import DiTConfig, VAEConfig
+from fastvideo.v1.configs.models.dits import StepVideoConfig
+from fastvideo.v1.configs.models.vaes import StepVideoVAEConfig
+from fastvideo.v1.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class StepVideoT2VConfig(PipelineConfig):
+    """Base configuration for StepVideo pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+    # DiT
+    dit_config: DiTConfig = field(default_factory=StepVideoConfig)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=StepVideoVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Video parameters
+    use_cpu_offload: bool = True
+
+    # Denoising stage
+    flow_shift: int = 13
+    timesteps_scale: bool = False
+    pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。"
+    neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "bf16"

fastvideo/v1/configs/quantization/__init__.py

@@ -1,3 +0,0 @@
-from fastvideo.v1.configs.quantization.base import QuantizationConfig
-
-__all__ = ["QuantizationConfig"]

fastvideo/v1/configs/quantization/base.py

@@ -1,6 +0,0 @@
-from dataclasses import dataclass
-
-
-@dataclass
-class QuantizationConfig:
-    pass

fastvideo/v1/configs/sample/base.py

@@ -78,6 +78,113 @@ class SamplingParam:
 
         return sampling_param
 
+    @staticmethod
+    def add_cli_args(parser: Any) -> Any:
+        """Add CLI arguments for SamplingParam fields"""
+        parser.add_argument(
+            "--prompt",
+            type=str,
+            default=SamplingParam.prompt,
+            help="Text prompt for video generation",
+        )
+        parser.add_argument(
+            "--negative-prompt",
+            type=str,
+            default=SamplingParam.negative_prompt,
+            help="Negative text prompt for video generation",
+        )
+        parser.add_argument(
+            "--prompt-path",
+            type=str,
+            default=SamplingParam.prompt_path,
+            help="Path to a text file containing the prompt",
+        )
+        parser.add_argument(
+            "--output-path",
+            type=str,
+            default=SamplingParam.output_path,
+            help="Path to save the generated video",
+        )
+        parser.add_argument(
+            "--num-videos-per-prompt",
+            type=int,
+            default=SamplingParam.num_videos_per_prompt,
+            help="Number of videos to generate per prompt",
+        )
+        parser.add_argument(
+            "--seed",
+            type=int,
+            default=SamplingParam.seed,
+            help="Random seed for generation",
+        )
+        parser.add_argument(
+            "--num-frames",
+            type=int,
+            default=SamplingParam.num_frames,
+            help="Number of frames to generate",
+        )
+        parser.add_argument(
+            "--height",
+            type=int,
+            default=SamplingParam.height,
+            help="Height of generated video",
+        )
+        parser.add_argument(
+            "--width",
+            type=int,
+            default=SamplingParam.width,
+            help="Width of generated video",
+        )
+        parser.add_argument(
+            "--fps",
+            type=int,
+            default=SamplingParam.fps,
+            help="Frames per second for saved video",
+        )
+        parser.add_argument(
+            "--num-inference-steps",
+            type=int,
+            default=SamplingParam.num_inference_steps,
+            help="Number of denoising steps",
+        )
+        parser.add_argument(
+            "--guidance-scale",
+            type=float,
+            default=SamplingParam.guidance_scale,
+            help="Classifier-free guidance scale",
+        )
+        parser.add_argument(
+            "--guidance-rescale",
+            type=float,
+            default=SamplingParam.guidance_rescale,
+            help="Guidance rescale factor",
+        )
+        parser.add_argument(
+            "--save-video",
+            action="store_true",
+            default=SamplingParam.save_video,
+            help="Whether to save the video to disk",
+        )
+        parser.add_argument(
+            "--no-save-video",
+            action="store_false",
+            dest="save_video",
+            help="Don't save the video to disk",
+        )
+        parser.add_argument(
+            "--return-frames",
+            action="store_true",
+            default=SamplingParam.return_frames,
+            help="Whether to return the raw frames",
+        )
+        parser.add_argument(
+            "--image-path",
+            type=str,
+            default=SamplingParam.image_path,
+            help="Path to input image for image-to-video generation",
+        )
+        return parser
+
 
 @dataclass
 class CacheParams:

fastvideo/v1/configs/sample/registry.py

@@ -3,6 +3,7 @@ from typing import Any, Callable, Dict, Optional
 
 from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                                  HunyuanSamplingParam)
+from fastvideo.v1.configs.sample.stepvideo import StepVideoT2VSamplingParam
 from fastvideo.v1.configs.sample.wan import (WanI2V_14B_480P_SamplingParam,
                                              WanI2V_14B_720P_SamplingParam,
                                              WanT2V_1_3B_SamplingParam,
@@ -20,6 +21,7 @@ SAMPLING_PARAM_REGISTRY: Dict[str, Any] = {
     "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,
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
     # Add other specific weight variants
 }
 
@@ -28,6 +30,7 @@ SAMPLING_PARAM_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(),
+    "stepvideo": lambda id: "stepvideo" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -38,6 +41,7 @@ SAMPLING_FALLBACK_PARAM: Dict[str, Any] = {
     "wanpipeline":
     WanT2V_1_3B_SamplingParam,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V_14B_480P_SamplingParam,
+    "stepvideo": StepVideoT2VSamplingParam
     # Other fallbacks by architecture
 }
 

fastvideo/v1/configs/sample/stepvideo.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+from fastvideo.v1.configs.sample.base import SamplingParam
+
+
+@dataclass
+class StepVideoT2VSamplingParam(SamplingParam):
+    # Video parameters
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 81
+
+    # Denoising stage
+    guidance_scale: float = 9.0
+    num_inference_steps: int = 50
+
+    # neg magic and pos magic
+    # pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。"
+    # neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"

fastvideo/v1/dataset/__init__.py

@@ -2,11 +2,12 @@ from torchvision import transforms
 from torchvision.transforms import Lambda
 from transformers import AutoTokenizer
 
-from fastvideo.dataset.t2v_datasets import T2V_dataset
-from fastvideo.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop
+from fastvideo.v1.dataset.t2v_datasets import T2V_dataset
+from fastvideo.v1.dataset.transform import (CenterCropResizeVideo, Normalize255,
+                                            TemporalRandomCrop)
 
 
-def getdataset(args):
+def getdataset(args, start_idx=0):
     temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
     norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
     resize_topcrop = [
@@ -25,15 +26,15 @@ def getdataset(args):
         norm_fun,
     ])
     # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
-    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name,
+                                              cache_dir=args.cache_dir)
     if args.dataset == "t2v":
-        return T2V_dataset(
-            args,
-            transform=transform,
-            temporal_sample=temporal_sample,
-            tokenizer=tokenizer,
-            transform_topcrop=transform_topcrop,
-        )
+        return T2V_dataset(args,
+                           transform=transform,
+                           temporal_sample=temporal_sample,
+                           tokenizer=tokenizer,
+                           transform_topcrop=transform_topcrop,
+                           start_idx=start_idx)
 
     raise NotImplementedError(args.dataset)
 
@@ -44,7 +45,7 @@ if __name__ == "__main__":
     from accelerate import Accelerator
     from tqdm import tqdm
 
-    from fastvideo.dataset.t2v_datasets import dataset_prog
+    from fastvideo.v1.dataset.t2v_datasets import dataset_prog
 
     args = type(
         "args",
@@ -63,7 +64,8 @@ if __name__ == "__main__":
             "interpolation_scale_h": 1,
             "interpolation_scale_w": 1,
             "cache_dir": "../cache_dir",
-            "image_data": "/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt",
+            "image_data":
+            "/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt",
             "video_data": "1",
             "train_fps": 24,
             "drop_short_ratio": 1.0,
@@ -80,7 +82,10 @@ if __name__ == "__main__":
     zero = 0
     for idx in tqdm(range(num)):
         image_data = dataset_prog.img_cap_list[idx]
-        caps = [i["cap"] if isinstance(i["cap"], list) else [i["cap"]] for i in image_data]
+        caps = [
+            i["cap"] if isinstance(i["cap"], list) else [i["cap"]]
+            for i in image_data
+        ]
         try:
             caps = [[random.choice(i)] for i in caps]
         except Exception as e:

fastvideo/v1/dataset/dataloader/schema.py

@@ -0,0 +1,44 @@
+# schema.py
+"""
+Unified data schema and format for saving and loading image/video data after
+preprocessing.
+
+It uses apache arrow in-memory format that can be consumed by modern data
+frameworks that can handle parquet or lance file.
+"""
+
+import pyarrow as pa
+
+pyarrow_schema = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Image/Video VAE latents ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("vae_latent_bytes", pa.binary()),
+    # e.g., [C, T, H, W] or [C, H, W]
+    pa.field("vae_latent_shape", pa.list_(pa.int64())),
+    # e.g., 'float32'
+    pa.field("vae_latent_dtype", pa.string()),
+    # --- Text encoder output tensor ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("text_embedding_bytes", pa.binary()),
+    # e.g., [SeqLen, Dim]
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    # e.g., 'bfloat16' or 'float32'
+    pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_attention_mask_bytes", pa.binary()),
+    # e.g., [SeqLen]
+    pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+    # e.g., 'bool' or 'int8'
+    pa.field("text_attention_mask_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # 'image' or 'video'
+    pa.field("width", pa.int64()),
+    pa.field("height", pa.int64()),
+    # -- Video-specific (can be null/default for images) ---
+    # Number of frames processed (e.g., 1 for image, N for video)
+    pa.field("num_frames", pa.int64()),
+    pa.field("duration_sec", pa.float64()),
+    pa.field("fps", pa.float64()),
+])

fastvideo/v1/dataset/latent_datasets.py

@@ -20,9 +20,11 @@ class LatentDataset(Dataset):
         self.datase_dir_path = os.path.dirname(json_path)
         self.video_dir = os.path.join(self.datase_dir_path, "video")
         self.latent_dir = os.path.join(self.datase_dir_path, "latent")
-        self.prompt_embed_dir = os.path.join(self.datase_dir_path, "prompt_embed")
-        self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask")
-        with open(self.json_path, "r") as f:
+        self.prompt_embed_dir = os.path.join(self.datase_dir_path,
+                                             "prompt_embed")
+        self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path,
+                                                      "prompt_attention_mask")
+        with open(self.json_path) as f:
             self.data_anno = json.load(f)
         # json.load(f) already keeps the order
         # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path'])
@@ -31,12 +33,16 @@ class LatentDataset(Dataset):
         self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32)
         # 256 zeros
         self.uncond_prompt_mask = torch.zeros(256).bool()
-        self.lengths = [data_item["length"] if "length" in data_item else 1 for data_item in self.data_anno]
+        self.lengths = [
+            data_item["length"] if "length" in data_item else 1
+            for data_item in self.data_anno
+        ]
 
     def __getitem__(self, idx):
         latent_file = self.data_anno[idx]["latent_path"]
         prompt_embed_file = self.data_anno[idx]["prompt_embed_path"]
-        prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"]
+        prompt_attention_mask_file = self.data_anno[idx][
+            "prompt_attention_mask"]
         # load
         latent = torch.load(
             os.path.join(self.latent_dir, latent_file),
@@ -54,7 +60,8 @@ class LatentDataset(Dataset):
                 weights_only=True,
             )
             prompt_attention_mask = torch.load(
-                os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file),
+                os.path.join(self.prompt_attention_mask_dir,
+                             prompt_attention_mask_file),
                 map_location="cpu",
                 weights_only=True,
             )
@@ -104,8 +111,12 @@ def latent_collate_function(batch):
 
 
 if __name__ == "__main__":
-    dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt", num_latent_t=28)
-    dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function)
+    dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt",
+                            num_latent_t=28)
+    dataloader = torch.utils.data.DataLoader(dataset,
+                                             batch_size=2,
+                                             shuffle=False,
+                                             collate_fn=latent_collate_function)
     for latent, prompt_embed, latent_attn_mask, prompt_attention_mask in dataloader:
         print(
             latent.shape,

fastvideo/v1/dataset/parquet_datasets.py

@@ -0,0 +1,418 @@
+import argparse
+import os
+import random
+import time
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+import tqdm
+from torch import distributed as dist
+from torch.utils.data import IterableDataset, get_worker_info
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+# Path to your dataset
+dataset_path = "/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/train/"
+
+
+class ParquetVideoTextDataset(IterableDataset):
+    """Efficient loader for video-text data from a directory of Parquet files."""
+
+    def __init__(self,
+                 path: str,
+                 batch_size: int = 1024,
+                 rank: int = 0,
+                 world_size: int = 1,
+                 cfg_rate: float = 0.0,
+                 num_latent_t: int = 2):
+        super().__init__()
+        self.path = str(path)
+        self.batch_size = batch_size
+        self.rank = rank
+        self.world_size = world_size
+        self.cfg_rate = cfg_rate
+        self.num_latent_t = num_latent_t
+
+        # Find all parquet files recursively
+        print(f"Scanning for parquet files in {self.path}")
+        self.parquet_files = []
+        for root, _, files in os.walk(self.path):
+            for file in files:
+                if file.endswith('.parquet'):
+                    self.parquet_files.append(os.path.join(root, file))
+        # Sort files for consistent ordering
+        self.parquet_files.sort()
+
+        # Distribute files among workers
+        # drop last unenven files
+        print(f"Total files: {len(self.parquet_files)}")
+        total_files = len(self.parquet_files)
+        base_count = total_files // world_size
+        extra_files = total_files % world_size
+
+        if rank < extra_files:
+            start_idx = rank * (base_count + 1)
+            end_idx = start_idx + base_count + 1
+        else:
+            start_idx = rank * base_count + extra_files
+            end_idx = start_idx + base_count
+
+        self.parquet_files = self.parquet_files[start_idx:end_idx]
+
+        print(f"Files assigned to rank {rank}: {len(self.parquet_files)}")
+        if len(self.parquet_files) > 0:
+            print(f"First file: {self.parquet_files[0]}")
+            print(f"Last file: {self.parquet_files[-1]}")
+
+        # Initialize current file index
+        self.current_file_idx = 0
+        self.current_reader = None
+        self.current_batches = None
+        self.total_samples = 0
+
+    def _open_next_file(self):
+        """Open the next parquet file for reading."""
+        num_workers = get_worker_info().num_workers
+        worker_id = get_worker_info().id
+        total_files = len(self.parquet_files)
+        base_count = total_files // num_workers
+        extra_files = total_files % num_workers
+
+        if worker_id < extra_files:
+            start_idx = worker_id * (base_count + 1)
+            end_idx = start_idx + base_count + 1
+        else:
+            start_idx = worker_id * base_count + extra_files
+            end_idx = start_idx + base_count
+
+        worker_parquet_files = self.parquet_files[start_idx:end_idx]
+        if self.current_file_idx >= len(worker_parquet_files):
+            print(
+                f"Rank {self.rank}, Worker {worker_id}: No more files to open (current_idx={self.current_file_idx}, total_files={len(worker_parquet_files)})"
+            )
+            return False
+
+        if self.current_reader is not None:
+            self.current_reader.close()
+
+        file_path = worker_parquet_files[self.current_file_idx]
+        print(
+            f"Rank {self.rank}, Worker {worker_id}: Opening file {self.current_file_idx + 1}/{len(worker_parquet_files)}: {file_path}"
+        )
+
+        try:
+            self.current_reader = pq.ParquetFile(file_path)
+            self.current_batches = self.current_reader.iter_batches(
+                batch_size=self.batch_size)
+            self.current_file_idx += 1
+            return True
+        except Exception as e:
+            print(f"Error opening file {file_path}: {str(e)}")
+            return False
+
+    def __iter__(self):
+        """Iterate over the dataset in a streaming fashion."""
+        print(f"Rank {self.rank}: Starting iteration")
+
+        # First try to open a file
+        if not self._open_next_file():
+            print(f"Rank {self.rank}: Failed to open first file")
+            return
+
+        while True:
+            try:
+                # Get next batch from current file
+                batch = next(self.current_batches)
+                batch_dict = batch.to_pydict()
+                processed = self._process_batch(batch_dict)
+
+                # Update sample count
+                batch_size = len(processed["latents"])
+                self.total_samples += batch_size
+
+                # Print progress
+                if self.total_samples % 1000 == 0:
+                    print(
+                        f"Rank {self.rank}: Processed {self.total_samples} samples"
+                    )
+
+                # Yield each item in the batch
+                for lat, emb, mask, info in zip(processed["latents"],
+                                                processed["embeddings"],
+                                                processed["masks"],
+                                                processed["info"]):
+                    if lat.numel() == 0:  # Split is validation
+                        yield lat, emb, mask, info
+                    else:
+                        yield lat[:, -self.num_latent_t:], emb, mask, info
+
+            except StopIteration:
+                # Current file is exhausted, try next file
+                print(
+                    f"Rank {self.rank}: Current file exhausted, trying next file"
+                )
+                self.current_batches = None
+                if not self._open_next_file():
+                    print(
+                        f"Rank {self.rank}: No more files to process. Total samples: {self.total_samples}"
+                    )
+                    break
+            except Exception as e:
+                print(f"Error processing batch: {str(e)}")
+                self.current_batches = None
+                if not self._open_next_file():
+                    print(
+                        f"Rank {self.rank}: Failed to open next file after error"
+                    )
+                    break
+
+        # Clean up
+        if self.current_reader is not None:
+            self.current_reader.close()
+
+    def _process_batch(self, batch):
+        """Process a PyArrow batch into tensors."""
+        out = {"lat": [], "emb": [], "msk": [], "info": []}
+
+        for i in range(len(batch["vae_latent_bytes"])):
+            vae_latent_bytes = batch["vae_latent_bytes"][i]
+            vae_latent_shape = batch["vae_latent_shape"][i]
+            text_embedding_bytes = batch["text_embedding_bytes"][i]
+            text_embedding_shape = batch["text_embedding_shape"][i]
+            text_attention_mask_bytes = batch["text_attention_mask_bytes"][i]
+            text_attention_mask_shape = batch["text_attention_mask_shape"][i]
+
+            # Process latent
+            if not vae_latent_shape:  # No VAE latent is stored. Split is validation
+                lat = np.array([])
+            else:
+                lat = np.frombuffer(vae_latent_bytes,
+                                    dtype=np.float32).reshape(vae_latent_shape)
+                # Make array writable
+                lat = np.copy(lat)
+
+            if random.random() < self.cfg_rate:
+                emb = np.zeros((512, 4096), dtype=np.float32)
+            else:
+                emb = np.frombuffer(
+                    text_embedding_bytes,
+                    dtype=np.float32).reshape(text_embedding_shape)
+                # Make array writable
+                emb = np.copy(emb)
+            if emb.shape[0] < 512:
+                padded_emb = np.zeros((512, emb.shape[1]), dtype=np.float32)
+                padded_emb[:emb.shape[0], :] = emb
+                emb = padded_emb
+            elif emb.shape[0] > 512:
+                emb = emb[:512, :]
+
+            # Process mask
+            if len(text_attention_mask_bytes) > 0 and len(
+                    text_attention_mask_shape) > 0:
+                msk = np.frombuffer(text_attention_mask_bytes,
+                                    dtype=np.uint8).astype(np.bool_)
+                msk = msk.reshape(1, -1)
+                # Make array writable
+                msk = np.copy(msk)
+                if msk.shape[1] < 512:
+                    padded_msk = np.zeros((1, 512), dtype=np.bool_)
+                    padded_msk[:, :msk.shape[1]] = msk
+                    msk = padded_msk
+                elif msk.shape[1] > 512:
+                    msk = msk[:, :512]
+            else:
+                msk = np.ones((1, 512), dtype=np.bool_)
+            # to string
+            file_name = str(batch["file_name"][i])
+            # Collect metadata
+            info = {
+                "width": batch["width"][i],
+                "height": batch["height"][i],
+                "num_frames": batch["num_frames"][i],
+                "duration_sec": batch["duration_sec"][i],
+                "fps": batch["fps"][i],
+                "file_name": batch["file_name"][i],
+                "caption": batch["caption"][i],
+            }
+
+            out["lat"].append(torch.from_numpy(lat))
+            out["emb"].append(torch.from_numpy(emb))
+            out["msk"].append(torch.from_numpy(msk))
+            out["info"].append(info)
+
+        return {
+            "latents": torch.stack(out["lat"]) if out["lat"] else None,
+            "embeddings": torch.stack(out["emb"]) if out["emb"] else None,
+            "masks": torch.stack(out["msk"]) if out["msk"] else None,
+            "info": out["info"]
+        }
+
+
+def bind_cpu_cores(local_rank, cpu_per_process=16):
+    """根据local_rank绑定固定cpu核。"""
+    start = local_rank * cpu_per_process
+    end = start + cpu_per_process
+    cores = list(range(start, end))
+    print(f"[Rank {local_rank}] Binding to CPU cores: {cores}")
+    os.sched_setaffinity(0, cores)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description='Benchmark Parquet dataset loading speed')
+    parser.add_argument('--path',
+                        type=str,
+                        default=dataset_path,
+                        help='Path to Parquet dataset')
+    parser.add_argument('--batch_size',
+                        type=int,
+                        default=4,
+                        help='Batch size for DataLoader')
+    parser.add_argument('--num_batches',
+                        type=int,
+                        default=100,
+                        help='Number of batches to benchmark')
+    parser.add_argument('--vae_debug', action="store_true")
+    args = parser.parse_args()
+
+    # Initialize distributed training
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    rank = int(os.environ.get("RANK", 0))
+
+    # Initialize CUDA device first
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+        device = torch.device(f"cuda:{local_rank}")
+    else:
+        device = torch.device("cpu")
+
+    # Initialize distributed training
+    if world_size > 1:
+        dist.init_process_group(backend="nccl",
+                                init_method="env://",
+                                world_size=world_size,
+                                rank=rank)
+        print(
+            f"Initialized process: rank={rank}, local_rank={local_rank}, world_size={world_size}, device={device}"
+        )
+
+    # Bind CPU cores after distributed initialization
+    # bind_cpu_cores(local_rank, cpu_per_process=16)
+
+    # Create dataset
+    dataset = ParquetVideoTextDataset(
+        args.path,
+        batch_size=args.batch_size,
+        rank=rank,
+        world_size=world_size,
+    )
+
+    # Create DataLoader with proper settings
+    dataloader = StatefulDataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=1,  # Reduce number of workers to avoid memory issues
+        prefetch_factor=2,
+        shuffle=False,
+        pin_memory=True,
+        drop_last=True)
+
+    # Example of how to load dataloader state
+    # if os.path.exists("/workspace/FastVideo/dataloader_state.pt"):
+    #     dataloader_state = torch.load("/workspace/FastVideo/dataloader_state.pt")
+    #     dataloader.load_state_dict(dataloader_state[rank])
+
+    # Warm-up with synchronization
+    if rank == 0:
+        print("Warming up...")
+    for i, (latents, embeddings, masks, infos) in enumerate(dataloader):
+        # Example of how to save dataloader state
+        # if i == 30:
+        #     dist.barrier()
+        #     local_data = {rank: dataloader.state_dict()}
+        #     gathered_data = [None] * world_size
+        #     dist.all_gather_object(gathered_data, local_data)
+        #     if rank == 0:
+        #         global_state_dict = {}
+        #         for d in gathered_data:
+        #             global_state_dict.update(d)
+        #         torch.save(global_state_dict, "dataloader_state.pt")
+        assert torch.sum(masks[0]).item() == torch.count_nonzero(
+            embeddings[0]).item() // 4096
+        if args.vae_debug:
+            from diffusers.utils import export_to_video
+            from diffusers.video_processor import VideoProcessor
+
+            from fastvideo.v1.configs.models.vaes import WanVAEConfig
+            from fastvideo.v1.fastvideo_args import FastVideoArgs
+            from fastvideo.v1.models.loader.component_loader import VAELoader
+            VAE_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/vae"
+            fastvideo_args = FastVideoArgs(
+                model_path=VAE_PATH,
+                vae_config=WanVAEConfig(load_encoder=False),
+                vae_precision="fp32")
+            fastvideo_args.device = device
+            vae_loader = VAELoader()
+            vae = vae_loader.load(model_path=VAE_PATH,
+                                  architecture="",
+                                  fastvideo_args=fastvideo_args)
+
+            videoprocessor = VideoProcessor(vae_scale_factor=8)
+
+            with torch.inference_mode():
+                video = vae.decode(latents[0].unsqueeze(0).to(device))
+                video = videoprocessor.postprocess_video(video)
+                video_path = os.path.join("/workspace/FastVideo/debug_videos",
+                                          infos["caption"][0][:50] + ".mp4")
+                export_to_video(video[0], video_path, fps=16)
+
+        # Move data to device
+        # latents = latents.to(device)
+        # embeddings = embeddings.to(device)
+
+    if world_size > 1:
+        dist.barrier()
+
+    # Benchmark
+    if rank == 0:
+        print(f"Benchmarking with batch_size={args.batch_size}")
+    start_time = time.time()
+    total_samples = 0
+    for i, (latents, embeddings, masks,
+            infos) in enumerate(tqdm.tqdm(dataloader, total=args.num_batches)):
+        if i >= args.num_batches:
+            break
+
+        # Move data to device
+        latents = latents.to(device)
+        embeddings = embeddings.to(device)
+
+        # Calculate actual batch size
+        batch_size = latents.size(0)
+        total_samples += batch_size
+
+        # Print progress only from rank 0
+        if rank == 0 and (i + 1) % 10 == 0:
+            elapsed = time.time() - start_time
+            samples_per_sec = total_samples / elapsed
+            print(
+                f"Batch {i+1}/{args.num_batches}, Speed: {samples_per_sec:.2f} samples/sec"
+            )
+
+    # Final statistics
+    if world_size > 1:
+        dist.barrier()
+
+    if rank == 0:
+        elapsed = time.time() - start_time
+        samples_per_sec = total_samples / elapsed
+
+        print("\nBenchmark Results:")
+        print(f"Total time: {elapsed:.2f} seconds")
+        print(f"Total samples: {total_samples}")
+        print(f"Average speed: {samples_per_sec:.2f} samples/sec")
+        print(f"Time per batch: {elapsed/args.num_batches*1000:.2f} ms")
+
+    if world_size > 1:
+        dist.destroy_process_group()

fastvideo/v1/dataset/t2v_datasets.py

@@ -46,7 +46,8 @@ class DataSetProg(metaclass=SingletonMeta):
 
         for i in range(self.num_workers):
             self.n_used_elements[i] = 0
-            per_worker = int(math.ceil(len(self.elements) / float(self.num_workers)))
+            per_worker = int(
+                math.ceil(len(self.elements) / float(self.num_workers)))
             start = i * per_worker
             end = min(start + per_worker, len(self.elements))
             self.worker_elements[i] = self.elements[start:end]
@@ -57,7 +58,9 @@ class DataSetProg(metaclass=SingletonMeta):
         else:
             worker_id = work_info.id
 
-        idx = self.worker_elements[worker_id][self.n_used_elements[worker_id] % len(self.worker_elements[worker_id])]
+        idx = self.worker_elements[worker_id][
+            self.n_used_elements[worker_id] %
+            len(self.worker_elements[worker_id])]
         self.n_used_elements[worker_id] += 1
         return idx
 
@@ -65,7 +68,10 @@ class DataSetProg(metaclass=SingletonMeta):
 dataset_prog = DataSetProg()
 
 
-def filter_resolution(h, w, max_h_div_w_ratio=17 / 16, min_h_div_w_ratio=8 / 16):
+def filter_resolution(h,
+                      w,
+                      max_h_div_w_ratio=17 / 16,
+                      min_h_div_w_ratio=8 / 16):
     if h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio:
         return True
     return False
@@ -73,7 +79,14 @@ def filter_resolution(h, w, max_h_div_w_ratio=17 / 16, min_h_div_w_ratio=8 / 16)
 
 class T2V_dataset(Dataset):
 
-    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
+    def __init__(self,
+                 args,
+                 transform,
+                 temporal_sample,
+                 tokenizer,
+                 transform_topcrop,
+                 start_idx=0):
+        self.start_idx = start_idx
         self.data = args.data_merge_path
         self.num_frames = args.num_frames
         self.train_fps = args.train_fps
@@ -102,7 +115,8 @@ class T2V_dataset(Dataset):
         self.lengths = self.sample_num_frames
 
         n_elements = len(cap_list)
-        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list, n_elements)
+        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list,
+                                  n_elements)
 
         print(f"video length: {len(dataset_prog.cap_list)}", flush=True)
 
@@ -129,7 +143,8 @@ class T2V_dataset(Dataset):
         video_path = dataset_prog.cap_list[idx]["path"]
         assert os.path.exists(video_path), f"file {video_path} do not exist!"
         frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
-        torchvision_video, _, metadata = torchvision.io.read_video(video_path, output_format="TCHW")
+        torchvision_video, _, metadata = torchvision.io.read_video(
+            video_path, output_format="TCHW")
         video = torchvision_video[frame_indices]
         video = self.transform(video)
         video = rearrange(video, "t c h w -> c t h w")
@@ -160,16 +175,17 @@ class T2V_dataset(Dataset):
         )
         input_ids = text_tokens_and_mask["input_ids"]
         cond_mask = text_tokens_and_mask["attention_mask"]
-        return dict(
-            pixel_values=video,
-            text=text,
-            input_ids=input_ids,
-            cond_mask=cond_mask,
-            path=video_path,
-        )
+        return dict(pixel_values=video,
+                    text=text,
+                    input_ids=input_ids,
+                    cond_mask=cond_mask,
+                    path=video_path,
+                    fps=dataset_prog.cap_list[idx]["fps"],
+                    duration=dataset_prog.cap_list[idx]["duration"])
 
     def get_image(self, idx):
-        image_data = dataset_prog.cap_list[idx]  # [{'path': path, 'cap': cap}, ...]
+        image_data = dataset_prog.cap_list[
+            idx]  # [{'path': path, 'cap': cap}, ...]
 
         image = Image.open(image_data["path"]).convert("RGB")  # [h, w, c]
         image = torch.from_numpy(np.array(image))  # [h, w, c]
@@ -178,13 +194,15 @@ class T2V_dataset(Dataset):
         #     h, w = i.shape[-2:]
         #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
 
-        image = (self.transform_topcrop(image) if "human_images" in image_data["path"] else self.transform(image)
+        image = (self.transform_topcrop(image) if "human_images"
+                 in image_data["path"] else self.transform(image)
                  )  #  [1 C H W] -> num_img [1 C H W]
         image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
 
         image = image.float() / 127.5 - 1.0
 
-        caps = (image_data["cap"] if isinstance(image_data["cap"], list) else [image_data["cap"]])
+        caps = (image_data["cap"]
+                if isinstance(image_data["cap"], list) else [image_data["cap"]])
         caps = [random.choice(caps)]
         text = caps
         input_ids, cond_mask = [], []
@@ -238,10 +256,12 @@ class T2V_dataset(Dataset):
                     cnt_no_resolution += 1
                     continue
                 else:
-                    if (resolution.get("height", None) is None or resolution.get("width", None) is None):
+                    if (resolution.get("height", None) is None
+                            or resolution.get("width", None) is None):
                         cnt_no_resolution += 1
                         continue
-                    height, width = i["resolution"]["height"], i["resolution"]["width"]
+                    height, width = i["resolution"]["height"], i["resolution"][
+                        "width"]
                     aspect = self.max_height / self.max_width
                     hw_aspect_thr = 1.5
                     is_pick = filter_resolution(
@@ -259,29 +279,34 @@ class T2V_dataset(Dataset):
                 i["num_frames"] = math.ceil(fps * duration)
                 # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
                 if i["num_frames"] / fps > self.video_length_tolerance_range * (
-                        self.num_frames / self.train_fps *
-                        self.speed_factor):  # too long video is not suitable for this training stage (self.num_frames)
+                        self.num_frames / self.train_fps * self.speed_factor
+                ):  # too long video is not suitable for this training stage (self.num_frames)
                     cnt_too_long += 1
                     continue
 
                 # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
                 frame_interval = fps / self.train_fps
                 start_frame_idx = 0
-                frame_indices = np.arange(start_frame_idx, i["num_frames"], frame_interval).astype(int)
+                frame_indices = np.arange(start_frame_idx, i["num_frames"],
+                                          frame_interval).astype(int)
 
                 # comment out it to enable dynamic frames training
-                if (len(frame_indices) < self.num_frames and random.random() < self.drop_short_ratio):
+                if (len(frame_indices) < self.num_frames
+                        and random.random() < self.drop_short_ratio):
                     cnt_too_short += 1
                     continue
 
                 #  too long video will be temporal-crop randomly
                 if len(frame_indices) > self.num_frames:
-                    begin_index, end_index = self.temporal_sample(len(frame_indices))
+                    begin_index, end_index = self.temporal_sample(
+                        len(frame_indices))
                     frame_indices = frame_indices[begin_index:end_index]
                     # frame_indices = frame_indices[:self.num_frames]  # head crop
                 i["sample_frame_index"] = frame_indices.tolist()
                 new_cap_list.append(i)
-                i["sample_num_frames"] = len(i["sample_frame_index"])  # will use in dataloader(group sampler)
+                i["sample_num_frames"] = len(
+                    i["sample_frame_index"]
+                )  # will use in dataloader(group sampler)
                 sample_num_frames.append(i["sample_num_frames"])
             elif path.endswith(".jpg"):  # image
                 cnt_img += 1
@@ -290,13 +315,15 @@ class T2V_dataset(Dataset):
                 sample_num_frames.append(i["sample_num_frames"])
             else:
                 raise NameError(
-                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image")
+                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image"
+                )
         # import ipdb;ipdb.set_trace()
         main_print(
             f"no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, "
             f"no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, "
             f"Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, "
-            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}")
+            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}"
+        )
         return new_cap_list, sample_num_frames
 
     def decord_read(self, path, frame_indices):
@@ -308,11 +335,14 @@ class T2V_dataset(Dataset):
 
     def read_jsons(self, data):
         cap_lists = []
-        with open(data, "r") as f:
-            folder_anno = [i.strip().split(",") for i in f.readlines() if len(i.strip()) > 0]
+        with open(data) as f:
+            folder_anno = [
+                i.strip().split(",") for i in f.readlines()
+                if len(i.strip()) > 0
+            ]
         print(folder_anno)
         for folder, anno in folder_anno:
-            with open(anno, "r") as f:
+            with open(anno) as f:
                 sub_list = json.load(f)
             for i in range(len(sub_list)):
                 sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
@@ -320,5 +350,5 @@ class T2V_dataset(Dataset):
         return cap_lists
 
     def get_cap_list(self):
-        cap_lists = self.read_jsons(self.data)
+        cap_lists = self.read_jsons(self.data)[self.start_idx:]
         return cap_lists

fastvideo/v1/dataset/transform.py

@@ -21,15 +21,19 @@ def center_crop_arr(pil_image, image_size):
     https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
     """
     while min(*pil_image.size) >= 2 * image_size:
-        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
+        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size),
+                                     resample=Image.BOX)
 
     scale = image_size / min(*pil_image.size)
-    pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
+    pil_image = pil_image.resize(tuple(
+        round(x * scale) for x in pil_image.size),
+                                 resample=Image.BICUBIC)
 
     arr = np.array(pil_image)
     crop_y = (arr.shape[0] - image_size) // 2
     crop_x = (arr.shape[1] - image_size) // 2
-    return Image.fromarray(arr[crop_y:crop_y + image_size, crop_x:crop_x + image_size])
+    return Image.fromarray(arr[crop_y:crop_y + image_size,
+                               crop_x:crop_x + image_size])
 
 
 def crop(clip, i, j, h, w):
@@ -44,7 +48,9 @@ def crop(clip, i, j, h, w):
 
 def resize(clip, target_size, interpolation_mode):
     if len(target_size) != 2:
-        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
+        raise ValueError(
+            f"target size should be tuple (height, width), instead got {target_size}"
+        )
     return torch.nn.functional.interpolate(
         clip,
         size=target_size,
@@ -56,7 +62,9 @@ def resize(clip, target_size, interpolation_mode):
 
 def resize_scale(clip, target_size, interpolation_mode):
     if len(target_size) != 2:
-        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
+        raise ValueError(
+            f"target size should be tuple (height, width), instead got {target_size}"
+        )
     H, W = clip.size(-2), clip.size(-1)
     scale_ = target_size[0] / min(H, W)
     return torch.nn.functional.interpolate(
@@ -166,7 +174,8 @@ def normalize_video(clip):
     """
     _is_tensor_video_clip(clip)
     if not clip.dtype == torch.uint8:
-        raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype))
+        raise TypeError("clip tensor should have data type uint8. Got %s" %
+                        str(clip.dtype))
     # return clip.float().permute(3, 0, 1, 2) / 255.0
     return clip.float() / 255.0
 
@@ -227,7 +236,9 @@ class RandomCropVideo:
         th, tw = self.size
 
         if h < th or w < tw:
-            raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}")
+            raise ValueError(
+                f"Required crop size {(th, tw)} is larger than input image size {(h, w)}"
+            )
 
         if w == tw and h == th:
             return 0, 0, h, w
@@ -301,7 +312,9 @@ class LongSideResizeVideo:
         else:
             h = int(h * self.size / w)
             w = self.size
-        resize_clip = resize(clip, target_size=(h, w), interpolation_mode=self.interpolation_mode)
+        resize_clip = resize(clip,
+                             target_size=(h, w),
+                             interpolation_mode=self.interpolation_mode)
         return resize_clip
 
     def __repr__(self) -> str:
@@ -321,7 +334,8 @@ class CenterCropResizeVideo:
         interpolation_mode="bilinear",
     ):
         if len(size) != 2:
-            raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            raise ValueError(
+                f"size should be tuple (height, width), instead got {size}")
         self.size = size
         self.top_crop = top_crop
         self.interpolation_mode = interpolation_mode
@@ -335,7 +349,10 @@ class CenterCropResizeVideo:
                 size is (T, C, crop_size, crop_size)
         """
         # clip_center_crop = center_crop_using_short_edge(clip)
-        clip_center_crop = center_crop_th_tw(clip, self.size[0], self.size[1], top_crop=self.top_crop)
+        clip_center_crop = center_crop_th_tw(clip,
+                                             self.size[0],
+                                             self.size[1],
+                                             top_crop=self.top_crop)
         # import ipdb;ipdb.set_trace()
         clip_center_crop_resize = resize(
             clip_center_crop,
@@ -361,7 +378,8 @@ class UCFCenterCropVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -376,7 +394,9 @@ class UCFCenterCropVideo:
             torch.tensor: scale resized / center cropped video clip.
                 size is (T, C, crop_size, crop_size)
         """
-        clip_resize = resize_scale(clip=clip, target_size=self.size, interpolation_mode=self.interpolation_mode)
+        clip_resize = resize_scale(clip=clip,
+                                   target_size=self.size,
+                                   interpolation_mode=self.interpolation_mode)
         clip_center_crop = center_crop(clip_resize, self.size)
         return clip_center_crop
 
@@ -396,7 +416,8 @@ class KineticsRandomCropResizeVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -405,7 +426,8 @@ class KineticsRandomCropResizeVideo:
 
     def __call__(self, clip):
         clip_random_crop = random_shift_crop(clip)
-        clip_resize = resize(clip_random_crop, self.size, self.interpolation_mode)
+        clip_resize = resize(clip_random_crop, self.size,
+                             self.interpolation_mode)
         return clip_resize
 
 
@@ -418,7 +440,8 @@ class CenterCropVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -514,7 +537,7 @@ class RandomHorizontalFlipVideo:
 #  ------------------------------------------------------------
 #  ---------------------  Sampling  ---------------------------
 #  ------------------------------------------------------------
-class TemporalRandomCrop(object):
+class TemporalRandomCrop:
     """Temporally crop the given frame indices at a random location.
 
     Args:
@@ -531,7 +554,7 @@ class TemporalRandomCrop(object):
         return begin_index, end_index
 
 
-class DynamicSampleDuration(object):
+class DynamicSampleDuration:
     """Temporally crop the given frame indices at a random location.
 
     Args:
@@ -545,7 +568,8 @@ class DynamicSampleDuration(object):
     def __call__(self, t, h, w):
         if self.extra_1:
             t = t - 1
-        truncate_t_list = list(range(t + 1))[t // 2:][::self.t_stride]  # need half at least
+        truncate_t_list = list(
+            range(t + 1))[t // 2:][::self.t_stride]  # need half at least
         truncate_t = random.choice(truncate_t_list)
         if self.extra_1:
             truncate_t = truncate_t + 1
@@ -560,14 +584,18 @@ if __name__ == "__main__":
     from torchvision import transforms
     from torchvision.utils import save_image
 
-    vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi", pts_unit="sec", output_format="TCHW")
+    vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi",
+                                           pts_unit="sec",
+                                           output_format="TCHW")
 
     trans = transforms.Compose([
         Normalize255(),
         RandomHorizontalFlipVideo(),
         UCFCenterCropVideo(512),
         # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5],
+                             std=[0.5, 0.5, 0.5],
+                             inplace=True),
     ])
 
     target_video_len = 32
@@ -582,7 +610,10 @@ if __name__ == "__main__":
     # print(start_frame_ind)
     # print(end_frame_ind)
     assert end_frame_ind - start_frame_ind >= target_video_len
-    frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, target_video_len, dtype=int)
+    frame_indice = np.linspace(start_frame_ind,
+                               end_frame_ind - 1,
+                               target_video_len,
+                               dtype=int)
     print(frame_indice)
 
     select_vframes = vframes[frame_indice]
@@ -593,11 +624,14 @@ if __name__ == "__main__":
     print(select_vframes_trans.shape)
     print(select_vframes_trans.dtype)
 
-    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * 255).to(dtype=torch.uint8)
+    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) *
+                                255).to(dtype=torch.uint8)
     print(select_vframes_trans_int.dtype)
     print(select_vframes_trans_int.permute(0, 2, 3, 1).shape)
 
-    io.write_video("./test.avi", select_vframes_trans_int.permute(0, 2, 3, 1), fps=8)
+    io.write_video("./test.avi",
+                   select_vframes_trans_int.permute(0, 2, 3, 1),
+                   fps=8)
 
     for i in range(target_video_len):
         save_image(

fastvideo/v1/dataset/upload_to_hf.py

@@ -0,0 +1,10 @@
+from huggingface_hub import HfApi, upload_folder
+
+api = HfApi()
+repo_id = "weizhou03/HD-Mixkit-Finetune-Wan"  # customize this
+api.create_repo(repo_id=repo_id, repo_type="dataset")
+
+upload_folder(repo_id=repo_id,
+              folder_path="/workspace/data/HD-Mixkit-Finetune-Wan",
+              repo_type="dataset",
+              path_in_repo="")

fastvideo/v1/default_configs/v1_inference_hunyuan_config.yaml

@@ -1,16 +0,0 @@
-num_gpus: 4
-model_path: FastVideo/FastHunyuan-diffusers
-master_port: 29503
-sp_size: 4 
-tp_size: 4 
-height: 720 
-width: 1280 
-num_frames: 125 
-num_inference_steps: 6 
-guidance_scale: 1 
-embedded_cfg_scale: 6 
-flow_shift: 17 
-prompt_path: ./assets/prompt.txt 
-seed: 1024 
-output_path: outputs_video/ 
-vae-sp: True

fastvideo/v1/distributed/__init__.py

@@ -5,7 +5,8 @@ from fastvideo.v1.distributed.parallel_state import (
     cleanup_dist_env_and_memory, get_sequence_model_parallel_rank,
     get_sequence_model_parallel_world_size, get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size, get_world_group,
-    init_distributed_environment, initialize_model_parallel)
+    init_distributed_environment, initialize_model_parallel,
+    model_parallel_is_initialized)
 from fastvideo.v1.distributed.utils import *
 
 __all__ = [
@@ -17,4 +18,5 @@ __all__ = [
     "get_tensor_model_parallel_world_size",
     "cleanup_dist_env_and_memory",
     "get_world_group",
+    "model_parallel_is_initialized",
 ]

fastvideo/v1/distributed/device_communicators/base_device_communicator.py

@@ -1,16 +1,182 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py
 
-from typing import Optional
+from typing import Any, Optional, Tuple
 
 import torch
 import torch.distributed as dist
-from torch.distributed import ProcessGroup
+from torch import Tensor
+from torch.distributed import ProcessGroup, ReduceOp
+
+
+class DistributedAutograd:
+    """Collection of autograd functions for distributed operations.
+    
+    This class provides custom autograd functions for distributed operations like all_reduce,
+    all_gather, and all_to_all. Each operation is implemented as a static inner class with
+    proper forward and backward implementations.
+    """
+
+    class AllReduce(torch.autograd.Function):
+        """Differentiable all_reduce operation.
+        
+        The gradient of all_reduce is another all_reduce operation since the operation
+        combines values from all ranks equally.
+        """
+
+        @staticmethod
+        def forward(ctx: Any,
+                    group: ProcessGroup,
+                    input_: Tensor,
+                    op: Optional[dist.ReduceOp] = None) -> Tensor:
+            ctx.group = group
+            ctx.op = op
+            output = input_.clone()
+            dist.all_reduce(output, group=group, op=op)
+            return output
+
+        @staticmethod
+        def backward(ctx: Any,
+                     grad_output: Tensor) -> Tuple[None, Tensor, None]:
+            grad_output = grad_output.clone()
+            dist.all_reduce(grad_output, group=ctx.group, op=ctx.op)
+            return None, grad_output, None
+
+    class AllGather(torch.autograd.Function):
+        """Differentiable all_gather operation.
+        
+        The operation gathers tensors from all ranks and concatenates them along a specified dimension.
+        The backward pass uses reduce_scatter to efficiently distribute gradients back to source ranks.
+        """
+
+        @staticmethod
+        def forward(ctx: Any, group: ProcessGroup, input_: Tensor,
+                    world_size: int, dim: int) -> Tensor:
+            ctx.group = group
+            ctx.world_size = world_size
+            ctx.dim = dim
+            ctx.input_shape = input_.shape
+
+            input_size = input_.size()
+            output_size = (input_size[0] * world_size, ) + input_size[1:]
+            output_tensor = torch.empty(output_size,
+                                        dtype=input_.dtype,
+                                        device=input_.device)
+
+            dist.all_gather_into_tensor(output_tensor, input_, group=group)
+
+            output_tensor = output_tensor.reshape((world_size, ) + input_size)
+            output_tensor = output_tensor.movedim(0, dim)
+            output_tensor = output_tensor.reshape(input_size[:dim] +
+                                                  (world_size *
+                                                   input_size[dim], ) +
+                                                  input_size[dim + 1:])
+            return output_tensor
+
+        @staticmethod
+        def backward(ctx: Any,
+                     grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
+            # Split the gradient tensor along the gathered dimension
+            dim_size = grad_output.size(ctx.dim) // ctx.world_size
+            grad_chunks = grad_output.reshape(grad_output.shape[:ctx.dim] +
+                                              (ctx.world_size, dim_size) +
+                                              grad_output.shape[ctx.dim + 1:])
+            grad_chunks = grad_chunks.movedim(ctx.dim, 0)
+
+            # Each rank only needs its corresponding gradient
+            grad_input = torch.empty(ctx.input_shape,
+                                     dtype=grad_output.dtype,
+                                     device=grad_output.device)
+            dist.reduce_scatter_tensor(grad_input,
+                                       grad_chunks.contiguous(),
+                                       group=ctx.group)
+
+            return None, grad_input, None, None
+
+    class AllToAll4D(torch.autograd.Function):
+        """Differentiable all_to_all operation specialized for 4D tensors.
+        
+        This operation is particularly useful for attention operations where we need to
+        redistribute data across ranks for efficient parallel processing.
+        
+        The operation supports two modes:
+        1. scatter_dim=2, gather_dim=1: Used for redistributing attention heads
+        2. scatter_dim=1, gather_dim=2: Used for redistributing sequence dimensions
+        """
+
+        @staticmethod
+        def forward(ctx: Any, group: ProcessGroup, input_: Tensor,
+                    world_size: int, scatter_dim: int,
+                    gather_dim: int) -> Tensor:
+            ctx.group = group
+            ctx.world_size = world_size
+            ctx.scatter_dim = scatter_dim
+            ctx.gather_dim = gather_dim
+
+            if world_size == 1:
+                return input_
+
+            assert input_.dim(
+            ) == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
+
+            if scatter_dim == 2 and gather_dim == 1:
+                bs, shard_seqlen, hc, hs = input_.shape
+                seqlen = shard_seqlen * world_size
+                shard_hc = hc // world_size
+
+                input_t = input_.reshape(bs, shard_seqlen, world_size, shard_hc,
+                                         hs).transpose(0, 2).contiguous()
+                output = torch.empty_like(input_t)
+
+                dist.all_to_all_single(output, input_t, group=group)
+
+                output = output.reshape(seqlen, bs, shard_hc,
+                                        hs).transpose(0, 1).contiguous()
+                output = output.reshape(bs, seqlen, shard_hc, hs)
+
+                return output
+            elif scatter_dim == 1 and gather_dim == 2:
+                bs, seqlen, shard_hc, hs = input_.shape
+                hc = shard_hc * world_size
+                shard_seqlen = seqlen // world_size
+
+                input_t = input_.reshape(bs, world_size, shard_seqlen, shard_hc,
+                                         hs)
+                input_t = input_t.transpose(0, 3).transpose(0, 1).contiguous()
+                input_t = input_t.reshape(world_size, shard_hc, shard_seqlen,
+                                          bs, hs)
+
+                output = torch.empty_like(input_t)
+                dist.all_to_all_single(output, input_t, group=group)
+
+                output = output.reshape(hc, shard_seqlen, bs, hs)
+                output = output.transpose(0, 2).contiguous()
+                output = output.reshape(bs, shard_seqlen, hc, hs)
+
+                return output
+            else:
+                raise RuntimeError(
+                    f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. "
+                    f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported."
+                )
+
+        @staticmethod
+        def backward(
+                ctx: Any,
+                grad_output: Tensor) -> Tuple[None, Tensor, None, None, None]:
+            if ctx.world_size == 1:
+                return None, grad_output, None, None, None
+
+            # For backward pass, we swap scatter_dim and gather_dim
+            output = DistributedAutograd.AllToAll4D.apply(
+                ctx.group, grad_output, ctx.world_size, ctx.gather_dim,
+                ctx.scatter_dim)
+            return None, output, None, None, None
 
 
 class DeviceCommunicatorBase:
     """
-    Base class for device-specific communicator.
+    Base class for device-specific communicator with autograd support.
     It can use the `cpu_group` to initialize the communicator.
     If the device has PyTorch integration (PyTorch can recognize its
     communication backend), the `device_group` will also be given.
@@ -33,35 +199,28 @@ class DeviceCommunicatorBase:
         self.rank_in_group = dist.get_group_rank(self.cpu_group,
                                                  self.global_rank)
 
-    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
-        dist.all_reduce(input_, group=self.device_group)
-        return input_
+    def all_reduce(self,
+                   input_: torch.Tensor,
+                   op: Optional[dist.ReduceOp] = ReduceOp.SUM) -> torch.Tensor:
+        """Performs an all_reduce operation with gradient support."""
+        return DistributedAutograd.AllReduce.apply(self.device_group, input_,
+                                                   op)
 
     def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        """Performs an all_gather operation with gradient support."""
         if dim < 0:
-            # Convert negative dim to positive.
             dim += input_.dim()
-        input_size = input_.size()
-        # NOTE: we have to use concat-style all-gather here,
-        # stack-style all-gather has compatibility issues with
-        # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
-        output_size = (input_size[0] * self.world_size, ) + input_size[1:]
-        # Allocate output tensor.
-        output_tensor = torch.empty(output_size,
-                                    dtype=input_.dtype,
-                                    device=input_.device)
-        # All-gather.
-        dist.all_gather_into_tensor(output_tensor,
-                                    input_,
-                                    group=self.device_group)
-        # Reshape
-        output_tensor = output_tensor.reshape((self.world_size, ) + input_size)
-        output_tensor = output_tensor.movedim(0, dim)
-        output_tensor = output_tensor.reshape(input_size[:dim] +
-                                              (self.world_size *
-                                               input_size[dim], ) +
-                                              input_size[dim + 1:])
-        return output_tensor
+        return DistributedAutograd.AllGather.apply(self.device_group, input_,
+                                                   self.world_size, dim)
+
+    def all_to_all_4D(self,
+                      input_: torch.Tensor,
+                      scatter_dim: int = 2,
+                      gather_dim: int = 1) -> torch.Tensor:
+        """Performs a 4D all-to-all operation with gradient support."""
+        return DistributedAutograd.AllToAll4D.apply(self.device_group, input_,
+                                                    self.world_size,
+                                                    scatter_dim, gather_dim)
 
     def gather(self,
                input_: torch.Tensor,
@@ -95,81 +254,6 @@ class DeviceCommunicatorBase:
             output_tensor = None
         return output_tensor
 
-    def all_to_all_4D(self,
-                      input_: torch.Tensor,
-                      scatter_dim: int = 2,
-                      gather_dim: int = 1) -> torch.Tensor:
-        """Specialized all-to-all operation for 4D tensors (e.g., for QKV matrices).
-        
-        Args:
-            input_ (torch.Tensor): 4D input tensor to be scattered and gathered.
-            scatter_dim (int, optional): Dimension along which to scatter. Defaults to 2.
-            gather_dim (int, optional): Dimension along which to gather. Defaults to 1.
-            
-        Returns:
-            torch.Tensor: Output tensor after all-to-all operation.
-        """
-        # Bypass the function if we are using only 1 GPU.
-        if self.world_size == 1:
-            return input_
-
-        assert input_.dim(
-        ) == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
-
-        if scatter_dim == 2 and gather_dim == 1:
-            # input: (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs)
-            bs, shard_seqlen, hc, hs = input_.shape
-            seqlen = shard_seqlen * self.world_size
-            shard_hc = hc // self.world_size
-
-            # Reshape and transpose for scattering
-            input_t = (input_.reshape(bs, shard_seqlen, self.world_size,
-                                      shard_hc, hs).transpose(0,
-                                                              2).contiguous())
-
-            output = torch.empty_like(input_t)
-
-            torch.distributed.all_to_all_single(output,
-                                                input_t,
-                                                group=self.device_group)
-            torch.cuda.synchronize()
-
-            # Reshape and transpose back
-            output = output.reshape(seqlen, bs, shard_hc,
-                                    hs).transpose(0, 1).contiguous().reshape(
-                                        bs, seqlen, shard_hc, hs)
-
-            return output
-
-        elif scatter_dim == 1 and gather_dim == 2:
-            # input: (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs)
-            bs, seqlen, shard_hc, hs = input_.shape
-            hc = shard_hc * self.world_size
-            shard_seqlen = seqlen // self.world_size
-
-            # Reshape and transpose for scattering
-            input_t = (input_.reshape(bs, self.world_size, shard_seqlen,
-                                      shard_hc, hs).transpose(0, 3).transpose(
-                                          0, 1).contiguous().reshape(
-                                              self.world_size, shard_hc,
-                                              shard_seqlen, bs, hs))
-            output = torch.empty_like(input_t)
-
-            torch.distributed.all_to_all_single(output,
-                                                input_t,
-                                                group=self.device_group)
-            torch.cuda.synchronize()
-
-            # Reshape and transpose back
-            output = output.reshape(hc, shard_seqlen, bs,
-                                    hs).transpose(0, 2).contiguous().reshape(
-                                        bs, shard_seqlen, hc, hs)
-
-            return output
-        else:
-            raise RuntimeError(
-                "scatter_dim must be 1 or 2 and gather_dim must be 1 or 2")
-
     def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
         """Sends a tensor to the destination rank in a non-blocking way"""
         """NOTE: `dst` is the local rank of the destination rank."""

fastvideo/v1/distributed/device_communicators/cuda_communicator.py

@@ -29,17 +29,19 @@ class CudaCommunicator(DeviceCommunicatorBase):
                 device=self.device,
             )
 
-    def all_reduce(self, input_):
+    def all_reduce(self,
+                   input_,
+                   op: Optional[torch.distributed.ReduceOp] = None):
         pynccl_comm = self.pynccl_comm
         assert pynccl_comm is not None
-        out = pynccl_comm.all_reduce(input_)
+        out = pynccl_comm.all_reduce(input_, op=op)
         if out is None:
             # fall back to the default all-reduce using PyTorch.
             # this usually happens during testing.
             # when we run the model, allreduce only happens for the TP
             # group, where we always have either custom allreduce or pynccl.
             out = input_.clone()
-            torch.distributed.all_reduce(out, group=self.device_group)
+            torch.distributed.all_reduce(out, group=self.device_group, op=op)
         return out
 
     def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:

fastvideo/v1/distributed/parallel_state.py

@@ -35,7 +35,7 @@ from unittest.mock import patch
 
 import torch
 import torch.distributed
-from torch.distributed import Backend, ProcessGroup
+from torch.distributed import Backend, ProcessGroup, ReduceOp
 
 import fastvideo.v1.envs as envs
 from fastvideo.v1.distributed.device_communicators.base_device_communicator import (
@@ -260,7 +260,11 @@ class GroupCoordinator:
         with torch.cuda.stream(stream):
             yield graph_capture_context
 
-    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
+    def all_reduce(
+        self,
+        input_: torch.Tensor,
+        op: Optional[torch.distributed.ReduceOp] = ReduceOp.SUM
+    ) -> torch.Tensor:
         """
         User-facing all-reduce function before we actually call the
         all-reduce operation.
@@ -283,10 +287,14 @@ class GroupCoordinator:
             return torch.ops.vllm.all_reduce(input_,
                                              group_name=self.unique_name)
         else:
-            return self._all_reduce_out_place(input_)
+            return self._all_reduce_out_place(input_, op=op)
 
-    def _all_reduce_out_place(self, input_: torch.Tensor) -> torch.Tensor:
-        return self.device_communicator.all_reduce(input_)
+    def _all_reduce_out_place(
+        self,
+        input_: torch.Tensor,
+        op: Optional[torch.distributed.ReduceOp] = ReduceOp.SUM
+    ) -> torch.Tensor:
+        return self.device_communicator.all_reduce(input_, op=op)
 
     def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
         world_size = self.world_size

fastvideo/v1/distributed/utils.py

@@ -94,14 +94,14 @@ class StatelessProcessGroup:
         key = f"send_to/{dst}/{self.send_dst_counter[dst]}"
         self.store.set(key, pickle.dumps(obj))
         self.send_dst_counter[dst] += 1
-        self.entries.append((key, time.time()))
+        self.entries.append((key, time.perf_counter()))
 
     def expire_data(self) -> None:
         """Expire data that is older than `data_expiration_seconds` seconds."""
         while self.entries:
             # check the oldest entry
             key, timestamp = self.entries[0]
-            if time.time() - timestamp > self.data_expiration_seconds:
+            if time.perf_counter() - timestamp > self.data_expiration_seconds:
                 self.store.delete_key(key)
                 self.entries.popleft()
             else:
@@ -125,7 +125,7 @@ class StatelessProcessGroup:
                    f"{self.broadcast_send_counter}")
             self.store.set(key, pickle.dumps(obj))
             self.broadcast_send_counter += 1
-            self.entries.append((key, time.time()))
+            self.entries.append((key, time.perf_counter()))
             return obj
         else:
             key = (f"broadcast_from/{src}/"

fastvideo/v1/entrypoints/cli/generate.py

@@ -2,13 +2,20 @@
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py
 
 import argparse
-from typing import List, cast
+import dataclasses
+import os
+from typing import Any, Dict, List, Optional, cast
 
-from fastvideo.v1.entrypoints.cli import utils
+from fastvideo import PipelineConfig, VideoGenerator
+from fastvideo.v1.configs.sample.base import SamplingParam
 from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
+from fastvideo.v1.entrypoints.cli.utils import RaiseNotImplementedAction
 from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import FlexibleArgumentParser
 
+logger = init_logger(__name__)
+
 
 class GenerateSubcommand(CLISubcommand):
     """The `generate` subcommand for the FastVideo CLI"""
@@ -16,45 +23,80 @@ class GenerateSubcommand(CLISubcommand):
     def __init__(self) -> None:
         self.name = "generate"
         super().__init__()
+        self.init_arg_names = self._get_init_arg_names()
+        self.generation_arg_names = self._get_generation_arg_names()
+
+    def _get_init_arg_names(self) -> List[str]:
+        """Get names of arguments for VideoGenerator initialization"""
+        return ["num_gpus", "tp_size", "sp_size", "model_path"]
+
+    def _get_generation_arg_names(self) -> List[str]:
+        """Get names of arguments for generate_video method"""
+        return [field.name for field in dataclasses.fields(SamplingParam)]
 
     def cmd(self, args: argparse.Namespace) -> None:
-        excluded_args = [
-            'subparser', 'config', 'num_gpus', 'master_port',
-            'dispatch_function'
-        ]
+        excluded_args = ['subparser', 'config', 'dispatch_function']
 
-        # Create a filtered dictionary of arguments
-        filtered_args = {
+        FastVideoArgs.from_cli_args(args)
+
+        provided_args = {}
+        for k, v in vars(args).items():
+            if (k not in excluded_args and v is not None
+                    and hasattr(args, '_provided') and k in args._provided):
+                provided_args[k] = v
+
+        if 'model_path' in vars(args) and args.model_path is not None:
+            provided_args['model_path'] = args.model_path
+
+        if 'prompt' in vars(args) and args.prompt is not None:
+            provided_args['prompt'] = args.prompt
+
+        merged_args = {**provided_args}
+
+        logger.info('CLI Args: %s', merged_args)
+
+        if 'model_path' not in merged_args or not merged_args['model_path']:
+            raise ValueError(
+                "model_path must be provided either in config file or via --model-path"
+            )
+
+        if 'prompt' not in merged_args or not merged_args['prompt']:
+            raise ValueError(
+                "prompt must be provided either in config file or via --prompt")
+
+        init_args = {
             k: v
-            for k, v in vars(args).items()
-            if k not in excluded_args and v is not None
+            for k, v in merged_args.items() if k in self.init_arg_names
+        }
+        generation_args = {
+            k: v
+            for k, v in merged_args.items() if k in self.generation_arg_names
         }
 
-        main_args = []
+        pipeline_config = PipelineConfig.from_pretrained(
+            merged_args['model_path'])
 
-        for key, value in filtered_args.items():
-            # Convert underscores to dashes in argument names
-            arg_name = f"--{key.replace('_', '-')}"
+        update_config_from_args(pipeline_config.dit_config, merged_args,
+                                "dit_config")
+        update_config_from_args(pipeline_config.vae_config, merged_args,
+                                "vae_config")
+        update_config_from_args(pipeline_config, merged_args)
 
-            # Handle boolean flags
-            if isinstance(value, bool):
-                if value:
-                    main_args.append(arg_name)
-            else:
-                main_args.append(arg_name)
-                main_args.append(str(value))
+        model_path = init_args.pop('model_path')
+        prompt = generation_args.pop('prompt')
 
-        utils.launch_distributed(args.num_gpus,
-                                 main_args,
-                                 master_port=args.master_port)
+        generator = VideoGenerator.from_pretrained(
+            model_path=model_path, **init_args, pipeline_config=pipeline_config)
+
+        generator.generate_video(prompt=prompt, **generation_args)
 
     def validate(self, args: argparse.Namespace) -> None:
+        """Validate the arguments for this command"""
         if args.num_gpus is not None and args.num_gpus <= 0:
             raise ValueError("Number of gpus must be positive")
 
-        if args.master_port is not None and (args.master_port < 1024
-                                             or args.master_port > 65535):
-            raise ValueError("Master port must be between 1024 and 65535")
+        if args.config and not os.path.exists(args.config):
+            raise ValueError(f"Config file not found: {args.config}")
 
     def subparser_init(
             self,
@@ -63,7 +105,7 @@ class GenerateSubcommand(CLISubcommand):
             "generate",
             help="Run inference on a model",
             usage=
-            "fastvideo generate --model-path MODEL_PATH_OR_ID --prompt PROMPT [OPTIONS]"
+            "fastvideo generate (--model-path MODEL_PATH_OR_ID --prompt PROMPT) | --config CONFIG_FILE [OPTIONS]"
         )
 
         generate_parser.add_argument(
@@ -71,17 +113,53 @@ class GenerateSubcommand(CLISubcommand):
             type=str,
             default='',
             required=False,
-            help="Read CLI options from a config YAML file.")
-
-        generate_parser.add_argument("--master-port",
-                                     type=int,
-                                     default=None,
-                                     help="Port for the master process")
+            help=
+            "Read CLI options from a config JSON or YAML file. If provided, --model-path and --prompt are optional."
+        )
 
         generate_parser = FastVideoArgs.add_cli_args(generate_parser)
+        generate_parser = SamplingParam.add_cli_args(generate_parser)
+
+        generate_parser.add_argument(
+            "--text-encoder-configs",
+            action=RaiseNotImplementedAction,
+            help=
+            "JSON array of text encoder configurations (NOT YET IMPLEMENTED)",
+        )
 
         return cast(FlexibleArgumentParser, generate_parser)
 
 
 def cmd_init() -> List[CLISubcommand]:
     return [GenerateSubcommand()]
+
+
+def update_config_from_args(config: Any,
+                            args_dict: Dict[str, Any],
+                            prefix: Optional[str] = None) -> None:
+    """
+    Update configuration object from arguments dictionary.
+    
+    Args:
+        config: The configuration object to update
+        args_dict: Dictionary containing arguments
+        prefix: Prefix for the configuration parameters in the args_dict.
+               If None, assumes direct attribute mapping without prefix.
+    """
+    # Handle top-level attributes (no prefix)
+    if prefix is None:
+        for key, value in args_dict.items():
+            if hasattr(config, key) and value is not None:
+                if key == "text_encoder_precisions" and isinstance(value, list):
+                    setattr(config, key, tuple(value))
+                else:
+                    setattr(config, key, value)
+        return
+
+    # Handle nested attributes with prefix
+    prefix_with_dot = f"{prefix}."
+    for key, value in args_dict.items():
+        if key.startswith(prefix_with_dot) and value is not None:
+            attr_name = key[len(prefix_with_dot):]
+            if hasattr(config, attr_name):
+                setattr(config, attr_name, value)

fastvideo/v1/entrypoints/cli/utils.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import argparse
 import os
 import subprocess
 import sys
@@ -10,6 +11,13 @@ from fastvideo.v1.logger import init_logger
 logger = init_logger(__name__)
 
 
+class RaiseNotImplementedAction(argparse.Action):
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        raise NotImplementedError(
+            f"The {option_string} option is not yet implemented")
+
+
 def launch_distributed(num_gpus: int,
                        args: List[str],
                        master_port: Optional[int] = None) -> int:

fastvideo/v1/entrypoints/data_preprocessor.py

@@ -0,0 +1,29 @@
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from fastvideo.v1.pipelines.wan.wan_latent_pipeline import WanLatentPipeline
+
+
+def main():
+    print("Starting data preprocessor")
+    pipeline = WanLatentPipeline.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+
+    train_dataset = getdataset(args)
+    sampler = DistributedSampler(train_dataset,
+                                 rank=local_rank,
+                                 num_replicas=world_size,
+                                 shuffle=True)
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+    )
+
+    for batch in train_dataloader:
+        pipeline(batch)
+
+
+if __name__ == "__main__":
+    main()

fastvideo/v1/entrypoints/video_generator.py

@@ -7,6 +7,7 @@ diffusion models.
 """
 
 import gc
+import math
 import os
 import time
 from typing import Any, Dict, List, Optional, Union
@@ -72,7 +73,6 @@ class VideoGenerator:
 
         Priority level: Default pipeline config < User's pipeline config < User's kwargs
         """
-
         config = None
         # 1. If users provide a pipeline config, it will override the default pipeline config
         if isinstance(pipeline_config, PipelineConfig):
@@ -181,12 +181,46 @@ class VideoGenerator:
                 f"height={sampling_param.height}, width={sampling_param.width}, "
                 f"num_frames={sampling_param.num_frames}")
 
-        if (
-                sampling_param.num_frames - 1
-        ) % fastvideo_args.vae_config.arch_config.temporal_compression_ratio != 0:
-            raise ValueError(
-                f"num_frames-1 must be a multiple of {fastvideo_args.vae_config.arch_config.temporal_compression_ratio}, got {sampling_param.num_frames}"
-            )
+        temporal_scale_factor = fastvideo_args.vae_config.arch_config.temporal_compression_ratio
+        num_frames = sampling_param.num_frames
+        num_gpus = fastvideo_args.num_gpus
+        use_temporal_scaling_frames = fastvideo_args.vae_config.use_temporal_scaling_frames
+
+        # Adjust number of frames based on number of GPUs
+        if use_temporal_scaling_frames:
+            orig_latent_num_frames = (num_frames -
+                                      1) // temporal_scale_factor + 1
+        else:  # stepvideo only
+            orig_latent_num_frames = sampling_param.num_frames // 17 * 3
+
+        if orig_latent_num_frames % fastvideo_args.num_gpus != 0:
+            # Adjust latent frames to be divisible by number of GPUs
+            if sampling_param.num_frames_round_down:
+                # Ensure we have at least 1 batch per GPU
+                new_latent_num_frames = max(
+                    1, (orig_latent_num_frames // num_gpus)) * num_gpus
+            else:
+                new_latent_num_frames = math.ceil(
+                    orig_latent_num_frames / num_gpus) * num_gpus
+
+            if use_temporal_scaling_frames:
+                # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor
+                new_num_frames = (new_latent_num_frames -
+                                  1) * temporal_scale_factor + 1
+            else:  # stepvideo only
+                # Find the least common multiple of 3 and num_gpus
+                divisor = math.lcm(3, num_gpus)
+                # Round up to the nearest multiple of this LCM
+                new_latent_num_frames = (
+                    (new_latent_num_frames + divisor - 1) // divisor) * divisor
+                # Convert back to actual frames using the StepVideo formula
+                new_num_frames = new_latent_num_frames // 3 * 17
+
+            logger.info(
+                "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
+                sampling_param.num_frames, new_num_frames,
+                fastvideo_args.num_gpus)
+            sampling_param.num_frames = new_num_frames
 
         # Calculate sizes
         target_height = align_to(sampling_param.height, 16)
@@ -225,11 +259,11 @@ class VideoGenerator:
         )
 
         # Run inference
-        start_time = time.time()
+        start_time = time.perf_counter()
         output_batch = self.executor.execute_forward(batch, fastvideo_args)
         samples = output_batch
 
-        gen_time = time.time() - start_time
+        gen_time = time.perf_counter() - start_time
         logger.info("Generated successfully in %.2f seconds", gen_time)
 
         # Process outputs

fastvideo/v1/fastvideo_args.py

@@ -10,7 +10,7 @@ from typing import Any, Callable, List, Optional, Tuple
 
 from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import FlexibleArgumentParser
+from fastvideo.v1.utils import FlexibleArgumentParser, StoreBoolean
 
 logger = init_logger(__name__)
 
@@ -68,8 +68,12 @@ class FastVideoArgs:
     image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)
 
     # Text encoder configuration
+    DEFAULT_TEXT_ENCODER_PRECISIONS = (
+        "fp16",
+        # "fp16",
+    )
     text_encoder_precisions: Tuple[str, ...] = field(
-        default_factory=lambda: ("fp16", ))
+        default_factory=lambda: FastVideoArgs.DEFAULT_TEXT_ENCODER_PRECISIONS)
     text_encoder_configs: Tuple[EncoderConfig, ...] = field(
         default_factory=lambda: (EncoderConfig(), ))
     preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
@@ -84,6 +88,11 @@ class FastVideoArgs:
     use_cpu_offload: bool = False
     disable_autocast: bool = False
 
+    # StepVideo specific parameters
+    pos_magic: Optional[str] = None
+    neg_magic: Optional[str] = None
+    timesteps_scale: Optional[bool] = None
+
     # Logging
     log_level: str = "info"
 
@@ -91,6 +100,10 @@ class FastVideoArgs:
     device_str: Optional[str] = None
     device = None
 
+    @property
+    def training_mode(self) -> bool:
+        return not self.inference_mode
+
     def __post_init__(self):
         pass
 
@@ -100,7 +113,6 @@ class FastVideoArgs:
         parser.add_argument(
             "--model-path",
             type=str,
-            required=True,
             help=
             "The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
         )
@@ -124,10 +136,17 @@ class FastVideoArgs:
             help="The distributed executor backend to use",
         )
 
+        parser.add_argument(
+            "--inference-mode",
+            action=StoreBoolean,
+            default=FastVideoArgs.inference_mode,
+            help="Whether to use inference mode",
+        )
+
         # HuggingFace specific parameters
         parser.add_argument(
             "--trust-remote-code",
-            action="store_true",
+            action=StoreBoolean,
             default=FastVideoArgs.trust_remote_code,
             help="Trust remote code when loading HuggingFace models",
         )
@@ -206,21 +225,21 @@ class FastVideoArgs:
         )
         parser.add_argument(
             "--vae-tiling",
-            action="store_true",
+            action=StoreBoolean,
             default=FastVideoArgs.vae_tiling,
             help="Enable VAE tiling",
         )
         parser.add_argument(
             "--vae-sp",
-            action="store_true",
+            action=StoreBoolean,
             help="Enable VAE spatial parallelism",
         )
 
         parser.add_argument(
-            "--text-encoder-precision",
+            "--text-encoder-precisions",
             nargs="+",
             type=str,
-            default=FastVideoArgs.text_encoder_precisions,
+            default=FastVideoArgs.DEFAULT_TEXT_ENCODER_PRECISIONS,
             choices=["fp32", "fp16", "bf16"],
             help="Precision for each text encoder",
         )
@@ -242,23 +261,42 @@ class FastVideoArgs:
         )
         parser.add_argument(
             "--enable-torch-compile",
-            action="store_true",
+            action=StoreBoolean,
             help=
             "Use torch.compile for speeding up STA inference without teacache",
         )
 
         parser.add_argument(
             "--use-cpu-offload",
-            action="store_true",
+            action=StoreBoolean,
             help="Use CPU offload for the model load",
         )
         parser.add_argument(
             "--disable-autocast",
-            action="store_true",
+            action=StoreBoolean,
             help=
             "Disable autocast for denoising loop and vae decoding in pipeline sampling",
         )
 
+        parser.add_argument(
+            "--pos_magic",
+            type=str,
+            default=FastVideoArgs.pos_magic,
+            help="Positive magic prompt for sampling",
+        )
+        parser.add_argument(
+            "--neg_magic",
+            type=str,
+            default=FastVideoArgs.neg_magic,
+            help="Negative magic prompt for sampling",
+        )
+        parser.add_argument(
+            "--timesteps_scale",
+            type=bool,
+            default=FastVideoArgs.timesteps_scale,
+            help="Bool for applying scheduler scale in set_timesteps",
+        )
+
         # Logging
         parser.add_argument(
             "--log-level",
@@ -267,6 +305,14 @@ class FastVideoArgs:
             help="The logging level of all loggers.",
         )
 
+        # Add VAE configuration arguments
+        from fastvideo.v1.configs.models.vaes.base import VAEConfig
+        VAEConfig.add_cli_args(parser)
+
+        # Add DiT configuration arguments
+        from fastvideo.v1.configs.models.dits.base import DiTConfig
+        DiTConfig.add_cli_args(parser)
+
         return parser
 
     @classmethod
@@ -388,3 +434,349 @@ def get_current_fastvideo_args() -> FastVideoArgs:
         # TODO(will): may need to handle this for CI.
         raise ValueError("Current fastvideo args is not set.")
     return _current_fastvideo_args
+
+
+@dataclasses.dataclass
+class TrainingArgs(FastVideoArgs):
+    data_path: str = ""
+    dataloader_num_workers: int = 0
+    num_height: int = 0
+    num_width: int = 0
+    num_frames: int = 0
+
+    train_batch_size: int = 0
+    num_latent_t: int = 0
+    group_frame: bool = False
+    group_resolution: bool = False
+
+    # text encoder & vae & diffusion model
+    pretrained_model_name_or_path: str = ""
+    dit_model_name_or_path: str = ""
+    cache_dir: str = ""
+
+    # diffusion setting
+    ema_decay: float = 0.0
+    ema_start_step: int = 0
+    cfg: float = 0.0
+    precondition_outputs: bool = False
+
+    # validation & logs
+    validation_prompt_dir: str = ""
+    validation_sampling_steps: str = ""
+    validation_guidance_scale: str = ""
+    validation_steps: float = 0.0
+    log_validation: bool = False
+    tracker_project_name: str = ""
+    # seed: int
+
+    # output
+    output_dir: str = ""
+    checkpoints_total_limit: int = 0
+    checkpointing_steps: int = 0
+    resume_from_checkpoint: str = ""
+    resume_from_lora_checkpoint: str = ""
+    logging_dir: str = ""
+
+    # optimizer & scheduler
+    num_train_epochs: int = 0
+    max_train_steps: int = 0
+    gradient_accumulation_steps: int = 0
+    learning_rate: float = 0.0
+    scale_lr: bool = False
+    lr_scheduler: str = ""
+    lr_warmup_steps: int = 0
+    max_grad_norm: float = 0.0
+    gradient_checkpointing: bool = False
+    selective_checkpointing: float = 0.0
+    allow_tf32: bool = False
+    mixed_precision: str = ""
+    use_cpu_offload: bool = False
+    # fp16_full_eval: bool
+    # fp16_backend: str
+    train_sp_batch_size: int = 0
+    use_lora: bool = False
+    lora_alpha: int = 0
+    lora_rank: int = 0
+    fsdp_sharding_startegy: str = ""
+
+    weighting_scheme: str = ""
+    logit_mean: float = 0.0
+    logit_std: float = 1.0
+    mode_scale: float = 0.0
+
+    # lr_scheduler
+    lr_scheduler: str = ""
+    num_euler_timesteps: int = 0
+    lr_num_cycles: int = 0
+    lr_power: float = 0.0
+    not_apply_cfg_solver: bool = False
+    distill_cfg: float = 0.0
+    scheduler_type: str = ""
+    linear_quadratic_threshold: float = 0.0
+    linear_range: float = 0.0
+    weight_decay: float = 0.0
+    use_ema: bool = False
+    multi_phased_distill_schedule: str = ""
+    pred_decay_weight: float = 0.0
+    pred_decay_type: str = ""
+    hunyuan_teacher_disable_cfg: bool = False
+
+    # master_weight_type
+    master_weight_type: str = ""
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
+        # Get all fields from the dataclass
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+
+        # Create a dictionary of attribute values, with defaults for missing attributes
+        kwargs = {}
+        for attr in attrs:
+            # Handle renamed attributes or those with multiple CLI names
+            if attr == 'tp_size' and hasattr(args, 'tensor_parallel_size'):
+                kwargs[attr] = args.tensor_parallel_size
+            elif attr == 'sp_size' and hasattr(args, 'sequence_parallel_size'):
+                kwargs[attr] = args.sequence_parallel_size
+            elif attr == 'flow_shift' and hasattr(args, 'shift'):
+                kwargs[attr] = args.shift
+            # Use getattr with default value from the dataclass for potentially missing attributes
+            else:
+                default_value = getattr(cls, attr, None)
+                kwargs[attr] = getattr(args, attr, default_value)
+
+        return cls(**kwargs)
+
+    @staticmethod
+    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
+        parser.add_argument("--data-path",
+                            type=str,
+                            required=True,
+                            help="Path to parquet files")
+        parser.add_argument("--dataloader-num-workers",
+                            type=int,
+                            required=True,
+                            help="Number of workers for dataloader")
+        parser.add_argument("--num-height",
+                            type=int,
+                            required=True,
+                            help="Number of heights")
+        parser.add_argument("--num-width",
+                            type=int,
+                            required=True,
+                            help="Number of widths")
+        parser.add_argument("--num-frames",
+                            type=int,
+                            required=True,
+                            help="Number of frames")
+
+        # Training batch and model configuration
+        parser.add_argument("--train-batch-size",
+                            type=int,
+                            required=True,
+                            help="Training batch size")
+        parser.add_argument("--num-latent-t",
+                            type=int,
+                            required=True,
+                            help="Number of latent time steps")
+        parser.add_argument("--group-frame",
+                            action=StoreBoolean,
+                            help="Whether to group frames during training")
+        parser.add_argument("--group-resolution",
+                            action=StoreBoolean,
+                            help="Whether to group resolutions during training")
+
+        # Model paths
+        parser.add_argument("--pretrained-model-name-or-path",
+                            type=str,
+                            required=True,
+                            help="Path to pretrained model or model name")
+        parser.add_argument("--dit-model-name-or-path",
+                            type=str,
+                            required=False,
+                            help="Path to DiT model or model name")
+        parser.add_argument("--cache-dir",
+                            type=str,
+                            help="Directory to cache models")
+
+        # Diffusion settings
+        parser.add_argument("--ema-decay",
+                            type=float,
+                            default=0.999,
+                            help="EMA decay rate")
+        parser.add_argument("--ema-start-step",
+                            type=int,
+                            default=0,
+                            help="Step to start EMA")
+        parser.add_argument("--cfg",
+                            type=float,
+                            help="Classifier-free guidance scale")
+        parser.add_argument(
+            "--precondition-outputs",
+            action=StoreBoolean,
+            help="Whether to precondition the outputs of the model")
+
+        # Validation and logging
+        parser.add_argument("--validation-prompt-dir",
+                            type=str,
+                            help="Directory containing validation prompts")
+        parser.add_argument("--validation-sampling-steps",
+                            type=str,
+                            help="Validation sampling steps")
+        parser.add_argument("--validation-guidance-scale",
+                            type=str,
+                            help="Validation guidance scale")
+        parser.add_argument("--validation-steps",
+                            type=float,
+                            help="Number of validation steps")
+        parser.add_argument("--log-validation",
+                            action=StoreBoolean,
+                            help="Whether to log validation results")
+        parser.add_argument("--tracker-project-name",
+                            type=str,
+                            help="Project name for tracking")
+
+        # Output configuration
+        parser.add_argument("--output-dir",
+                            type=str,
+                            required=True,
+                            help="Output directory for checkpoints and logs")
+        parser.add_argument("--checkpoints-total-limit",
+                            type=int,
+                            help="Maximum number of checkpoints to keep")
+        parser.add_argument("--checkpointing-steps",
+                            type=int,
+                            help="Steps between checkpoints")
+        parser.add_argument("--resume-from-checkpoint",
+                            type=str,
+                            help="Path to checkpoint to resume from")
+        parser.add_argument("--resume-from-lora-checkpoint",
+                            type=str,
+                            help="Path to LoRA checkpoint to resume from")
+        parser.add_argument("--logging-dir",
+                            type=str,
+                            help="Directory for logging")
+
+        # Training configuration
+        parser.add_argument("--num-train-epochs",
+                            type=int,
+                            help="Number of training epochs")
+        parser.add_argument("--max-train-steps",
+                            type=int,
+                            help="Maximum number of training steps")
+        parser.add_argument("--gradient-accumulation-steps",
+                            type=int,
+                            help="Number of steps to accumulate gradients")
+        parser.add_argument("--learning-rate",
+                            type=float,
+                            required=True,
+                            help="Learning rate")
+        parser.add_argument("--scale-lr",
+                            action=StoreBoolean,
+                            help="Whether to scale learning rate")
+        parser.add_argument("--lr-scheduler",
+                            type=str,
+                            default="constant",
+                            help="Learning rate scheduler type")
+        parser.add_argument("--lr-warmup-steps",
+                            type=int,
+                            default=10,
+                            help="Number of warmup steps for learning rate")
+        parser.add_argument("--max-grad-norm",
+                            type=float,
+                            help="Maximum gradient norm")
+        parser.add_argument("--gradient-checkpointing",
+                            action=StoreBoolean,
+                            help="Whether to use gradient checkpointing")
+        parser.add_argument("--selective-checkpointing",
+                            type=float,
+                            help="Selective checkpointing threshold")
+        parser.add_argument("--allow-tf32",
+                            action=StoreBoolean,
+                            help="Whether to allow TF32")
+        parser.add_argument("--mixed-precision",
+                            type=str,
+                            help="Mixed precision training type")
+        parser.add_argument("--train-sp-batch-size",
+                            type=int,
+                            help="Training spatial parallelism batch size")
+
+        # LoRA configuration
+        parser.add_argument("--use-lora",
+                            action=StoreBoolean,
+                            help="Whether to use LoRA")
+        parser.add_argument("--lora-alpha",
+                            type=int,
+                            help="LoRA alpha parameter")
+        parser.add_argument("--lora-rank", type=int, help="LoRA rank")
+        parser.add_argument("--fsdp-sharding-strategy",
+                            type=str,
+                            help="FSDP sharding strategy")
+
+        parser.add_argument(
+            "--weighting_scheme",
+            type=str,
+            default="uniform",
+            choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"],
+        )
+        parser.add_argument(
+            "--logit_mean",
+            type=float,
+            default=0.0,
+            help="mean to use when using the `'logit_normal'` weighting scheme.",
+        )
+        parser.add_argument(
+            "--logit_std",
+            type=float,
+            default=1.0,
+            help="std to use when using the `'logit_normal'` weighting scheme.",
+        )
+        parser.add_argument(
+            "--mode_scale",
+            type=float,
+            default=1.29,
+            help=
+            "Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
+        )
+
+        # Additional training parameters
+        parser.add_argument("--num-euler-timesteps",
+                            type=int,
+                            help="Number of Euler timesteps")
+        parser.add_argument("--lr-num-cycles",
+                            type=int,
+                            help="Number of learning rate cycles")
+        parser.add_argument("--lr-power",
+                            type=float,
+                            help="Learning rate power")
+        parser.add_argument("--not-apply-cfg-solver",
+                            action=StoreBoolean,
+                            help="Whether to not apply CFG solver")
+        parser.add_argument("--distill-cfg",
+                            type=float,
+                            help="Distillation CFG scale")
+        parser.add_argument("--scheduler-type", type=str, help="Scheduler type")
+        parser.add_argument("--linear-quadratic-threshold",
+                            type=float,
+                            help="Linear quadratic threshold")
+        parser.add_argument("--linear-range", type=float, help="Linear range")
+        parser.add_argument("--weight-decay", type=float, help="Weight decay")
+        parser.add_argument("--use-ema",
+                            action=StoreBoolean,
+                            help="Whether to use EMA")
+        parser.add_argument("--multi-phased-distill-schedule",
+                            type=str,
+                            help="Multi-phased distillation schedule")
+        parser.add_argument("--pred-decay-weight",
+                            type=float,
+                            help="Prediction decay weight")
+        parser.add_argument("--pred-decay-type",
+                            type=str,
+                            help="Prediction decay type")
+        parser.add_argument("--hunyuan-teacher-disable-cfg",
+                            action=StoreBoolean,
+                            help="Whether to disable CFG for Hunyuan teacher")
+        parser.add_argument("--master-weight-type",
+                            type=str,
+                            help="Master weight type")
+
+        return parser

fastvideo/v1/inference_engine.py

@@ -191,7 +191,7 @@ class InferenceEngine:
         # ========================================================================
         # Pipeline inference
         # ========================================================================
-        start_time = time.time()
+        start_time = time.perf_counter()
         samples = self.pipeline.forward(
             batch=batch,
             fastvideo_args=fastvideo_args,
@@ -201,7 +201,7 @@ class InferenceEngine:
         out_dict["samples"] = samples
         out_dict["prompts"] = prompt
 
-        gen_time = time.time() - start_time
+        gen_time = time.perf_counter() - start_time
         logger.info("Success, time: %s", gen_time)
 
         return out_dict

fastvideo/v1/layers/activation.py

@@ -9,7 +9,6 @@ import torch.nn.functional as F
 
 # TODO (will): remove this dependency
 from fastvideo.v1.layers.custom_op import CustomOp
-from fastvideo.v1.platforms import current_platform
 
 
 @CustomOp.register("silu_and_mul")
@@ -25,21 +24,12 @@ class SiluAndMul(CustomOp):
 
     def __init__(self) -> None:
         super().__init__()
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            self.op = torch.ops._C.silu_and_mul
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         d = x.shape[-1] // 2
         return F.silu(x[..., :d]) * x[..., d:]
 
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = (x.shape[:-1] + (d, ))
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x)
-        return out
-
 
 @CustomOp.register("gelu_and_mul")
 class GeluAndMul(CustomOp):
@@ -57,24 +47,12 @@ class GeluAndMul(CustomOp):
         self.approximate = approximate
         if approximate not in ("none", "tanh"):
             raise ValueError(f"Unknown approximate mode: {approximate}")
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            if approximate == "none":
-                self.op = torch.ops._C.gelu_and_mul
-            elif approximate == "tanh":
-                self.op = torch.ops._C.gelu_tanh_and_mul
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         d = x.shape[-1] // 2
         return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:]
 
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = (x.shape[:-1] + (d, ))
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x)
-        return out
-
     def extra_repr(self) -> str:
         return f'approximate={repr(self.approximate)}'
 
@@ -84,8 +62,6 @@ class NewGELU(CustomOp):
 
     def __init__(self):
         super().__init__()
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            self.op = torch.ops._C.gelu_new
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
@@ -93,32 +69,17 @@ class NewGELU(CustomOp):
         return 0.5 * x * (1.0 + torch.tanh(c *
                                            (x + 0.044715 * torch.pow(x, 3.0))))
 
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        self.op(out, x)
-        return out
-
-    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
-        return self.op(x)
-
 
 @CustomOp.register("quick_gelu")
 class QuickGELU(CustomOp):
     # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90
     def __init__(self):
         super().__init__()
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            self.op = torch.ops._C.gelu_quick
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         return x * torch.sigmoid(1.702 * x)
 
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        self.op(out, x)
-        return out
-
 
 _ACTIVATION_REGISTRY = {
     "gelu": nn.GELU,

fastvideo/v1/layers/custom_op.py

@@ -50,6 +50,10 @@ class CustomOp(nn.Module):
         return self.forward_native(*args, **kwargs)
 
     def dispatch_forward(self) -> Callable:
+        # FIXME(will): for now, we always use the native implementation, since
+        # forward_cuda is using vllm's custom ops and it doesn't support
+        # backwards. We should add our own custom ops that support backwards.
+        return self.forward_native
         # NOTE(woosuk): Here we assume that vLLM was built for only one
         # specific backend. Currently, we do not support dynamic dispatching.
         enabled = self.enabled()

fastvideo/v1/layers/layernorm.py

@@ -75,33 +75,6 @@ class RMSNorm(CustomOp):
         else:
             return x, residual
 
-    def forward_cuda(
-        self,
-        x: torch.Tensor,
-        residual: Optional[torch.Tensor] = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        if self.variance_size_override is not None:
-            return self.forward_native(x, residual)
-
-        from vllm import _custom_ops as ops
-
-        if residual is not None:
-            ops.fused_add_rms_norm(
-                x,
-                residual,
-                self.weight.data,
-                self.variance_epsilon,
-            )
-            return x, residual
-        out = torch.empty_like(x)
-        ops.rms_norm(
-            out,
-            x,
-            self.weight.data,
-            self.variance_epsilon,
-        )
-        return out
-
     def extra_repr(self) -> str:
         s = f"hidden_size={self.weight.data.size(0)}"
         s += f", eps={self.variance_epsilon}"
@@ -173,7 +146,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
         # Apply normalization
         normalized = self.norm(residual_output)
         # Apply scale and shift
-        modulated = normalized * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        modulated = normalized * (1.0 + scale) + shift
         return modulated, residual_output
 
 
@@ -209,4 +182,4 @@ class LayerNormScaleShift(nn.Module):
                 scale: torch.Tensor) -> torch.Tensor:
         """Apply ln followed by scale and shift in a single fused operation."""
         normalized = self.norm(x)
-        return normalized * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        return normalized * (1.0 + scale) + shift

fastvideo/v1/layers/linear.py

@@ -7,16 +7,14 @@ from typing import Optional, Union
 import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
-# TODO(will): remove this import by copying the definition from vLLM then
-# manually import each quantization method we want to use. Refer to SGLang
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig, QuantizeMethodBase)
 
 from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank,
                                       get_tensor_model_parallel_world_size,
                                       split_tensor_along_last_dim,
                                       tensor_model_parallel_all_gather,
                                       tensor_model_parallel_all_reduce)
+from fastvideo.v1.layers.quantization.base_config import (QuantizationConfig,
+                                                          QuantizeMethodBase)
 from fastvideo.v1.logger import init_logger
 # yapf: disable
 from fastvideo.v1.models.parameter import (BasevLLMParameter,
@@ -545,19 +543,21 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         tp_size = get_tensor_model_parallel_world_size()
 
         if isinstance(param, BlockQuantScaleParameter):
-            from vllm.model_executor.layers.quantization.fp8 import (
-                Fp8LinearMethod, Fp8MoEMethod)
-            assert self.quant_method is not None
-            assert isinstance(self.quant_method,
-                              (Fp8LinearMethod, Fp8MoEMethod))
-            weight_block_size = self.quant_method.quant_config.weight_block_size
-            assert weight_block_size is not None
-            block_n, _ = weight_block_size[0], weight_block_size[1]
-            shard_offset = (
-                (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
-                block_n) // tp_size
-            shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
-                          block_n // tp_size)
+            raise NotImplementedError("FP8 is not implemented yet")
+            # FIXME(will): add fp8 support
+            # from vllm.model_executor.layers.quantization.fp8 import (
+            #     Fp8LinearMethod, Fp8MoEMethod)
+            # assert self.quant_method is not None
+            # assert isinstance(self.quant_method,
+            #                   (Fp8LinearMethod, Fp8MoEMethod))
+            # weight_block_size = self.quant_method.quant_config.weight_block_size
+            # assert weight_block_size is not None
+            # block_n, _ = weight_block_size[0], weight_block_size[1]
+            # shard_offset = (
+            #     (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
+            #     block_n) // tp_size
+            # shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
+            #               block_n // tp_size)
         else:
             shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
             shard_size = self.output_sizes[loaded_shard_id] // tp_size

fastvideo/v1/layers/quantization/__init__.py

@@ -0,0 +1,65 @@
+from typing import Literal, get_args
+
+from fastvideo.v1.layers.quantization.base_config import QuantizationConfig
+
+QuantizationMethods = Literal[None]
+
+QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
+
+# The customized quantization methods which will be added to this dict.
+_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {}
+
+
+def register_quantization_config(quantization: str):
+    """Register a customized vllm quantization config.
+
+    When a quantization method is not supported by vllm, you can register a customized
+    quantization config to support it.
+
+    Args:
+        quantization (str): The quantization method name.
+
+    Examples:
+        >>> from fastvideo.v1.layers.quantization import register_quantization_config
+        >>> from fastvideo.v1.layers.quantization import get_quantization_config
+        >>> from fastvideo.v1.layers.quantization.base_config import QuantizationConfig
+        >>>
+        >>> @register_quantization_config("my_quant")
+        ... class MyQuantConfig(QuantizationConfig):
+        ...     pass
+        >>>
+        >>> get_quantization_config("my_quant")
+        <class 'MyQuantConfig'>
+    """  # noqa: E501
+
+    def _wrapper(quant_config_cls):
+        if quantization in QUANTIZATION_METHODS:
+            raise ValueError(
+                f"The quantization method `{quantization}` is already exists.")
+        if not issubclass(quant_config_cls, QuantizationConfig):
+            raise ValueError("The quantization config must be a subclass of "
+                             "`QuantizationConfig`.")
+        _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
+        QUANTIZATION_METHODS.append(quantization)
+        return quant_config_cls
+
+    return _wrapper
+
+
+def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
+    if quantization not in QUANTIZATION_METHODS:
+        raise ValueError(f"Invalid quantization method: {quantization}")
+
+    method_to_config: dict[str, type[QuantizationConfig]] = {}
+    # Update the `method_to_config` with customized quantization methods.
+    method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
+
+    return method_to_config[quantization]
+
+
+all = [
+    "QuantizationMethods",
+    "QuantizationConfig",
+    "get_quantization_config",
+    "QUANTIZATION_METHODS",
+]

fastvideo/v1/layers/quantization/base_config.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py
+
+import inspect
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from fastvideo.v1.layers.quantization import QuantizationMethods
+else:
+    QuantizationMethods = str
+
+
+class QuantizeMethodBase(ABC):
+    """Base class for different quantized methods."""
+
+    @abstractmethod
+    def create_weights(self, layer: torch.nn.Module, *weight_args,
+                       **extra_weight_attrs):
+        """Create weights for a layer.
+
+        The weights will be set as attributes of the layer."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
+        """Apply the weights in layer to the input tensor.
+
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+    # Not required functions
+    def embedding(self, layer: torch.nn.Module, *args,
+                  **kwargs) -> torch.Tensor:
+        """Gather embeddings in the layer based on indices in the input tensor.
+
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        """Process the weight after loading.
+
+        This can be used for example, to transpose weights for computation.
+        """
+        return
+
+
+def method_has_implemented_embedding(
+        method_class: type[QuantizeMethodBase]) -> bool:
+    """
+    Not all quant methods have embedding implemented, so we need to check that
+    it exists for our given method. We check this by making sure the function
+    has been changed from the base implementation.
+    """
+    base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding",
+                                            None)
+    class_embedding = inspect.getattr_static(method_class, "embedding", None)
+
+    return (class_embedding is not None
+            and class_embedding is not base_embedding)
+
+
+class QuantizationConfig(ABC):
+    """Base class for quantization configs."""
+
+    def __init__(self):
+        super().__init__()
+        # mapping is updated by models as they initialize
+        self.packed_modules_mapping: dict[str, list[str]] = dict()
+
+    @abstractmethod
+    def get_name(self) -> QuantizationMethods:
+        """Name of the quantization method."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        """List of supported activation dtypes."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def get_min_capability(cls) -> int:
+        """Minimum GPU capability to support the quantization method.
+
+        E.g., 70 for Volta, 75 for Turing, 80 for Ampere.
+        This requirement is due to the custom CUDA kernels used by the
+        quantization method.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_config_filenames() -> list[str]:
+        """List of filenames to search for in the model directory."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig":
+        """Create a config class from the model's quantization config."""
+        raise NotImplementedError
+
+    @classmethod
+    def override_quantization_method(
+            cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:
+        """
+           Detects if this quantization method can support a given checkpoint
+           format by overriding the user specified quantization method -- 
+           this method should only be overwritten by subclasses in exceptional 
+           circumstances
+        """
+        return None
+
+    @staticmethod
+    def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any:
+        """Get a value from the model's quantization config."""
+        for key in keys:
+            if key in config:
+                return config[key]
+        raise ValueError(f"Cannot find any of {keys} in the model's "
+                         "quantization config.")
+
+    @staticmethod
+    def get_from_keys_or(config: dict[str, Any], keys: list[str],
+                         default: Any) -> Any:
+        """Get a optional value from the model's quantization config."""
+        try:
+            return QuantizationConfig.get_from_keys(config, keys)
+        except ValueError:
+            return default
+
+    @abstractmethod
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional[QuantizeMethodBase]:
+        """Get the quantize method to use for the quantized layer.
+        
+        Args:
+            layer: The layer for the quant method.
+            prefix: The full name of the layer in the state dict
+        Returns:
+            The quantize method. None if the given layer doesn't support quant
+            method.
+        """
+        raise NotImplementedError
+
+    def get_cache_scale(self, name: str) -> Optional[str]:
+        return None

fastvideo/v1/layers/rotary_embedding.py

@@ -152,28 +152,6 @@ class RotaryEmbedding(CustomOp):
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
-    def forward_cuda(
-        self,
-        positions: torch.Tensor,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        offsets: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        from vllm import _custom_ops as ops
-
-        self.cos_sin_cache = self.cos_sin_cache.to(query.device,
-                                                   dtype=query.dtype)
-        # ops.rotary_embedding()/batched_rotary_embedding()
-        # are in-place operations that update the query and key tensors.
-        if offsets is not None:
-            ops.batched_rotary_embedding(positions, query, key, self.head_size,
-                                         self.cos_sin_cache, self.is_neox_style,
-                                         self.rotary_dim, offsets)
-        else:
-            ops.rotary_embedding(positions, query, key, self.head_size,
-                                 self.cos_sin_cache, self.is_neox_style)
-        return query, key
-
     def extra_repr(self) -> str:
         s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
         s += f", max_position_embeddings={self.max_position_embeddings}"

fastvideo/v1/layers/vocab_parallel_embedding.py

@@ -6,12 +6,12 @@ from typing import List, Optional, Sequence, Tuple
 import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter, UninitializedParameter
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
 
 from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank,
                                       get_tensor_model_parallel_world_size,
                                       tensor_model_parallel_all_reduce)
+from fastvideo.v1.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
 from fastvideo.v1.models.parameter import BasevLLMParameter
 from fastvideo.v1.models.utils import set_weight_attrs
 from fastvideo.v1.platforms import current_platform

fastvideo/v1/models/dits/base.py

@@ -33,9 +33,11 @@ class BaseDiT(nn.Module, ABC):
                     f"Subclasses of BaseDiT must define '{attr}' class variable"
                 )
 
-    def __init__(self, config: DiTConfig, **kwargs) -> None:
+    def __init__(self, config: DiTConfig, hf_config: dict[str, Any],
+                 **kwargs) -> None:
         super().__init__()
         self.config = config
+        self.hf_config = hf_config
         if not self.supported_attention_backends:
             raise ValueError(
                 f"Subclass {self.__class__.__name__} must define _supported_attention_backends"
@@ -95,8 +97,6 @@ class CachableDiT(BaseDiT):
         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
@@ -106,7 +106,9 @@ class CachableDiT(BaseDiT):
         else:
             self.accumulated_rel_l1_distance = 0
             self.previous_modulated_input = None
-            self.previous_residual = None
+            self.previous_resiual = None
+        self.previous_e0_even: torch.Tensor | None = None
+        self.previous_e0_odd: torch.Tensor | None = None
 
     def maybe_cache_states(self, hidden_states: torch.Tensor,
                            original_hidden_states: torch.Tensor) -> None:

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -442,8 +442,8 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
     )._supported_attention_backends
     _param_names_mapping = HunyuanVideoConfig()._param_names_mapping
 
-    def __init__(self, config: HunyuanVideoConfig):
-        super().__init__(config=config)
+    def __init__(self, config: HunyuanVideoConfig, hf_config: dict[str, Any]):
+        super().__init__(config=config, hf_config=hf_config)
 
         self.patch_size = [
             config.patch_size_t, config.patch_size, config.patch_size
@@ -562,8 +562,7 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
         """
         forward_context = get_forward_context()
         forward_batch = forward_context.forward_batch
-        assert forward_batch is not None
-        enable_teacache = forward_batch.enable_teacache
+        enable_teacache = forward_batch is not None and forward_batch.enable_teacache
 
         if guidance is None:
             guidance = torch.tensor([6016.0],
@@ -661,7 +660,8 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
 
         forward_context = get_forward_context()
         forward_batch = forward_context.forward_batch
-        assert forward_batch is not None
+        if forward_batch is None:
+            return False
         current_timestep = forward_context.current_timestep
         enable_teacache = forward_batch.enable_teacache
 

fastvideo/v1/models/dits/stepvideo.py

@@ -0,0 +1,683 @@
+# Copyright 2025 StepFun Inc. All Rights Reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# ==============================================================================
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+from einops import rearrange, repeat
+from torch import nn
+
+from fastvideo.v1.attention import DistributedAttention, LocalAttention
+from fastvideo.v1.configs.models.dits import StepVideoConfig
+from fastvideo.v1.distributed.parallel_state import (
+    get_sequence_model_parallel_world_size)
+from fastvideo.v1.layers.layernorm import LayerNormScaleShift
+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 TimestepEmbedder
+from fastvideo.v1.models.dits.base import BaseDiT
+from fastvideo.v1.platforms import _Backend
+
+
+class PatchEmbed2D(nn.Module):
+    """2D Image to Patch Embedding
+
+    Image to Patch Embedding using Conv2d
+
+    A convolution based approach to patchifying a 2D image w/ embedding projection.
+
+    Based on the impl in https://github.com/google-research/vision_transformer
+
+    Hacked together by / Copyright 2020 Ross Wightman
+
+    Remove the _assert function in forward function to be compatible with multi-resolution images.
+    """
+
+    def __init__(self,
+                 patch_size=16,
+                 in_chans=3,
+                 embed_dim=768,
+                 norm_layer=None,
+                 flatten=True,
+                 bias=True,
+                 dtype=None,
+                 prefix: str = ""):
+        super().__init__()
+        # Convert patch_size to 2-tuple
+        if isinstance(patch_size, (list, tuple)):
+            if len(patch_size) == 1:
+                patch_size = (patch_size[0], patch_size[0])
+        else:
+            patch_size = (patch_size, patch_size)
+
+        self.patch_size = patch_size
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans,
+                              embed_dim,
+                              kernel_size=patch_size,
+                              stride=patch_size,
+                              bias=bias,
+                              dtype=dtype)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+class StepVideoRMSNorm(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        elementwise_affine=True,
+        eps: float = 1e-6,
+        device=None,
+        dtype=None,
+    ):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.eps = eps
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
+
+    def _norm(self, x) -> torch.Tensor:
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        if hasattr(self, "weight"):
+            output = output * self.weight
+        return output
+
+
+class SelfAttention(nn.Module):
+
+    def __init__(self,
+                 hidden_dim,
+                 head_dim,
+                 rope_split: Tuple[int, int, int] = (64, 32, 32),
+                 bias: bool = False,
+                 with_rope: bool = True,
+                 with_qk_norm: bool = True,
+                 attn_type: str = "torch",
+                 supported_attention_backends=(_Backend.FLASH_ATTN,
+                                               _Backend.TORCH_SDPA)):
+        super().__init__()
+        self.head_dim = head_dim
+        self.hidden_dim = hidden_dim
+        self.rope_split = list(rope_split)
+        self.n_heads = hidden_dim // head_dim
+
+        self.wqkv = ReplicatedLinear(hidden_dim, hidden_dim * 3, bias=bias)
+        self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)
+
+        self.with_rope = with_rope
+        self.with_qk_norm = with_qk_norm
+        if self.with_qk_norm:
+            self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
+            self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
+
+        # self.core_attention = self.attn_processor(attn_type=attn_type)
+        self.parallel = attn_type == 'parallel'
+        self.attn = DistributedAttention(
+            num_heads=self.n_heads,
+            head_size=head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends)
+
+    def _apply_rope(self, x: torch.Tensor, cos: torch.Tensor,
+                    sin: torch.Tensor):
+        """
+        x:   [B, S, H, D]
+        cos: [S, D/2]  where D = head_dim = sum(self.rope_split)
+        sin: [S, D/2]
+        returns x with rotary applied exactly as v0 did
+        """
+        B, S, H, D = x.shape
+        # 1) split cos/sin per chunk
+        half_splits = [c // 2
+                       for c in self.rope_split]  # [32,16,16] for [64,32,32]
+        cos_splits = cos.split(half_splits, dim=1)
+        sin_splits = sin.split(half_splits, dim=1)
+
+        outs = []
+        idx = 0
+        for (chunk_size, cos_i, sin_i) in zip(self.rope_split, cos_splits,
+                                              sin_splits):
+            # slice the corresponding channels
+            x_chunk = x[..., idx:idx + chunk_size]  # [B,S,H,chunk_size]
+            idx += chunk_size
+
+            # flatten to [S, B*H, chunk_size]
+            x_flat = rearrange(x_chunk, 'b s h d -> s (b h) d')
+
+            # apply rotary on *that* chunk
+            out_flat = _apply_rotary_emb(x_flat,
+                                         cos_i,
+                                         sin_i,
+                                         is_neox_style=True)
+
+            # restore [B,S,H,chunk_size]
+            out = rearrange(out_flat, 's (b h) d -> b s h d', b=B, h=H)
+            outs.append(out)
+
+        # concatenate back to [B,S,H,D]
+        return torch.cat(outs, dim=-1)
+
+    def forward(self,
+                x,
+                cu_seqlens=None,
+                max_seqlen=None,
+                rope_positions=None,
+                cos_sin=None,
+                attn_mask=None,
+                mask_strategy=None):
+
+        B, S, _ = x.shape
+        xqkv, _ = self.wqkv(x)
+        xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim)
+        q, k, v = torch.split(xqkv, [self.head_dim] * 3, dim=-1)  # [B,S,H,D]
+
+        if self.with_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.with_rope:
+            if rope_positions is not None:
+                F, Ht, W = rope_positions
+            assert F * Ht * W == S, "rope_positions mismatches sequence length"
+
+            cos, sin = cos_sin
+            cos = cos.to(x.device, dtype=x.dtype)
+            sin = sin.to(x.device, dtype=x.dtype)
+
+            q = self._apply_rope(q, cos, sin)
+            k = self._apply_rope(k, cos, sin)
+
+        output, _ = self.attn(q, k, v)  # [B,heads,S,D]
+
+        output = rearrange(output, 'b s h d -> b s (h d)')
+        output, _ = self.wo(output)
+
+        return output
+
+
+class CrossAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_dim,
+        head_dim,
+        bias=False,
+        with_qk_norm=True,
+        supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+    ) -> None:
+        super().__init__()
+        self.head_dim = head_dim
+        self.n_heads = hidden_dim // head_dim
+
+        self.wq = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)
+        self.wkv = ReplicatedLinear(hidden_dim, hidden_dim * 2, bias=bias)
+        self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)
+
+        self.with_qk_norm = with_qk_norm
+        if self.with_qk_norm:
+            self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
+            self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
+
+        self.attn = LocalAttention(
+            num_heads=self.n_heads,
+            head_size=head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends)
+
+    def forward(self,
+                x: torch.Tensor,
+                encoder_hidden_states: torch.Tensor,
+                attn_mask=None) -> torch.Tensor:
+
+        xq, _ = self.wq(x)
+        xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim)
+
+        xkv, _ = self.wkv(encoder_hidden_states)
+        xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim)
+
+        xk, xv = torch.split(xkv, [self.head_dim] * 2,
+                             dim=-1)  ## seq_len, n, dim
+
+        if self.with_qk_norm:
+            xq = self.q_norm(xq)
+            xk = self.k_norm(xk)
+
+        output = self.attn(xq, xk, xv)
+
+        output = rearrange(output, 'b s h d -> b s (h d)')
+        output, _ = self.wo(output)
+
+        return output
+
+
+class AdaLayerNormSingle(nn.Module):
+    r"""
+        Norm layer adaptive layer norm single (adaLN-single).
+
+        As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
+
+        Parameters:
+            embedding_dim (`int`): The size of each embedding vector.
+            use_additional_conditions (`bool`): To use additional conditions for normalization or not.
+    """
+
+    def __init__(self, embedding_dim: int, time_step_rescale=1000):
+        super().__init__()
+
+        self.emb = TimestepEmbedder(embedding_dim)
+
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(embedding_dim,
+                                       6 * embedding_dim,
+                                       bias=True)
+
+        self.time_step_rescale = time_step_rescale  ## timestep usually in [0, 1], we rescale it to [0,1000] for stability
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        embedded_timestep = self.emb(timestep * self.time_step_rescale)
+
+        out, _ = self.linear(self.silu(embedded_timestep))
+
+        return out, embedded_timestep
+
+
+class StepVideoTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        num_embeds_ada_norm (:
+            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
+        attention_bias (:
+            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+        only_cross_attention (`bool`, *optional*):
+            Whether to use only cross-attention layers. In this case two cross attention layers are used.
+        double_self_attention (`bool`, *optional*):
+            Whether to use two self-attention layers. In this case no cross attention layers are used.
+        upcast_attention (`bool`, *optional*):
+            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
+        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
+        final_dropout (`bool` *optional*, defaults to False):
+            Whether to apply a final dropout after the last feed-forward layer.
+        positional_embeddings (`str`, *optional*, defaults to `None`):
+            The type of positional embeddings to apply to.
+        num_positional_embeddings (`int`, *optional*, defaults to `None`):
+            The maximum number of positional embeddings to apply.
+    """
+
+    def __init__(self,
+                 dim: int,
+                 attention_head_dim: int,
+                 norm_eps: float = 1e-5,
+                 ff_inner_dim: Optional[int] = None,
+                 ff_bias: bool = False,
+                 attention_type: str = 'torch'):
+        super().__init__()
+        self.dim = dim
+        self.norm1 = LayerNormScaleShift(dim,
+                                         norm_type="layer",
+                                         elementwise_affine=True,
+                                         eps=norm_eps)
+        self.attn1 = SelfAttention(
+            dim,
+            attention_head_dim,
+            bias=False,
+            with_rope=True,
+            with_qk_norm=True,
+        )
+
+        self.norm2 = LayerNormScaleShift(dim,
+                                         norm_type="layer",
+                                         elementwise_affine=True,
+                                         eps=norm_eps)
+        self.attn2 = CrossAttention(dim,
+                                    attention_head_dim,
+                                    bias=False,
+                                    with_qk_norm=True)
+
+        self.ff = MLP(input_dim=dim,
+                      mlp_hidden_dim=dim *
+                      4 if ff_inner_dim is None else ff_inner_dim,
+                      act_type="gelu_pytorch_tanh",
+                      bias=ff_bias)
+
+        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+    @torch.no_grad()
+    def forward(self,
+                q: torch.Tensor,
+                kv: torch.Tensor,
+                t_expand: torch.LongTensor,
+                attn_mask=None,
+                rope_positions: Optional[list] = None,
+                cos_sin=None,
+                mask_strategy=None) -> torch.Tensor:
+
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            torch.clone(chunk)
+            for chunk in (self.scale_shift_table[None] +
+                          t_expand.reshape(-1, 6, self.dim)).chunk(6, dim=1))
+
+        scale_shift_q = self.norm1(q,
+                                   scale=scale_msa.squeeze(1),
+                                   shift=shift_msa.squeeze(1))
+
+        attn_q = self.attn1(scale_shift_q,
+                            rope_positions=rope_positions,
+                            cos_sin=cos_sin,
+                            mask_strategy=mask_strategy)
+
+        q = attn_q * gate_msa + q
+
+        attn_q = self.attn2(q, kv, attn_mask)
+
+        q = attn_q + q
+
+        scale_shift_q = self.norm2(q,
+                                   scale=scale_mlp.squeeze(1),
+                                   shift=shift_mlp.squeeze(1))
+
+        ff_output = self.ff(scale_shift_q)
+
+        q = ff_output * gate_mlp + q
+
+        return q
+
+
+class StepVideoModel(BaseDiT):
+    # (Optional) Keep the same attribute for compatibility with splitting, etc.
+    _fsdp_shard_conditions = [
+        lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit(),
+        # lambda n, m: "pos_embed" in n  # If needed for the patch embedding.
+    ]
+    _param_names_mapping = StepVideoConfig()._param_names_mapping
+    _supported_attention_backends = StepVideoConfig(
+    )._supported_attention_backends
+
+    def __init__(self, config: StepVideoConfig, hf_config: dict[str,
+                                                                Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_dim = config.attention_head_dim
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels
+        self.num_layers = config.num_layers
+        self.dropout = config.dropout
+        self.patch_size = config.patch_size
+        self.norm_type = config.norm_type
+        self.norm_elementwise_affine = config.norm_elementwise_affine
+        self.norm_eps = config.norm_eps
+        self.use_additional_conditions = config.use_additional_conditions
+        self.caption_channels = config.caption_channels
+        self.attention_type = config.attention_type
+        self.num_channels_latents = config.num_channels_latents
+        # Compute inner dimension.
+        self.hidden_size = config.hidden_size
+
+        # Image/video patch embedding.
+        self.pos_embed = PatchEmbed2D(
+            patch_size=self.patch_size,
+            in_chans=self.in_channels,
+            embed_dim=self.hidden_size,
+        )
+
+        self._rope_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
+        # Transformer blocks.
+        self.transformer_blocks = nn.ModuleList([
+            StepVideoTransformerBlock(
+                dim=self.hidden_size,
+                attention_head_dim=self.attention_head_dim,
+                attention_type=self.attention_type)
+            for _ in range(self.num_layers)
+        ])
+
+        # Output blocks.
+        self.norm_out = LayerNormScaleShift(
+            self.hidden_size,
+            norm_type="layer",
+            eps=self.norm_eps,
+            elementwise_affine=self.norm_elementwise_affine)
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, self.hidden_size) / (self.hidden_size**0.5))
+        self.proj_out = ReplicatedLinear(
+            self.hidden_size,
+            self.patch_size * self.patch_size * self.out_channels)
+        # Time modulation via adaptive layer norm.
+        self.adaln_single = AdaLayerNormSingle(self.hidden_size)
+
+        # Set up caption conditioning.
+        if isinstance(self.caption_channels, int):
+            caption_channel = self.caption_channels
+        else:
+            caption_channel, clip_channel = self.caption_channels
+            self.clip_projection = ReplicatedLinear(clip_channel,
+                                                    self.hidden_size)
+        self.caption_norm = nn.LayerNorm(
+            caption_channel,
+            eps=self.norm_eps,
+            elementwise_affine=self.norm_elementwise_affine)
+        self.caption_projection = MLP(input_dim=caption_channel,
+                                      mlp_hidden_dim=self.hidden_size,
+                                      act_type="gelu_pytorch_tanh")
+
+        # Flag to indicate if using parallel attention.
+        self.parallel = (self.attention_type == "parallel")
+
+        self.__post_init__()
+
+    def patchfy(self, hidden_states) -> torch.Tensor:
+        hidden_states = rearrange(hidden_states, 'b f c h w -> (b f) c h w')
+        hidden_states = self.pos_embed(hidden_states)
+        return hidden_states
+
+    def prepare_attn_mask(self, encoder_attention_mask, encoder_hidden_states,
+                          q_seqlen) -> Tuple[torch.Tensor, torch.Tensor]:
+        kv_seqlens = encoder_attention_mask.sum(dim=1).int()
+        mask = torch.zeros([len(kv_seqlens), q_seqlen,
+                            max(kv_seqlens)],
+                           dtype=torch.bool,
+                           device=encoder_attention_mask.device)
+        encoder_hidden_states = encoder_hidden_states[:, :max(kv_seqlens)]
+        for i, kv_len in enumerate(kv_seqlens):
+            mask[i, :, :kv_len] = 1
+        return encoder_hidden_states, mask
+
+    def block_forward(self,
+                      hidden_states,
+                      encoder_hidden_states=None,
+                      t_expand=None,
+                      rope_positions=None,
+                      cos_sin=None,
+                      attn_mask=None,
+                      parallel=True,
+                      mask_strategy=None) -> torch.Tensor:
+
+        for i, block in enumerate(self.transformer_blocks):
+            hidden_states = block(hidden_states,
+                                  encoder_hidden_states,
+                                  t_expand=t_expand,
+                                  attn_mask=attn_mask,
+                                  rope_positions=rope_positions,
+                                  cos_sin=cos_sin,
+                                  mask_strategy=mask_strategy[i])
+
+        return hidden_states
+
+    def _get_rope(self, rope_positions: tuple[int, int, int],
+                  dtype: torch.dtype, device: torch.device):
+        F, Ht, W = rope_positions
+        key = (F, Ht, W, dtype)
+        if key not in self._rope_cache:
+            cos, sin = get_rotary_pos_embed(
+                rope_sizes=(F * get_sequence_model_parallel_world_size(), Ht,
+                            W),
+                hidden_size=self.hidden_size,
+                heads_num=self.hidden_size // self.attention_head_dim,
+                rope_dim_list=(64, 32, 32),  # same split you used
+                rope_theta=1.0e4,
+                dtype=torch.float32  # build once in fp32
+            )
+            # move & cast once
+            self._rope_cache[key] = (cos.to(device, dtype=dtype),
+                                     sin.to(device, dtype=dtype))
+        return self._rope_cache[key]
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        t_expand: Optional[torch.LongTensor] = None,
+        encoder_hidden_states_2: Optional[torch.Tensor] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        fps: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+        mask_strategy=None,
+        guidance=None,
+    ):
+        assert hidden_states.ndim == 5
+        "hidden_states's shape should be (bsz, f, ch, h ,w)"
+        frame = hidden_states.shape[2]
+        hidden_states = rearrange(hidden_states,
+                                  'b c f h w -> b f c h w',
+                                  f=frame)
+        if mask_strategy is None:
+            mask_strategy = [None, None]
+        bsz, frame, _, height, width = hidden_states.shape
+        height, width = height // self.patch_size, width // self.patch_size
+
+        hidden_states = self.patchfy(hidden_states)
+        len_frame = hidden_states.shape[1]
+
+        t_expand, embedded_timestep = self.adaln_single(t_expand)
+        encoder_hidden_states = self.caption_projection(
+            self.caption_norm(encoder_hidden_states))
+
+        if encoder_hidden_states_2 is not None and hasattr(
+                self, 'clip_projection'):
+            clip_embedding, _ = self.clip_projection(encoder_hidden_states_2)
+            encoder_hidden_states = torch.cat(
+                [clip_embedding, encoder_hidden_states], dim=1)
+
+        hidden_states = rearrange(hidden_states,
+                                  '(b f) l d->  b (f l) d',
+                                  b=bsz,
+                                  f=frame,
+                                  l=len_frame).contiguous()
+        encoder_hidden_states, attn_mask = self.prepare_attn_mask(
+            encoder_attention_mask,
+            encoder_hidden_states,
+            q_seqlen=frame * len_frame)
+
+        cos_sin = self._get_rope((frame, height, width), hidden_states.dtype,
+                                 hidden_states.device)
+
+        hidden_states = self.block_forward(
+            hidden_states,
+            encoder_hidden_states,
+            t_expand=t_expand,
+            rope_positions=[frame, height, width],
+            cos_sin=cos_sin,
+            attn_mask=attn_mask,
+            parallel=self.parallel,
+            mask_strategy=mask_strategy)
+
+        hidden_states = rearrange(hidden_states,
+                                  'b (f l) d -> (b f) l d',
+                                  b=bsz,
+                                  f=frame,
+                                  l=len_frame)
+
+        embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d',
+                                   f=frame).contiguous()
+
+        shift, scale = (self.scale_shift_table[None] +
+                        embedded_timestep[:, None]).chunk(2, dim=1)
+        hidden_states = self.norm_out(hidden_states,
+                                      shift=shift.squeeze(1),
+                                      scale=scale.squeeze(1))
+        # Modulation
+        hidden_states, _ = self.proj_out(hidden_states)
+
+        # unpatchify
+        hidden_states = hidden_states.reshape(shape=(-1, height, width,
+                                                     self.patch_size,
+                                                     self.patch_size,
+                                                     self.out_channels))
+
+        hidden_states = rearrange(hidden_states, 'n h w p q c -> n c h p w q')
+        output = hidden_states.reshape(shape=(-1, self.out_channels,
+                                              height * self.patch_size,
+                                              width * self.patch_size))
+
+        output = rearrange(output, '(b f) c h w -> b c f h w', f=frame)
+        return output

fastvideo/v1/models/dits/wanvideo.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
-from typing import List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -298,7 +298,7 @@ class WanTransformerBlock(nn.Module):
             hidden_states = hidden_states.squeeze(1)
         bs, seq_length, _ = hidden_states.shape
         orig_dtype = hidden_states.dtype
-        assert orig_dtype != torch.float32
+        # assert orig_dtype != torch.float32
         e = self.scale_shift_table + temb.float()
         shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
             6, dim=1)
@@ -360,8 +360,9 @@ class WanTransformer3DModel(CachableDiT):
     )._supported_attention_backends
     _param_names_mapping = WanVideoConfig()._param_names_mapping
 
-    def __init__(self, config: WanVideoConfig) -> None:
-        super().__init__(config=config)
+    def __init__(self, config: WanVideoConfig, hf_config: dict[str,
+                                                               Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
 
         inner_dim = config.num_attention_heads * config.attention_head_dim
         self.hidden_size = config.hidden_size
@@ -424,8 +425,7 @@ class WanTransformer3DModel(CachableDiT):
                 guidance=None,
                 **kwargs) -> torch.Tensor:
         forward_batch = get_forward_context().forward_batch
-        assert forward_batch is not None
-        enable_teacache = forward_batch.enable_teacache
+        enable_teacache = forward_batch is not None and forward_batch.enable_teacache
 
         orig_dtype = hidden_states.dtype
         if not isinstance(encoder_hidden_states, torch.Tensor):
@@ -525,8 +525,7 @@ class WanTransformer3DModel(CachableDiT):
 
         forward_context = get_forward_context()
         forward_batch = forward_context.forward_batch
-        assert forward_batch is not None
-        if not forward_batch.enable_teacache:
+        if forward_batch is None or not forward_batch.enable_teacache:
             return False
         teacache_params = forward_batch.teacache_params
         assert teacache_params is not None, "teacache_params is not initialized"

fastvideo/v1/models/encoders/bert.py

@@ -0,0 +1,40 @@
+# type: ignore
+import os
+
+import torch
+import torch.nn as nn
+from transformers import BertModel, BertTokenizer
+
+
+class HunyuanClip(nn.Module):
+    """
+        Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
+        hunyuan's clip used BertModel and BertTokenizer, so we copy it.
+    """
+
+    def __init__(self, model_dir, max_length=77):
+        super().__init__()
+
+        self.max_length = max_length
+        self.tokenizer = BertTokenizer.from_pretrained(
+            os.path.join(model_dir, 'tokenizer'))
+        self.text_encoder = BertModel.from_pretrained(
+            os.path.join(model_dir, 'clip_text_encoder'))
+
+    @torch.no_grad
+    def forward(self, prompts, with_mask=True):
+        self.device = next(self.text_encoder.parameters()).device
+        text_inputs = self.tokenizer(
+            prompts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        prompt_embeds = self.text_encoder(
+            text_inputs.input_ids.to(self.device),
+            attention_mask=text_inputs.attention_mask.to(self.device)
+            if with_mask else None,
+        )
+        return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output

fastvideo/v1/models/encoders/clip.py

@@ -13,12 +13,12 @@ from fastvideo.v1.attention import LocalAttention
 from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
                                                   CLIPTextConfig,
                                                   CLIPVisionConfig)
-from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.distributed import (divide,
                                       get_tensor_model_parallel_world_size)
 from fastvideo.v1.layers.activation import get_act_fn
 from fastvideo.v1.layers.linear import (ColumnParallelLinear, QKVParallelLinear,
                                         RowParallelLinear)
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.encoders.base import ImageEncoder, TextEncoder
 from fastvideo.v1.models.encoders.vision import resolve_visual_encoder_outputs

fastvideo/v1/models/encoders/llama.py

@@ -32,12 +32,12 @@ from torch import nn
 from fastvideo.v1.attention import LocalAttention
 # from ..utils import (extract_layer_index)
 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
 from fastvideo.v1.layers.layernorm import RMSNorm
 from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
                                         QKVParallelLinear, RowParallelLinear)
+from fastvideo.v1.layers.quantization import QuantizationConfig
 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 TextEncoder

fastvideo/v1/models/encoders/stepllm.py

@@ -0,0 +1,584 @@
+# type: ignore
+# Copyright 2025 StepFun Inc. All Rights Reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# ==============================================================================
+import os
+from functools import wraps
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from transformers.modeling_utils import PretrainedConfig, PreTrainedModel
+
+from fastvideo.v1.models.dits.stepvideo import StepVideoRMSNorm
+
+
+class EmptyInitOnDevice(torch.overrides.TorchFunctionMode):
+
+    def __init__(self, device=None):
+        self.device = device
+
+    def __torch_function__(self, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+        if getattr(func, '__module__', None) == 'torch.nn.init':
+            if 'tensor' in kwargs:
+                return kwargs['tensor']
+            else:
+                return args[0]
+        if self.device is not None and func in torch.utils._device._device_constructors(
+        ) and kwargs.get('device') is None:
+            kwargs['device'] = self.device
+        return func(*args, **kwargs)
+
+
+def with_empty_init(func):
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        with EmptyInitOnDevice('cpu'):
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+class LLaMaEmbedding(nn.Module):
+    """Language model embeddings.
+
+    Arguments:
+        hidden_size: hidden size
+        vocab_size: vocabulary size
+        max_sequence_length: maximum size of sequence. This
+                             is used for positional embedding
+        embedding_dropout_prob: dropout probability for embeddings
+        init_method: weight initialization method
+        num_tokentypes: size of the token-type embeddings. 0 value
+                        will ignore this embedding
+    """
+
+    def __init__(
+        self,
+        cfg,
+    ):
+        super().__init__()
+        self.hidden_size = cfg.hidden_size
+        self.params_dtype = cfg.params_dtype
+        self.fp32_residual_connection = cfg.fp32_residual_connection
+        self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32
+        self.word_embeddings = torch.nn.Embedding(
+            cfg.padded_vocab_size,
+            self.hidden_size,
+        )
+        self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout)
+
+    def forward(self, input_ids):
+        # Embeddings.
+        if self.embedding_weights_in_fp32:
+            self.word_embeddings = self.word_embeddings.to(torch.float32)
+        embeddings = self.word_embeddings(input_ids)
+        if self.embedding_weights_in_fp32:
+            embeddings = embeddings.to(self.params_dtype)
+            self.word_embeddings = self.word_embeddings.to(self.params_dtype)
+
+        # Data format change to avoid explicit transposes : [b s h] --> [s b h].
+        embeddings = embeddings.transpose(0, 1).contiguous()
+
+        # If the input flag for fp32 residual connection is set, convert for float.
+        if self.fp32_residual_connection:
+            embeddings = embeddings.float()
+
+        # Dropout.
+        embeddings = self.embedding_dropout(embeddings)
+
+        return embeddings
+
+
+class StepChatTokenizer:
+    """Step Chat Tokenizer"""
+
+    def __init__(
+            self,
+            model_file,
+            name="StepChatTokenizer",
+            bot_token="<|BOT|>",  # Begin of Turn
+            eot_token="<|EOT|>",  # End of Turn
+            call_start_token="<|CALL_START|>",  # Call Start
+            call_end_token="<|CALL_END|>",  # Call End
+            think_start_token="<|THINK_START|>",  # Think Start
+            think_end_token="<|THINK_END|>",  # Think End
+            mask_start_token="<|MASK_1e69f|>",  # Mask start
+            mask_end_token="<|UNMASK_1e69f|>",  # Mask end
+    ):
+        import sentencepiece
+
+        self._tokenizer = sentencepiece.SentencePieceProcessor(
+            model_file=model_file)
+
+        self._vocab = {}
+        self._inv_vocab = {}
+
+        self._special_tokens = {}
+        self._inv_special_tokens = {}
+
+        self._t5_tokens = []
+
+        for idx in range(self._tokenizer.get_piece_size()):
+            text = self._tokenizer.id_to_piece(idx)
+            self._inv_vocab[idx] = text
+            self._vocab[text] = idx
+
+            if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown(
+                    idx):
+                self._special_tokens[text] = idx
+                self._inv_special_tokens[idx] = text
+
+        self._unk_id = self._tokenizer.unk_id()
+        self._bos_id = self._tokenizer.bos_id()
+        self._eos_id = self._tokenizer.eos_id()
+
+        for token in [
+                bot_token, eot_token, call_start_token, call_end_token,
+                think_start_token, think_end_token
+        ]:
+            assert token in self._vocab, f"Token '{token}' not found in tokenizer"
+            assert token in self._special_tokens, f"Token '{token}' is not a special token"
+
+        for token in [mask_start_token, mask_end_token]:
+            assert token in self._vocab, f"Token '{token}' not found in tokenizer"
+
+        self._bot_id = self._tokenizer.piece_to_id(bot_token)
+        self._eot_id = self._tokenizer.piece_to_id(eot_token)
+        self._call_start_id = self._tokenizer.piece_to_id(call_start_token)
+        self._call_end_id = self._tokenizer.piece_to_id(call_end_token)
+        self._think_start_id = self._tokenizer.piece_to_id(think_start_token)
+        self._think_end_id = self._tokenizer.piece_to_id(think_end_token)
+        self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token)
+        self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token)
+
+        self._underline_id = self._tokenizer.piece_to_id("\u2581")
+
+    @property
+    def vocab(self):
+        return self._vocab
+
+    @property
+    def inv_vocab(self):
+        return self._inv_vocab
+
+    @property
+    def vocab_size(self):
+        return self._tokenizer.vocab_size()
+
+    def tokenize(self, text: str) -> List[int]:
+        return self._tokenizer.encode_as_ids(text)
+
+    def detokenize(self, token_ids: List[int]) -> str:
+        return self._tokenizer.decode_ids(token_ids)
+
+
+class Tokens:
+
+    def __init__(self, input_ids, cu_input_ids, attention_mask, cu_seqlens,
+                 max_seq_len) -> None:
+        self.input_ids = input_ids
+        self.attention_mask = attention_mask
+        self.cu_input_ids = cu_input_ids
+        self.cu_seqlens = cu_seqlens
+        self.max_seq_len = max_seq_len
+
+    def to(self, device):
+        self.input_ids = self.input_ids.to(device)
+        self.attention_mask = self.attention_mask.to(device)
+        self.cu_input_ids = self.cu_input_ids.to(device)
+        self.cu_seqlens = self.cu_seqlens.to(device)
+        return self
+
+
+class Wrapped_StepChatTokenizer(StepChatTokenizer):
+
+    def __call__(self,
+                 text,
+                 max_length=320,
+                 padding="max_length",
+                 truncation=True,
+                 return_tensors="pt"):
+        # [bos, ..., eos, pad, pad, ..., pad]
+        self.BOS = 1
+        self.EOS = 2
+        self.PAD = 2
+        out_tokens = []
+        attn_mask = []
+        if len(text) == 0:
+            part_tokens = [self.BOS] + [self.EOS]
+            valid_size = len(part_tokens)
+            if len(part_tokens) < max_length:
+                part_tokens += [self.PAD] * (max_length - valid_size)
+            out_tokens.append(part_tokens)
+            attn_mask.append([1] * valid_size + [0] * (max_length - valid_size))
+        else:
+            for part in text:
+                part_tokens = self.tokenize(part)
+                part_tokens = part_tokens[:(max_length -
+                                            2)]  # leave 2 space for bos and eos
+                part_tokens = [self.BOS] + part_tokens + [self.EOS]
+                valid_size = len(part_tokens)
+                if len(part_tokens) < max_length:
+                    part_tokens += [self.PAD] * (max_length - valid_size)
+                out_tokens.append(part_tokens)
+                attn_mask.append([1] * valid_size + [0] *
+                                 (max_length - valid_size))
+
+        out_tokens = torch.tensor(out_tokens, dtype=torch.long)
+        attn_mask = torch.tensor(attn_mask, dtype=torch.long)
+
+        # padding y based on tp size
+        padded_len = 0
+        padded_flag = False
+        if padded_len > 0:
+            padded_flag = True
+        if padded_flag:
+            pad_tokens = torch.tensor([[self.PAD] * max_length],
+                                      device=out_tokens.device)
+            pad_attn_mask = torch.tensor([[1] * padded_len + [0] *
+                                          (max_length - padded_len)],
+                                         device=attn_mask.device)
+            out_tokens = torch.cat([out_tokens, pad_tokens], dim=0)
+            attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0)
+
+        # cu_seqlens
+        cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0)
+        seqlen = attn_mask.sum(dim=1).tolist()
+        cu_seqlens = torch.cumsum(torch.tensor([0] + seqlen),
+                                  0).to(device=out_tokens.device,
+                                        dtype=torch.int32)
+        max_seq_len = max(seqlen)
+        return Tokens(out_tokens, cu_out_tokens, attn_mask, cu_seqlens,
+                      max_seq_len)
+
+
+def flash_attn_func(q,
+                    k,
+                    v,
+                    dropout_p=0.0,
+                    softmax_scale=None,
+                    causal=True,
+                    return_attn_probs=False,
+                    tp_group_rank=0,
+                    tp_group_size=1):
+    softmax_scale = q.size(-1)**(
+        -0.5) if softmax_scale is None else softmax_scale
+    return torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale,
+                                 causal, return_attn_probs, None, tp_group_rank,
+                                 tp_group_size)[0]
+
+
+class FlashSelfAttention(torch.nn.Module):
+
+    def __init__(
+        self,
+        attention_dropout=0.0,
+    ):
+        super().__init__()
+        self.dropout_p = attention_dropout
+
+    def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None):
+        if cu_seqlens is None:
+            output = flash_attn_func(q, k, v, dropout_p=self.dropout_p)
+        else:
+            raise ValueError('cu_seqlens is not supported!')
+
+        return output
+
+
+def safediv(n, d):
+    q, r = divmod(n, d)
+    assert r == 0
+    return q
+
+
+class MultiQueryAttention(nn.Module):
+
+    def __init__(self, cfg, layer_id=None):
+        super().__init__()
+
+        self.head_dim = cfg.hidden_size // cfg.num_attention_heads
+        self.max_seq_len = cfg.seq_length
+        self.use_flash_attention = cfg.use_flash_attn
+        assert self.use_flash_attention, 'FlashAttention is required!'
+
+        self.n_groups = cfg.num_attention_groups
+        self.tp_size = 1
+        self.n_local_heads = cfg.num_attention_heads
+        self.n_local_groups = self.n_groups
+
+        self.wqkv = nn.Linear(
+            cfg.hidden_size,
+            cfg.hidden_size + self.head_dim * 2 * self.n_groups,
+            bias=False,
+        )
+        self.wo = nn.Linear(
+            cfg.hidden_size,
+            cfg.hidden_size,
+            bias=False,
+        )
+
+        # assert self.use_flash_attention, 'non-Flash attention not supported yet.'
+        self.core_attention = FlashSelfAttention(
+            attention_dropout=cfg.attention_dropout)
+        # self.core_attention = LocalAttention(
+        #     num_heads = self.n_local_heads,
+        #     head_size = self.head_dim,
+        #     # num_kv_heads  = self.n_local_groups,
+        #     casual = True,
+        #     supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA], # RIVER TODO
+        # )
+        self.layer_id = layer_id
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        cu_seqlens: Optional[torch.Tensor],
+        max_seq_len: Optional[torch.Tensor],
+    ):
+        seqlen, bsz, dim = x.shape
+        xqkv = self.wqkv(x)
+
+        xq, xkv = torch.split(
+            xqkv,
+            (dim // self.tp_size,
+             self.head_dim * 2 * self.n_groups // self.tp_size),
+            dim=-1,
+        )
+
+        # gather on 1st dimension
+        xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim)
+        xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim)
+        xk, xv = xkv.chunk(2, -1)
+
+        # rotary embedding + flash attn
+        xq = rearrange(xq, "s b h d -> b s h d")
+        xk = rearrange(xk, "s b h d -> b s h d")
+        xv = rearrange(xv, "s b h d -> b s h d")
+
+        # q_per_kv = self.n_local_heads // self.n_local_groups
+        # if q_per_kv > 1:
+        #     b, s, h, d = xk.size()
+        #     if h == 1:
+        #         xk = xk.expand(b, s, q_per_kv, d)
+        #         xv = xv.expand(b, s, q_per_kv, d)
+        #     else:
+        #         ''' To cover the cases where h > 1, we have
+        #             the following implementation, which is equivalent to:
+        #                 xk = xk.repeat_interleave(q_per_kv, dim=-2)
+        #                 xv = xv.repeat_interleave(q_per_kv, dim=-2)
+        #             but can avoid calling aten::item() that involves cpu.
+        #         '''
+        #         idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten()
+        #         xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous()
+        #         xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous()
+        if self.use_flash_attention:
+            output = self.core_attention(xq, xk, xv)
+            # reduce-scatter only support first dimension now
+            output = rearrange(output, "b s h d -> s b (h d)").contiguous()
+        else:
+            xq, xk, xv = [
+                rearrange(x, "b s ... -> s b ...").contiguous()
+                for x in (xq, xk, xv)
+            ]
+            output = self.core_attention(xq, xk, xv)  #, mask)
+        output = self.wo(output)
+        return output
+
+
+class FeedForward(nn.Module):
+
+    def __init__(
+        self,
+        cfg,
+        dim: int,
+        hidden_dim: int,
+        layer_id: int,
+        multiple_of: int = 256,
+    ):
+        super().__init__()
+
+        hidden_dim = multiple_of * (
+            (hidden_dim + multiple_of - 1) // multiple_of)
+
+        def swiglu(x):
+            x = torch.chunk(x, 2, dim=-1)
+            return F.silu(x[0]) * x[1]
+
+        self.swiglu = swiglu
+
+        self.w1 = nn.Linear(
+            dim,
+            2 * hidden_dim,
+            bias=False,
+        )
+        self.w2 = nn.Linear(
+            hidden_dim,
+            dim,
+            bias=False,
+        )
+
+    def forward(self, x):
+        x = self.swiglu(self.w1(x))
+        output = self.w2(x)
+        return output
+
+
+class TransformerBlock(nn.Module):
+
+    def __init__(self, cfg, layer_id: int):
+        super().__init__()
+
+        self.n_heads = cfg.num_attention_heads
+        self.dim = cfg.hidden_size
+        self.head_dim = cfg.hidden_size // cfg.num_attention_heads
+        self.attention = MultiQueryAttention(
+            cfg,
+            layer_id=layer_id,
+        )
+
+        self.feed_forward = FeedForward(
+            cfg,
+            dim=cfg.hidden_size,
+            hidden_dim=cfg.ffn_hidden_size,
+            layer_id=layer_id,
+        )
+        self.layer_id = layer_id
+        self.attention_norm = StepVideoRMSNorm(
+            cfg.hidden_size,
+            eps=cfg.layernorm_epsilon,
+        )
+        self.ffn_norm = StepVideoRMSNorm(
+            cfg.hidden_size,
+            eps=cfg.layernorm_epsilon,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        cu_seqlens: Optional[torch.Tensor],
+        max_seq_len: Optional[torch.Tensor],
+    ):
+        residual = self.attention.forward(self.attention_norm(x), mask,
+                                          cu_seqlens, max_seq_len)
+        h = x + residual
+        ffn_res = self.feed_forward.forward(self.ffn_norm(h))
+        out = h + ffn_res
+        return out
+
+
+class Transformer(nn.Module):
+
+    def __init__(
+        self,
+        config,
+        max_seq_size=8192,
+    ):
+        super().__init__()
+        self.num_layers = config.num_layers
+        self.layers = self._build_layers(config)
+
+    def _build_layers(self, config):
+        layers = torch.nn.ModuleList()
+        for layer_id in range(self.num_layers):
+            layers.append(TransformerBlock(
+                config,
+                layer_id=layer_id + 1,
+            ))
+        return layers
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        cu_seqlens=None,
+        max_seq_len=None,
+    ):
+
+        if max_seq_len is not None and not isinstance(max_seq_len,
+                                                      torch.Tensor):
+            max_seq_len = torch.tensor(max_seq_len,
+                                       dtype=torch.int32,
+                                       device="cpu")
+
+        for lid, layer in enumerate(self.layers):
+            hidden_states = layer(
+                hidden_states,
+                attention_mask,
+                cu_seqlens,
+                max_seq_len,
+            )
+        return hidden_states
+
+
+class Step1Model(PreTrainedModel):
+    config_class = PretrainedConfig
+
+    @with_empty_init
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__(config)
+        self.tok_embeddings = LLaMaEmbedding(config)
+        self.transformer = Transformer(config)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+    ):
+
+        hidden_states = self.tok_embeddings(input_ids)
+
+        hidden_states = self.transformer(
+            hidden_states,
+            attention_mask,
+        )
+        return hidden_states
+
+
+class STEP1TextEncoder(torch.nn.Module):
+
+    def __init__(self, model_dir, max_length=320):
+        super().__init__()
+        self.max_length = max_length
+        self.text_tokenizer = Wrapped_StepChatTokenizer(
+            os.path.join(model_dir, 'step1_chat_tokenizer.model'))
+        text_encoder = Step1Model.from_pretrained(model_dir)
+        self.text_encoder = text_encoder.eval().to(torch.bfloat16)
+
+    @torch.no_grad
+    def forward(self, prompts, with_mask=True, max_length=None):
+        self.device = next(self.text_encoder.parameters()).device
+
+        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.bfloat16):
+            if type(prompts) is str:
+                prompts = [prompts]
+            txt_tokens = self.text_tokenizer(prompts,
+                                             max_length=max_length
+                                             or self.max_length,
+                                             padding="max_length",
+                                             truncation=True,
+                                             return_tensors="pt")
+            y = self.text_encoder(txt_tokens.input_ids.to(self.device),
+                                  attention_mask=txt_tokens.attention_mask.to(
+                                      self.device) if with_mask else None)
+            y_mask = txt_tokens.attention_mask
+        return y.transpose(0, 1), y_mask

fastvideo/v1/models/encoders/t5.py

@@ -28,13 +28,13 @@ import torch.nn.functional as F
 from torch import nn
 
 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)
 from fastvideo.v1.layers.activation import get_act_fn
 from fastvideo.v1.layers.layernorm import RMSNorm
 from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
                                         QKVParallelLinear, RowParallelLinear)
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from fastvideo.v1.models.encoders.base import TextEncoder
 from fastvideo.v1.models.loader.weight_utils import default_weight_loader

fastvideo/v1/models/loader/component_loader.py

@@ -6,6 +6,7 @@ import json
 import os
 import time
 from abc import ABC, abstractmethod
+from copy import deepcopy
 from typing import Any, Generator, Iterable, List, Optional, Tuple, cast
 
 import torch
@@ -366,6 +367,7 @@ class TransformerLoader(ComponentLoader):
              fastvideo_args: FastVideoArgs):
         """Load the transformer based on the model path, architecture, and inference args."""
         config = get_diffusers_config(model=model_path)
+        hf_config = deepcopy(config)
         cls_name = config.pop("_class_name")
         if cls_name is None:
             raise ValueError(
@@ -392,13 +394,30 @@ class TransformerLoader(ComponentLoader):
         default_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
 
         # Load the model using FSDP loader
-        logger.info("Loading model from %s", cls_name)
+        logger.info("Loading model from %s, default_dtype: %s", cls_name, default_dtype)
+        # model = load_fsdp_model(model_cls=model_cls,
+        #                         init_params={
+        #                             "config": dit_config,
+        #                             "hf_config": hf_config
+        #                         },
+        #                         weight_dir_list=safetensors_list,
+        #                         device=fastvideo_args.device,
+        #                         cpu_offload=fastvideo_args.use_cpu_offload,
+        #                         default_dtype=default_dtype)
         model = load_fsdp_model(model_cls=model_cls,
-                                init_params={"config": dit_config},
+                                init_params={
+                                    "config": dit_config,
+                                    "hf_config": hf_config
+                                },
                                 weight_dir_list=safetensors_list,
                                 device=fastvideo_args.device,
                                 cpu_offload=fastvideo_args.use_cpu_offload,
-                                default_dtype=default_dtype)
+                                default_dtype=default_dtype,
+                                # TODO(will): make these configurable
+                                param_dtype=torch.bfloat16,
+                                reduce_dtype=torch.float32,
+                                output_dtype=None,
+                                )
         if fastvideo_args.enable_torch_compile:
             logger.info("Torch Compile enabled for DiT")
             for n, m in reversed(list(model.named_modules())):
@@ -446,7 +465,8 @@ class SchedulerLoader(ComponentLoader):
         scheduler = scheduler_cls(**config)
         if fastvideo_args.flow_shift is not None:
             scheduler.set_shift(fastvideo_args.flow_shift)
-
+        if fastvideo_args.timesteps_scale is not None:
+            scheduler.set_timesteps_scale(fastvideo_args.timesteps_scale)
         return scheduler
 
 

fastvideo/v1/models/loader/fsdp_load.py

@@ -14,13 +14,16 @@ from typing import (Any, Callable, DefaultDict, Dict, Generator, Hashable, List,
 import torch
 from torch import nn
 from torch.distributed import DeviceMesh, init_device_mesh
-from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard
+from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard, MixedPrecisionPolicy
 from torch.distributed._tensor import distribute_tensor
 from torch.nn.modules.module import _IncompatibleKeys
 
 from fastvideo.v1.distributed.parallel_state import (
     get_sequence_model_parallel_world_size)
 from fastvideo.v1.models.loader.weight_utils import safetensors_weights_iterator
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
 
 
 # TODO(PY): move this to utils elsewhere
@@ -86,16 +89,29 @@ def get_param_names_mapping(
 
 
 # TODO(PY): add compile option
+    # param_dtype: torch.dtype,
+    # reduce_dtype: torch.dtype,
+    # output_dtype: torch.dtype,
+    # pp_enabled: bool = False,
+    # cpu_offload: bool = False,
 def load_fsdp_model(
     model_cls: Type[nn.Module],
     init_params: Dict[str, Any],
     weight_dir_list: List[str],
     device: torch.device,
+    default_dtype: torch.dtype,
+    param_dtype: torch.dtype,
+    reduce_dtype: torch.dtype,
     cpu_offload: bool = False,
-    default_dtype: Optional[torch.dtype] = torch.bfloat16,
+    output_dtype: Optional[torch.dtype] = None,
 ) -> torch.nn.Module:
+
+    mp_policy = MixedPrecisionPolicy(param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=True)
+
+    # with set_default_dtype(default_dtype), torch.device("meta"):
     with set_default_dtype(default_dtype), torch.device("meta"):
         model = model_cls(**init_params)
+
     device_mesh = init_device_mesh(
         "cuda",
         mesh_shape=(get_sequence_model_parallel_world_size(), ),
@@ -104,6 +120,7 @@ def load_fsdp_model(
     shard_model(model,
                 cpu_offload=cpu_offload,
                 reshard_after_forward=True,
+                mp_policy=mp_policy,
                 dp_mesh=device_mesh["dp"])
     weight_iterator = safetensors_weights_iterator(weight_dir_list)
     param_names_mapping_fn = get_param_names_mapping(model._param_names_mapping)
@@ -121,6 +138,7 @@ def load_fsdp_model(
                 f"Unexpected param or buffer {n} on meta device.")
     for p in model.parameters():
         p.requires_grad = False
+    # set_state_dict(model, StateDictType.LOCAL_STATE_DICT)
     return model
 
 
@@ -129,6 +147,7 @@ def shard_model(
     *,
     cpu_offload: bool,
     reshard_after_forward: bool = True,
+    mp_policy: Optional[MixedPrecisionPolicy] = None,
     dp_mesh: Optional[DeviceMesh] = None,
 ) -> None:
     """
@@ -156,14 +175,17 @@ def shard_model(
     """
     fsdp_kwargs = {
         "reshard_after_forward": reshard_after_forward,
-        "mesh": dp_mesh
+        "mesh": dp_mesh,
+        "mp_policy": mp_policy,
     }
     if cpu_offload:
         fsdp_kwargs["offload_policy"] = CPUOffloadPolicy()
 
-    # Shard the model with FSDP, iterating in reverse to start with
+    # iterating in reverse to start with
     # lowest-level modules first
     num_layers_sharded = 0
+    # TODO(will): don't reshard after forward for the last layer to save on the
+    # all-gather that will immediately happen Shard the model with FSDP,
     for n, m in reversed(list(model.named_modules())):
         if any([
                 shard_condition(n, m)
@@ -210,6 +232,10 @@ def load_fsdp_model_from_full_model_state_dict(
         NotImplementedError: If got FSDP with more than 1D.
     """
     meta_sharded_sd = model.state_dict()
+    # s = fully_shard.state(model)
+    # logger.info(f"type(s): {type(s)}")
+    # logger.info(f"s: {s}")
+    # import pdb; pdb.set_trace()
 
     sharded_sd = {}
     to_merge_params: DefaultDict[Hashable, Dict[Any, Any]] = defaultdict(dict)

fastvideo/v1/models/registry.py

@@ -23,6 +23,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "HunyuanVideoTransformer3DModel":
     ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+    "StepVideoModel": ("dits", "stepvideo", "StepVideoModel")
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -34,6 +35,8 @@ _TEXT_ENCODER_MODELS = {
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "LlamaModel": ("encoders", "llama", "LlamaModel"),
     "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"),
+    "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"),
+    "BertModel": ("encoders", "clip", "CLIPTextModel"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -45,6 +48,7 @@ _VAE_MODELS = {
     "AutoencoderKLHunyuanVideo":
     ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
     "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),
+    "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo")
 }
 
 _SCHEDULERS = {

fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -153,9 +153,11 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             sigmas = 1 - sigmas
 
         self.sigmas = sigmas
-        self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(
-            dtype=torch.float32, device=device)
-
+        if not getattr(self.config, "timesteps_scale", True):
+            self.timesteps = sigmas[:-1]  # for stepvideo
+        else:
+            self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(
+                dtype=torch.float32, device=device)
         # Reset step index
         self._step_index = None
 
@@ -178,6 +180,9 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     def set_shift(self, shift: float) -> None:
         self.config.shift = shift
 
+    def set_timesteps_scale(self, timesteps_scale: bool) -> None:
+        self.config.timesteps_scale = timesteps_scale
+
     def _init_step_index(self, timestep) -> None:
         if self.begin_index is None:
             if isinstance(timestep, torch.Tensor):

fastvideo/v1/pipelines/composed_pipeline_base.py

@@ -5,19 +5,25 @@ Base class for composed pipelines.
 This module defines the base class for pipelines that are composed of multiple stages.
 """
 
+import argparse
 import os
 from abc import ABC, abstractmethod
 from copy import deepcopy
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Dict, List, Optional, Union, cast
 
 import torch
 
-from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.pipelines import (PipelineConfig,
+                                            get_pipeline_config_cls_for_name)
+from fastvideo.v1.distributed import (init_distributed_environment,
+                                      initialize_model_parallel,
+                                      model_parallel_is_initialized)
+from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.loader.component_loader import PipelineComponentLoader
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages import PipelineStage
-from fastvideo.v1.utils import (maybe_download_model,
+from fastvideo.v1.utils import (maybe_download_model, shallow_asdict,
                                 verify_model_config_and_directory)
 
 logger = init_logger(__name__)
@@ -39,15 +45,20 @@ class ComposedPipelineBase(ABC):
     def __init__(self,
                  model_path: str,
                  fastvideo_args: FastVideoArgs,
-                 config: Optional[Dict[str, Any]] = None):
+                 config: Optional[Dict[str, Any]] = None,
+                 required_config_modules: Optional[List[str]] = None):
         """
         Initialize the pipeline. After __init__, the pipeline should be ready to
         use. The pipeline should be stateless and not hold any batch state.
         """
+        self.fastvideo_args = fastvideo_args
         self.model_path = model_path
         self._stages: List[PipelineStage] = []
         self._stage_name_mapping: Dict[str, PipelineStage] = {}
 
+        if required_config_modules is not None:
+            self._required_config_modules = required_config_modules
+
         if self._required_config_modules is None:
             raise NotImplementedError(
                 "Subclass must set _required_config_modules")
@@ -59,16 +70,135 @@ class ComposedPipelineBase(ABC):
         else:
             self.config = config
 
+        self.maybe_init_distributed_environment(fastvideo_args)
+
         # Load modules directly in initialization
         logger.info("Loading pipeline modules...")
         self.modules = self.load_modules(fastvideo_args)
 
+        if fastvideo_args.training_mode:
+            if fastvideo_args.log_validation:
+                self.initialize_validation_pipeline(fastvideo_args)
+            self.initialize_training_pipeline(fastvideo_args)
+
         self.initialize_pipeline(fastvideo_args)
 
-        logger.info("Creating pipeline stages...")
-        self.create_pipeline_stages(fastvideo_args)
+        # logger.info("Creating pipeline stages...")
+        # self.create_pipeline_stages(fastvideo_args)
 
-    def get_module(self, module_name: str) -> Any:
+        if fastvideo_args.training_mode:
+            logger.info("Creating training pipeline stages...")
+            self.create_training_stages(fastvideo_args)
+        else:
+            logger.info("Creating pipeline stages...")
+            self.create_pipeline_stages(fastvideo_args)
+
+    def initialize_training_pipeline(self, fastvideo_args: FastVideoArgs):
+        raise NotImplementedError(
+            "if training_mode is True, the pipeline must implement this method")
+
+    def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
+        raise NotImplementedError(
+            "if log_validation is True, the pipeline must implement this method"
+        )
+
+    @classmethod
+    def from_pretrained(cls,
+                        model_path: str,
+                        device: Optional[str] = None,
+                        torch_dtype: Optional[torch.dtype] = None,
+                        pipeline_config: Optional[
+                            Union[str
+                                  | PipelineConfig]] = None,
+                        args: Optional[argparse.Namespace] = None,
+                        required_config_modules: Optional[List[str]] = None,
+                        **kwargs) -> "ComposedPipelineBase":
+        config = None
+        # 1. If users provide a pipeline config, it will override the default pipeline config
+        if isinstance(pipeline_config, PipelineConfig):
+            config = pipeline_config
+        else:
+            config_cls = get_pipeline_config_cls_for_name(model_path)
+            if config_cls is not None:
+                config = config_cls()
+                if isinstance(pipeline_config, str):
+                    config.load_from_json(pipeline_config)
+
+        # 2. If users also provide some kwargs, it will override the pipeline config.
+        # The user kwargs shouldn't contain model config parameters!
+        if config is None:
+            logger.warning("No config found for model %s, using default config",
+                           model_path)
+            config_args = kwargs
+        else:
+            config_args = shallow_asdict(config)
+            config_args.update(kwargs)
+
+        if args.inference_mode:
+            fastvideo_args = FastVideoArgs(model_path=model_path,
+                                           device_str=device or "cuda" if
+                                           torch.cuda.is_available() else "cpu",
+                                           **config_args)
+
+            fastvideo_args.model_path = model_path
+            fastvideo_args.device_str = device or "cuda" if torch.cuda.is_available(
+            ) else "cpu"
+            for key, value in config_args.items():
+                setattr(fastvideo_args, key, value)
+        else:
+            assert args is not None, "args must be provided for training mode"
+            fastvideo_args = TrainingArgs.from_cli_args(args)
+            # TODO(will): fix this so that its not so ugly
+            fastvideo_args.model_path = model_path
+            fastvideo_args.device_str = device or "cuda" if torch.cuda.is_available(
+            ) else "cpu"
+            for key, value in config_args.items():
+                setattr(fastvideo_args, key, value)
+
+            # we use cpu offload for training 
+            fastvideo_args.use_cpu_offload = False
+            # make sure we are in training mode
+            fastvideo_args.inference_mode = False
+            # we hijack the precision to be the master weight type so that the
+            # model is loaded with the correct precision. Subsequently we will
+            # use FSDP2's MixedPrecisionPolicy to set the precision for the
+            # fwd, bwd, and other operations' precision.
+            fastvideo_args.precision = fastvideo_args.master_weight_type
+            assert fastvideo_args.precision == 'fp32', 'only fp32 is supported for training'
+
+        fastvideo_args.check_fastvideo_args()
+
+        logger.info(f"fastvideo_args in from_pretrained: {fastvideo_args}")
+
+        return cls(model_path,
+                   fastvideo_args,
+                   required_config_modules=required_config_modules)
+
+    def maybe_init_distributed_environment(self, fastvideo_args: FastVideoArgs):
+        if model_parallel_is_initialized():
+            return
+        local_rank = int(os.environ.get("LOCAL_RANK", -1))
+        world_size = int(os.environ.get("WORLD_SIZE", -1))
+        rank = int(os.environ.get("RANK", -1))
+
+        if local_rank == -1 or world_size == -1 or rank == -1:
+            raise ValueError(
+                "Local rank, world size, and rank must be set. Use torchrun to launch the script."
+            )
+
+        torch.cuda.set_device(local_rank)
+        init_distributed_environment(world_size=world_size,
+                                     rank=rank,
+                                     local_rank=local_rank)
+        initialize_model_parallel(
+            tensor_model_parallel_size=fastvideo_args.tp_size,
+            sequence_model_parallel_size=fastvideo_args.sp_size)
+        device = torch.device(f"cuda:{local_rank}")
+        fastvideo_args.device = device
+
+    def get_module(self, module_name: str, default_value: Any = None) -> Any:
+        if module_name not in self.modules:
+            return default_value
         return self.modules[module_name]
 
     def add_module(self, module_name: str, module: Any):
@@ -114,6 +244,19 @@ class ComposedPipelineBase(ABC):
         """
         raise NotImplementedError
 
+    # @abstractmethod
+    # def create_validation_stages(self, fastvideo_args: FastVideoArgs):
+    #     """
+    #     Create the validation pipeline stages.
+    #     """
+    #     raise NotImplementedError
+
+    def create_training_stages(self, fastvideo_args: FastVideoArgs):
+        """
+        Create the training pipeline stages.
+        """
+        raise NotImplementedError
+
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
         """
         Initialize the pipeline.
@@ -136,19 +279,21 @@ class ComposedPipelineBase(ABC):
             modules_config
         ) > 1, "model_index.json must contain at least one pipeline module"
 
-        required_modules = [
-            "vae", "text_encoder", "transformer", "scheduler", "tokenizer"
-        ]
-        for module_name in required_modules:
+        for module_name in self.required_config_modules:
             if module_name not in modules_config:
                 raise ValueError(
                     f"model_index.json must contain a {module_name} module")
-        logger.info("Diffusers config passed sanity checks")
 
         # all the component models used by the pipeline
+        required_modules = self.required_config_modules
+        logger.info("Loading required modules: %s", required_modules)
+
         modules = {}
         for module_name, (transformers_or_diffusers,
                           architecture) in modules_config.items():
+            if module_name not in required_modules:
+                logger.info("Skipping module %s", module_name)
+                continue
             component_model_path = os.path.join(self.model_path, module_name)
             module = PipelineComponentLoader.load_module(
                 module_name=module_name,
@@ -164,7 +309,6 @@ class ComposedPipelineBase(ABC):
                 logger.warning("Overwriting module %s", module_name)
             modules[module_name] = module
 
-        required_modules = self.required_config_modules
         # Check if all required modules were loaded
         for module_name in required_modules:
             if module_name not in modules or modules[module_name] is None:
@@ -198,7 +342,7 @@ class ComposedPipelineBase(ABC):
         # Execute each stage
         logger.info("Running pipeline stages: %s",
                     self._stage_name_mapping.keys())
-        logger.info("Batch: %s", batch)
+        # logger.info("Batch: %s", batch)
         for stage in self.stages:
             batch = stage(batch, fastvideo_args)
 

fastvideo/v1/pipelines/distill_pipeline.py

@@ -0,0 +1,303 @@
+import os
+import sys
+import time
+from collections import deque
+
+import torch
+from tqdm.auto import tqdm
+
+# import torch.distributed as dist
+import wandb
+from fastvideo.utils.checkpoint import save_checkpoint
+from fastvideo.utils.communications import sp_parallel_dataloader_wrapper
+from fastvideo.utils.validation import log_validation
+from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines import ComposedPipelineBase
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage,
+                                           DenoisingStage, InputValidationStage,
+                                           LatentPreparationStage,
+                                           TextEncodingStage,
+                                           TimestepPreparationStage)
+
+logger = init_logger(__name__)
+
+
+class WanTrainingPipeline(ComposedPipelineBase):  # == distill_one_step
+    _required_config_modules = ["scheduler", "transformer"]
+
+    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")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ):
+        device = fastvideo_args.device
+        local_rank = int(os.environ.get("LOCAL_RANK", -1))
+        rank = int(os.environ.get("RANK", -1))
+        assert rank != -1
+        assert local_rank != -1
+        sp_group = get_sp_group()
+        world_size = sp_group.world_size
+        rank = sp_group.rank
+        args = fastvideo_args
+        transformer = self.get_module("transformer")
+        teacher_transformer = self.get_module("teacher_transformer")
+        ema_transformer = None
+        assert not fastvideo_args.use_ema, "ema is not supported now"
+        assert teacher_transformer is not None
+        assert transformer is not None
+        train_dataset = self.train_dataset
+        train_dataloader = self.train_dataloader
+        init_steps = self.init_steps
+        lr_scheduler = self.lr_scheduler
+        optimizer = self.optimizer
+        noise_scheduler = self.noise_scheduler
+        solver = self.solver
+        noise_random_generator = None
+        uncond_prompt_embed = self.uncond_prompt_embed
+        uncond_prompt_mask = self.uncond_prompt_mask
+
+        # Train!
+        total_batch_size = (world_size * args.gradient_accumulation_steps /
+                            args.sp_size * args.train_sp_batch_size)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %s", len(train_dataset))
+        logger.info("  Dataloader size = %s", len(train_dataloader))
+        logger.info("  Num Epochs = %s", args.num_train_epochs)
+        logger.info("  Resume training from step %s", init_steps)
+        logger.info("  Instantaneous batch size per device = %s",
+                    args.train_batch_size)
+        logger.info(
+            "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
+            total_batch_size)
+        logger.info("  Gradient Accumulation steps = %s",
+                    args.gradient_accumulation_steps)
+        logger.info("  Total optimization steps = %s", args.max_train_steps)
+        logger.info(
+            "  Total training parameters per FSDP shard = %s B",
+            sum(p.numel()
+                for p in transformer.parameters() if p.requires_grad) / 1e9)
+        # print dtype
+        logger.info("  Master weight dtype: %s",
+                    transformer.parameters().__next__().dtype)
+
+        # Potentially load in the weights and states from a previous save
+        if args.resume_from_checkpoint:
+            assert NotImplementedError(
+                "resume_from_checkpoint is not supported now.")
+            # TODO
+
+        progress_bar = tqdm(
+            range(0, args.max_train_steps),
+            initial=init_steps,
+            desc="Steps",
+            # Only show the progress bar once on each machine.
+            disable=local_rank > 0,
+        )
+
+        loader = sp_parallel_dataloader_wrapper(
+            train_dataloader,
+            device,
+            args.train_batch_size,
+            args.sp_size,
+            args.train_sp_batch_size,
+        )
+        # loader = self.get_module("train_dataloader")
+
+        step_times = deque(maxlen=100)
+
+        # todo future
+        for i in range(init_steps):
+            next(loader)
+
+        # log_validation(args, transformer, device,
+        #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
+        def get_num_phases(multi_phased_distill_schedule, step):
+            # step-phase,step-phase
+            multi_phases = multi_phased_distill_schedule.split(",")
+            phase = multi_phases[-1].split("-")[-1]
+            for step_phases in multi_phases:
+                phase_step, phase = step_phases.split("-")
+                if step <= int(phase_step):
+                    return int(phase)
+            return phase
+
+        for step in range(init_steps + 1, args.max_train_steps + 1):
+            start_time = time.time()
+            assert args.multi_phased_distill_schedule is not None
+            num_phases = get_num_phases(args.multi_phased_distill_schedule,
+                                        step)
+
+            loss, grad_norm, pred_norm = self.distill_one_step(
+                transformer,
+                args.model_type,
+                teacher_transformer,
+                ema_transformer,
+                optimizer,
+                lr_scheduler,
+                loader,
+                noise_scheduler,
+                solver,
+                noise_random_generator,
+                args.gradient_accumulation_steps,
+                args.sp_size,
+                args.max_grad_norm,
+                uncond_prompt_embed,
+                uncond_prompt_mask,
+                args.num_euler_timesteps,
+                num_phases,
+                args.not_apply_cfg_solver,
+                args.distill_cfg,
+                args.ema_decay,
+                args.pred_decay_weight,
+                args.pred_decay_type,
+                args.hunyuan_teacher_disable_cfg,
+            )
+
+            step_time = time.time() - start_time
+            step_times.append(step_time)
+            avg_step_time = sum(step_times) / len(step_times)
+
+            progress_bar.set_postfix({
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+                "phases": num_phases,
+            })
+            progress_bar.update(1)
+            if rank <= 0:
+                wandb.log(
+                    {
+                        "train_loss":
+                        loss,
+                        "learning_rate":
+                        lr_scheduler.get_last_lr()[0],
+                        "step_time":
+                        step_time,
+                        "avg_step_time":
+                        avg_step_time,
+                        "grad_norm":
+                        grad_norm,
+                        "pred_fro_norm":
+                        pred_norm["fro"],  # codespell:ignore
+                        "pred_largest_singular_value":
+                        pred_norm["largest singular value"],
+                        "pred_absolute_mean":
+                        pred_norm["absolute mean"],
+                        "pred_absolute_max":
+                        pred_norm["absolute max"],
+                    },
+                    step=step,
+                )
+            if step % args.checkpointing_steps == 0:
+                if args.use_lora:
+                    # Save LoRA weights
+                    raise NotImplementedError("lora is not supported now")
+                    # save_lora_checkpoint(transformer, optimizer, rank,
+                    #                      args.output_dir, step)
+                else:
+                    # Your existing checkpoint saving code
+                    if args.use_ema:
+                        raise NotImplementedError("ema is not supported now")
+                        save_checkpoint(ema_transformer, rank, args.output_dir,
+                                        step)
+                    else:
+                        save_checkpoint(transformer, rank, args.output_dir,
+                                        step)
+
+                sp_group.barrier()
+            if args.log_validation and step % args.validation_steps == 0:
+                log_validation(
+                    args,
+                    transformer,
+                    device,
+                    torch.bfloat16,
+                    step,
+                    scheduler_type=args.scheduler_type,
+                    shift=args.shift,
+                    num_euler_timesteps=args.num_euler_timesteps,
+                    linear_quadratic_threshold=args.linear_quadratic_threshold,
+                    linear_range=args.linear_range,
+                    ema=False,
+                )
+                if args.use_ema:
+                    log_validation(
+                        args,
+                        ema_transformer,
+                        device,
+                        torch.bfloat16,
+                        step,
+                        scheduler_type=args.scheduler_type,
+                        shift=args.shift,
+                        num_euler_timesteps=args.num_euler_timesteps,
+                        linear_quadratic_threshold=args.
+                        linear_quadratic_threshold,
+                        linear_range=args.linear_range,
+                        ema=True,
+                    )
+
+        if args.use_lora:
+            raise NotImplementedError("lora is not supported now")
+            save_lora_checkpoint(transformer, optimizer, rank, args.output_dir,
+                                 args.max_train_steps)
+        else:
+            save_checkpoint(transformer, rank, args.output_dir,
+                            args.max_train_steps)
+
+        if get_sp_group():
+            cleanup_dist_env_and_memory()
+
+
+def main(args):
+    logger.info("Starting training pipeline...")
+    pipeline = WanTrainingPipeline.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", args=args)
+    args = pipeline.fastvideo_args
+    pipeline.forward(None, args)
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.v1.fastvideo_args import TrainingArgs
+    from fastvideo.v1.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    print(args)
+    main(args)