update sfwan ssim

Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

.buildkite/pipeline.yml

@@ -104,6 +104,18 @@ steps:
                   - TEST_TYPE=distillation_dmd
                 agents:
                   queue: "default"
+            - path:
+                - "fastvideo/training/*self_forcing_distillation_pipeline.py"
+                - "fastvideo/tests/training/self-forcing/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Self-Forcing Tests"
+                env:
+                  - TEST_TYPE=self_forcing
+                agents:
+                  queue: "default"
             - path:
                 - "fastvideo/**"
                 - "pyproject.toml"
@@ -198,4 +210,15 @@ steps:
                 env:  
                   - TEST_TYPE=inference_vmoba
                 agents:
-                  queue: "default"
+                  queue: "default"
+            - path:
+                - "fastvideo/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Unit Tests"
+                env:
+                  - TEST_TYPE=unit_test
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -110,6 +110,10 @@ case "$TEST_TYPE" in
         MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_distill_dmd_tests"
         ;;
         # run_inference_tests_vmoba
+    "self_forcing")
+        log "Running self-forcing tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_self_forcing_tests"
+        ;;
     "inference_vmoba")
         log "Running V-MoBA inference tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_vmoba"
@@ -118,6 +122,10 @@ case "$TEST_TYPE" in
         log "Running V-MoBA precision tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_vmoba"
         ;;
+    "unit_test")
+        log "Running unit tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_unit_test"
+        ;;
     *)
         log "Error: Unknown test type: $TEST_TYPE"
         exit 1

.github/workflows/pr-test.yml

@@ -62,8 +62,8 @@ on:
         required: false
         default: false
         type: boolean
-      run_nightly_test:
-        description: "Run nightly-test"
+      run_unit_test:
+        description: "Run unit-test"
         required: false
         default: false
         type: boolean
@@ -93,6 +93,7 @@ jobs:
       inference-test-STA: ${{ steps.filter.outputs.inference-test-STA }}
       precision-test-STA: ${{ steps.filter.outputs.precision-test-STA }}
       precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }}
+      unit-test: ${{ steps.filter.outputs.unit-test }}
     steps:
       - uses: actions/checkout@v4
       - uses: dorny/paths-filter@v3
@@ -102,6 +103,8 @@ jobs:
             # Define reusable path patterns
             common-paths: &common-paths
               - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.10'
+              - 'docker/Dockerfile.python3.11'
               - 'docker/Dockerfile.python3.12'
             sta-kernel-paths: &sta-kernel-paths
               - 'csrc/attn/sliding_tile_attn/**'
@@ -155,6 +158,9 @@ jobs:
             precision-test-VSA:
               - *common-paths
               - *vsa-kernel-paths
+            unit-test:
+              - 'fastvideo/**'
+              - *common-paths
 
   encoder-test:
     needs: change-filter
@@ -333,23 +339,42 @@ jobs:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
 
-  nightly-test:
+  unit-test:
+    needs: change-filter
     if: >-
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
+      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.unit-test == 'true') || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_unit_test == 'true')
     uses: ./.github/workflows/runpod-test.yml
     with:
-      job_id: "nightly-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 4
+      job_id: "unit-test"
+      gpu_type: "NVIDIA L40S"
+      gpu_count: 1
       volume_size: 100
       disk_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/dataset/ -vs && pytest ./fastvideo/workflow/ -vs"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+
+  # nightly-test:
+  #   if: >-
+  #     (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
+  #   uses: ./.github/workflows/runpod-test.yml
+  #   with:
+  #     job_id: "nightly-test"
+  #     gpu_type: "NVIDIA A40"
+  #     gpu_count: 4
+  #     volume_size: 100
+  #     disk_size: 100
+  #     image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+  #     test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+  #     timeout_minutes: 30
+  #   secrets:
+  #     RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+  #     RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
+  #     WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
 
   runpod-cleanup:
     # Add other jobs to this list as you create them

.gitignore

@@ -64,4 +64,6 @@ docs/source/distillation/examples/
 !docs/source/_static/images/**/*.png
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif
-dmd_t2v_output/
+
+dmd_t2v_output/
+preprocess_output_text/

.pre-commit-config.yaml

@@ -12,9 +12,6 @@ exclude: |
     scripts/.*|
     fastvideo/data_preprocess/.*|
     fastvideo/dataset/.*|
-    fastvideo/distill/.*|
-    fastvideo/distill\.py|
-    fastvideo/distill_adv\.py|
     fastvideo/models/.*|
     fastvideo/sample/.*|
     fastvideo/train\.py|
@@ -44,10 +41,10 @@ repos:
   - id: codespell
     additional_dependencies: ['tomli']
     args: ['--toml', 'pyproject.toml']
-- repo: https://github.com/PyCQA/isort
-  rev: 6.0.1
-  hooks:
-  - id: isort
+# - repo: https://github.com/PyCQA/isort
+#   rev: 6.0.1
+#   hooks:
+#   - id: isort
 - repo: https://github.com/jackdewinter/pymarkdown
   rev: v0.9.30
   hooks:

README.md

@@ -7,7 +7,7 @@
 FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
 
 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</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/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/S7HLCSTh" target="_blank"> <b> WeChat </b> </a> |
+    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</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/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/q46BbX6" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">
@@ -155,8 +155,8 @@ If you find FastVideo useful, please considering citing our work:
 }
 
 @article{zhang2025vsa,
-  title={VSA: Faster Video Diffusion with Trainable Sparse Attention},
-  author={Zhang, Peiyuan and Huang, Haofeng and Chen, Yongqi and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
+  title={Vsa: Faster video diffusion with trainable sparse attention},
+  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
   journal={arXiv preprint arXiv:2505.13389},
   year={2025}
 }

csrc/attn/vmoba_attn/setup.py

@@ -20,5 +20,7 @@ setup(
         "License :: OSI Approved :: Apache Software License",
     ],
     python_requires='>=3.12',
-    install_requires=[]
+    install_requires=[
+        "flash-attn >= 2.7.1",
+    ]
 )

csrc/attn/vmoba_attn/vmoba/vmoba.py

@@ -6,8 +6,16 @@ import time
 import os
 import torch
 from typing import Tuple
-from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
-from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+try:
+    from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
+    from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+except ImportError:
+    def _unsupported(*args, **kwargs):
+        raise ImportError("flash-attn is not installed. Please install it, e.g., `pip install flash-attn`.")
+    _flash_attn_varlen_forward = _unsupported
+    _flash_attn_varlen_backward = _unsupported
+    flash_attn_varlen_func = _unsupported
+    
 from functools import lru_cache
 from einops import rearrange
 

docs/source/inference/optimizations.md

@@ -1,32 +1,40 @@
 (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)
+  - [Sage Attention 3](#optimizations-sage3)
 
 - 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`
 - Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`
 - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
+- Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE`
 
 ### 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
@@ -34,6 +42,7 @@ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
 ```
 
 #### 2. In CLI
+
 You can also set the environment variable on the command line:
 
 ```bash
@@ -41,6 +50,7 @@ FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
 ```
 
 (optimizations-flash)=
+
 ### Flash Attention
 
 **`FLASH_ATTN`**
@@ -57,7 +67,7 @@ And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://git
 git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention
 
 cd hopper
-pip install ninja 
+pip install ninja
 python setup.py install
 ```
 
@@ -66,7 +76,9 @@ FastVideo will automatically detect and use `FA3` if it is installed when using
 :::
 
 (optimizations-sta)=
+
 ### Sliding Tile Attention
+
 **`SLIDING_TILE_ATTN`**
 
 ```bash
@@ -76,7 +88,9 @@ pip install st_attn==0.0.4
 Please see [this page](#sta-installation) for more installation instructions.
 
 (optimizations-vsa)=
+
 ### Video Sparse Attention
+
 **`VIDEO_SPARSE_ATTN`**
 
 ```bash
@@ -87,19 +101,45 @@ python setup_vsa.py install
 Please see [this page](#vsa-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 
+cd sageattention
 python setup.py install  # or pip install -e .
 ```
 
+(optimizations-sage3)=
+
+### Sage Attention 3
+
+**`SAGE_ATTN_THREE`**
+
+[SageAttention 3](https://huggingface.co/jt-zhang/SageAttention3) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements.
+
+#### Hardware Requirements
+
+- RTX5090
+
+#### Installation
+
+Note that Sage Attention 3 requires `python>=3.13`, `torch>=2.8.0`, `CUDA >=12.8`. If you are using `uv` and using `torch==2.8.0` make sure that `sentencepiece==0.2.1` in the pyproject.toml file.
+
+To use Sage Attention 3 in FastVideo, first get access to the SageAttention3 code, then move `sageattn/` and `setup.py` to the directory `fastvideo/attention/backends`, then install from using:
+
+```bash
+python setup.py install
+```
+
 (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?

examples/dataset/vidprom/README.md

@@ -0,0 +1,9 @@
+# VidProm Dataset
+
+From [Self-Forcing](https://github.com/gdhe17/Self-Forcing) repository.
+
+## Download the dataset
+
+```bash
+./download_dataset.sh
+```

examples/dataset/vidprom/download_dataset.sh

@@ -0,0 +1,3 @@
+#! /bin/bash
+
+huggingface-cli download gdhe17/Self-Forcing vidprom_filtered_extended.txt --local-dir prompts

examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh

@@ -1,21 +1,15 @@
 #!/bin/bash
-#SBATCH --job-name=wl_t2v
+#SBATCH --job-name=t2v
 #SBATCH --partition=main
-#SBATCH --nodes=4
-#SBATCH --ntasks=4
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
 #SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
+#SBATCH --gres=gpu:1
 #SBATCH --cpus-per-task=128
 #SBATCH --mem=1440G
-#SBATCH --output=dmd_t2v_output/sf.out
-#SBATCH --error=dmd_t2v_output/sf.err
+#SBATCH --output=dmd_t2v_output/t2v_%j.out
+#SBATCH --error=dmd_t2v_output/t2v_%j.err
 #SBATCH --exclusive
-set -e -x
-
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-fv
 
 # Basic Info
 export NCCL_P2P_DISABLE=1
@@ -24,82 +18,71 @@ export TORCH_NCCL_ENABLE_MONITORING=0
 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
 export MASTER_PORT=29503
 export TOKENIZERS_PARALLELISM=false
-export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
+export WANDB_API_KEY=your_wandb_api_key
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 
 # Configs
-NUM_GPUS=8
+NUM_GPUS=1
 
 # Model paths for Self-Forcing DMD distillation:
 GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
 FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
 
-# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
-DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
-VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
-VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+DATA_DIR=your_data_dir
+VALIDATION_DATASET_FILE=your_validation_data_dir
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
-# Training arguments
 training_args=(
-  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
-  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan_t2v_finetune"
-  # --use_sf_wan
-  # --sf_ode_init_path "checkpoints/ode_init.pt"
-  --max_train_steps 6000
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  
+  --output_dir your_output_dir
+  --max_train_steps 4000
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
   --num_latent_t 21
   --num_height 480
   --num_width 832
-  --num_frames 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
   --enable_gradient_checkpointing_type "full"
   --log_visualization
   --simulate_generator_forward
+  --num_frames 81
   --num_frame_per_block 3  # Frame generation block size for self-forcing
   --enable_gradient_masking
   --gradient_mask_last_n_frames 21
 )
 
-# Parallel arguments
 parallel_args=(
-  --num_gpus 32 # 64
+  --num_gpus $NUM_GPUS # 64
   --sp_size 1
   --tp_size 1
-  --hsdp_replicate_dim 32 # 64
-  --hsdp_shard_dim 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_GPUS
 )
 
-# Model arguments
 model_args=(
   --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
   --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
-  --generator_model_path $GENERATOR_MODEL_PATH
   --real_score_model_path $REAL_SCORE_MODEL_PATH
   --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
-# Dataset arguments
 dataset_args=(
   --data_path "$DATA_DIR"
   --dataloader_num_workers 4
 )
 
-# Validation arguments
 validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 10
+  --validation_steps 50
   --validation_sampling_steps "4"
   --validation_guidance_scale "6.0" # not used for dmd inference
 )
 
-# Optimizer arguments
 optimizer_args=(
   --learning_rate 1e-5
   --mixed_precision "bf16"
@@ -110,7 +93,6 @@ optimizer_args=(
   --max_grad_norm 1.0
 )
 
-# Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
   --checkpoints_total_limit 3
@@ -121,11 +103,9 @@ miscellaneous_args=(
   --use_ema True
   --ema_decay 0.99
   --ema_start_step 100
-  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
-  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/FastVideo2/warp_vidprom_8b16k_test_warp_1e-5/checkpoint-2000/transformer/diffusion_pytorch_model.safetensors"
+  --init_weights_from_safetensors your_ode_init_weights_path
 )
 
-# Self-forcing DMD arguments
 dmd_args=(
   --dmd_denoising_steps '1000,750,500,250'
   --min_timestep_ratio 0.02
@@ -137,20 +117,16 @@ dmd_args=(
   --warp_denoising_step
 )
 
-# Self-forcing specific arguments
 self_forcing_args=(
   --independent_first_frame False  # Whether to treat first frame independently
-  --same_step_across_blocks False  # Whether to use same denoising step across all blocks
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
   --last_step_only False  # Whether to only use the last denoising step
   --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
-  --validate_cache_structure False  # Set to True for debugging KV cache issues
 )
 
-srun torchrun \
---nnodes $SLURM_JOB_NUM_NODES \
---node_rank $SLURM_PROCID \
---rdzv_backend=c10d \
---rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+torchrun \
+--nnodes 1 \
+--master_port $MASTER_PORT \
 --nproc_per_node $NUM_GPUS \
     fastvideo/training/wan_self_forcing_distillation_pipeline.py \
     "${parallel_args[@]}" \
@@ -161,4 +137,4 @@ srun torchrun \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
     "${dmd_args[@]}" \
-    "${self_forcing_args[@]}"
+    "${self_forcing_args[@]}"

examples/distill/SFWan2.1-T2V/download_dataset.sh

@@ -1,3 +1,3 @@
 #!/bin/bash
 
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/distill/SFWan2.1-T2V/preprocess_data.sh

@@ -21,4 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
     --samples_per_file 8 \
     --flush_frequency 8 \
     --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 
+    --preprocess_task "t2v" 

examples/distill/SFWan2.2-A14B/distill_dmd.sh

@@ -0,0 +1,157 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=4
+#SBATCH --ntasks=4
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_t2v_output/t2v_%j.out
+#SBATCH --error=dmd_t2v_output/t2v_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+# Configs
+NUM_GPUS=8
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Updated to Wan2.2
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
+  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan2.2_t2v_finetune"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  --simulate_generator_forward
+  # --log_visualization
+  --num_frames 81
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
+)
+
+parallel_args=(
+  --num_gpus 32 # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim 32
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 20
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 100
+)
+
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 3.0
+  --fake_score_learning_rate 8e-6
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm

@@ -39,15 +39,18 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # Training arguments
 training_args=(
   --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir"checkpoints/wan_t2v_finetune"
+  --output_dir $OUTPUT_DIR
   --max_train_steps 4000
   --train_batch_size 1
   --train_sp_batch_size 1
@@ -72,6 +75,8 @@ parallel_args=(
 model_args=(
   --model_path $MODEL_PATH
   --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
 # Dataset arguments
@@ -91,7 +96,7 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 2e-6
   --mixed_precision "bf16"
   --training_state_checkpointing_steps 500
   --weight_only_checkpointing_steps 500
@@ -134,4 +139,4 @@ srun torchrun \
     "${optimizer_args[@]}" \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm

@@ -39,15 +39,18 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # Training arguments
 training_args=(
   --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir "checkpoints/wan_t2v_finetune"
+  --output_dir "$OUTPUT_DIR"
   --max_train_steps 4000
   --train_batch_size 1
   --train_sp_batch_size 1
@@ -72,6 +75,8 @@ parallel_args=(
 model_args=(
   --model_path $MODEL_PATH
   --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
 # Dataset arguments
@@ -91,7 +96,7 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 2e-6
   --mixed_precision "bf16"
   --training_state_checkpointing_steps 500
   --weight_only_checkpointing_steps 500
@@ -134,4 +139,4 @@ srun torchrun \
     "${optimizer_args[@]}" \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm

@@ -39,15 +39,18 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # Training arguments
 training_args=(
   --tracker_project_name wan_t2v_distill_dmd
-  --output_dir "checkpoints/wan_t2v_finetune"
+  --output_dir "$OUTPUT_DIR"
   --max_train_steps 4000
   --train_batch_size 1
   --train_sp_batch_size 1
@@ -72,6 +75,8 @@ parallel_args=(
 model_args=(
   --model_path $MODEL_PATH
   --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
 # Dataset arguments
@@ -91,7 +96,7 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 2e-6
   --mixed_precision "bf16"
   --training_state_checkpointing_steps 500
   --weight_only_checkpointing_steps 500
@@ -133,4 +138,4 @@ srun torchrun \
     "${optimizer_args[@]}" \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/download_dataset.sh

@@ -1,3 +1,3 @@
 #!/bin/bash
 
-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"

examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh

@@ -40,15 +40,18 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DIR=your_validation_path  #(example:validation_64.json)
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # Training arguments
 training_args=(
   --tracker_project_name Wan_distillation
-  --output_dir "your_output_dir"
+  --output_dir "$OUTPUT_DIR"
   --max_train_steps 4000
   --train_batch_size 1
   --train_sp_batch_size 1
@@ -73,6 +76,8 @@ parallel_args=(
 model_args=(
   --model_path $MODEL_PATH
   --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
 # Dataset arguments
@@ -92,11 +97,11 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-5
+  --learning_rate 4e-6
   --lr_scheduler "cosine_with_min_lr"
   --min_lr_ratio 0.5
   --lr_warmup_steps 100
-  --fake_score_learning_rate 1e-5
+  --fake_score_learning_rate 2e-6
   --fake_score_lr_scheduler "cosine_with_min_lr"
   --mixed_precision "bf16"
   --training_state_checkpointing_steps 500
@@ -141,4 +146,4 @@ srun torchrun \
     "${optimizer_args[@]}" \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"

examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh

@@ -40,6 +40,8 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DIR=your_validation_path  #(example:validation_64.json)
 # export CUDA_VISIBLE_DEVICES=4,5
@@ -73,6 +75,8 @@ parallel_args=(
 model_args=(
   --model_path $MODEL_PATH
   --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
 # Dataset arguments
@@ -92,11 +96,11 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-5
+  --learning_rate 4e-6
   --lr_scheduler "cosine_with_min_lr"
   --min_lr_ratio 0.5
   --lr_warmup_steps 100
-  --fake_score_learning_rate 1e-5
+  --fake_score_learning_rate 2e-6
   --fake_score_lr_scheduler "cosine_with_min_lr"
   --mixed_precision "bf16"
   --training_state_checkpointing_steps 500
@@ -142,4 +146,4 @@ srun torchrun \
     "${optimizer_args[@]}" \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"

examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json

@@ -55,6 +55,462 @@
       "height": 448,
       "width": 832,
       "num_frames": 61
+    },
+    {
+      "caption": "In an industrial setting, a person leans casually against a railing, exuding a sense of confidence and composure. They are wearing a striking outfit, consisting of a vibrant, patterned jacket over a simple white crop top, creating a bold contrast. The atmosphere is infused with warm, ambient lighting that casts soft shadows on the concrete walls and metallic surfaces. Intricate wiring and pipes form an intricate backdrop, enhancing the urban aesthetic. Their relaxed posture and direct, engaging gaze suggest a sense of ease in this industrial environment. This scene encapsulates a blend of modern fashion and gritty, urban architecture, creating a visually compelling narrative.",
+      "video_path": "Fashion/mixkit-portrait-of-a-hipster-woman-walking-down-a-stairs-1297_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man is energetically stretching in an open-air setting, surrounded by rows of vibrant red seats that suggest an amphitheater or outdoor venue. He wears a sleeveless black shirt layered with a hooded vest, emphasizing his athletic build as he engages in a warm-up routine. Behind him, the striking modern architecture of the building features geometric panels, with large sections of glass and overlapping metallic beams creating a dynamic backdrop. The scene captures the contrast between his focused movements and the static, bold design of the structure, while the surrounding greenery adds a touch of nature to the environment. The overall atmosphere is one of preparation and anticipation, with the man appearing determined and ready for an upcoming event or performance.",
+      "video_path": "Sport/mixkit-man-doing-arm-stretches-595_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young woman is seated on the floor in front of a plush, beige tufted couch, fully engrossed in sorting through a stack of papers. Her dark hair falls loosely past her shoulders, and she wears a green plaid shirt, contributing to the casual yet focused atmosphere. She gently places the papers onto a small round white table, occasionally lifting individual sheets to examine them more closely. Her expression shifts subtly, reflecting concentration and contemplation as she processes the information on the pages. Two small, round nested tables hold her documents, along with a small plant in a gray pot, adding a touch of greenery to the scene. The background features a dark paneled wall, creating a contrasting backdrop for the light-colored furniture. The setting is tranquil and organized, the couch and tables arranged symmetrically, conveying a sense of harmony. A calculator rests on the smaller table, hinting at a task involving calculations or budgeting.",
+      "video_path": "Woman/mixkit-frustrated-woman-throws-paperwork-on-the-floor-4526_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A heavily rusted metal gate stands firmly locked, with two vertical bars joined by a thick, old chain that loops elegantly around them. The chain's texture is coarse and rugged, its surface reflecting varying shades of orange and brown, indicative of years exposed to the elements. At the heart of the chain, a black iron padlock, slightly worn yet imposing, secures the gate, its curves and edges smooth against the aged links. The gate's metalwork is outlined by a backdrop of soft, blurred greenery, suggesting a serene and isolated location beyond the barrier. Tall trees rise in the distance, their trunks and leaves creating a lush, forest-like setting that contrasts with the gate's severe rust. A pathway leads away from the gate, its surface uneven with patches of moss and weathered stone visible in the soft focus, inviting yet inaccessible. The ambiance is quiet and mysterious, with a sense of abandonment hanging subtly in the air, evoking curiosity about what lies beyond. Shadows play across the gate, cast by branches swaying gently in the breeze, adding to the dynamic interaction of light and texture. This scene, rich in detail and atmosphere, captures the viewer's imagination, evoking both the allure of the forbidden and the beauty of decay.",
+      "video_path": "forest/mixkit-rusty-fence-with-a-chain-of-a-property-in-nature-5294_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a serene and softly lit yoga studio, three individuals engage in a yoga session, each performing an upward-facing stretch. The central figure is a woman with shoulder-length brown hair, dressed in a light cropped top and green leggings, her posture reflecting grace and concentration. To her right, another participant, a woman in a purple outfit, mirrors the pose with equal poise. On her left, a person with a bun focuses intently, supported slightly by yoga blocks beneath their hands. The warm-colored wooden floor contrasts soothingly with the soft pastel mural on the back wall, featuring an abstract design and partial visage of a serene face. Natural light floods the space from a large window on the right, where lush greens peek through, adding an element of tranquility. In the corner of the room, a collection of meditation instruments, including a gong and a Buddha statue, subtly frame the peaceful setting. The mood is calm yet focused, as all three participants are deeply engaged in their practice. The scene combines elements of balance, harmony, and a shared journey towards mindfulness. This depiction captures the essence of a yoga session that blends personal growth with collective experience.",
+      "video_path": "People/mixkit-small-group-of-people-doing-yoga-together-43730_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the deep blue expanse of the ocean, two dolphins glide effortlessly, their sleek bodies reflecting the sunlight filtering through the water. The prominent shadows and caustics create a shimmering effect on their skin, capturing the beauty of their natural habitat. Each dolphin moves with a fluid grace, occasionally interacting with gentle nudges, showcasing their playful and social nature. The scene is vibrant and dynamic, with the clear blue background accentuating the dolphins' movements, making it an ideal subject for AI recreation.",
+      "video_path": "sea/mixkit-dolphins-underwater-4133_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a young woman stands against a vibrant graffiti-covered wall, deeply engrossed in her smartphone. Her expression reflects a mix of focus and subtle satisfaction as she interacts with the screen. She wears a black floral-patterned top, which contrasts with the bright, abstract shapes and bold colors of the mural behind her. As she continues to engage with her phone, a series of like count notifications appear on the screen, indicating a growing online appreciation. The wall behind her features a striking mix of geometric and organic shapes, including swirls of teal, orange, and black, with large humanoid figures in a pop-art style. Her long, light-brown hair frames her face, adding a calm, composed aura amidst the lively backdrop. The video captures a blend of contemporary digital interaction and expressive urban art, creating a dynamic yet harmonious scene.",
+      "video_path": "Girl/mixkit-girl-looking-at-the-likes-in-her-post-4914_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young mother and her baby sit comfortably on a bed, surrounded by an inviting, cozy atmosphere. The woman, wearing a sleeveless top and jeans, is gently engaging with the baby, who is dressed in an adorable animal-print onesie. The child is seated on the bed with colorful toys scattered around, including a plush toy and a board book. The warm glow from a hanging lamp casts a soft light on them, enhancing the serene environment. Pillows are propped up against the headboard, providing a cushioned backdrop as the mother leans slightly over to interact with the baby. A small bottle is visible beside her, suggesting a nurturing setting. Her hand gestures animatedly as she holds up a soft, white cushion with red and blue accents, likely stimulating the baby\u2019s curiosity. Their shared moment is filled with affection and joy, a perfect snapshot of familial bonding.",
+      "video_path": "Baby/mixkit-loving-mother-and-her-baby-playing-with-soft-toys-49966_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young girl with long brown hair sits at a round wooden table, engrossed in working on her laptop. The laptop screen is a vivid green, suggesting a green screen effect is in use. To her left, a doll dressed in a yellow and white outfit is casually laid on top of some books, adding a playful and innocent touch to the scene. The setting is cozy, with sheer curtains in the background allowing soft natural light to spill into the room. The girl's posture and focused attention on the laptop suggest she is either playing a game or learning something new. This serene and domestic atmosphere is complemented by the slight blur of a dark couch in the foreground, framing the focused activity of the child.",
+      "video_path": "Girl/mixkit-little-girl-doing-homework-on-a-laptop-4757_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "An expansive view of a calm bay reveals a fleet of sailboats, each anchored in a regimented line stretching toward the horizon. The water is a serene blue, reflecting the soft hues of the early morning sky. A gentle breeze is indicated by the subtle ripples trailing behind the boats, while a single, larger vessel cuts a distinct path, leaving a graceful wake in its journey to the open sea. On one side, a cluster of modern high-rise buildings stands, contrasting against the natural simplicity of the water, suggesting a blend of urban and marine life. The distant shoreline is barely visible, softened by the atmospheric perspective, giving a sense of endless waters meeting the sky. The overall mood is peaceful and orderly, with the boats appearing almost as sentinels guarding the expanse of the tranquil bay.",
+      "video_path": "beach/mixkit-flying-backwards-over-the-sea-near-a-coast-50187_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a person is standing in the center of a dark, featureless space, illuminated by a spotlight that emphasizes their presence. The individual is dressed in a traditional martial arts uniform, known as a gi, which is predominantly white with a black belt tied around the waist, indicating a high level of expertise. The background remains pitch black, creating a stark contrast with the brightly lit figure, ensuring complete focus on them. The person's expression is serious and focused, reflecting a deep sense of discipline and concentration. Their hands move gracefully, transitioning through various martial arts stances, demonstrating practiced skill and fluidity. The uniform's crisp fabric folds and subtly reflects the light, further highlighting each precise movement. Despite the simplicity of the environment, the scene is dynamic, with each motion capturing the essence of martial arts practice. The video effectively conveys a sense of calm strength and mastery, making it ideal for an AI to recreate with attention to posture, lighting, and attire.",
+      "video_path": "Sport/mixkit-karate-fighter-bowing-to-the-front-49706_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a dimly lit room bathed in a mix of neon purple and blue lights, a focused individual is seated in a gaming chair. She wears a white hoodie and large headphones with cat ears that glow softly, creating a striking silhouette. Her hands rest on a keyboard, typing swiftly as she concentrates intently on the screen in front of her. The atmosphere exudes a sense of intensity and immersion, with the soft-colored lighting enhancing the futuristic vibe. Her long hair cascades down her shoulders, adding a touch of elegance to the otherwise tech-centric setting. The overall scene captures the essence of a dedicated gamer deeply engaged in her virtual world.",
+      "video_path": "earth/mixkit-a-young-woman-wearing-headphones-with-rgb-lights-suddenly-gets-51621_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "Inside a dimly-lit bus, five individuals are seated along the rows of worn seats, each subtly illuminated by the colorful lights emanating from overhead. On the left, a woman sits with a relaxed posture, her curly hair accented by a patterned scarf, wearing a plaid outfit paired with bright neon socks. Next to her, a person clad in a denim jacket appears deep in thought, resting their head on a hand. Further back, another figure in a bucket hat and oversized yellow attire gazes across the aisle, evoking a sense of introspection. The atmosphere is enriched by the soft glow of red and green lights, bathing the bus interior in an almost surreal ambiance, creating a compelling tableau of urban life.",
+      "video_path": "Music/mixkit-conceptual-urban-fashion-42581_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "An aerial view captures two tennis players on a court, with one dressed in white on the left and another in red on the right. They are mid-game, each poised for action with rackets in hand, accentuated by their strategic positioning at opposite baselines. The court itself is a stark, deep blue, bordered by the vibrant green of the surrounding area, with a dark central net dividing the space. Long shadows stretch dramatically across the ground, suggesting a late afternoon setting. The subtly textured surface of the court contrasts with the crisp, white lines marking its boundaries and sections. This scene creates a vivid, balanced composition, highlighting both the competitive tension and serene atmosphere of the game.",
+      "video_path": "People/mixkit-two-people-playing-tennis-aerial-view-880_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a vibrant, dreamlike setting, a lone figure moves energetically against a backdrop of deep blue and purple hues, casting emotive shadows that ripple with dynamic motion. The figure, almost obscured by a smeared effect, suggests a rhythmic dance or a passionate performance, arms blurred as they sweep through colorful, streaked lighting. A neon glow accentuates their form, particularly highlighting the face which is abstractly illuminated in bursts of orange and red, suggesting intense emotional expression. The scene is dominated by two primary elements \u2013 the figure\u2019s motion and the dramatic lighting, creating a synergy of human emotion and visual spectacle. Swirling trails of light seem to intertwine with the figure, like a visual symphony of movement and color that floods the space. The lighting changes, casting intricate patterns on the figure and the surrounding space, giving the impression of a kaleidoscope in motion. Despite the blurred and abstract portrayal, there is a sense of focus conveyed through the figure\u2019s intent movements, akin to a conductor orchestrating a visual and auditory performance. The environment resonates with an electric energy, suggesting a seamless fusion of art and technology. As the visual drama unfolds, the scene invites viewers to lose themselves in the abstract dance and the play of vivid luminance.",
+      "video_path": "Music/mixkit-dancer-dancing-with-a-light-bar-in-his-hands-42221_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a brightly lit studio, a photographer wearing a denim jacket focuses intently, capturing shots with a professional camera. Facing him, a model stands gracefully, adjusting her long, flowing hair with delicate movements. The scene is characterized by strong contrasts; the model's soft pink attire and gentle gestures complement the rugged, precise demeanor of the photographer. Positioned against a minimalist backdrop, the pair work seamlessly, with the camera\u2019s lens pointed directly at the model, capturing her elegance. The soft, diffused lighting casts a gentle glow on both subjects, creating an airy and ethereal atmosphere perfect for a high-fashion photo shoot.",
+      "video_path": "Fashion/mixkit-professional-photo-session-with-a-young-female-model-41621_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video showcases a serene, expansive landscape covered with a variety of trees dotting the hills. The hills gently slope across the frame, with patches of dry grass contrasting against the lush green foliage. Tall trees with dense canopies stand elegantly, casting soft shadows on the ground below. The sunlight bathes the entire scene, highlighting the varied textures of the leaves and terrain. Gaps between the trees reveal a narrow dirt path meandering through the hills, suggesting a sense of quiet solitude. The undulating hills extend into the distance, creating depth and a calming sense of vast space. The verdant hues of the leaves contrast with the earthy tones of the hills, enhancing the visual richness. In the background, a faint outline of distant hills can be seen, blurred softly by the atmospheric perspective. This tranquil setting could be efficiently recreated in a virtual environment by focusing on its layered composition, color palette, and natural textures.",
+      "video_path": "forest/mixkit-aerial-panorama-of-a-sunny-mountain-landscape-40846_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A bustling ski slope comes alive with skiers descending a pristine, snow-covered hill, surrounded by towering, snow-draped evergreens. Several figures stand atop the slope, silhouetted against a clear blue sky, preparing to embark on their ski run. The chair lift on the right continuously drops off eager adventurers, adding to the excitement at the hilltop. Each skier, clad in colorful winter gear, carves distinct paths into the textured snow as they weave their way down. The interplay of sunlight and shadows accentuates the myriad tracks etched into the slope, creating a dynamic visual rhythm. The scene captures a vibrant winter wonderland, full of action and the thrill of a perfect ski day.",
+      "video_path": "Car/mixkit-skiers-on-a-snowy-slope-3327_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The scene unfolds within a dimly lit bus, where three young individuals are seated, each absorbed in their unique world. To the left, a person with tied-back hair rests their head on their hand, dressed casually in a jacket and jeans, projecting a relaxed demeanor. Central to the frame is another individual, sitting upright with intense focus, donning a plaid blazer and oversize hoops, enhancing their confident presence. The muted green and red lighting casts an atmospheric glow, adding depth and intrigue to the setting. On the right, a person in a bucket hat and striped shirt leans back, appearing contemplative as they adjust their hat with a nonchalant gesture. The interplay of light and shadow highlights their expressions, creating an intimate and cinematic ambiance. Together, these figures form a cohesive tableau, capturing a moment of introspection amid a bustling yet serene urban environment.",
+      "video_path": "City/mixkit-three-models-posing-to-the-lens-while-on-board-a-42575_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A determined climber is scaling a massive rock face, showcasing exceptional strength and skill. The person, clad in a teal shirt and dark pants, climbs with precision, their movements measured and deliberate. They are secured by climbing gear, which includes ropes and a harness, emphasizing their commitment to safety. The rugged texture of the sandy-colored rock provides an imposing backdrop, adding drama and scale to the climb. In the distance, other large rock formations and sparse vegetation can be seen under a bright, overcast sky, contributing to the natural and adventurous atmosphere. The scene captures a moment of focus and challenge, highlighting the climber's tenacity and the breathtaking environment.",
+      "video_path": "Sport/mixkit-alpinist-climbing-a-huge-rock-in-a-desert-43306_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A woman stands confidently in front of a large array of solar panels, her navy blue jumpsuit contrasting against the lush green grass beneath her feet. Her expression is calm and focused, eyes facing directly ahead, suggesting a deep connection to the subject matter\u2014renewable energy. The sunlight bathes the scene in warm hues, casting gentle shadows and highlighting the geometric precision of the solar panels' grid-like structure. The background reveals a blend of nature and technology, as the panels are anchored on a grassy slope with foliage on the left side of the frame. This composition captures a harmonious blend of human innovation and environmental consciousness, accentuated by the serene outdoor setting.",
+      "video_path": "Business/mixkit-woman-standing-in-front-of-a-solar-panel-4880_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, two people are working at a wooden desk, using an iMac computer. One person, wearing a white knit sweater, is using the apple wireless mouse with their right hand, while their left hand rests on the sleek white keyboard. Their movements are smooth yet intentional, suggesting they are focused on a task on the computer screen. The monitor displays a well-organized array of files and folders, hinting at a task that involves detailed organization or detailed data navigation. The second person, only subtly visible, sits closely by and appears to observe or assist, creating a collaborative atmosphere. Their presence adds a quiet dynamic to the scene, as if they are ready to provide input or guidance. Sticky notes with handwritten notes are attached to the monitor\u2019s stand, adding a touch of personal organization amidst the digital workspace. The focus on the keyboard and mouse emphasizes a streamlined workflow, indicative of a productive work environment. The overall ambiance is calm and focuses on teamwork, technology, and efficient workspace management.",
+      "video_path": "People/mixkit-person-with-glasses-working-on-a-desktop-computer-3248_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man stands in front of a modern glass facade, taking off a dark hoodie to reveal his gray tank top underneath. His arms are lifted high as he maneuvers the hoodie over his head, showcasing a fluid motion that conveys a sense of calm and routine. The lighting highlights the contours of his muscles, emphasizing a combination of strength and quiet determination. Behind him, the reflective surface of the glass panels provides a subtle backdrop, enhancing the focus on his focused and serene demeanor.",
+      "video_path": "Sport/mixkit-man-puts-on-sleeveless-hoodie-603_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video displays a captivating dance of fiery orange flames against a stark black background, creating an intense visual contrast. The flames twist and intertwine, forming symmetrical, swirling patterns that expand and contract rhythmically across the frame. Each fiery tendril seems to be alive, moving with an almost hypnotic fluidity that captures the viewer's attention. The illumination from the flames casts subtle shadows, enhancing the depth and texture of the scene. Overall, the dynamic movement and vibrant color palette create an atmosphere of both beauty and power.",
+      "video_path": "fire/mixkit-two-orange-flames-on-black-background-685_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In this scene, a person is seated in a dimly lit room, possibly a recording studio, holding several drumsticks in their hands. The individual's face is partially obscured by sunglasses, adding a touch of mystery to their demeanor. They are wearing a colorful, patterned shirt with a mix of orange and blue tones that stands out against the darker background. The person appears focused and engaged with the drumsticks, their hands prominently displayed. The ambient light casts warm, soft shadows, emphasizing the texture and colors of their shirt and the wooden drumsticks. The room features wooden paneling, which complements the overall cozy, music-centric setting of the scene. The use of perspective centers on the drumsticks, highlighting the importance of rhythm and music in the captured moment.",
+      "video_path": "Music/mixkit-drummer-stretching-before-playing-42783_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man is casually sitting on a sofa, engrossed in his meal and entertainment. He is holding a TV remote in one hand while reaching for food with the other, indicating a laid-back, comfortable evening. The table before him is filled with takeout containers, revealing a variety of appetizers and dishes, suggestive of a casual dining experience at home. The background is defined by colorful patterned cushions, adding a cozy, homey feel to the scene. Warm, ambient lighting highlights the relaxed atmosphere, casting soft shadows that contribute to the intimate setting. In this moment, he takes a bite of a sandwich, comfortably balancing his attention between food and whatever is playing on the screen.",
+      "video_path": "Man/mixkit-man-watching-tv-and-eating-fast-food-26089_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The scene opens to a breathtaking view of a tranquil ocean horizon at dusk, displaying a vibrant tapestry of oranges, pinks, and purples as the sun sets. In the foreground, tall, swaying palm trees frame the scene, their silhouettes stark against the colorful sky. The ocean itself shimmers with reflections of the sunset, creating a peaceful, almost ethereal atmosphere. A small boat can be seen in the distance, centered on the horizon, adding a sense of scale and solitude to the scene. The waves gently lap the shore, creating faint patterns on the sandy beach, which stretches across the foreground. Above, the sky is dotted with scattered clouds that catch the last light of the day, enhancing the drama and beauty of the scene. The overall mood is serene and contemplative, capturing a perfect moment of nature\u2019s grandeur.",
+      "video_path": "beach/mixkit-sunset-with-sailing-boats-2166_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man sits hunched on a couch, the weight of emotions clearly visible on his posture. He wears a simple, gray t-shirt, and his head is bowed, resting in his hands, which cover most of his face, obscuring his features. The gentle light filtering through sheer curtains in the background casts a soft glow upon him, emphasizing the contrast between his static form and the hazy brightness behind. His elbows rest upon his knees, suggesting a posture of deep contemplation or distress. The simplicity of the room, with its muted colors, highlights the focus on the man's internal struggle. Delicate detailing on the fabric of his shirt adds texture, enhancing the scene's realism. Subtle changes in the natural light indicate the passage of time, as the man remains unmoving, absorbed in thought. This intimate moment captures a profound vulnerability, making the scene universally relatable and poignant.",
+      "video_path": "Man/mixkit-worried-and-sad-man-with-his-head-down-4701_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A pair of hands, belonging to an unseen figure, carefully unrolls a large sheet of crisp, white paper on a dark wooden table. The lighting is warm, casting a gentle glow that highlights the textures of the paper and the wood grain of the table. As the paper unfurls, the edges reveal the faint beginnings of a colorful map printed on its surface. The arms, clad in a casual gray T-shirt, suggest a relaxed and focused task at hand. Each motion is deliberate, with fingers deftly guiding the paper, ensuring it lays flat without creases. In the background, a hint of a red curtain can be seen, adding a touch of color and depth to the setting. The composition of the scene emphasizes the contrast between the bright paper and the rich tones of the surroundings. This serene and methodical action evokes a sense of exploration and preparation.",
+      "video_path": "Man/mixkit-unrolling-a-world-map-on-a-table-21626_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young woman sits on a vibrant green seat inside a bus, illuminated by the soft glow of pink and blue lights. Her outfit is a striking mix of colors: a neon pink top paired with a jacket featuring dark sleeves, and jeans that provide a neutral contrast. She wears large, hoop earrings that catch the light as she moves slightly, exuding an air of cool confidence. Her gaze is directed thoughtfully to the side, suggesting contemplation or daydreaming during her commute. The metallic pole beside her adds a geometric element to the composition, reflecting the kaleidoscope of neon hues. The background is a clean, futuristic white, serving as a blank canvas that amplifies the neon atmosphere. Her relaxed posture and the modern bus setting create a scene that captures a blend of urban life and personal introspection.",
+      "video_path": "City/mixkit-fashion-model-posing-on-a-bus-42578_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A silver SUV drives along a winding, snow-covered mountain road, with dense pine trees blanketed in snow lining both sides. The scene is serene, with the vehicle moving smoothly, possibly on a winter journey or vacation. As the SUV disappears around the bend, another, darker SUV follows, creating a sense of motion and perspective on the snow-dusted asphalt. The towering, snow-laden rock formation to the right contrasts with the dark green of the pines, highlighting the peacefulness of the wintry landscape.",
+      "video_path": "Car/mixkit-curve-on-a-snowy-forest-road-3317_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video showcases a vibrant urban skyline during twilight, with towering buildings reflecting the warm hues of the setting sun. A series of tall, cylindrical structures dominate the foreground, adjacent to a complex of industrial equipment and grids. The scene includes modern high-rise buildings with glass exteriors, capturing the evolving architecture of a bustling cityscape. A prominent structure labeled \"CITY OF AUSTIN POWER PLANT\" stands out, highlighting the industrial theme amidst the urban backdrop. The soft glow of city lights begins to pierce the approaching dusk, creating an inviting yet dynamic atmosphere. Shadows cast by the buildings add depth and contrast, emphasizing their massive scale and intricate designs. The overall composition is balanced between the natural light of the sunset and the artificial illumination of the city, offering a compelling visual narrative.",
+      "video_path": "Car/mixkit-slow-air-travel-in-reverse-over-a-big-city-49841_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the scene, a striking architectural structure dominates the view, bathed in a soft, ambient light. The enormous yellow arches serve as the centerpiece, drawing the eye upwards with their majestic curves and towering presence. The smooth, clean surfaces of the structure reflect the light, highlighting the texture and depth of the architecture. In the foreground, blurred streaks of headlights and taillights suggest the motion of vehicles passing by, adding dynamic energy to the otherwise still scene. The contrast between the fast-moving lights and the static arches creates a balanced composition. To the left, a lone streetlamp and a small tree provide a touch of nature and urban elements against the monumental backdrop. The night sky subtly peeks through the gaps in the structure, hinting at a clear, calm evening. Shadows from the arches create patterns on the ground, adding an intricate detail to the scene. Overall, the combination of light, shadow, and movement makes for a dramatic and visually captivating moment.",
+      "video_path": "Car/mixkit-a-fast-timelapse-of-the-street-with-a-monumental-yellow-50993_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A tranquil marina comes into full view under the golden hues of a setting sun. A collection of gleaming yachts and boats are neatly moored, their reflections shimmering softly on the gentle water. The sun's low position casts elongated shadows over the bustling harbor scene, while rolling hillsides surround the distant cityscape. The skyline is interspersed with modern buildings and clusters of residences, adding layers to the vibrant community. At the center, a broad wooden pier juts confidently into the harbor, extending an invitation for leisurely strolls. To the left, various shops and colorful structures line the waterfront, indicating a vibrant coastal economy. The entire atmosphere exudes a serene yet lively charm, balancing the hustle of maritime activity with the peacefulness of the encroaching dusk. It's a scene of calm anticipation, as if the whole place holds its breath before the night's events unfold.",
+      "video_path": "beach/mixkit-harbor-on-a-tourist-coast-with-many-boats-and-yachts-40077_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video features a confident individual standing atop a structure against a clear blue sky, exuding a sense of freedom and style. The person is clad in a striking yellow button-up shirt tied at the waist, and beneath it, they wear a simple white top that adds to their relaxed yet stylish appearance. Completing the ensemble are high-waisted white jeans paired with a black belt, adding a touch of contrast. Around their neck is a bold red scarf, providing a splash of color and an air of vintage flair. The person's sunglasses, tinted in yellow, reflect the sunlight and contribute to the overall cool and composed demeanor. Their hair is styled elegantly, pulled back with headphones resting over the ears, suggesting they are immersed in music. One hand casually grazes the headphones, while the other rests gently on the railing, grounding the individual in the moment. The scene is an effortless blend of fashion and tranquility, capturing the spirit of sunny, carefree days.",
+      "video_path": "Music/mixkit-standing-woman-listening-to-music-460_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A ballerina gracefully spins and moves across a pink-hued studio, her poised figure accentuated by a shimmering white tutu and bodice. The background, a continuous wash of soft pink, provides a serene and ethereal atmosphere, emphasizing her fluid movements. Her arms extend with elegance, highlighting the delicacy and precision of her ballet pose, while her focused expression adds intensity to the scene. The subtle details of her costume, combined with the pink monochromatic ambiance, create a dreamlike spectacle, ideal for an AI to envision a oneiric dance setting.",
+      "video_path": "Dance/mixkit-portrait-of-a-ballerina-spinning-with-pink-background-40163_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The scene unfolds with two human figures in the distance, making their way through a serene meadow, thick with tall golden grass swaying gently in the breeze. The sun hangs low in the sky, casting a soft, diffused glow that illuminates the landscape with a warm, ethereal light. These figures, clad in hiking gear, move deliberately, suggesting they're either embarking on or concluding a journey. Their silhouettes contrast against the lush greenery of the surrounding trees, whose branches reach out, framing the horizon. The play of light and shadow among the trees creates a quilt of textures, with each leaf catching a hint of the sun's dying rays. This tranquil setting evokes a sense of calm and adventure, capturing the quintessential beauty of nature\u2019s landscape.",
+      "video_path": "People/mixkit-landscape-in-nature-while-two-people-are-jogging-44348_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A large cargo ship is docked at an industrial port, its white superstructure contrasting with the deep green and yellow of its deck. The foreground is dominated by the calm, deep blue waters of the harbor, which reflect the vessel\u2019s imposing presence. Surrounding the ship, a series of industrial buildings and storage facilities are visible, hinting at the bustling activity of the port. The deck is intricately detailed, featuring an array of pipes, equipment, and railings, showcasing the ship's functionality and purpose. In the background, a paved area with green patches and a few parked vehicles adds to the busy, industrious atmosphere of the scene.",
+      "video_path": "sea/mixkit-empty-cargo-ship-waiting-at-the-port-4209_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A lone climber ascends a towering rock face, clad in a pink shirt and gray pants, displaying a determined and focused expression. The climber navigates the rugged surface, where the texture of the rock is peppered with natural pockets and crevices that offer handholds and footholds. Sunlight casts soft shadows across the cliff, highlighting the intricate patterns and the climber\u2019s strategic movements. The cliff looms high, with sparse vegetation breaking the monotony of the stone, while distant rocky formations form a dramatic backdrop against the clear blue sky. The climber\u2019s gear, including a harness and chalk bag, underscores the adventure and challenge woven into this majestic, vertical journey.",
+      "video_path": "Sport/mixkit-mountaineer-girl-climbing-a-steep-rocky-mountain-41089_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A person is seen in a close-up shot, skillfully adjusting the tuning pegs of a guitar, showcasing a focused and practiced hand. The image is in black and white, highlighting the contrast between the textures of the instrument and the clothing. The individual's shirt, visible in the background, adds a soft, subtle texture, while the dark tones of the guitar neck create depth in the scene. This composition captures a moment of concentration and finesse, perfect for recreating an intimate musical setting.",
+      "video_path": "Music/mixkit-guitarist-playing-so-inspired-black-and-white-shot-44178_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A musician is playing a large brass instrument with the words \"Brass Band\" clearly visible on its bell. The scene is set against a vibrant yellow backdrop, casting a warm glow on the subject. The musician wears a dark cap and a matching suit, adding a formal touch to his attire. He is deeply focused on his performance, with the instrument's intricate tubing adding complexity to the visual composition. The lighting creates dramatic shadows and highlights, emphasizing the musician's expression and the instrument's metallic sheen. This harmonious blend of color and form captures the essence of a live brass band performance.",
+      "video_path": "Music/mixkit-musician-playing-the-trombone-while-dancing-43752_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a lone musician stands gracefully in front of a grand cathedral, playing an accordion while surrounded by the lively water display of a central fountain. Dressed in a casual ensemble, he wears a light-colored shirt, dark pants, and a flat cap that gives him a vintage charm. His posture is relaxed, yet engaged, as he sways gently in rhythm with the music, casting soft shadows on the cobblestone steps beneath him. The backdrop features the cathedral's towering twin spires, with intricate stonework that casts a rich, historical aura around the scene. Sunlight bathes the entire setting, enhancing the golden hues of the cathedral facade and creating a halo-like effect around the musician. The fountain's water jets splash playfully, catching glimmers of light and adding a dynamic element to the tranquil atmosphere. The scene captures a harmonious blend of architectural majesty and human creativity, framed by the clear, azure sky that extends infinitely above. It's a vivid depiction of solitude and artistry, set against a timeless urban landscape.",
+      "video_path": "Music/mixkit-man-plays-an-accordion-in-front-of-a-fountain-630_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the tranquil video, a person sits in a meditative pose on a gentle hillside, silhouetted against the dawning sky. The person is facing the breathtaking sunrise, with their back slightly turned to the viewer, wearing a simple, light-colored shirt. Their right hand rests on their knee, fingers relaxed in a common meditation mudra, symbolizing calmness and peace. The sky, a stunning blend of soft oranges and deep purples, gradually brightens, casting a warm glow over the lush, green landscape. To the left, the outlines of distant urban buildings can be seen against the horizon, adding a contrast between nature and city life. A river reflecting the sky's colors meanders through the scene, lending a serene, flowing dynamic to the landscape. Trees rise and fall gently across the terrain, their leaves rustling only faintly in the morning breeze. The person remains still and focused, embodying a moment of mindfulness and connection with nature. This visual captures a harmonious balance, evoking a sense of tranquility and introspection.",
+      "video_path": "City/mixkit-girl-meditating-in-yoga-pose-at-sunset-4803_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A serene landscape video captures a breathtaking panoramic view of a vast valley covered in a gentle mist. The undulating hills are lush with dense greenery, their rich foliage creating a vibrant border on the left side of the frame. The mist weaves through the landscape like a soft, ethereal blanket, lending a dream-like quality to the scene. In the distance, several mountain peaks emerge, their dark outlines contrasting against the pale blue sky. A few faint, wispy clouds drift lazily across the horizon, complementing the tranquil atmosphere. The sunlight filters through the haze, casting a warm glow and highlighting different textures of the flora. The overall mood is calm and contemplative, inviting the viewer to pause and appreciate nature's untouched beauty. The composition emphasizes depth and expansiveness, drawing attention to the harmony between earth and sky. This captivating scene embodies tranquility, offering a perfect backdrop for meditation or relaxation.",
+      "video_path": "forest/mixkit-flying-over-a-hill-with-a-view-of-the-surrounding-49743_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In this scene, a bearded individual is intently focused on their smartphone, with the sun setting in the background, casting a warm glow across the cityscape. The person, partially visible, is wearing a dark, buttoned shirt that contrasts with the golden hue of the sunset. Their hands are holding the smartphone delicately but purposefully, reflecting a sense of engagement and focus on the screen. The sunlight creates a striking lens flare effect, enhancing the dramatic atmosphere of the moment as it glimmers off the phone\u2019s surface. The surrounding environment hints at an elevated vantage point, providing a panoramic view of the urban landscape below.",
+      "video_path": "City/mixkit-guy-texting-at-sunset-265_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In an expansive, industrial space defined by towering columns and high ceilings, a solitary figure takes center stage. The person, dressed in dark, fitted clothing, assumes a powerful, dynamic stance with one leg bent forward and both arms outstretched in a horizontal arc. Framing this pose are intense flames that engulf their arms, creating a striking visual contrast against the muted tones of the room. The fire forms a brilliant halo of orange and yellow, casting flickering shadows on the weathered walls and worn, tiled floor. This interplay between light and dark showcases the dancer's poise and agility, as they maintain balance amidst the intense heat. Windows line the background, their panes dimly illuminated by the daylight filtering in, adding depth and perspective to the scene. The entire performance evokes a sense of raw energy and elemental mastery, as the figure continues to manipulate the fire in a seamless, mesmerizing display.",
+      "video_path": "fire/mixkit-expert-juggler-doing-tricks-with-a-stick-with-fire-43663_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man is playing the violin, focused intently on his music. His fingers gracefully dance along the strings, flawlessly executing each note. He holds the violin close to his chin with a sense of familiarity and expertise. The rich, warm tones of the violin reflect in the soft lighting of the room. He wears a dark shirt, and a subtle necklace rests against his chest, adding a personal touch to his attire. The bow moves smoothly across the strings, producing a melody that seems to fill the space with emotion. His expression is one of concentration and passion, immersing himself fully in the performance. The background is softly blurred, bringing the violin's intricate craftsmanship and his precise movements into sharp focus. This serene and intimate moment captures the essence of his musical artistry.",
+      "video_path": "Music/mixkit-fiddler-playing-a-song-639_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the dimly lit parking garage, two figures engage in an impromptu game of soccer. The first person, wearing a light grey shirt and black pants with three white stripes, skillfully maneuvers the ball with precise footwork. The ground is slick with patches of water, reflecting the vibrant neon lights above. A second figure, clad in dark clothing, stands poised in the background, ready to intercept. The space is defined by stark yellow lines and orange safety bollards, adding structure to the chaotic energy of the scene. The soccer ball glides smoothly across the wet floor, kicking up droplets as it passes. Despite the muted colors of the environment, the players' movements are dynamic and full of life. Their shadowy silhouettes dance with the reflecting light, creating a mesmerizing visual interplay. The atmosphere is charged with focus and camaraderie, encapsulating the essence of a late-night urban soccer experience.",
+      "video_path": "Sport/mixkit-player-making-skillful-play-in-a-street-soccer-game-43504_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A lone climber is seen scaling a towering vertical rock face, demonstrating remarkable strength and focus. Dressed in a light-colored shirt and jeans, the climber grips the stone tightly, navigating the rough textures and crevices with precision. The sheer cliff is massive, exhibiting a range of natural hues from light tan to deep gray, accentuating the climber's figure against the vast rocky backdrop. Surrounding the cliff, scattered greenery and rugged terrain provide a sense of wilderness and isolation. The scene portrays a daring ascension requiring concentration and skill, capturing the essence of human endeavor against nature's formidable beauty.",
+      "video_path": "Sport/mixkit-skilled-mountaineer-climbing-a-gigantic-mountain-41083_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In this serene landscape, a lush meadow stretches across the foreground, dotted with vibrant yellow wildflowers swaying gently in the breeze. A towering tree stands majestically on the right side, its branches reaching wide under the bright blue sky filled with fluffy white clouds. On the left, dense trees form a natural corridor leading to the horizon, suggesting a sense of journey and possibility. The richness of the green grass contrasts beautifully with the golden hue of the distant fields, creating a harmonious palette of nature\u2019s colors. The play of light and shadow adds depth and dimension, evoking a tranquil, inviting atmosphere. It's a scene where nature\u2019s beauty simply commands attention, offering a perfect escape into tranquility.",
+      "video_path": "sky/mixkit-countryside-meadow-4075_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A solitary boat glides across the expansive, tranquil expanse of a serene lake. The vessel leaves a gentle wake behind, creating delicate ripples across the mirror-like surface. The water appears a rich shade of teal, seamlessly blending with the sky at the horizon. Silhouettes of distant trees are faintly visible, creating a picturesque backdrop that enhances the solitary journey of the boat. The sky is a calm gradient, shifting from soft oranges near the shore to the pale blues above. In the distance, a few slender poles emerge from the water, remnants of an old structure or natural formation. The mood of the scene is one of peace and solitude, with the boat journeying steadily through the quiet landscape. There is a sense of endless possibilities as the boat moves toward the unseen beyond the frame. The simplicity and stillness of the scene invite contemplation and reflection, encapsulating a perfect moment of quietude on the water.",
+      "video_path": "mountain/mixkit-motorboat-on-a-large-lake-with-turquoise-blue-waters-4996_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a cozy, dimly lit caf\u00e9, a woman sits alone at a rustic wooden table, fully engrossed in her reading. Her dark, wavy hair frames her face as she leans forward over an open book, suggesting deep focus and contemplation. The caf\u00e9\u2019s ambiance is warm, with hanging pendant lights casting a soft glow over the wooden shelves lined with jars and coffee paraphernalia in the background. A small cup of coffee rests just within her reach, alongside a glass dome encasing a solitary pastry, adding a touch of tranquility to the scene. Her casual attire, a denim jacket over a simple shirt, complements the laid-back, comfortable setting of the caf\u00e9. The contrast between her concentrated expression and the bustling, yet subdued caf\u00e9 atmosphere creates a harmonious, serene visual. The overall composition captures a quiet moment of introspection amidst the gentle hum of caf\u00e9 life.",
+      "video_path": "Woman/mixkit-woman-drinking-coffee-in-a-cafe-223_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a vast, deserted landscape under the night sky, a solitary figure stands at a small music setup, illuminated by strategically placed lights. The person is engrossed in playing a keyboard, with various electronic equipment surrounding them, casting soft glows of orange and blue hues across the scene. To the left, a large circular light adds a dramatic focal point, highlighting the intense contrast between the darkness and the lit performance area. This setup, with its minimalistic design and strategic lighting, creates a captivating and easily recognizable scene that merges the serene, expansive backdrop with an intimate, focused music performance.",
+      "video_path": "Music/mixkit-talented-dj-playing-in-a-lonely-desert-42414_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a bustling urban scene, cars zoom past a weathered building, their blurred motion a testament to the city\u2019s lively pace. The building, with its faded yellow and brown facade, boasts graffiti that speaks of both art and decay, framing the scene with an air of urban grit. A solitary figure stands slightly to the side, clad casually in a gray top and mustard trousers, gazing into the street, seemingly detached from the surrounding flurry. The motion of the traffic creates a dynamic contrast against the static backdrop, emphasizing the relentless movement of the city. As the video progresses, a bright yellow taxi appears, slowing down as it approaches the figure, adding a pop of color to the desaturated hues of the environment. The interaction suggests a routine, a possibly daily exchange between the driver and the pedestrian, hinting at the rhythms of city life. Overhead, a soft, overcast sky casts a diffused light, lending the scene a subdued, timeless quality. Small elements, like the vertical pole cutting through the frame and the distant chatter of urban sounds, complete this vivid tableau of urban existence.",
+      "video_path": "Car/mixkit-morning-in-the-street-time-lapse-1648_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young woman sits on a curb in a tranquil park, basking in the golden hue of the setting sun. Beside her, a collie dog rests calmly, its fur illuminated by the warm sunlight, creating a serene glow. The woman's hand gently strokes the dog's back, highlighting the bond and affection between them. Tall trees surround the pair, casting elongated shadows on the leaf-laden ground, adding to the peaceful and intimate ambiance of the scene.",
+      "video_path": "Pets/mixkit-a-woman-pets-a-dog-in-a-park-1562_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a grand, majestic elephant stands in an open, sunlit field, its massive form dominating the scene. The elephant's skin is a tapestry of earthy tones, with rough, textured wrinkles that add character to its already imposing presence. Its trunk, a powerful and flexible appendage, moves gently, swaying as the elephant possibly enjoys the warmth of the day. The background is a blur of greenery, suggesting a lively environment filled with trees and shrubs that provide a natural habitat. Light plays on the elephant's skin, highlighting patches of dust and dirt that give it an authentic wilderness look. The scene captures the tranquility and majesty of this gentle giant in its natural surroundings.",
+      "video_path": "Zoo/mixkit-wet-elephant-in-the-savanna-3663_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a fluffy dog with brown patches is intently engaged with a bright red toy shaped like a fire hydrant, which has a yellow and orange rope attached. The dog's body is relaxed as it lies on a plain white background, concentrating on nudging and playfully biting the toy. Its ears perk up slightly with curiosity, and its eyes are fixated on the toy, suggesting a scene of focused playfulness. The neutral tones of the dog's fur contrast starkly against the vivid red of the toy, creating a visually striking moment.",
+      "video_path": "Pets/mixkit-a-cute-border-collie-dog-play-with-a-fire-street-50662_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
     }
   ]
 }

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh

@@ -14,26 +14,29 @@ export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
 # Configs
 NUM_GPUS=1
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" 
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # Training arguments
 training_args=(
   --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir="checkpoints/wan_t2v_finetune"
-  --max_train_steps=4000
-  --train_batch_size=1
+  --output_dir "$OUTPUT_DIR"
+  --max_train_steps 4000
+  --train_batch_size 1
   --train_sp_batch_size 1
-  --gradient_accumulation_steps=1
+  --gradient_accumulation_steps 1
   --num_latent_t 31
   --num_height 704
   --num_width 1280
   --num_frames 121
   --enable_gradient_checkpointing_type "full"
-  --training_state_checkpointing_steps=500
-  --weight_only_checkpointing_steps=500
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
 )
 
 # Parallel arguments
@@ -49,6 +52,8 @@ parallel_args=(
 model_args=(
   --model_path $MODEL_PATH
   --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
 # Dataset arguments
@@ -68,8 +73,8 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate=1e-5
-  --mixed_precision="bf16"
+  --learning_rate 2e-6
+  --mixed_precision "bf16"
   --weight_decay 0.01
   --max_grad_norm 1.0
 )
@@ -107,4 +112,4 @@ torchrun \
     "${optimizer_args[@]}" \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh

@@ -14,6 +14,8 @@ export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
 # Configs
 NUM_GPUS=1
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
@@ -51,6 +53,8 @@ parallel_args=(
 model_args=(
   --model_path $MODEL_PATH
   --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
 # Dataset arguments
@@ -109,4 +113,4 @@ torchrun \
     "${optimizer_args[@]}" \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/download_dataset.sh

@@ -1,3 +1,3 @@
 #!/bin/bash
 
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/preprocess_wan_data_ti2v_5b.sh

@@ -21,4 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
     --samples_per_file 8 \
     --flush_frequency 8 \
     --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 
+    --preprocess_task "t2v" 

examples/inference/basic/basic.py

@@ -9,9 +9,9 @@ def main():
     # If a local path is provided, FastVideo will make a best effort
     # attempt to identify the optimal arguments.
     generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-14B-Diffusers",
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         # FastVideo will automatically handle distributed setup
-        num_gpus=4,
+        num_gpus=1,
         use_fsdp_inference=True,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
@@ -25,7 +25,9 @@ def main():
     # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
     # Generate videos with the same simple API, regardless of GPU count
     prompt = (
-        "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open."
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
     )
     video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
     # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
@@ -33,7 +35,11 @@ def main():
     # Generate another video with a different prompt, without reloading the
     # model!
     prompt2 = (
-        "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object")
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
     video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
 
 

examples/inference/basic/basic_self_forcing_causal.py

@@ -16,7 +16,6 @@ def main():
         use_fsdp_inference=True,
         text_encoder_cpu_offload=False,
         dit_cpu_offload=False,
-        use_sf_wan=True,
     )
 
     sampling_param = SamplingParam.from_pretrained(model_name)

examples/training/consistency_finetune/causal_ode_init/crush_smol_prompts.txt

@@ -0,0 +1,47 @@
+A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.
+The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object.
+The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign.
+A red toy car is being crushed by a large hydraulic press, which is flattening objects as if they were under a hydraulic press.
+A large, cylindrical object is seen pressing down on a small orange ball, causing it to flatten as if it were under a hydraulic press. The background features a green wall with yellow and red warning signs.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is shown compressing a wooden object, which shatters into small pieces. The background features a green wall with a yellow sign displaying a lightning bolt.
+A large metal cylinder is seen descending, flattening objects as if they were under a hydraulic press. The cylinder compresses a stack of matches and boxes, causing them to crumble into small pieces. The scene is set against a green background with yellow and red signs.
+A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out.
+The video shows a metal press flattening objects as if they were under a hydraulic press. The press is pressing down on a pile of colorful gummy candies, squishing them into a pile of squiggly shapes. The press is made of metal and has a large base, and the gummy candies are of various colors, including red, green, and orange. The background is a green wall, and the press is placed on a metal surface.
+A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces.
+The video shows a stack of colorful sponges being flattened as if they were under a hydraulic press. The sponges, which are pink, white, blue, and green, are compressed into a smaller size, demonstrating the press's power. The background features a green wall with a yellow and red sign, adding context to the setting.
+A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, leaving a pile of debris around it.
+A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press.
+The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.
+The video shows a close-up of a metal cylinder pressing down on a yellow object, which is being flattened as if it were under a hydraulic press. The cylinder is positioned above the object, and the force is causing the object to compress and spread out, creating a visible deformation. The background is blurred, focusing attention on the action of the cylinder and the object being flattened.
+A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.
+The video shows a hydraulic press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing two colorful objects that resemble sandwiches. The press is yellow and black striped, and the objects being flattened are placed on a metal plate. The background is green, and the press is moving down, compressing the objects.
+The scene shows a metal press with a yellow and black striped pattern, holding a container filled with chocolate. A metal cylinder is descending, flattening the chocolate as if it were under a hydraulic press. The background is a green wall, and the press is mounted on a sturdy metal frame.
+The video shows a colorful sponge being flattened as if it were under a hydraulic press, with the sponge being compressed and eventually flattened into a thin layer.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is pressing down on a stack of wooden blocks, causing them to crumble and break apart. The press is black and yellow striped, and the wooden blocks are small and rectangular. The background is green, and the press is sitting on a metal table.
+A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces.
+The video shows a stack of colorful sponges being flattened by a large, cylindrical object, which appears to be a hydraulic press. The sponges, which are pink, blue, white, and green, are compressed into a single layer, demonstrating the press's powerful force. The background features a green wall with a yellow and red sign, adding context to the industrial setting.
+A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, demonstrating the immense pressure applied by the cylinder.
+A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press. The popcorn is crushed and scattered around the base of the cylinder, creating a satisfying visual effect.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is composed of a large, cylindrical metal cylinder with yellow and black stripes, and a metal base. The objects being flattened are two cylindrical blocks of cotton candy, one pink and one blue. The press is positioned on a metal table, and the background features a green wall with a yellow and red sign.
+The video shows a large orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.
+The video shows a cylindrical object being pressed down onto a flat surface, causing the objects beneath it to be flattened as if they were under a hydraulic press. The objects being flattened appear to be yellow and are being crushed into a pile of debris. The background is a greenish-gray color, and the surface on which the objects are being flattened is metallic and shiny.
+A green and blue object with a spiky texture is being flattened by a large, cylindrical metal press, demonstrating its resilience and durability.
+The video shows a stack of caramelized sugar cubes being flattened as if they were under a hydraulic press, resulting in a messy pile of broken sugar on the table.
+A large metal cylinder is seen pressing down on a pile of colorful jelly beans, flattening them as if they were under a hydraulic press.
+The video shows a machine with a yellow and black striped cylinder pressing down on a stack of colorful sponges, flattening them as if they were under a hydraulic press. The machine is situated in a green-walled room with warning signs in the background.
+The video shows a machine with a yellow and black striped cylinder, which is pressing down on two colorful objects, flattening them as if they were under a hydraulic press. The machine appears to be in a workshop or industrial setting, with a green wall in the background. The objects being flattened are green and orange, and the machine is covered in dirt and grime, indicating it has been used frequently.
+The video shows a large, industrial press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing a pile of pink objects into a pile of crumbs. The press is large and metallic, with a yellow and black striped pattern on its side. The background is a green wall with a yellow warning sign.
+The video shows a pink, sparkly ball being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press.
+A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and segments.
+The video shows a machine with a yellow and black striped cylinder, which is flattening objects as if they were under a hydraulic press. The machine is pressing down on two colorful objects, causing them to compress and flatten. The background is a green wall, and the machine appears to be in a workshop or industrial setting.
+The video shows a large, yellow and black striped cylinder flattening objects as if they were under a hydraulic press. The objects being flattened are pink and are being crushed into small pieces. The background is a green wall with a yellow sign.
+The video shows a machine with a yellow and black striped cylinder pressing down on two colorful objects, which are flattened as if they were under a hydraulic press. The machine is positioned on a metal platform, and the background is a green wall.
+A green cube is being compressed by a hydraulic press, which flattens the object as if it were under a hydraulic press. The press is shown in action, with the cube being squeezed into a smaller shape.
+A pink, sparkly ball is being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press.
+A red cabbage is being crushed by a hydraulic press, which flattens the objects as if they were under a hydraulic press. The press is shown in action, compressing the cabbage into a smaller, more compact form.
+A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and pulp.
+A large metal press is shown compressing a stack of burgers, causing them to be flattened and crushed into a pile of ground meat.
+A pizza is being crushed by a hydraulic press, causing the toppings to spread out and the crust to crumble.
+A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.
+A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.
+A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.

examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh

@@ -0,0 +1,94 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_crush_smol"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "wan_ode_init_crush_smol"
+  --max_train_steps 6000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --warp_denoising_step
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR" 
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 6e-6
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/consistency_finetune/causal_ode_init/finetune_ode_init_16k_vidprom.slurm

@@ -0,0 +1,136 @@
+#!/bin/bash
+#SBATCH --job-name=2e6B8_16kFV_ode_vidprom
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=ode_vidprom16k/ode_vidprom8b16k_2e-6.out
+#SBATCH --error=ode_vidprom16k/ode_vidprom8b16k_2e-6.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate your-conda-env
+
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_API_KEY=your-wandb-api-key
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+
+MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="your-data-dir"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/causal_ode_init/validation.json"
+OUTPUT_DIR="your-output-dir"
+INIT_WEIGHTS_FROM_SAFETENSORS="your-init-weights-from-safetensors" # bidirectional weights from Wan2.1-T2V-1.3B-Diffusers
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir $OUTPUT_DIR
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "vidprom_8b16k_ode_init_2e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --warp_denoising_step
+  --log_visualization
+  --max_train_steps 6001
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 81
+  --dmd_denoising_steps "1000,750,500,250"
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim $NUM_GPUS
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR" 
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+  --init_weights_from_safetensors $INIT_WEIGHTS_FROM_SAFETENSORS
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-6
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 500
+  --training_state_checkpointing_steps 500
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+srun torchrun \
+  --nnodes $SLURM_JOB_NUM_NODES \
+  --nproc_per_node $NUM_GPUS \
+  --node_rank $SLURM_PROCID \
+  --rdzv_backend=c10d \
+  --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/consistency_finetune/causal_ode_init/preprocess_wan_data_crush_smol.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="$(dirname "$0")/crush_smol_prompts.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 1 \
+    --seed 42 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 81 \
+    --flow_shift 5.0 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --train_fps 16 \
+    --samples_per_file 8 \
+    --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "ode_trajectory" 

examples/training/consistency_finetune/causal_ode_init/validation.json

@@ -0,0 +1,76 @@
+{
+  "data": [
+    {
+      "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "Elon Musk, dressed in a sleek white spacesuit with a reflective visor, walks confidently across the lunar surface. His posture is upright, and he moves steadily with purpose. The moon's rocky terrain and scattered boulders surround him, casting shadows under the dim sunlight. The background shows vast stretches of the moon's barren landscape with craters and dust clouds kicked up by his boots. The scene captures a wide shot, emphasizing the vastness and desolation of the lunar environment. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "In a dynamic action-packed sequence set in the Marvel multiverse, Spider-Man and Venom engage in an intense battle. Spider-Man, in his classic red and blue suit, swings and dodges venomous attacks from the black symbiote-covered Venom. Both characters display a range of acrobatic moves and powerful strikes. The environment is a chaotic urban landscape with crumbling buildings and neon lights, reflecting the multiversal theme. The camera captures the epic fight from various angles, including wide shots to show the scale of destruction and close-ups to highlight their fierce expressions and physical combat. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A warm, family-oriented scene depicting a father getting ready to leave the house to buy milk. The father, a middle-aged man with a kind face and a casual outfit, picks up a jacket from the coat rack. His posture is upright as he bends down slightly to put on his shoes. In the background, there are glimpses of a cozy living room with a family photograph on the wall. The camera focuses closely on the father, capturing his gentle smile and reassuring nod towards the camera before he opens the front door and steps outside. Static medium close-up shot. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "Close-up shot of a man with a prosthetic hand that functions as a rocket launcher. He looks at his new hand with a mix of amazement and concern, his facial expression showing a blend of curiosity and apprehension. The prosthetic hand is sleek and metallic, with intricate details that resemble a high-tech weapon. The background is a dimly lit laboratory with various scientific equipment and monitors displaying data. The man stands in a relaxed posture, his other hand resting on his hip, as he inspects his new limb. The scene is rendered in a realistic sci-fi style, emphasizing the futuristic technology and the man's emotional response to his new appendage. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "Realistic CCTV footage style, Kim Taehyung from the band BTS is involved in a drug deal, caught on camera. Kim Taehyung appears nervous and cautious, wearing casual clothing typical of a public space. He exchanges items discreetly with another person, who is partially obscured. Both individuals maintain a watchful demeanor, occasionally glancing around to ensure no one is watching them. The lighting is dim, with flickering fluorescent lights casting shadows on their faces. The background shows a typical urban setting with blurred figures moving in the distance. Static camera angle, medium close-up shot focusing on the interaction between Taehyung and the other individual. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "Photorealistic studio setup with professional lighting, showcasing detailed cubic dissections of experimental plastic and felt-like materials on a pristine white background. Each cube reveals intricate layers and textures of the materials, emphasizing their unique properties. The scene has a shallow depth of field initially, then slowly pulls out to reveal the full arrangement of cubes, maintaining a wide depth of field throughout the transition. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}

examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh

@@ -61,7 +61,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 2e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 2000
+  --weight_only_checkpointing_steps 2000
+  --training_state_checkpointing_steps 2000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm

@@ -95,7 +95,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 1e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm

@@ -93,7 +93,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 1e-6
   --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
   --weight_decay 0.01
   --max_grad_norm 1.0
 )

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm

@@ -0,0 +1,134 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=8
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=VSA_t2v_output/t2v_%j.out
+#SBATCH --error=VSA_t2v_output/t2v_%j.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate your_env
+
+# Basic Info
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR=your_data_dir
+VALIDATION_DATASET_FILE=your_validation_dataset_file
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name wan_t2v_VSA
+  --output_dir "checkpoints/wan_t2v_finetune_VSA"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 81
+  # --enable_gradient_checkpointing_type "full" # if OOM enable this
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 64
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file $VALIDATION_DATASET_FILE
+  --validation_steps 200
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "5.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-6
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
+  --weight_decay 0.01
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --dit_precision "fp32"
+  --ema_start_step 0
+  --flow_shift 1
+  --seed 1000
+)
+
+# VSA arguments
+vsa_args=(
+  --VSA_decay_rate 0.03 \
+  --VSA_decay_interval_steps 50 \
+  --VSA_sparsity 0.9 \
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/wan_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${vsa_args[@]}"  

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm

@@ -91,9 +91,10 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 1e-6
   --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
   --weight_decay 0.01
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh

@@ -61,7 +61,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 1e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm

@@ -95,7 +95,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 1e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh

@@ -61,7 +61,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 1e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh

@@ -61,7 +61,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 5e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm

@@ -92,7 +92,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 5e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 400
+  --weight_only_checkpointing_steps 400
+  --training_state_checkpointing_steps 400
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh

@@ -61,7 +61,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 5e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 400
+  --weight_only_checkpointing_steps 400
+  --training_state_checkpointing_steps 400
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json

@@ -28,4 +28,4 @@
       "num_frames": 77
     }
   ]
-}
+}

fastvideo/attention/backends/sage_attn3.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from fastvideo.attention.backends.sageattn.api import sageattn_blackwell
+
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttention3Backend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SAGE_ATTN_THREE"
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttention3Impl"]:
+        return SageAttention3Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+    # @staticmethod
+    # def get_metadata_cls() -> Type["AttentionMetadata"]:
+    #     return FlashAttentionMetadata
+
+
+class SageAttention3Impl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        output = sageattn_blackwell(query, key, value, is_causal=self.causal)
+        output = output.transpose(1, 2)
+        return output

fastvideo/attention/backends/vmoba.py

@@ -5,7 +5,6 @@ from dataclasses import dataclass
 
 import torch
 from einops import rearrange
-from flash_attn.bert_padding import pad_input
 
 from csrc.attn.vmoba_attn.vmoba import (moba_attn_varlen, process_moba_input,
                                         process_moba_output)
@@ -134,6 +133,8 @@ class VMOBAAttentionImpl(AttentionImpl):
                  **extra_impl_args) -> None:
         self.prefix = prefix
         self.layer_idx = self._get_layer_idx(prefix)
+        from flash_attn.bert_padding import pad_input
+        self.pad_input = pad_input
 
     def _get_layer_idx(self, prefix: str) -> int | None:
         match = re.search(r"blocks\.(\d+)", prefix)
@@ -169,7 +170,6 @@ class VMOBAAttentionImpl(AttentionImpl):
             moba_chunk_size = attn_metadata.st_chunk_size
             moba_topk = attn_metadata.st_topk
 
-        # torch.distributed.breakpoint()
         query, chunk_size = process_moba_input(query,
                                                attn_metadata.patch_resolution,
                                                moba_chunk_size)
@@ -205,8 +205,8 @@ class VMOBAAttentionImpl(AttentionImpl):
             simsum_threshold=attn_metadata.moba_threshold,
             threshold_type=attn_metadata.moba_threshold_type,
         )
-        hidden_states = pad_input(hidden_states, indices_q, batch_size,
-                                  sequence_length)
+        hidden_states = self.pad_input(hidden_states, indices_q, batch_size,
+                                       sequence_length)
         hidden_states = process_moba_output(hidden_states,
                                             attn_metadata.patch_resolution,
                                             moba_chunk_size)

fastvideo/configs/models/dits/base.py

@@ -15,18 +15,16 @@ class DiTArchConfig(ArchConfig):
     reverse_param_names_mapping: dict = field(default_factory=dict)
     lora_param_names_mapping: dict = field(default_factory=dict)
     _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
-        AttentionBackendEnum.SLIDING_TILE_ATTN,
-        AttentionBackendEnum.SAGE_ATTN,
-        AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA,
-        AttentionBackendEnum.VIDEO_SPARSE_ATTN,
-        AttentionBackendEnum.VMOBA_ATTN,
-    )
+        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
+        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
+        AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN,
+        AttentionBackendEnum.SAGE_ATTN_THREE)
 
     hidden_size: int = 0
     num_attention_heads: int = 0
     num_channels_latents: int = 0
     exclude_lora_layers: list[str] = field(default_factory=list)
+    boundary_ratio: float | None = None
 
     def __post_init__(self) -> None:
         if not self._compile_conditions:

fastvideo/configs/models/dits/wanvideo.py

@@ -92,6 +92,9 @@ class WanVideoArchConfig(DiTArchConfig):
     pos_embed_seq_len: int | None = None
     exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
 
+    # Wan MoE
+    boundary_ratio: float | None = None
+
     # Causal Wan
     local_attn_size: int = -1  # Window size for temporal local attention (-1 indicates global attention)
     sink_size: int = 0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache

fastvideo/configs/pipelines/__init__.py

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

fastvideo/configs/pipelines/base.py

@@ -87,6 +87,7 @@ class PipelineConfig:
 
     # Wan2.2 TI2V parameters
     ti2v_task: bool = False
+    boundary_ratio: float | None = None
 
     # Compilation
     # enable_torch_compile: bool = False

fastvideo/configs/pipelines/registry.py

@@ -11,9 +11,9 @@ from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 # isort: off
 from fastvideo.configs.pipelines.wan import (
     FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
-    SelfForcingWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config,
-    Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig,
-    WanT2V720PConfig)
+    Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
+    WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
+    SelfForcingWanT2V480PConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -48,6 +48,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "wanpipeline": lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
     "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
+    "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
     "stepvideo": lambda id: "stepvideo" in id.lower(),
     # Add other pipeline architecture detectors
 }
@@ -60,6 +61,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
     "wandmdpipeline": FastWan2_1_T2V_480P_Config,
+    "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
     "stepvideo": StepVideoT2VConfig
     # Other fallbacks by architecture
 }

fastvideo/configs/pipelines/wan.py

@@ -54,6 +54,9 @@ class WanT2V480PConfig(PipelineConfig):
     text_encoder_precisions: tuple[str, ...] = field(
         default_factory=lambda: ("fp32", ))
 
+    # self-forcing params
+    warp_denoising_step: bool = True
+
     # WanConfig-specific added parameters
 
     def __post_init__(self):
@@ -82,7 +85,7 @@ class WanI2V480PConfig(WanT2V480PConfig):
         default_factory=CLIPVisionConfig)
     image_encoder_precision: str = "fp32"
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
 
@@ -108,19 +111,17 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 757, 522])
 
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True
-
 
 @dataclass
 class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
     flow_shift: float | None = 5.0
     ti2v_task: bool = True
+    expand_timesteps: bool = True
 
     def __post_init__(self) -> None:
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
+        self.dit_config.expand_timesteps = self.expand_timesteps
 
 
 @dataclass
@@ -132,12 +133,26 @@ class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config):
 
 @dataclass
 class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
-    pass
+    flow_shift: float | None = 12.0
+    boundary_ratio: float | None = 0.875
+
+    # self-forcing params
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000, 750, 500, 250])
+    warp_denoising_step: bool = True
+
+    def __post_init__(self) -> None:
+        self.dit_config.boundary_ratio = self.boundary_ratio
 
 
 @dataclass
-class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
-    pass
+class Wan2_2_I2V_A14B_Config(WanI2V480PConfig):
+    flow_shift: float | None = 5.0
+    boundary_ratio: float | None = 0.900
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.dit_config.boundary_ratio = self.boundary_ratio
 
 
 # =============================================

fastvideo/configs/sample/base.py

@@ -40,6 +40,7 @@ class SamplingParam:
     num_inference_steps: int = 50
     guidance_scale: float = 1.0
     guidance_rescale: float = 0.0
+    boundary_ratio: float | None = None
 
     # TeaCache parameters
     enable_teacache: bool = False
@@ -47,6 +48,8 @@ class SamplingParam:
     # Misc
     save_video: bool = True
     return_frames: bool = False
+    return_trajectory_latents: bool = False  # returns all latents for each timestep
+    return_trajectory_decoded: bool = False  # returns decoded latents for each timestep
 
     def __post_init__(self) -> None:
         self.data_type = "video" if self.num_frames > 1 else "image"
@@ -167,6 +170,12 @@ class SamplingParam:
             default=SamplingParam.guidance_rescale,
             help="Guidance rescale factor",
         )
+        parser.add_argument(
+            "--boundary-ratio",
+            type=float,
+            default=SamplingParam.boundary_ratio,
+            help="Boundary timestep ratio",
+        )
         parser.add_argument(
             "--save-video",
             action="store_true",
@@ -198,6 +207,18 @@ class SamplingParam:
             help=
             "Path to a JSON file containing V-MoBA specific configurations.",
         )
+        parser.add_argument(
+            "--return-trajectory-latents",
+            action="store_true",
+            default=SamplingParam.return_trajectory_latents,
+            help="Whether to return the trajectory",
+        )
+        parser.add_argument(
+            "--return-trajectory-decoded",
+            action="store_true",
+            default=SamplingParam.return_trajectory_decoded,
+            help="Whether to return the decoded trajectory",
+        )
         return parser
 
 

fastvideo/configs/sample/wan.py

@@ -144,18 +144,22 @@ class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParam):
 
 @dataclass
 class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
-    guidance_scale: float = 4.0
-    guidance_scale_2: float = 3.0
+    guidance_scale: float = 4.0  # high_noise
+    guidance_scale_2: float = 3.0  # low_noise
     num_inference_steps: int = 40
     fps: int = 16
+    # NOTE(will): default boundary timestep is tracked by PipelineConfig, but
+    # can be overridden during sampling
 
 
 @dataclass
 class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
-    guidance_scale: float = 3.5
-    guidance_scale_2: float = 3.5
+    guidance_scale: float = 3.5  # high_noise
+    guidance_scale_2: float = 3.5  # low_noise
     num_inference_steps: int = 40
     fps: int = 16
+    # NOTE(will): default boundary timestep is tracked by PipelineConfig, but
+    # can be overridden during sampling
 
 
 # =============================================

fastvideo/dataset/__init__.py

@@ -4,7 +4,7 @@ from torchvision.transforms import Lambda
 
 from fastvideo.dataset.parquet_dataset_map_style import (
     build_parquet_map_style_dataloader)
-from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset
+from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, TextDataset
 from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                          TemporalRandomCrop)
 from fastvideo.dataset.validation_dataset import ValidationDataset
@@ -37,9 +37,15 @@ def getdataset(args) -> VideoCaptionMergedDataset:
                                      temporal_sample=temporal_sample,
                                      transform_topcrop=transform_topcrop,
                                      seed=args.seed)
+                                    
+
+def gettextdataset(args) -> TextDataset:
+    return TextDataset(data_merge_path=args.data_merge_path,
+                       args=args,
+                       seed=args.seed)
 
 
 __all__ = [
     "build_parquet_map_style_dataloader", "ValidationDataset",
-    "VideoCaptionMergedDataset"
+    "VideoCaptionMergedDataset", "TextDataset"
 ]

fastvideo/dataset/dataloader/parquet_io.py

@@ -0,0 +1,264 @@
+"""
+Utilities for converting preprocessing records (dicts) into Arrow tables and
+writing Parquet datasets in fixed-size chunks.
+
+This module centralizes table construction and Parquet file writing so
+pipelines only need to define their PyArrow schema and produce per-sample
+record dictionaries.
+
+Key APIs:
+- records_to_table(records, schema): Safely convert a list of dictionaries into
+  a pa.Table, casting to the provided schema.
+- ParquetDatasetWriter: Buffer tables and flush to a directory as multiple
+  Parquet files with a fixed number of rows per file. Uses temporary files and
+  atomic rename to avoid partially written outputs.
+"""
+
+from __future__ import annotations
+
+import multiprocessing
+import os
+from concurrent.futures import ProcessPoolExecutor
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+
+def records_to_table(records: list[dict[str, Any]], schema: pa.Schema) -> pa.Table:
+    """Build a PyArrow table from Python record dicts using an explicit schema.
+
+    Arrow will cast values to the target schema when possible (e.g., promoting
+    Python ints/floats to pa.int64/pa.float64), eliminating hand-written per-
+    field array construction.
+
+    Args:
+        records: List of dictionaries, each representing one row. Keys must
+            match schema field names.
+        schema: Target PyArrow schema. Controls field names and types.
+
+    Returns:
+        pa.Table: In-memory table matching the provided schema. If ``records``
+        is empty, returns an empty table with the given schema.
+    """
+    if not records:
+        return pa.table({}, schema=schema)
+    return pa.Table.from_pylist(records, schema=schema)
+
+
+class ParquetDatasetWriter:
+    """Accumulate tables and flush them to a Parquet directory in fixed-size chunks.
+
+    Behavior:
+    - Writes files under worker-specific subdirectories for parallelism.
+    - Uses temporary files and atomic rename to avoid partial files being left
+      behind on failure.
+    - Only full chunks of ``samples_per_file`` rows are written on each flush;
+      any remainder rows are re-buffered for the next flush.
+
+    Note:
+    - Instances are not meant to be shared across processes. Create one writer
+      per process if using multiprocessing.
+    """
+
+    def __init__(self, out_dir: str, samples_per_file: int, compression: str = "zstd") -> None:
+        """Initialize the dataset writer.
+
+        Args:
+            out_dir: Output directory where Parquet files will be written.
+            samples_per_file: Fixed number of rows per Parquet file.
+            compression: Compression codec passed to ``pyarrow.parquet.write_table``
+                (e.g., ``"zstd"``, ``"snappy"``, ``"gzip"``).
+        """
+        self.out_dir = out_dir
+        self.samples_per_file = max(int(samples_per_file), 1)
+        self.compression = compression
+        os.makedirs(self.out_dir, exist_ok=True)
+        self._tables: list[pa.Table] = []
+
+    def append_table(self, table: pa.Table) -> None:
+        """Append a non-empty table to the internal buffer.
+
+        Args:
+            table: A ``pa.Table`` to buffer. Empty or ``None`` tables are ignored.
+        """
+        if table is None or len(table) == 0:
+            return
+        self._tables.append(table)
+
+    def _combine(self) -> pa.Table | None:
+        """Combine all buffered tables into a single table, if any.
+
+        Returns:
+            A concatenated table, a single table if only one was buffered, or
+            ``None`` if no tables are buffered.
+        """
+        if not self._tables:
+            return None
+        if len(self._tables) == 1:
+            return self._tables[0]
+        return pa.concat_tables(self._tables, promote_options='none')
+
+    def flush(self, num_workers: int | None = None, write_remainder: bool = False) -> int:
+        """Write accumulated tables to disk and clear the written portion.
+
+        Only complete chunks of size ``samples_per_file`` are written. Any
+        remainder rows are kept buffered for the next flush.
+
+        Args:
+            num_workers: Optional override for the number of parallel workers
+                used to write chunks. Defaults to ``min(cpu_count, chunks)``.
+            write_remainder: If True, also write any leftover rows (< samples_per_file)
+                as a final small Parquet file (useful for the last flush at the
+                end of preprocessing).
+
+        Returns:
+            int: Number of rows successfully written in this flush call.
+        """
+        combined = self._combine()
+        self._tables = []
+        if combined is None or len(combined) == 0:
+            return 0
+
+        num_samples = len(combined)
+        total_chunks = num_samples // self.samples_per_file
+        if total_chunks == 0:
+            if not write_remainder:
+                # Not enough to form a full chunk; keep buffered for next round
+                # Re-buffer and return 0 written
+                self._tables = [combined]
+                return 0
+            # Last flush: write the small remainder as a final file in worker_0
+            worker_dir = os.path.join(self.out_dir, "worker_0")
+            os.makedirs(worker_dir, exist_ok=True)
+            # Determine next index
+            num_parquets = 0
+            for _, _, files in os.walk(worker_dir):
+                for file in files:
+                    if file.endswith('.parquet'):
+                        num_parquets += 1
+            chunk_path = os.path.join(worker_dir, f"data_chunk_{num_parquets}.parquet")
+            temp_path = chunk_path + '.tmp'
+            pq.write_table(combined, temp_path, compression=self.compression)
+            if os.path.exists(chunk_path):
+                os.remove(chunk_path)
+            os.rename(temp_path, chunk_path)
+            return num_samples
+
+        # Only write full chunks; keep remainder for next flush
+        written_rows = total_chunks * self.samples_per_file
+        remainder = num_samples - written_rows
+
+        table_to_write = combined.slice(0, written_rows)
+        remainder_table = combined.slice(written_rows, remainder) if remainder > 0 else None
+        if remainder_table is not None and len(remainder_table) > 0:
+            if write_remainder:
+                # Write the remainder as a final small file (worker_0)
+                worker_dir = os.path.join(self.out_dir, "worker_0")
+                os.makedirs(worker_dir, exist_ok=True)
+                num_parquets = 0
+                for _, _, files in os.walk(worker_dir):
+                    for file in files:
+                        if file.endswith('.parquet'):
+                            num_parquets += 1
+                remainder_path = os.path.join(worker_dir,
+                                              f"data_chunk_{num_parquets}.parquet")
+                temp_path = remainder_path + '.tmp'
+                pq.write_table(remainder_table,
+                               temp_path,
+                               compression=self.compression)
+                if os.path.exists(remainder_path):
+                    os.remove(remainder_path)
+                os.rename(temp_path, remainder_path)
+            else:
+                self._tables = [remainder_table]
+
+        # Parallel write by chunk ranges
+        if num_workers is None:
+            num_workers = min(multiprocessing.cpu_count(), max(total_chunks, 1))
+        num_workers = max(int(num_workers), 1)
+        chunks_per_worker = (total_chunks + num_workers - 1) // num_workers
+
+        work_ranges: list[tuple[int, int, pa.Table, int, str, int, str]] = []
+        for worker_id in range(num_workers):
+            start_chunk = worker_id * chunks_per_worker
+            end_chunk = min((worker_id + 1) * chunks_per_worker, total_chunks)
+            if start_chunk < end_chunk:
+                work_ranges.append(
+                    (
+                        start_chunk,
+                        end_chunk,
+                        table_to_write,
+                        worker_id,
+                        self.out_dir,
+                        self.samples_per_file,
+                        self.compression,
+                    )
+                )
+
+        written_total = 0
+        if len(work_ranges) == 1:
+            written_total += _process_chunk_range(work_ranges[0])
+            return written_total
+
+        with ProcessPoolExecutor(max_workers=num_workers) as executor:
+            futures = [executor.submit(_process_chunk_range, args) for args in work_ranges]
+            for f in futures:
+                written_total += f.result()
+        return written_total + (len(remainder_table) if write_remainder and remainder_table is not None else 0)
+
+
+def _process_chunk_range(args: Any) -> int:
+    """Worker function to write a contiguous range of chunk files.
+
+    Args:
+        args: Tuple containing
+            - start_chunk (int): inclusive start chunk index
+            - end_chunk (int): exclusive end chunk index
+            - table (pa.Table): concatenated table containing all rows to write
+            - worker_id (int): numeric worker identifier
+            - output_dir (str): base output directory
+            - samples_per_file (int): rows per chunk file
+            - compression (str): compression codec for Parquet
+
+    Returns:
+        int: Total number of rows written by this worker.
+    """
+    start_chunk, end_chunk, table, worker_id, output_dir, samples_per_file, compression = args
+    total_written = 0
+    num_samples = len(table)
+
+    worker_dir = os.path.join(output_dir, f"worker_{worker_id}")
+    os.makedirs(worker_dir, exist_ok=True)
+
+    # Offset to continue numbering if files exist
+    num_parquets = 0
+    for root, _, files in os.walk(worker_dir):
+        for file in files:
+            if file.endswith('.parquet'):
+                num_parquets += 1
+
+    for i in range(start_chunk, end_chunk):
+        start_sample = i * samples_per_file
+        end_sample = min((i + 1) * samples_per_file, num_samples)
+        if end_sample <= start_sample:
+            continue
+        chunk = table.slice(start_sample, end_sample - start_sample)
+
+        chunk_path = os.path.join(worker_dir, f"data_chunk_{i + num_parquets}.parquet")
+        temp_path = chunk_path + '.tmp'
+        try:
+            pq.write_table(chunk, temp_path, compression=compression)
+            if os.path.exists(chunk_path):
+                os.remove(chunk_path)
+            os.rename(temp_path, chunk_path)
+            total_written += len(chunk)
+        except Exception:
+            if os.path.exists(temp_path):
+                os.remove(temp_path)
+            raise
+
+    return total_written
+
+
+

fastvideo/dataset/dataloader/record_schema.py

@@ -1,5 +1,7 @@
 from typing import Any
 
+import numpy as np
+
 from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
 
 
@@ -120,3 +122,69 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
             })
 
     return records
+
+
+def ode_text_only_record_creator(
+        video_name: str, text_embedding: np.ndarray, caption: str,
+        trajectory_latents: np.ndarray,
+        trajectory_timesteps: np.ndarray) -> dict[str, Any]:
+    """Create a text-only ODE trajectory record matching pyarrow_schema_ode_trajectory_text_only.
+
+    Args:
+        video_name: Base name/id for the sample (without extension).
+        text_embedding: Text encoder output array [SeqLen, Dim].
+        caption: Original text prompt.
+        trajectory_latents: Collected trajectory latents array.
+        trajectory_timesteps: Collected timesteps array.
+
+    Returns:
+        dict suitable for records_to_table(…, pyarrow_schema_ode_trajectory_text_only)
+    """
+    assert trajectory_latents is not None, "trajectory_latents is required"
+    assert trajectory_timesteps is not None, "trajectory_timesteps is required"
+
+    record = {
+        "id": f"text_{video_name}",
+        "text_embedding_bytes": text_embedding.tobytes(),
+        "text_embedding_shape": list(text_embedding.shape),
+        "text_embedding_dtype": str(text_embedding.dtype),
+        "file_name": video_name,
+        "caption": caption,
+        "media_type": "text",
+    }
+
+    record.update({
+        "trajectory_latents_bytes": trajectory_latents.tobytes(),
+        "trajectory_latents_shape": list(trajectory_latents.shape),
+        "trajectory_latents_dtype": str(trajectory_latents.dtype),
+    })
+
+    record.update({
+        "trajectory_timesteps_bytes": trajectory_timesteps.tobytes(),
+        "trajectory_timesteps_shape": list(trajectory_timesteps.shape),
+        "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype),
+    })
+
+    return record
+
+
+def text_only_record_creator(text_name: str, text_embedding: np.ndarray,
+                             caption: str) -> dict[str, Any]:
+    """Create a text-only record matching pyarrow_schema_text_only.
+
+    Args:
+        text_name: Base id/name for the text sample.
+        text_embedding: Text encoder output array [SeqLen, Dim].
+        caption: Original text prompt.
+
+    Returns:
+        dict suitable for records_to_table(…, pyarrow_schema_text_only)
+    """
+    record = {
+        "id": f"text_{text_name}",
+        "text_embedding_bytes": text_embedding.tobytes(),
+        "text_embedding_shape": list(text_embedding.shape),
+        "text_embedding_dtype": str(text_embedding.dtype),
+        "caption": caption,
+    }
+    return record

fastvideo/dataset/dataloader/schema.py

@@ -50,6 +50,7 @@ pyarrow_schema_i2v = pa.schema([
     pa.field("fps", pa.float64()),
 ])
 
+
 pyarrow_schema_t2v = pa.schema([
     pa.field("id", pa.string()),
     # --- Image/Video VAE latents ---
@@ -78,3 +79,40 @@ pyarrow_schema_t2v = pa.schema([
     pa.field("duration_sec", pa.float64()),
     pa.field("fps", pa.float64()),
 ])
+
+
+pyarrow_schema_ode_trajectory_text_only = pa.schema([
+    pa.field("id", 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()),
+    # --- ODE Trajectory ---
+    pa.field("trajectory_latents_bytes", pa.binary()),
+    pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_latents_dtype", pa.string()),
+    pa.field("trajectory_timesteps_bytes", pa.binary()),
+    pa.field("trajectory_timesteps_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_timesteps_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # Always 'text' for text-only
+])
+
+
+pyarrow_schema_text_only = pa.schema([
+    pa.field("id", 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()),
+    # --- Metadata ---
+    pa.field("caption", pa.string()),
+])

fastvideo/dataset/preprocessing_datasets.py

@@ -628,3 +628,134 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         """Load state dict from checkpoint."""
         self.processed_batches = state_dict["processed_batches"]
+
+
+class TextDataset(torch.utils.data.IterableDataset,
+                  torch.distributed.checkpoint.stateful.Stateful):
+    """
+    Text-only dataset for processing prompts from a simple text file.
+    
+    Assumes that data_merge_path is a text file with one prompt per line:
+    A cat playing with a ball
+    A dog running in the park
+    A person cooking dinner
+    ...
+    
+    This dataset processes text data through text encoding stages only.
+    """
+
+    def __init__(self,
+                 data_merge_path: str,
+                 args,
+                 start_idx: int = 0,
+                 seed: int = 42):
+        self.data_merge_path = data_merge_path
+        self.start_idx = start_idx
+        self.args = args
+        self.seed = seed
+
+        # Initialize tokenizer
+        tokenizer_path = os.path.join(args.model_path, "tokenizer")
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
+                                                  cache_dir=args.cache_dir)
+
+        # Initialize text encoding stage
+        self.text_encoding_stage = TextEncodingStage(
+            tokenizer=tokenizer,
+            text_max_length=args.text_max_length,
+            cfg_rate=getattr(args, 'training_cfg_rate', 0.0),
+            seed=self.seed)
+
+        # Process text data
+        self.processed_batches = self._process_text_data()
+
+    def _load_text_data(self) -> list[str]:
+        """Load text prompts from file."""
+        prompts = []
+        with open(self.data_merge_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                if line:  # Skip empty lines
+                    prompts.append(line)
+        
+        logger.info(f"Loaded {len(prompts)} text prompts from {self.data_merge_path}")
+        return prompts
+
+    def _process_text_data(self) -> list[PreprocessBatch]:
+        """Process the text prompts through text encoding stage."""
+        raw_prompts = self._load_text_data()
+        processed_batches = []
+
+        for idx, prompt in enumerate(raw_prompts):
+            # Create a text-only batch with dummy path
+            batch = PreprocessBatch(
+                path=f"text_prompt_{idx}",
+                cap=[prompt],  # TextEncodingStage expects a list
+                resolution=None,
+                fps=None,
+                duration=None,
+                num_frames=0,
+                sample_frame_index=None,
+                sample_num_frames=0
+            )
+            
+            processed_batches.append(batch)
+
+        logger.info(f"Processed {len(processed_batches)} text batches")
+        return processed_batches
+
+    def __iter__(self):
+        """Iterator for the dataset."""
+        # Set up distributed sampling if needed
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            rank = torch.distributed.get_rank()
+            world_size = torch.distributed.get_world_size()
+        else:
+            rank = 0
+            world_size = 1
+
+        # Calculate chunk for this rank
+        total_items = len(self.processed_batches)
+        items_per_rank = math.ceil(total_items / world_size)
+        start_idx = rank * items_per_rank + self.start_idx
+        end_idx = min(start_idx + items_per_rank, total_items)
+
+        # Yield items for this rank
+        for idx in range(start_idx, end_idx):
+            if idx < len(self.processed_batches):
+                yield self._get_item(idx)
+
+    def _get_item(self, idx: int) -> dict:
+        """Get a single processed text item."""
+        batch = self.processed_batches[idx]
+
+        # Apply text encoding stage
+        batch = self.text_encoding_stage.process(batch)
+
+        # Build result dictionary for text-only processing with required schema fields
+        result = {
+            "text": batch.text,
+            "input_ids": batch.input_ids,
+            "cond_mask": batch.cond_mask,
+            "path": batch.path,
+            # Required schema fields for ODE trajectory processing
+            "id": f"text_{idx}",
+            "file_name": batch.path,
+            "caption": batch.text,
+            "media_type": "text",
+            "width": 1,
+            "height": 1,
+            "num_frames": 0,
+            "duration_sec": 0.0,
+            "fps": 0.0,
+        }
+
+        return result
+
+    def state_dict(self) -> dict[str, Any]:
+        """Return state dict for checkpointing."""
+        return {"processed_batches": self.processed_batches}
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """Load state dict from checkpoint."""
+        self.processed_batches = state_dict["processed_batches"]

fastvideo/dataset/utils.py

@@ -5,7 +5,7 @@ import numpy as np
 import torch
 
 
-def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
+def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Pad or crop an embedding [L, D] to exactly padding_length tokens.
     Return:

fastvideo/entrypoints/video_generator.py

@@ -344,6 +344,9 @@ class VideoGenerator:
                 "size": (target_height, target_width, batch.num_frames),
                 "generation_time": gen_time,
                 "logging_info": logging_info,
+                "trajectory": output_batch.trajectory_latents,
+                "trajectory_timesteps": output_batch.trajectory_timesteps,
+                "trajectory_decoded": output_batch.trajectory_decoded,
             }
 
     def set_lora_adapter(self,

fastvideo/envs.py

@@ -175,6 +175,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention
     # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention
     # - "SAGE_ATTN": use Sage Attention
+    # - "SAGE_ATTN_THREE": use Sage Attention 3
     "FASTVIDEO_ATTENTION_BACKEND":
     lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
 

fastvideo/fastvideo_args.py

@@ -133,6 +133,7 @@ class FastVideoArgs:
 
     # Compilation
     enable_torch_compile: bool = False
+    torch_compile_kwargs: dict[str, Any] = field(default_factory=dict)
 
     disable_autocast: bool = False
 
@@ -158,15 +159,15 @@ class FastVideoArgs:
         "transformer": True,
         "vae": True,
     })
+    override_transformer_cls_name: str | None = None
+    init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading
+    init_weights_from_safetensors_2: str = ""  # path to safetensors file for initial weight loading for transformer_2
 
     # # DMD parameters
     # dmd_denoising_steps: List[int] | None = field(default=None)
 
     # MoE parameters used by Wan2.2
-    boundary_ratio: float | None = None
-
-    # XXX
-    use_sf_wan: bool = False # force self-forcing Wan model for both distillation and validation
+    boundary_ratio: float | None = 0.875
 
     @property
     def training_mode(self) -> bool:
@@ -194,12 +195,6 @@ class FastVideoArgs:
             help=
             "The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
         )
-        parser.add_argument(
-            "--use-sf-wan",
-            action=StoreBoolean,
-            default=FastVideoArgs.use_sf_wan,
-            help="Use self-forcing Wan model for both distillation and validation",
-        )
         parser.add_argument(
             "--model-dir",
             type=str,
@@ -338,6 +333,13 @@ class FastVideoArgs:
             help="Use torch.compile to speed up DiT inference." +
             "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)",
         )
+        parser.add_argument(
+            "--torch-compile-kwargs",
+            type=str,
+            default=None,
+            help=
+            "JSON string of kwargs to pass to torch.compile. Example: '{\"backend\":\"inductor\",\"mode\":\"reduce-overhead\"}'",
+        )
 
         parser.add_argument(
             "--dit-cpu-offload",
@@ -405,6 +407,20 @@ class FastVideoArgs:
             default=FastVideoArgs.enable_stage_verification,
             help="Enable input/output verification for pipeline stages",
         )
+        parser.add_argument(
+            "--override-transformer-cls-name",
+            type=str,
+            default=FastVideoArgs.override_transformer_cls_name,
+            help="Override transformer cls name",
+        )
+        parser.add_argument(
+            "--init-weights-from-safetensors",
+            type=str,
+            help="Path to safetensors file for initial weight loading")
+        parser.add_argument(
+            "--init-weights-from-safetensors-2",
+            type=str,
+            help="Path to safetensors file for initial weight loading")
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 
@@ -433,6 +449,21 @@ class FastVideoArgs:
                 mode_value = getattr(args, attr, FastVideoArgs.mode.value)
                 kwargs['mode'] = ExecutionMode.from_string(
                     mode_value) if isinstance(mode_value, str) else mode_value
+            elif attr == 'torch_compile_kwargs':
+                # Parse JSON string for torch.compile kwargs
+                torch_compile_kwargs_str = getattr(args, 'torch_compile_kwargs',
+                                                   None)
+                if torch_compile_kwargs_str:
+                    try:
+                        import json
+                        kwargs['torch_compile_kwargs'] = json.loads(
+                            torch_compile_kwargs_str)
+                    except json.JSONDecodeError as e:
+                        raise ValueError(
+                            f"Invalid JSON for torch_compile_kwargs: {e}"
+                        ) from e
+                else:
+                    kwargs['torch_compile_kwargs'] = {}
             elif attr == 'workload_type':
                 # Convert string to WorkloadType enum
                 workload_type_value = getattr(args, 'workload_type',
@@ -612,10 +643,8 @@ class TrainingArgs(FastVideoArgs):
 
     # text encoder & vae & diffusion model
     pretrained_model_name_or_path: str = ""
-    dit_model_name_or_path: str = ""
 
     # DMD model paths - separate paths for each network
-    generator_model_path: str = ""  # path for generator (student) model
     real_score_model_path: str = ""  # path for real score (teacher) model
     fake_score_model_path: str = ""  # path for fake score (critic) model
 
@@ -639,9 +668,7 @@ class TrainingArgs(FastVideoArgs):
     # output
     output_dir: str = ""
     checkpoints_total_limit: int = 0
-    checkpointing_steps: int = 0
     resume_from_checkpoint: str = ""  # specify the checkpoint folder to resume from
-    init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading
 
     # optimizer & scheduler
     num_train_epochs: int = 0
@@ -713,11 +740,9 @@ class TrainingArgs(FastVideoArgs):
     independent_first_frame: bool = False
     enable_gradient_masking: bool = True
     gradient_mask_last_n_frames: int = 21
-    validate_cache_structure: bool = False  # Debug flag for cache validation
     same_step_across_blocks: bool = False  # Use same exit timestep for all blocks
     last_step_only: bool = False  # Only use the last timestep for training
     context_noise: int = 0  # Context noise level for cache updates
-    sf_ode_init_path: str = ""  # Path to ODE init weights for self-forcing model
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -888,9 +913,6 @@ class TrainingArgs(FastVideoArgs):
         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(
             "--training-state-checkpointing-steps",
             type=int,
@@ -903,10 +925,6 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--resume-from-checkpoint",
                             type=str,
                             help="Path to checkpoint to resume from")
-        parser.add_argument(
-            "--init-weights-from-safetensors",
-            type=str,
-            help="Path to safetensors file for initial weight loading")
         parser.add_argument("--logging-dir",
                             type=str,
                             help="Directory for logging")
@@ -1143,16 +1161,10 @@ class TrainingArgs(FastVideoArgs):
             "--last-step-only",
             action=StoreBoolean,
             help="Whether to only use the last timestep for training")
-        parser.add_argument(
-            "--context-noise",
-            type=int,
-            default=TrainingArgs.context_noise,
-            help="Context noise level for cache updates")
-        parser.add_argument(
-            "--sf-ode-init-path",
-            type=str,
-            default=TrainingArgs.sf_ode_init_path,
-            help="Path to ODE init weights for self-forcing model")
+        parser.add_argument("--context-noise",
+                            type=int,
+                            default=TrainingArgs.context_noise,
+                            help="Context noise level for cache updates")
 
         return parser
 
@@ -1160,4 +1172,4 @@ class TrainingArgs(FastVideoArgs):
 def parse_int_list(value: str) -> list[int]:
     if not value:
         return []
-    return [int(x.strip()) for x in value.split(",")]
+    return [int(x.strip()) for x in value.split(",")]

fastvideo/layers/lora/linear.py

@@ -77,9 +77,11 @@ class BaseLayerWithLoRA(nn.Module):
             lora_A = self.lora_A.to_local()
 
         if not self.merged and not self.disable_lora:
-            delta = x @ (
-                self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
-                @ self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)))
+            lora_A_sliced = self.slice_lora_a_weights(
+                lora_A.to(x, non_blocking=True))
+            lora_B_sliced = self.slice_lora_b_weights(
+                lora_B.to(x, non_blocking=True))
+            delta = x @ lora_A_sliced.T @ lora_B_sliced.T
             if self.lora_alpha != self.lora_rank:
                 delta = delta * (
                     self.lora_alpha / self.lora_rank  # type: ignore

fastvideo/models/dits/causal_wanvideo.py

@@ -147,17 +147,11 @@ class CausalWanSelfAttention(nn.Module):
                 # Assign new keys/values directly up to current_end
                 local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
                 local_start_index = local_end_index - num_new_tokens
-                # kv_cache["k"] = kv_cache["k"].detach()
-                # kv_cache["v"] = kv_cache["v"].detach()
-                logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None)
-                s, e = local_start_index, local_end_index
-                k = kv_cache["k"]
-                kv_cache["k"] = torch.cat([k[:, :s], roped_key, k[:, e:]], dim=1)
-
-                v0 = kv_cache["v"]
-                kv_cache["v"] = torch.cat([v0[:, :s], v, v0[:, e:]], dim=1)
-                # kv_cache["k"][:, local_start_index:local_end_index] = roped_key
-                # kv_cache["v"][:, local_start_index:local_end_index] = v
+                kv_cache["k"] = kv_cache["k"].detach()
+                kv_cache["v"] = kv_cache["v"].detach()
+                # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None)
+                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                kv_cache["v"][:, local_start_index:local_end_index] = v
             x = self.attn(
                 roped_query,
                 kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index],
@@ -185,7 +179,7 @@ class CausalWanTransformerBlock(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
         self.to_q = ReplicatedLinear(dim, dim, bias=True)
         self.to_k = ReplicatedLinear(dim, dim, bias=True)
         self.to_v = ReplicatedLinear(dim, dim, bias=True)
@@ -218,8 +212,7 @@ class CausalWanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=True,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            dtype=torch.float32)
 
         # 2. Cross-attention
         # Only T2V for now
@@ -232,8 +225,7 @@ class CausalWanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=False,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -258,29 +250,29 @@ class CausalWanTransformerBlock(nn.Module):
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
         num_frames = temb.shape[1]
-        frame_seqlen = hidden_states.shape[1] // num_frames
+        frame_seqlen = hidden_states.shape[1] // num_frames    
         bs, seq_length, _ = hidden_states.shape
         orig_dtype = hidden_states.dtype
         # assert orig_dtype != torch.float32
-        e = self.scale_shift_table + temb.float()
+        e = self.scale_shift_table + temb
         # e.shape: [batch_size, num_frames, 6, inner_dim]
         assert e.shape == (bs, num_frames, 6, self.hidden_dim)
         shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
             6, dim=2)
         # *_msa.shape: [batch_size, num_frames, 1, inner_dim]
-        assert shift_msa.dtype == torch.float32
+        # assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
-                        (1 + scale_msa) + shift_msa).flatten(1, 2).to(orig_dtype)
+        norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
+                        (1 + scale_msa) + shift_msa).flatten(1, 2)
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
 
         if self.norm_q is not None:
-            query = self.norm_q(query)
+            query = self.norm_q.forward_native(query)
         if self.norm_k is not None:
-            key = self.norm_k(key)
+            key = self.norm_k.forward_native(key)
 
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
@@ -294,8 +286,6 @@ class CausalWanTransformerBlock(nn.Module):
         null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
         norm_hidden_states, hidden_states = self.self_attn_residual_norm(
             hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
 
         # 2. Cross-attention
         attn_output = self.attn2(norm_hidden_states,
@@ -304,13 +294,10 @@ class CausalWanTransformerBlock(nn.Module):
                                  crossattn_cache=crossattn_cache)
         norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
             hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
         hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
-        hidden_states = hidden_states.to(orig_dtype)
 
         return hidden_states
 
@@ -373,8 +360,7 @@ class CausalWanTransformer3DModel(BaseDiT):
                                             norm_type="layer",
                                             eps=config.eps,
                                             elementwise_affine=False,
-                                            dtype=torch.float32,
-                                            compute_dtype=torch.float32)
+                                            dtype=torch.float32)
         self.proj_out = nn.Linear(
             inner_dim, config.out_channels * math.prod(config.patch_size))
         self.scale_shift_table = nn.Parameter(
@@ -384,7 +370,8 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         # Causal-specific
         self.block_mask = None
-        self.num_frame_per_block = 1
+        self.num_frame_per_block = config.arch_config.num_frames_per_block
+        assert self.num_frame_per_block <= 3
         self.independent_first_frame = False
 
         self.__post_init__()
@@ -496,12 +483,16 @@ class CausalWanTransformer3DModel(BaseDiT):
         )
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos.float(),
-                     freqs_sin.float()) if freqs_cos is not None else None
+        freqs_cis = (freqs_cos,
+                     freqs_sin) if freqs_cos is not None else None
 
         hidden_states = self.patch_embedding(hidden_states)
+        grid_sizes = torch.stack(
+            [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)])
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
+        encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
+
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
                         timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
         timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
@@ -548,14 +539,9 @@ class CausalWanTransformer3DModel(BaseDiT):
         hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
 
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
-        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
-        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+        output = self.unpatchify(hidden_states, grid_sizes)
 
-        return output
+        return torch.stack(output)
 
     def _forward_train(self,
                 hidden_states: torch.Tensor,
@@ -596,8 +582,8 @@ class CausalWanTransformer3DModel(BaseDiT):
         )
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos.float(),
-                     freqs_sin.float()) if freqs_cos is not None else None
+        freqs_cis = (freqs_cos,
+                     freqs_sin) if freqs_cos is not None else None
 
         # Construct blockwise causal attn mask
         if self.block_mask is None:
@@ -610,8 +596,12 @@ class CausalWanTransformer3DModel(BaseDiT):
             )
 
         hidden_states = self.patch_embedding(hidden_states)
+        grid_sizes = torch.stack(
+            [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)])
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
+        encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
+
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
                         timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
         timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
@@ -646,14 +636,9 @@ class CausalWanTransformer3DModel(BaseDiT):
         hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
 
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
-        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
-        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+        output = self.unpatchify(hidden_states, grid_sizes)
 
-        return output
+        return torch.stack(output)
 
     def forward(
         self,
@@ -664,3 +649,30 @@ class CausalWanTransformer3DModel(BaseDiT):
             return self._forward_inference(*args, **kwargs)
         else:
             return self._forward_train(*args, **kwargs)
+
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+
+
+        Returns:
+            Tensor:
+                Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_channels
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = u.permute(6, 0, 3, 1, 4, 2, 5)
+            # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous())
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out

fastvideo/models/loader/component_loader.py

@@ -416,6 +416,11 @@ class TransformerLoader(ComponentLoader):
                 "Model config does not contain a _class_name attribute. "
                 "Only diffusers format is supported.")
 
+        logger.info("transformer cls_name: %s", cls_name)
+        if fastvideo_args.override_transformer_cls_name is not None:
+            cls_name = fastvideo_args.override_transformer_cls_name
+            logger.info("Overriding transformer cls_name to %s", cls_name)
+
         fastvideo_args.model_paths["transformer"] = model_path
 
         # Config from Diffusers supersedes fastvideo's model config
@@ -433,15 +438,21 @@ class TransformerLoader(ComponentLoader):
         # Check if we should use custom initialization weights
         custom_weights_path = getattr(fastvideo_args, 'init_weights_from_safetensors', None)
         use_custom_weights = (custom_weights_path and os.path.exists(custom_weights_path) and 
-                            fastvideo_args.training_mode and 
                             not hasattr(fastvideo_args, '_loading_teacher_critic_model'))
 
         if use_custom_weights:
-            logger.info("Using custom initialization weights from: %s", custom_weights_path)
-            safetensors_list = [custom_weights_path]
+            if 'transformer_2' in model_path:
+                custom_weights_path = getattr(fastvideo_args, 'init_weights_from_safetensors_2', None)
+            assert custom_weights_path is not None, "Custom initialization weights must be provided"
+            if os.path.isdir(custom_weights_path):
+                safetensors_list = glob.glob(
+                    os.path.join(str(custom_weights_path), "*.safetensors"))
+            else:
+                assert custom_weights_path.endswith(".safetensors"), "Custom initialization weights must be a safetensors file"
+                safetensors_list = [custom_weights_path]
 
-        logger.info("Loading model from %s safetensors files in %s",
-                    len(safetensors_list), model_path)
+        logger.info("Loading model from %s safetensors files: %s",
+                    len(safetensors_list), safetensors_list)
 
         default_dtype = PRECISION_TO_TYPE[
             fastvideo_args.pipeline_config.dit_precision]
@@ -464,18 +475,20 @@ class TransformerLoader(ComponentLoader):
             pin_cpu_memory=fastvideo_args.pin_cpu_memory,
             fsdp_inference=fastvideo_args.use_fsdp_inference,
             # TODO(will): make these configurable
+            default_dtype=default_dtype,
             param_dtype=torch.bfloat16,
             reduce_dtype=torch.float32,
             output_dtype=None,
-            training_mode=fastvideo_args.training_mode)
+            training_mode=fastvideo_args.training_mode,
+            enable_torch_compile=fastvideo_args.enable_torch_compile,
+            torch_compile_kwargs=fastvideo_args.torch_compile_kwargs)
 
 
         total_params = sum(p.numel() for p in model.parameters())
         logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
 
-        dtypes = set(param.dtype for param in model.parameters())
-        if len(dtypes) > 1:
-            model = model.to(default_dtype)
+        assert next(model.parameters()).dtype == default_dtype, "Model dtype does not match default dtype"
+
         model = model.eval()
         return model
 

fastvideo/models/loader/fsdp_load.py

@@ -54,7 +54,7 @@ def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
         torch.set_default_dtype(old_dtype)
 
 
-# TODO(PY): add compile option
+# Supports optional torch.compile for FSDP-wrapped models during training
 def maybe_load_fsdp_model(
     model_cls: type[nn.Module],
     init_params: dict[str, Any],
@@ -62,6 +62,7 @@ def maybe_load_fsdp_model(
     device: torch.device,
     hsdp_replicate_dim: int,
     hsdp_shard_dim: int,
+    default_dtype: torch.dtype,
     param_dtype: torch.dtype,
     reduce_dtype: torch.dtype,
     cpu_offload: bool = False,
@@ -69,6 +70,8 @@ def maybe_load_fsdp_model(
     output_dtype: torch.dtype | None = None,
     training_mode: bool = True,
     pin_cpu_memory: bool = True,
+    enable_torch_compile: bool = False,
+    torch_compile_kwargs: dict[str, Any] | None = None,
 ) -> torch.nn.Module:
     """
     Load the model with FSDP if is training, else load the model without FSDP.
@@ -87,7 +90,8 @@ def maybe_load_fsdp_model(
         mp_policy=mp_policy,
     )
 
-    with set_default_dtype(param_dtype), torch.device("meta"):
+    logger.info("Loading model with default_dtype: %s", default_dtype)
+    with set_default_dtype(default_dtype), torch.device("meta"):
         model = model_cls(**init_params)
 
     # Check if we should use FSDP
@@ -125,7 +129,7 @@ def maybe_load_fsdp_model(
         model,
         weight_iterator,
         device,
-        param_dtype,
+        default_dtype,
         strict=True,
         cpu_offload=cpu_offload,
         param_names_mapping=param_names_mapping_fn,
@@ -137,6 +141,14 @@ def maybe_load_fsdp_model(
         # Avoid unintended computation graph accumulation during inference
         if isinstance(p, torch.nn.Parameter):
             p.requires_grad = False
+
+    compile_in_loader = enable_torch_compile and training_mode
+    if compile_in_loader:
+        compile_kwargs = torch_compile_kwargs or {}
+        logger.info("Enabling torch.compile for FSDP training module with kwargs=%s",
+                    compile_kwargs)
+        model = torch.compile(model, **compile_kwargs)
+        logger.info("torch.compile enabled for %s", type(model).__name__)
     return model
 
 

fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py

@@ -64,8 +64,15 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
     def step(self, model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final=False, return_dict=False, **kwargs):
         if timestep.ndim == 2:
             timestep = timestep.flatten(0, 1)
+        elif timestep.ndim == 0:
+            # handles the case where timestep is a scalar, this occurs when we
+            # use this scheduler for ODE trajectory
+            timestep = timestep.unsqueeze(0)
+
         self.sigmas = self.sigmas.to(model_output.device)
         self.timesteps = self.timesteps.to(model_output.device)
+        timestep = timestep.to(model_output.device)
+
         timestep_id = torch.argmin(
             (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
         sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)

fastvideo/models/utils.py

@@ -171,10 +171,10 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
     # timestep shape should be [B]
     dtype = pred_noise.dtype
     device = pred_noise.device
-    pred_noise = pred_noise.float().to(device)
-    noise_input_latent = noise_input_latent.float().to(device)
-    sigmas = scheduler.sigmas.float().to(device)
-    timesteps = scheduler.timesteps.float().to(device)
+    pred_noise = pred_noise.double().to(device)
+    noise_input_latent = noise_input_latent.double().to(device)
+    sigmas = scheduler.sigmas.double().to(device)
+    timesteps = scheduler.timesteps.double().to(device)
     timestep_id = torch.argmin(
         (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
     sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)

fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py

@@ -7,8 +7,6 @@ This module wires the causal DMD denoising stage into the modular pipeline.
 
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler)
 from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
 
 # isort: off
@@ -19,14 +17,8 @@ from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
                                         TextEncodingStage)
 # isort: on
 
-import torch
-
-from fastvideo.sf_utils.wan_wrapper import WanDiffusionWrapper
-
 logger = init_logger(__name__)
 
-from fastvideo.distributed import get_local_torch_device
-
 
 class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
 
@@ -36,21 +28,6 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
 
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
         """Set up pipeline stages with proper dependency injection."""
-        if fastvideo_args.use_sf_wan:
-            # timestep shift is 5.0 for self-forcing Wan model
-            # see https://github.com/guandeh17/Self-Forcing/blob/33593df3e81fa3ec10239271dd2c100facac6de1/configs/self_forcing_dmd.yaml#L50
-            config = self.get_module("transformer").config
-            if not isinstance(self.modules["transformer"], WanDiffusionWrapper):
-                sf_transformer = WanDiffusionWrapper(
-                    model_name="Wan2.1-T2V-1.3B", timestep_shift=5.0, is_causal=True, config=config)
-                del self.modules["transformer"]
-                state_dict = torch.load('checkpoints/self_forcing_dmd.pt')
-                sf_transformer.load_state_dict(state_dict['generator_ema'])
-                sf_transformer.to(get_local_torch_device())
-                self.modules["transformer"] = sf_transformer
-                logger.info("Using self-forcing Wan model for DMD inference")
-            else:
-                logger.info("transformer is already a WanDiffusionWrapper")
 
         self.add_stage(stage_name="input_validation_stage",
                        stage=InputValidationStage())
@@ -72,6 +49,7 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
         self.add_stage(stage_name="denoising_stage",
                        stage=CausalDMDDenosingStage(
                            transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2", None),
                            scheduler=self.get_module("scheduler")))
 
         self.add_stage(stage_name="decoding_stage",

fastvideo/pipelines/composed_pipeline_base.py

@@ -40,7 +40,7 @@ class ComposedPipelineBase(ABC):
     _extra_config_module_map: dict[str, str] = {}
     training_args: TrainingArgs | None = None
     fastvideo_args: FastVideoArgs | TrainingArgs | None = None
-    modules: dict[str, torch.nn.Module] = {}
+    modules: dict[str, Any] = {}
     post_init_called: bool = False
 
     # TODO(will): args should support both inference args and training args
@@ -99,9 +99,30 @@ class ComposedPipelineBase(ABC):
 
         self.initialize_pipeline(self.fastvideo_args)
         if self.fastvideo_args.enable_torch_compile:
-            self.modules["transformer"] = torch.compile(
-                self.modules["transformer"])
-            logger.info("Torch Compile enabled for DiT")
+            transformer_module = self.modules["transformer"]
+            if self.fastvideo_args.training_mode:
+                logger.info(
+                    "Torch Compile enabled via FSDP loader for training; skipping additional pipeline compile"
+                )
+            else:
+                fsdp_module_cls = None
+                try:
+                    from torch.distributed.fsdp import FSDPModule  # type: ignore
+                    fsdp_module_cls = FSDPModule
+                except Exception:  # pragma: no cover - FSDP not always available
+                    fsdp_module_cls = None
+                if fsdp_module_cls is not None and isinstance(
+                        transformer_module, fsdp_module_cls):
+                    logger.info(
+                        "Transformer is already FSDP-wrapped; skipping torch.compile in pipeline"
+                    )
+                else:
+                    compile_kwargs = self.fastvideo_args.torch_compile_kwargs or {}
+                    logger.info("Enabling torch.compile for DiT with kwargs=%s",
+                                compile_kwargs)
+                    self.modules["transformer"] = torch.compile(
+                        transformer_module, **compile_kwargs)
+                    logger.info("Torch Compile enabled for DiT")
 
         if not self.fastvideo_args.training_mode:
             logger.info("Creating pipeline stages...")
@@ -237,20 +258,19 @@ class ComposedPipelineBase(ABC):
         # remove keys that are not pipeline modules
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
-        # @TODO(Wei): Temporary hack
         if "boundary_ratio" in model_index and model_index[
                 "boundary_ratio"] is not None:
             logger.info(
                 "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..."
             )
             self.required_config_modules.append("transformer_2")
-            if fastvideo_args.boundary_ratio is None:
-                logger.info(
-                    "MoE pipeline detected. Setting boundary ratio to %s",
-                    model_index["boundary_ratio"])
-                fastvideo_args.boundary_ratio = model_index["boundary_ratio"]
+            logger.info("MoE pipeline detected. Setting boundary ratio to %s",
+                        model_index["boundary_ratio"])
+            fastvideo_args.pipeline_config.dit_config.boundary_ratio = model_index[
+                "boundary_ratio"]
 
         model_index.pop("boundary_ratio", None)
+        # used by Wan2.2 ti2v
         model_index.pop("expand_timesteps", None)
 
         # some sanity checks
@@ -283,8 +303,8 @@ class ComposedPipelineBase(ABC):
                           architecture) in model_index.items():
             if transformers_or_diffusers is None:
                 logger.warning(
-                    "Module in model_index.json has null value, removing from required_config_modules"
-                )
+                    "Module %s in model_index.json has null value, removing from required_config_modules",
+                    module_name)
                 if module_name in self.required_config_modules:
                     self.required_config_modules.remove(module_name)
                 continue

fastvideo/pipelines/pipeline_batch_info.py

@@ -129,6 +129,7 @@ class ForwardBatch:
     timesteps: torch.Tensor | None = None
     timestep: torch.Tensor | float | int | None = None
     step_index: int | None = None
+    boundary_ratio: float | None = None
 
     # Scheduler parameters
     num_inference_steps: int = 50
@@ -147,7 +148,12 @@ class ForwardBatch:
     modules: dict[str, Any] = field(default_factory=dict)
 
     # Final output (after pipeline completion)
-    output: Any = None
+    output: torch.Tensor | None = None
+    return_trajectory_latents: bool = False
+    return_trajectory_decoded: bool = False
+    trajectory_timesteps: list[torch.Tensor] | None = None
+    trajectory_latents: torch.Tensor | None = None
+    trajectory_decoded: list[torch.Tensor] | None = None
 
     # Extra parameters that might be needed by specific pipeline implementations
     extra: dict[str, Any] = field(default_factory=dict)
@@ -206,6 +212,10 @@ class TrainingBatch:
     infos: list[dict[str, Any]] | None = None
     mask_lat_size: torch.Tensor | None = None
 
+    # ODE trajectory supervision
+    trajectory_latents: torch.Tensor | None = None
+    trajectory_timesteps: torch.Tensor | None = None
+
     # Transformer inputs
     noisy_model_input: torch.Tensor | None = None
     timesteps: torch.Tensor | None = None
@@ -236,6 +246,7 @@ class TrainingBatch:
     fake_score_loss: float = 0.0
 
     dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
+    latent_vis_dict: dict[str, Any] = field(default_factory=dict)
     fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
 
 

fastvideo/pipelines/pipeline_registry.py

@@ -205,7 +205,7 @@ def import_pipeline_classes(
         except ImportError as e:
             raise ImportError(
                 f"Could not import {pipeline_type_package_name} when importing pipeline classes: {e}"
-            ) from e
+            ) from None
 
         type_to_arch_to_pipeline_dict[pipeline_type_str] = arch_to_pipeline_dict
 

fastvideo/pipelines/preprocess/preprocess_pipeline_base.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-import multiprocessing
 import os
-from concurrent.futures import ProcessPoolExecutor
 from typing import Any
 
 import numpy as np
@@ -12,6 +10,8 @@ from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from fastvideo.dataset import getdataset
+from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter,
+                                                     records_to_table)
 from fastvideo.dataset.preprocessing_datasets import PreprocessBatch
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -54,10 +54,14 @@ class BasePreprocessPipeline(ComposedPipelineBase):
         """Get additional features specific to the pipeline type. Override in subclasses."""
         return {}
 
-    def get_schema_fields(self) -> list[str]:
-        """Get the schema fields for the pipeline type. Override in subclasses."""
+    def get_pyarrow_schema(self) -> pa.Schema:
+        """Return the PyArrow schema for this pipeline. Must be overridden."""
         raise NotImplementedError
 
+    def get_schema_fields(self) -> list[str]:
+        """Get the schema fields for the pipeline type."""
+        return [f.name for f in self.get_pyarrow_schema()]
+
     def create_record_for_schema(self,
                                  preprocess_batch: PreprocessBatch,
                                  schema: pa.Schema,
@@ -400,166 +404,22 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                 batch_data.append(record)
 
             if batch_data:
-                # Add progress bar for writing to Parquet dataset
                 write_pbar = tqdm(total=1,
                                   desc="Writing to Parquet dataset",
                                   unit="batch")
-                # Convert batch data to PyArrow arrays
-                arrays = []
-                for field in self.get_schema_fields():
-                    if field.endswith('_bytes'):
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data],
-                                     type=pa.binary()))
-                    elif field.endswith('_shape'):
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data],
-                                     type=pa.list_(pa.int32())))
-                    elif field in ['width', 'height', 'num_frames']:
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data],
-                                     type=pa.int32()))
-                    elif field in ['duration_sec', 'fps']:
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data],
-                                     type=pa.float32()))
-                    else:
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data]))
-
-                table = pa.Table.from_arrays(arrays,
-                                             names=self.get_schema_fields())
+                table = records_to_table(batch_data, self.get_pyarrow_schema())
                 write_pbar.update(1)
                 write_pbar.close()
 
-                # Store the table in a list for later processing
-                if not hasattr(self, 'all_tables'):
-                    self.all_tables = []
-                self.all_tables.append(table)
-
+                if not hasattr(self, 'dataset_writer'):
+                    self.dataset_writer = ParquetDatasetWriter(
+                        out_dir=combined_parquet_dir,
+                        samples_per_file=args.samples_per_file,
+                    )
+                self.dataset_writer.append_table(table)
                 logger.info("Collected batch with %s samples", len(table))
 
             if num_processed_samples >= args.flush_frequency:
-                self._flush_tables(num_processed_samples, args,
-                                   combined_parquet_dir)
+                written = self.dataset_writer.flush()
+                logger.info("Flushed %s samples to parquet", written)
                 num_processed_samples = 0
-                self.all_tables = []
-
-    def _flush_tables(self, num_processed_samples: int, args,
-                      combined_parquet_dir: str):
-        """Flush collected tables to disk."""
-        assert hasattr(self, 'all_tables') and self.all_tables
-        print(f"Combining {len(self.all_tables)} batches...")
-        combined_table = pa.concat_tables(self.all_tables)
-        assert len(combined_table) == num_processed_samples
-        print(f"Total samples collected: {len(combined_table)}")
-
-        # Calculate total number of chunks needed, discarding remainder
-        total_chunks = max(num_processed_samples // args.samples_per_file, 1)
-
-        print(f"Fixed samples per parquet file: {args.samples_per_file}")
-        print(f"Total number of parquet files: {total_chunks}")
-        print(
-            f"Total samples to be processed: {total_chunks * args.samples_per_file} (discarding {num_processed_samples % args.samples_per_file} samples)"
-        )
-
-        # Split work among processes
-        num_workers = int(min(multiprocessing.cpu_count(), total_chunks))
-        chunks_per_worker = (total_chunks + num_workers - 1) // num_workers
-
-        print(f"Using {num_workers} workers to process {total_chunks} chunks")
-        logger.info("Chunks per worker: %s", chunks_per_worker)
-
-        # Prepare work ranges
-        work_ranges = []
-        for i in range(num_workers):
-            start_idx = i * chunks_per_worker
-            end_idx = min((i + 1) * chunks_per_worker, total_chunks)
-            if start_idx < total_chunks:
-                work_ranges.append(
-                    (start_idx, end_idx, combined_table, i,
-                     combined_parquet_dir, args.samples_per_file))
-
-        total_written = 0
-        failed_ranges = []
-        with ProcessPoolExecutor(max_workers=num_workers) as executor:
-            futures = {
-                executor.submit(self.process_chunk_range, work_range):
-                work_range
-                for work_range in work_ranges
-            }
-            for future in tqdm(futures, desc="Processing chunks"):
-                try:
-                    written = future.result()
-                    total_written += written
-                    logger.info("Processed chunk with %s samples", written)
-                except Exception as e:
-                    work_range = futures[future]
-                    failed_ranges.append(work_range)
-                    logger.error("Failed to process range %s-%s: %s",
-                                 work_range[0], work_range[1], str(e))
-
-        # Retry failed ranges sequentially
-        if failed_ranges:
-            logger.warning("Retrying %s failed ranges sequentially",
-                           len(failed_ranges))
-            for work_range in failed_ranges:
-                try:
-                    total_written += self.process_chunk_range(work_range)
-                except Exception as e:
-                    logger.error(
-                        "Failed to process range %s-%s after retry: %s",
-                        work_range[0], work_range[1], str(e))
-
-        logger.info("Total samples written: %s", total_written)
-
-    @staticmethod
-    def process_chunk_range(args: Any) -> int:
-        start_idx, end_idx, table, worker_id, output_dir, samples_per_file = args
-        try:
-            total_written = 0
-            num_samples = len(table)
-
-            # Create worker-specific subdirectory
-            worker_dir = os.path.join(output_dir, f"worker_{worker_id}")
-            os.makedirs(worker_dir, exist_ok=True)
-
-            # Check how many files there are already in the dir, and update i accordingly
-            num_parquets = 0
-            for root, _, files in os.walk(worker_dir):
-                for file in files:
-                    if file.endswith('.parquet'):
-                        num_parquets += 1
-
-            for i in range(start_idx, end_idx):
-                start_sample = i * samples_per_file
-                end_sample = min((i + 1) * samples_per_file, num_samples)
-                chunk = table.slice(start_sample, end_sample - start_sample)
-
-                # Create chunk file in worker's directory
-                chunk_path = os.path.join(
-                    worker_dir, f"data_chunk_{i + num_parquets}.parquet")
-                temp_path = chunk_path + '.tmp'
-
-                try:
-                    # Write to temporary file
-                    pq.write_table(chunk, temp_path, compression='zstd')
-
-                    # Rename temporary file to final file
-                    if os.path.exists(chunk_path):
-                        os.remove(
-                            chunk_path)  # Remove existing file if it exists
-                    os.rename(temp_path, chunk_path)
-
-                    total_written += len(chunk)
-                except Exception as e:
-                    # Clean up temporary file if it exists
-                    if os.path.exists(temp_path):
-                        os.remove(temp_path)
-                    raise e
-
-            return total_written
-        except Exception as e:
-            logger.error("Error processing chunks %s-%s for worker %s: %s",
-                         start_idx, end_idx, worker_id, str(e))
-            raise

fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py

@@ -40,9 +40,9 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
                            image_processor=self.get_module("image_processor"),
                        ))
 
-    def get_schema_fields(self) -> list[str]:
-        """Get the schema fields for I2V pipeline."""
-        return [f.name for f in pyarrow_schema_i2v]
+    def get_pyarrow_schema(self):
+        """Return the PyArrow schema for I2V pipeline."""
+        return pyarrow_schema_i2v
 
     def get_extra_features(self, valid_data: dict[str, Any],
                            fastvideo_args: FastVideoArgs) -> dict[str, Any]:

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -0,0 +1,323 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+ODE Trajectory Data Preprocessing pipeline implementation.
+
+This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline
+using the modular pipeline architecture.
+
+Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772
+"""
+
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import pyarrow as pa
+import torch
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm import tqdm
+
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset import gettextdataset
+from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter,
+                                                     records_to_table)
+from fastvideo.dataset.dataloader.record_schema import (
+    ode_text_only_record_creator)
+from fastvideo.dataset.dataloader.schema import (
+    pyarrow_schema_ode_trajectory_text_only)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (
+    SelfForcingFlowMatchScheduler)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
+    BasePreprocessPipeline)
+from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage)
+from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict
+
+logger = init_logger(__name__)
+
+
+class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
+    """ODE Trajectory preprocessing pipeline implementation."""
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+    preprocess_dataloader: StatefulDataLoader
+    preprocess_loader_iter: Iterator[dict[str, Any]]
+    pbar: Any
+    num_processed_samples: int
+
+    def get_pyarrow_schema(self) -> pa.Schema:
+        """Return the PyArrow schema for ODE Trajectory pipeline."""
+        return pyarrow_schema_ode_trajectory_text_only
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        assert fastvideo_args.pipeline_config.flow_shift == 5
+        self.modules["scheduler"] = SelfForcingFlowMatchScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift,
+            sigma_min=0.0,
+            extra_one_step=True)
+        self.modules["scheduler"].set_timesteps(num_inference_steps=48,
+                                                denoising_strength=1.0)
+
+        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="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", None)))
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler"),
+                           pipeline=self,
+                       ))
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+    def preprocess_text_and_trajectory(self, fastvideo_args: FastVideoArgs,
+                                       args):
+        """Preprocess text-only data and generate trajectory information."""
+
+        for batch_idx, data in enumerate(self.pbar):
+            if data is None:
+                continue
+
+            with torch.inference_mode():
+                # For text-only processing, we only need text data
+                # Filter out samples without text
+                valid_indices = []
+                for i, text in enumerate(data["text"]):
+                    if text and text.strip():  # Check if text is not empty
+                        valid_indices.append(i)
+                self.num_processed_samples += len(valid_indices)
+
+                if not valid_indices:
+                    continue
+
+                # Create new batch with only valid samples (text-only)
+                valid_data = {
+                    "text": [data["text"][i] for i in valid_indices],
+                    "path": [data["path"][i] for i in valid_indices],
+                }
+
+                # Add fps and duration if available in data
+                if "fps" in data:
+                    valid_data["fps"] = [data["fps"][i] for i in valid_indices]
+                if "duration" in data:
+                    valid_data["duration"] = [
+                        data["duration"][i] for i in valid_indices
+                    ]
+
+                batch_captions = valid_data["text"]
+                # Encode text using the standalone TextEncodingStage API
+                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                    batch_captions,
+                    fastvideo_args,
+                    encoder_index=[0],
+                    return_attention_mask=True,
+                )
+                prompt_embeds = prompt_embeds_list[0]
+                prompt_attention_masks = prompt_masks_list[0]
+                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
+
+                sampling_params = SamplingParam.from_pretrained(args.model_path)
+
+                # encode negative prompt for trajectory collection
+                if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None:
+                    negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                        sampling_params.negative_prompt,
+                        fastvideo_args,
+                        encoder_index=[0],
+                        return_attention_mask=True,
+                    )
+                    negative_prompt_embed = negative_prompt_embeds_list[0][0]
+                    negative_prompt_attention_mask = negative_prompt_masks_list[
+                        0][0]
+                else:
+                    negative_prompt_embed = None
+                    negative_prompt_attention_mask = None
+
+                trajectory_latents = []
+                trajectory_timesteps = []
+                trajectory_decoded = []
+
+                for i, (prompt_embed, prompt_attention_mask) in enumerate(
+                        zip(prompt_embeds, prompt_attention_masks,
+                            strict=False)):
+                    prompt_embed = prompt_embed.unsqueeze(0)
+                    prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
+
+                    # Collect the trajectory data (text-to-video generation)
+                    batch = ForwardBatch(**shallow_asdict(sampling_params), )
+                    batch.prompt_embeds = [prompt_embed]
+                    batch.prompt_attention_mask = [prompt_attention_mask]
+                    batch.negative_prompt_embeds = [negative_prompt_embed]
+                    batch.negative_attention_mask = [
+                        negative_prompt_attention_mask
+                    ]
+                    batch.num_inference_steps = 48
+                    batch.return_trajectory_latents = True
+                    # Enabling this will save the decoded trajectory videos.
+                    # Used for debugging.
+                    batch.return_trajectory_decoded = False
+                    batch.height = args.max_height
+                    batch.width = args.max_width
+                    batch.fps = args.train_fps
+                    batch.guidance_scale = 6.0
+                    batch.do_classifier_free_guidance = True
+
+                    result_batch = self.input_validation_stage(
+                        batch, fastvideo_args)
+                    result_batch = self.timestep_preparation_stage(
+                        batch, fastvideo_args)
+                    result_batch = self.latent_preparation_stage(
+                        result_batch, fastvideo_args)
+                    result_batch = self.denoising_stage(result_batch,
+                                                        fastvideo_args)
+                    result_batch = self.decoding_stage(result_batch,
+                                                       fastvideo_args)
+
+                    trajectory_latents.append(
+                        result_batch.trajectory_latents.cpu())
+                    trajectory_timesteps.append(
+                        result_batch.trajectory_timesteps.cpu())
+                    trajectory_decoded.append(result_batch.trajectory_decoded)
+
+                # Prepare extra features for text-only processing
+                extra_features = {
+                    "trajectory_latents": trajectory_latents,
+                    "trajectory_timesteps": trajectory_timesteps
+                }
+
+                if batch.return_trajectory_decoded:
+                    for i, decoded_frames in enumerate(trajectory_decoded):
+                        for j, decoded_frame in enumerate(decoded_frames):
+                            save_decoded_latents_as_video(
+                                decoded_frame,
+                                f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
+                                args.train_fps)
+
+                # Prepare batch data for Parquet dataset
+                batch_data: list[dict[str, Any]] = []
+
+                # Add progress bar for saving outputs
+                save_pbar = tqdm(enumerate(valid_data["path"]),
+                                 desc="Saving outputs",
+                                 unit="item",
+                                 leave=False)
+
+                for idx, video_path in save_pbar:
+                    video_name = os.path.basename(video_path).split(".")[0]
+
+                    # Convert tensors to numpy arrays
+                    text_embedding = prompt_embeds[idx].cpu().numpy()
+
+                    # Get extra features for this sample
+                    sample_extra_features = {}
+                    if extra_features:
+                        for key, value in extra_features.items():
+                            if isinstance(value, torch.Tensor):
+                                sample_extra_features[key] = value[idx].cpu(
+                                ).numpy()
+                            else:
+                                assert isinstance(value, list)
+                                if isinstance(value[idx], torch.Tensor):
+                                    sample_extra_features[key] = value[idx].cpu(
+                                    ).float().numpy()
+                                else:
+                                    sample_extra_features[key] = value[idx]
+
+                    # Create record for Parquet dataset (text-only ODE schema)
+                    record: dict[str, Any] = ode_text_only_record_creator(
+                        video_name=video_name,
+                        text_embedding=text_embedding,
+                        caption=valid_data["text"][idx],
+                        trajectory_latents=sample_extra_features[
+                            "trajectory_latents"],
+                        trajectory_timesteps=sample_extra_features[
+                            "trajectory_timesteps"],
+                    )
+                    batch_data.append(record)
+
+                if batch_data:
+                    write_pbar = tqdm(total=1,
+                                      desc="Writing to Parquet dataset",
+                                      unit="batch")
+                    table = records_to_table(batch_data,
+                                             self.get_pyarrow_schema())
+                    write_pbar.update(1)
+                    write_pbar.close()
+
+                    if not hasattr(self, 'dataset_writer'):
+                        self.dataset_writer = ParquetDatasetWriter(
+                            out_dir=self.combined_parquet_dir,
+                            samples_per_file=args.samples_per_file,
+                        )
+                    self.dataset_writer.append_table(table)
+
+                    logger.info("Collected batch with %s samples", len(table))
+
+                if self.num_processed_samples >= args.flush_frequency:
+                    written = self.dataset_writer.flush()
+                    logger.info("Flushed %s samples to parquet", written)
+                    self.num_processed_samples = 0
+
+        # Final flush for any remaining samples
+        if hasattr(self, 'dataset_writer'):
+            written = self.dataset_writer.flush(write_remainder=True)
+            if written:
+                logger.info("Final flush wrote %s samples", written)
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
+        if not self.post_init_called:
+            self.post_init()
+
+        self.local_rank = int(os.getenv("RANK", 0))
+        os.makedirs(args.output_dir, exist_ok=True)
+        # Create directory for combined data
+        self.combined_parquet_dir = os.path.join(args.output_dir,
+                                                 "combined_parquet_dataset")
+        os.makedirs(self.combined_parquet_dir, exist_ok=True)
+
+        # Loading dataset
+        train_dataset = gettextdataset(args)
+
+        self.preprocess_dataloader = DataLoader(
+            train_dataset,
+            batch_size=args.preprocess_video_batch_size,
+            num_workers=args.dataloader_num_workers,
+        )
+
+        self.preprocess_loader_iter = iter(self.preprocess_dataloader)
+
+        self.num_processed_samples = 0
+        # Add progress bar for video preprocessing
+        self.pbar = tqdm(self.preprocess_loader_iter,
+                         desc="Processing videos",
+                         unit="batch",
+                         disable=self.local_rank != 0)
+
+        # Initialize class variables for data sharing
+        self.video_data: dict[str, Any] = {}  # Store video metadata and paths
+        self.latent_data: dict[str, Any] = {}  # Store latent tensors
+        self.preprocess_text_and_trajectory(fastvideo_args, args)
+
+
+EntryClass = PreprocessPipeline_ODE_Trajectory

fastvideo/pipelines/preprocess/preprocess_pipeline_t2v.py

@@ -15,9 +15,9 @@ class PreprocessPipeline_T2V(BasePreprocessPipeline):
 
     _required_config_modules = ["text_encoder", "tokenizer", "vae"]
 
-    def get_schema_fields(self):
-        """Get the schema fields for T2V pipeline."""
-        return [f.name for f in pyarrow_schema_t2v]
+    def get_pyarrow_schema(self):
+        """Return the PyArrow schema for T2V pipeline."""
+        return pyarrow_schema_t2v
 
 
 EntryClass = PreprocessPipeline_T2V

fastvideo/pipelines/preprocess/preprocess_pipeline_text.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Text-only Data Preprocessing pipeline implementation.
+
+This module contains an implementation of the Text-only Data Preprocessing pipeline
+using the modular pipeline architecture, based on the ODE Trajectory preprocessing.
+"""
+
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import torch
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm import tqdm
+
+from fastvideo.dataset import gettextdataset
+from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter,
+                                                     records_to_table)
+from fastvideo.dataset.dataloader.record_schema import text_only_record_creator
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_text_only
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
+    BasePreprocessPipeline)
+from fastvideo.pipelines.stages import TextEncodingStage
+
+logger = init_logger(__name__)
+
+
+class PreprocessPipeline_Text(BasePreprocessPipeline):
+    """Text-only preprocessing pipeline implementation."""
+
+    _required_config_modules = ["text_encoder", "tokenizer"]
+    preprocess_dataloader: StatefulDataLoader
+    preprocess_loader_iter: Iterator[dict[str, Any]]
+    pbar: Any
+    num_processed_samples: int = 0
+
+    def get_pyarrow_schema(self):
+        """Return the PyArrow schema for text-only pipeline."""
+        return pyarrow_schema_text_only
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+    def preprocess_text_only(self, fastvideo_args: FastVideoArgs, args):
+        """Preprocess text-only data."""
+
+        for batch_idx, data in enumerate(self.pbar):
+            if data is None:
+                continue
+
+            with torch.inference_mode():
+                # For text-only processing, we only need text data
+                # Filter out samples without text
+                valid_indices = []
+                for i, text in enumerate(data["text"]):
+                    if text and text.strip():  # Check if text is not empty
+                        valid_indices.append(i)
+                self.num_processed_samples += len(valid_indices)
+
+                if not valid_indices:
+                    continue
+
+                # Create new batch with only valid samples (text-only)
+                valid_data = {
+                    "text": [data["text"][i] for i in valid_indices],
+                    "path": [data["path"][i] for i in valid_indices],
+                }
+
+                batch_captions = valid_data["text"]
+                # Encode text using the standalone TextEncodingStage API
+                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                    batch_captions,
+                    fastvideo_args,
+                    encoder_index=[0],
+                    return_attention_mask=True,
+                )
+                prompt_embeds = prompt_embeds_list[0]
+                prompt_attention_masks = prompt_masks_list[0]
+                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
+
+                logger.info("===== prompt_embeds: %s", prompt_embeds.shape)
+                logger.info("===== prompt_attention_masks: %s",
+                            prompt_attention_masks.shape)
+
+                # Prepare batch data for Parquet dataset
+                batch_data = []
+
+                # Add progress bar for saving outputs
+                save_pbar = tqdm(enumerate(valid_data["path"]),
+                                 desc="Saving outputs",
+                                 unit="item",
+                                 leave=False)
+
+                for idx, text_path in save_pbar:
+                    text_name = os.path.basename(text_path).split(".")[0]
+
+                    # Convert tensors to numpy arrays
+                    text_embedding = prompt_embeds[idx].cpu().numpy()
+
+                    # Create record for Parquet dataset (text-only schema)
+                    record = text_only_record_creator(
+                        text_name=text_name,
+                        text_embedding=text_embedding,
+                        caption=valid_data["text"][idx],
+                    )
+                    batch_data.append(record)
+
+                if batch_data:
+                    write_pbar = tqdm(total=1,
+                                      desc="Writing to Parquet dataset",
+                                      unit="batch")
+                    table = records_to_table(batch_data,
+                                             pyarrow_schema_text_only)
+                    write_pbar.update(1)
+                    write_pbar.close()
+
+                    if not hasattr(self, 'dataset_writer'):
+                        self.dataset_writer = ParquetDatasetWriter(
+                            out_dir=self.combined_parquet_dir,
+                            samples_per_file=args.samples_per_file,
+                        )
+                    self.dataset_writer.append_table(table)
+
+                    logger.info("Collected batch with %s samples", len(table))
+
+                if self.num_processed_samples >= args.flush_frequency:
+                    written = self.dataset_writer.flush()
+                    logger.info("Flushed %s samples to parquet", written)
+                    self.num_processed_samples = 0
+
+        # Final flush for any remaining samples
+        if hasattr(self, 'dataset_writer'):
+            written = self.dataset_writer.flush(write_remainder=True)
+            if written:
+                logger.info("Final flush wrote %s samples", written)
+
+    # Text-only record creation moved to fastvideo.dataset.dataloader.record_schema
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
+        if not self.post_init_called:
+            self.post_init()
+
+        self.local_rank = int(os.getenv("RANK", 0))
+        os.makedirs(args.output_dir, exist_ok=True)
+        # Create directory for combined data
+        self.combined_parquet_dir = os.path.join(args.output_dir,
+                                                 "combined_parquet_dataset")
+        os.makedirs(self.combined_parquet_dir, exist_ok=True)
+
+        # Loading text dataset
+        train_dataset = gettextdataset(args)
+
+        self.preprocess_dataloader = DataLoader(
+            train_dataset,
+            batch_size=args.preprocess_video_batch_size,
+            num_workers=args.dataloader_num_workers,
+        )
+
+        self.preprocess_loader_iter = iter(self.preprocess_dataloader)
+
+        self.num_processed_samples = 0
+        # Add progress bar for text preprocessing
+        self.pbar = tqdm(self.preprocess_loader_iter,
+                         desc="Processing text",
+                         unit="batch",
+                         disable=self.local_rank != 0)
+
+        # Initialize class variables for data sharing
+        self.text_data: dict[str, Any] = {}  # Store text metadata and paths
+
+        self.preprocess_text_only(fastvideo_args, args)
+
+
+EntryClass = PreprocessPipeline_Text

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+from typing import Any
 
 from fastvideo import PipelineConfig
 from fastvideo.configs.models.vaes import WanVAEConfig
@@ -9,8 +10,12 @@ from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import (
     PreprocessPipeline_I2V)
+from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import (
+    PreprocessPipeline_ODE_Trajectory)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import (
     PreprocessPipeline_T2V)
+from fastvideo.pipelines.preprocess.preprocess_pipeline_text import (
+    PreprocessPipeline_Text)
 from fastvideo.utils import maybe_download_model
 
 logger = init_logger(__name__)
@@ -21,12 +26,22 @@ def main(args) -> None:
     maybe_init_distributed_environment_and_model_parallel(1, 1)
     num_gpus = int(os.environ["WORLD_SIZE"])
     assert num_gpus == 1, "Only support 1 GPU"
+
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
-    kwargs = {
-        "vae_precision": "fp32",
-        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
-    }
+
+    kwargs: dict[str, Any] = {}
+    if args.preprocess_task == "text_only":
+        kwargs = {
+            "text_encoder_cpu_offload": False,
+        }
+    else:
+        # Full config for video/image processing
+        kwargs = {
+            "vae_precision": "fp32",
+            "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
+        }
     pipeline_config.update_config_from_dict(kwargs)
+
     fastvideo_args = FastVideoArgs(
         model_path=args.model_path,
         num_gpus=get_world_size(),
@@ -35,7 +50,23 @@ def main(args) -> None:
         text_encoder_cpu_offload=False,
         pipeline_config=pipeline_config,
     )
-    PreprocessPipeline = PreprocessPipeline_I2V if args.preprocess_task == "i2v" else PreprocessPipeline_T2V
+    if args.preprocess_task == "t2v":
+        PreprocessPipeline = PreprocessPipeline_T2V
+    elif args.preprocess_task == "i2v":
+        PreprocessPipeline = PreprocessPipeline_I2V
+    elif args.preprocess_task == "text_only":
+        PreprocessPipeline = PreprocessPipeline_Text
+    elif args.preprocess_task == "ode_trajectory":
+        assert args.flow_shift is not None, "flow_shift is required for ode_trajectory"
+        fastvideo_args.pipeline_config.flow_shift = args.flow_shift
+        PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
+    else:
+        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. "
+                         f"Valid options: t2v, i2v, ode_trajectory, text_only")
+
+    logger.info("Preprocess task: %s using %s", args.preprocess_task,
+                PreprocessPipeline.__name__)
+
     pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 
@@ -74,7 +105,12 @@ if __name__ == "__main__":
     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("--preprocess_task", type=str, default="t2v")
+    parser.add_argument("--flow_shift", type=float, default=None)
+    parser.add_argument("--preprocess_task",
+                        type=str,
+                        default="t2v",
+                        choices=["t2v", "i2v", "text_only", "ode_trajectory"],
+                        help="Type of preprocessing task to run")
     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)

fastvideo/pipelines/stages/causal_denoising.py

@@ -34,13 +34,15 @@ class CausalDMDDenosingStage(DenoisingStage):
     Denoising stage for causal diffusion.
     """
 
-    def __init__(self, transformer, scheduler) -> None:
-        super().__init__(transformer, scheduler)
+    def __init__(self, transformer, scheduler, transformer_2=None) -> None:
+        super().__init__(transformer, scheduler, transformer_2)
         # KV and cross-attention cache state (initialized on first forward)
+        self.transformer = transformer
+        self.transformer_2 = transformer_2
         self.kv_cache1: list | None = None
         self.crossattn_cache: list | None = None
         # Model-dependent constants (aligned with causal_inference.py assumptions)
-        self.num_transformer_blocks = self.transformer.config.arch_config.num_layers
+        self.num_transformer_blocks = len(self.transformer.blocks)
         self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block
         self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames
 
@@ -65,21 +67,18 @@ class CausalDMDDenosingStage(DenoisingStage):
             -1] * self.transformer.config.arch_config.patch_size[-2]
         self.frame_seq_length = latent_seq_length // patch_ratio
         # TODO(will): make this a parameter once we add i2v support
-        independent_first_frame = self.transformer.independent_first_frame
-
+        independent_first_frame = self.transformer.independent_first_frame if hasattr(
+            self.transformer, 'independent_first_frame') else False
         # Timesteps for DMD
         timesteps = torch.tensor(
             fastvideo_args.pipeline_config.dmd_denoising_steps,
             dtype=torch.long).cpu()
-
         if fastvideo_args.pipeline_config.warp_denoising_step:
-            logger.info("Warping timesteps...")
             scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(),
                                              torch.tensor([0],
                                                           dtype=torch.float32)))
             timesteps = scheduler_timesteps[1000 - timesteps]
         timesteps = timesteps.to(get_local_torch_device())
-        logger.info("Using timesteps: %s", timesteps)
 
         # Image kwargs (kept empty unless caller provides compatible args)
         image_kwargs: dict = {}
@@ -150,7 +149,6 @@ class CausalDMDDenosingStage(DenoisingStage):
                 with torch.autocast(device_type="cuda",
                                     dtype=target_dtype,
                                     enabled=autocast_enabled):
-                    assert False, "image_first_btchw is not supported"
                     _ = self.transformer(
                         image_first_btchw,
                         prompt_embeds,
@@ -177,7 +175,6 @@ class CausalDMDDenosingStage(DenoisingStage):
                 with torch.autocast(device_type="cuda",
                                     dtype=target_dtype,
                                     enabled=autocast_enabled):
-                    assert False, "ref_btchw is not supported"
                     _ = self.transformer(
                         ref_btchw,
                         prompt_embeds,
@@ -225,6 +222,10 @@ class CausalDMDDenosingStage(DenoisingStage):
                 video_raw_latent_shape = noise_latents_btchw.shape
 
                 for i, t_cur in enumerate(timesteps):
+                    if self.transformer_2 is not None and fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None and t_cur < fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps:
+                        current_model = self.transformer_2
+                    else:
+                        current_model = self.transformer
                     # Copy for pred conversion
                     noise_latents = noise_latents_btchw.clone()
                     latent_model_input = current_latents.to(target_dtype)
@@ -275,10 +276,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                             (latent_model_input.shape[0], 1),
                             device=latent_model_input.device,
                             dtype=torch.long)
-                        # if fastvideo_args.use_sf_wan:
-                        #     # SF wan wrapper requires BTCHW input
-                        #     latent_model_input = latent_model_input.permute(0, 2, 1, 3, 4)
-                        pred_noise_btchw = self.transformer(
+                        pred_noise_btchw = current_model(
                             latent_model_input,
                             prompt_embeds,
                             t_expanded_noise,
@@ -290,14 +288,6 @@ class CausalDMDDenosingStage(DenoisingStage):
                             **image_kwargs,
                             **pos_cond_kwargs,
                         ).permute(0, 2, 1, 3, 4)
-                        # if fastvideo_args.use_sf_wan:
-                        #     flow_pred, x0_pred = pred_noise_btchw
-                        #     logger.info(f"flow_pred.shape: {flow_pred.shape}")
-                        #     logger.info(f"x0_pred.shape: {x0_pred.shape}")
-                        #     # SF wan wrapper requires BTCHW output
-                        #     pred_noise_btchw = flow_pred
-                        # else:
-                        #     pred_noise_btchw = pred_noise_btchw.permute(0, 2, 1, 3, 4)
 
                     # Convert pred noise to pred video with FM Euler scheduler utilities
                     pred_video_btchw = pred_noise_to_pred_video(
@@ -351,10 +341,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                                         attn_metadata=attn_metadata,
                                         forward_batch=batch):
                     t_expanded_context = t_context.unsqueeze(1)
-                    # if fastvideo_args.use_sf_wan:
-                        # SF wan wrapper requires BTCHW input
-                        # context_bcthw = context_bcthw.permute(0, 2, 1, 3, 4)
-                    _ = self.transformer(
+                    _ = current_model(
                         context_bcthw,
                         prompt_embeds,
                         t_expanded_context,
@@ -365,10 +352,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                         start_frame=start_index,
                         **image_kwargs,
                         **pos_cond_kwargs,
-                    ).permute(0, 2, 1, 3, 4)
-                    # if fastvideo_args.use_sf_wan:
-                        # SF wan wrapper requires BTCHW output
-                        # context_bcthw = context_bcthw.permute(0, 2, 1, 3, 4)
+                    )
                 start_index += current_num_frames
 
         batch.latents = latents

fastvideo/pipelines/stages/decoding.py

@@ -50,6 +50,63 @@ class DecodingStage(PipelineStage):
         result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
         return result
 
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor,
+               fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        """
+        Decode latent representations into pixel space using VAE.
+        
+        Args:
+            latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)
+            fastvideo_args: Configuration containing:
+                - disable_autocast: Whether to disable automatic mixed precision (default: False)
+                - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16")
+                - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency
+            
+        Returns:
+            Decoded video tensor with shape (batch, channels, frames, height, width), 
+            normalized to [0, 1] range and moved to CPU as float32
+        """
+        self.vae = self.vae.to(get_local_torch_device())
+        latents = latents.to(get_local_torch_device())
+
+        # Setup VAE precision
+        vae_dtype = PRECISION_TO_TYPE[
+            fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        if isinstance(self.vae.scaling_factor, torch.Tensor):
+            latents = latents / self.vae.scaling_factor.to(
+                latents.device, latents.dtype)
+        else:
+            latents = latents / self.vae.scaling_factor
+
+        # Apply shifting if needed
+        if (hasattr(self.vae, "shift_factor")
+                and self.vae.shift_factor is not None):
+            if isinstance(self.vae.shift_factor, torch.Tensor):
+                latents += self.vae.shift_factor.to(latents.device,
+                                                    latents.dtype)
+            else:
+                latents += self.vae.shift_factor
+
+        # Decode latents
+        with torch.autocast(device_type="cuda",
+                            dtype=vae_dtype,
+                            enabled=vae_autocast_enabled):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            # if fastvideo_args.vae_sp:
+            #     self.vae.enable_parallel()
+            if not vae_autocast_enabled:
+                latents = latents.to(vae_dtype)
+            image = self.vae.decode(latents)
+
+        # Normalize image to [0, 1] range
+        image = (image / 2 + 0.5).clamp(0, 1)
+        return image
+
     @torch.no_grad()
     def forward(
         self,
@@ -59,13 +116,28 @@ class DecodingStage(PipelineStage):
         """
         Decode latent representations into pixel space.
         
+        This method processes the batch through the VAE decoder, converting latent
+        representations to pixel-space video/images. It also optionally decodes
+        trajectory latents for visualization purposes.
+        
         Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
+            batch: The current batch containing:
+                - latents: Tensor to decode (batch, channels, frames, height_latents, width_latents)
+                - return_trajectory_decoded (optional): Flag to decode trajectory latents
+                - trajectory_latents (optional): Latents at different timesteps
+                - trajectory_timesteps (optional): Corresponding timesteps
+            fastvideo_args: Configuration containing:
+                - output_type: "latent" to skip decoding, otherwise decode to pixels
+                - vae_cpu_offload: Whether to offload VAE to CPU after decoding
+                - model_loaded: Track VAE loading state
+                - model_paths: Path to VAE model if loading needed
             
         Returns:
-            The batch with decoded outputs.
+            Modified batch with:
+                - output: Decoded frames (batch, channels, frames, height, width) as CPU float32
+                - trajectory_decoded (if requested): List of decoded frames per timestep
         """
+        # load vae if not already loaded (used for memory constrained devices)
         pipeline = self.pipeline() if self.pipeline else None
         if not fastvideo_args.model_loaded["vae"]:
             loader = VAELoader()
@@ -75,58 +147,29 @@ class DecodingStage(PipelineStage):
                 pipeline.add_module("vae", self.vae)
             fastvideo_args.model_loaded["vae"] = True
 
-        self.vae = self.vae.to(get_local_torch_device())
-
-        latents = batch.latents
-        # TODO(will): remove this once we add input/output validation for stages
-        if latents is None:
-            raise ValueError("Latents must be provided")
-
-        # Skip decoding if output type is latent
         if fastvideo_args.output_type == "latent":
-            image = latents
+            frames = batch.latents
         else:
-            # Setup VAE precision
-            vae_dtype = PRECISION_TO_TYPE[
-                fastvideo_args.pipeline_config.vae_precision]
-            vae_autocast_enabled = (vae_dtype != torch.float32
-                                    ) and not fastvideo_args.disable_autocast
+            frames = self.decode(batch.latents, fastvideo_args)
 
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                latents = latents / self.vae.scaling_factor.to(
-                    latents.device, latents.dtype)
-            else:
-                latents = latents / self.vae.scaling_factor
-
-            # Apply shifting if needed
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    latents += self.vae.shift_factor.to(latents.device,
-                                                        latents.dtype)
-                else:
-                    latents += self.vae.shift_factor
-
-            # Decode latents
-            with torch.autocast(device_type="cuda",
-                                dtype=vae_dtype,
-                                enabled=vae_autocast_enabled):
-                if fastvideo_args.pipeline_config.vae_tiling:
-                    self.vae.enable_tiling()
-                # if fastvideo_args.vae_sp:
-                #     self.vae.enable_parallel()
-                if not vae_autocast_enabled:
-                    latents = latents.to(vae_dtype)
-                image = self.vae.decode(latents)
-
-        # Normalize image to [0, 1] range
-        image = (image / 2 + 0.5).clamp(0, 1)
+        # decode trajectory latents if needed
+        if batch.return_trajectory_decoded:
+            batch.trajectory_decoded = []
+            assert batch.trajectory_latents is not None, "batch should have trajectory latents"
+            for idx in range(batch.trajectory_latents.shape[1]):
+                # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
+                cur_latent = batch.trajectory_latents[:, idx, :, :, :, :]
+                cur_timestep = batch.trajectory_timesteps[idx]
+                logger.info("decoding trajectory latent for timestep: %s",
+                            cur_timestep)
+                decoded_frames = self.decode(cur_latent, fastvideo_args)
+                batch.trajectory_decoded.append(decoded_frames.cpu().float())
 
         # Convert to CPU float32 for compatibility
-        image = image.cpu().float()
+        frames = frames.cpu().float()
 
         # Update batch with decoded image
-        batch.output = image
+        batch.output = frames
 
         # Offload models if needed
         if hasattr(self, 'maybe_free_model_hooks'):

fastvideo/pipelines/stages/denoising.py

@@ -85,8 +85,9 @@ class DenoisingStage(PipelineStage):
                 AttentionBackendEnum.SLIDING_TILE_ATTN,
                 AttentionBackendEnum.VIDEO_SPARSE_ATTN,
                 AttentionBackendEnum.VMOBA_ATTN,
-                AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA
-            )  # hack
+                AttentionBackendEnum.FLASH_ATTN,
+                AttentionBackendEnum.TORCH_SDPA,
+                AttentionBackendEnum.SAGE_ATTN_THREE)  # hack
         )
 
     def forward(
@@ -204,8 +205,14 @@ class DenoisingStage(PipelineStage):
                 neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
 
         # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
-        if fastvideo_args.boundary_ratio is not None:
-            boundary_timestep = fastvideo_args.boundary_ratio * self.scheduler.num_train_timesteps
+        boundary_ratio = fastvideo_args.pipeline_config.dit_config.boundary_ratio
+        if batch.boundary_ratio is not None:
+            logger.info("Overriding boundary ratio from %s to %s",
+                        boundary_ratio, batch.boundary_ratio)
+            boundary_ratio = batch.boundary_ratio
+
+        if boundary_ratio is not None:
+            boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
         else:
             boundary_timestep = None
         latent_model_input = latents.to(target_dtype)
@@ -247,6 +254,10 @@ class DenoisingStage(PipelineStage):
                                                              patch_size[2])
             seq_len = int(math.ceil(seq_len / sp_world_size)) * sp_world_size
 
+        # Initialize lists for ODE trajectory
+        trajectory_timesteps: list[torch.Tensor] = []
+        trajectory_latents: list[torch.Tensor] = []
+
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -427,6 +438,11 @@ class DenoisingStage(PipelineStage):
                         latents = (1. - mask2[0]) * z + mask2[0] * latents
                         # latents = latents.unsqueeze(0)
 
+                # save trajectory latents if needed
+                if batch.return_trajectory_latents:
+                    trajectory_timesteps.append(t)
+                    trajectory_latents.append(latents)
+
                 # Update progress bar
                 if i == len(timesteps) - 1 or (
                     (i + 1) > num_warmup_steps and
@@ -434,9 +450,28 @@ class DenoisingStage(PipelineStage):
                         and progress_bar is not None):
                     progress_bar.update()
 
+        # Gather results if using sequence parallelism
+        trajectory_tensor: torch.Tensor | None = None
+        if trajectory_latents:
+            trajectory_tensor = torch.stack(trajectory_latents, dim=1)
+            trajectory_timesteps_tensor = torch.stack(trajectory_timesteps,
+                                                      dim=0)
+        else:
+            trajectory_tensor = None
+            trajectory_timesteps_tensor = None
+
         # Gather results if using sequence parallelism
         if sp_group:
             latents = sequence_model_parallel_all_gather(latents, dim=2)
+            if batch.return_trajectory_latents:
+                trajectory_tensor = trajectory_tensor.to(
+                    get_local_torch_device())
+                trajectory_tensor = sequence_model_parallel_all_gather(
+                    trajectory_tensor, dim=3)
+
+        if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
+            batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
+            batch.trajectory_latents = trajectory_tensor.cpu()
 
         # Update batch with final latents
         batch.latents = latents

fastvideo/platforms/cuda.py

@@ -115,6 +115,7 @@ class CudaPlatformBase(Platform):
 
         logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s",
                     envs.FASTVIDEO_ATTENTION_BACKEND)
+        logger.info("Selected backend: %s", selected_backend)
         if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN:
             try:
                 from st_attn import sliding_tile_attention  # noqa: F401
@@ -144,6 +145,20 @@ class CudaPlatformBase(Platform):
                 logger.info(
                     "Sage Attention backend is not installed. Fall back to Flash Attention."
                 )
+        elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE:
+            try:
+                from fastvideo.attention.backends.sageattn.api import sageattn_blackwell  # noqa: F401
+
+                from fastvideo.attention.backends.sage_attn3 import (  # noqa: F401
+                    SageAttention3Backend)
+                logger.info("Using Sage Attention 3 backend.")
+
+                return "fastvideo.attention.backends.sage_attn3.SageAttention3Backend"
+            except ImportError as e:
+                logger.info(e)
+                logger.info(
+                    "Sage Attention 3 backend is not installed. Fall back to Flash Attention."
+                )
         elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
             try:
                 from vsa import block_sparse_attn  # noqa: F401

fastvideo/platforms/interface.py

@@ -18,6 +18,7 @@ class AttentionBackendEnum(enum.Enum):
     SLIDING_TILE_ATTN = enum.auto()
     TORCH_SDPA = enum.auto()
     SAGE_ATTN = enum.auto()
+    SAGE_ATTN_THREE = enum.auto()
     VIDEO_SPARSE_ATTN = enum.auto()
     VMOBA_ATTN = enum.auto()
     NO_ATTENTION = enum.auto()

fastvideo/sf_utils/dataset.py

@@ -1,220 +0,0 @@
-from utils.lmdb import get_array_shape_from_lmdb, retrieve_row_from_lmdb
-from torch.utils.data import Dataset
-import numpy as np
-import torch
-import lmdb
-import json
-from pathlib import Path
-from PIL import Image
-import os
-
-
-class TextDataset(Dataset):
-    def __init__(self, prompt_path, extended_prompt_path=None):
-        with open(prompt_path, encoding="utf-8") as f:
-            self.prompt_list = [line.rstrip() for line in f]
-
-        if extended_prompt_path is not None:
-            with open(extended_prompt_path, encoding="utf-8") as f:
-                self.extended_prompt_list = [line.rstrip() for line in f]
-            assert len(self.extended_prompt_list) == len(self.prompt_list)
-        else:
-            self.extended_prompt_list = None
-
-    def __len__(self):
-        return len(self.prompt_list)
-
-    def __getitem__(self, idx):
-        batch = {
-            "prompts": self.prompt_list[idx],
-            "idx": idx,
-        }
-        if self.extended_prompt_list is not None:
-            batch["extended_prompts"] = self.extended_prompt_list[idx]
-        return batch
-
-
-class ODERegressionLMDBDataset(Dataset):
-    def __init__(self, data_path: str, max_pair: int = int(1e8)):
-        self.env = lmdb.open(data_path, readonly=True,
-                             lock=False, readahead=False, meminit=False)
-
-        self.latents_shape = get_array_shape_from_lmdb(self.env, 'latents')
-        self.max_pair = max_pair
-
-    def __len__(self):
-        return min(self.latents_shape[0], self.max_pair)
-
-    def __getitem__(self, idx):
-        """
-        Outputs:
-            - prompts: List of Strings
-            - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
-        """
-        latents = retrieve_row_from_lmdb(
-            self.env,
-            "latents", np.float16, idx, shape=self.latents_shape[1:]
-        )
-
-        if len(latents.shape) == 4:
-            latents = latents[None, ...]
-
-        prompts = retrieve_row_from_lmdb(
-            self.env,
-            "prompts", str, idx
-        )
-        return {
-            "prompts": prompts,
-            "ode_latent": torch.tensor(latents, dtype=torch.float32)
-        }
-
-
-class ShardingLMDBDataset(Dataset):
-    def __init__(self, data_path: str, max_pair: int = int(1e8)):
-        self.envs = []
-        self.index = []
-
-        for fname in sorted(os.listdir(data_path)):
-            path = os.path.join(data_path, fname)
-            env = lmdb.open(path,
-                            readonly=True,
-                            lock=False,
-                            readahead=False,
-                            meminit=False)
-            self.envs.append(env)
-
-        self.latents_shape = [None] * len(self.envs)
-        for shard_id, env in enumerate(self.envs):
-            self.latents_shape[shard_id] = get_array_shape_from_lmdb(env, 'latents')
-            for local_i in range(self.latents_shape[shard_id][0]):
-                self.index.append((shard_id, local_i))
-
-            # print("shard_id ", shard_id, " local_i ", local_i)
-
-        self.max_pair = max_pair
-
-    def __len__(self):
-        return len(self.index)
-
-    def __getitem__(self, idx):
-        """
-            Outputs:
-                - prompts: List of Strings
-                - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
-        """
-        shard_id, local_idx = self.index[idx]
-
-        latents = retrieve_row_from_lmdb(
-            self.envs[shard_id],
-            "latents", np.float16, local_idx,
-            shape=self.latents_shape[shard_id][1:]
-        )
-
-        if len(latents.shape) == 4:
-            latents = latents[None, ...]
-
-        prompts = retrieve_row_from_lmdb(
-            self.envs[shard_id],
-            "prompts", str, local_idx
-        )
-
-        return {
-            "prompts": prompts,
-            "ode_latent": torch.tensor(latents, dtype=torch.float32)
-        }
-
-
-class TextImagePairDataset(Dataset):
-    def __init__(
-        self,
-        data_dir,
-        transform=None,
-        eval_first_n=-1,
-        pad_to_multiple_of=None
-    ):
-        """
-        Args:
-            data_dir (str): Path to the directory containing:
-                - target_crop_info_*.json (metadata file)
-                - */ (subdirectory containing images with matching aspect ratio)
-            transform (callable, optional): Optional transform to be applied on the image
-        """
-        self.transform = transform
-        data_dir = Path(data_dir)
-
-        # Find the metadata JSON file
-        metadata_files = list(data_dir.glob('target_crop_info_*.json'))
-        if not metadata_files:
-            raise FileNotFoundError(f"No metadata file found in {data_dir}")
-        if len(metadata_files) > 1:
-            raise ValueError(f"Multiple metadata files found in {data_dir}")
-
-        metadata_path = metadata_files[0]
-        # Extract aspect ratio from metadata filename (e.g. target_crop_info_26-15.json -> 26-15)
-        aspect_ratio = metadata_path.stem.split('_')[-1]
-
-        # Use aspect ratio subfolder for images
-        self.image_dir = data_dir / aspect_ratio
-        if not self.image_dir.exists():
-            raise FileNotFoundError(f"Image directory not found: {self.image_dir}")
-
-        # Load metadata
-        with open(metadata_path, 'r') as f:
-            self.metadata = json.load(f)
-
-        eval_first_n = eval_first_n if eval_first_n != -1 else len(self.metadata)
-        self.metadata = self.metadata[:eval_first_n]
-
-        # Verify all images exist
-        for item in self.metadata:
-            image_path = self.image_dir / item['file_name']
-            if not image_path.exists():
-                raise FileNotFoundError(f"Image not found: {image_path}")
-
-        self.dummy_prompt = "DUMMY PROMPT"
-        self.pre_pad_len = len(self.metadata)
-        if pad_to_multiple_of is not None and len(self.metadata) % pad_to_multiple_of != 0:
-            # Duplicate the last entry
-            self.metadata += [self.metadata[-1]] * (
-                pad_to_multiple_of - len(self.metadata) % pad_to_multiple_of
-            )
-
-    def __len__(self):
-        return len(self.metadata)
-
-    def __getitem__(self, idx):
-        """
-        Returns:
-            dict: A dictionary containing:
-                - image: PIL Image
-                - caption: str
-                - target_bbox: list of int [x1, y1, x2, y2]
-                - target_ratio: str
-                - type: str
-                - origin_size: tuple of int (width, height)
-        """
-        item = self.metadata[idx]
-
-        # Load image
-        image_path = self.image_dir / item['file_name']
-        image = Image.open(image_path).convert('RGB')
-
-        # Apply transform if specified
-        if self.transform:
-            image = self.transform(image)
-
-        return {
-            'image': image,
-            'prompts': item['caption'],
-            'target_bbox': item['target_crop']['target_bbox'],
-            'target_ratio': item['target_crop']['target_ratio'],
-            'type': item['type'],
-            'origin_size': (item['origin_width'], item['origin_height']),
-            'idx': idx
-        }
-
-
-def cycle(dl):
-    while True:
-        for data in dl:
-            yield data

fastvideo/sf_utils/distributed.py

@@ -1,125 +0,0 @@
-from datetime import timedelta
-from functools import partial
-import os
-import torch
-import torch.distributed as dist
-from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType
-from torch.distributed.fsdp.api import CPUOffload
-from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
-
-
-def fsdp_state_dict(model):
-    fsdp_fullstate_save_policy = FullStateDictConfig(
-        offload_to_cpu=True, rank0_only=True
-    )
-    with FSDP.state_dict_type(
-        model, StateDictType.FULL_STATE_DICT, fsdp_fullstate_save_policy
-    ):
-        checkpoint = model.state_dict()
-
-    return checkpoint
-
-
-def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", min_num_params=int(5e7), transformer_module=None, cpu_offload=False):
-    if mixed_precision:
-        mixed_precision_policy = MixedPrecision(
-            param_dtype=torch.bfloat16,
-            reduce_dtype=torch.float32,
-            buffer_dtype=torch.float32,
-            cast_forward_inputs=False
-        )
-    else:
-        mixed_precision_policy = None
-
-    if wrap_strategy == "transformer":
-        auto_wrap_policy = partial(
-            transformer_auto_wrap_policy,
-            transformer_layer_cls=transformer_module
-        )
-    elif wrap_strategy == "size":
-        auto_wrap_policy = partial(
-            size_based_auto_wrap_policy,
-            min_num_params=min_num_params
-        )
-    else:
-        raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")
-
-    os.environ["NCCL_CROSS_NIC"] = "1"
-
-    sharding_strategy = {
-        "full": ShardingStrategy.FULL_SHARD,
-        "hybrid_full": ShardingStrategy.HYBRID_SHARD,
-        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
-        "no_shard": ShardingStrategy.NO_SHARD,
-    }[sharding_strategy]
-
-    module = FSDP(
-        module,
-        auto_wrap_policy=auto_wrap_policy,
-        sharding_strategy=sharding_strategy,
-        mixed_precision=mixed_precision_policy,
-        device_id=torch.cuda.current_device(),
-        limit_all_gathers=True,
-        use_orig_params=True,
-        cpu_offload=CPUOffload(offload_params=cpu_offload),
-        sync_module_states=False  # Load ckpt on rank 0 and sync to other ranks
-    )
-    return module
-
-
-def barrier():
-    if dist.is_initialized():
-        dist.barrier()
-
-
-def launch_distributed_job(backend: str = "nccl"):
-    rank = int(os.environ["RANK"])
-    local_rank = int(os.environ["LOCAL_RANK"])
-    world_size = int(os.environ["WORLD_SIZE"])
-    host = os.environ["MASTER_ADDR"]
-    port = int(os.environ["MASTER_PORT"])
-
-    if ":" in host:  # IPv6
-        init_method = f"tcp://[{host}]:{port}"
-    else:  # IPv4
-        init_method = f"tcp://{host}:{port}"
-    dist.init_process_group(rank=rank, world_size=world_size, backend=backend,
-                            init_method=init_method, timeout=timedelta(minutes=30))
-    torch.cuda.set_device(local_rank)
-
-
-class EMA_FSDP:
-    def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
-        self.decay = decay
-        self.shadow = {}
-        self._init_shadow(fsdp_module)
-
-    @torch.no_grad()
-    def _init_shadow(self, fsdp_module):
-        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-        with FSDP.summon_full_params(fsdp_module, writeback=False):
-            for n, p in fsdp_module.module.named_parameters():
-                self.shadow[n] = p.detach().clone().float().cpu()
-
-    @torch.no_grad()
-    def update(self, fsdp_module):
-        d = self.decay
-        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-        with FSDP.summon_full_params(fsdp_module, writeback=False):
-            for n, p in fsdp_module.module.named_parameters():
-                self.shadow[n].mul_(d).add_(p.detach().float().cpu(), alpha=1. - d)
-
-    # Optional helpers ---------------------------------------------------
-    def state_dict(self):
-        return self.shadow            # picklable
-
-    def load_state_dict(self, sd):
-        self.shadow = {k: v.clone() for k, v in sd.items()}
-
-    def copy_to(self, fsdp_module):
-        # load EMA weights into an (unwrapped) copy of the generator
-        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-        with FSDP.summon_full_params(fsdp_module, writeback=True):
-            for n, p in fsdp_module.module.named_parameters():
-                if n in self.shadow:
-                    p.data.copy_(self.shadow[n].to(p.dtype, device=p.device))

fastvideo/sf_utils/lmdb.py

@@ -1,72 +0,0 @@
-import numpy as np
-
-
-def get_array_shape_from_lmdb(env, array_name):
-    with env.begin() as txn:
-        image_shape = txn.get(f"{array_name}_shape".encode()).decode()
-        image_shape = tuple(map(int, image_shape.split()))
-    return image_shape
-
-
-def store_arrays_to_lmdb(env, arrays_dict, start_index=0):
-    """
-    Store rows of multiple numpy arrays in a single LMDB.
-    Each row is stored separately with a naming convention.
-    """
-    with env.begin(write=True) as txn:
-        for array_name, array in arrays_dict.items():
-            for i, row in enumerate(array):
-                # Convert row to bytes
-                if isinstance(row, str):
-                    row_bytes = row.encode()
-                else:
-                    row_bytes = row.tobytes()
-
-                data_key = f'{array_name}_{start_index + i}_data'.encode()
-
-                txn.put(data_key, row_bytes)
-
-
-def process_data_dict(data_dict, seen_prompts):
-    output_dict = {}
-
-    all_videos = []
-    all_prompts = []
-    for prompt, video in data_dict.items():
-        if prompt in seen_prompts:
-            continue
-        else:
-            seen_prompts.add(prompt)
-
-        video = video.half().numpy()
-        all_videos.append(video)
-        all_prompts.append(prompt)
-
-    if len(all_videos) == 0:
-        return {"latents": np.array([]), "prompts": np.array([])}
-
-    all_videos = np.concatenate(all_videos, axis=0)
-
-    output_dict['latents'] = all_videos
-    output_dict['prompts'] = np.array(all_prompts)
-
-    return output_dict
-
-
-def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape=None):
-    """
-    Retrieve a specific row from a specific array in the LMDB.
-    """
-    data_key = f'{array_name}_{row_index}_data'.encode()
-
-    with lmdb_env.begin() as txn:
-        row_bytes = txn.get(data_key)
-
-    if dtype == str:
-        array = row_bytes.decode()
-    else:
-        array = np.frombuffer(row_bytes, dtype=dtype)
-
-    if shape is not None and len(shape) > 0:
-        array = array.reshape(shape)
-    return array

fastvideo/sf_utils/loss.py

@@ -1,81 +0,0 @@
-from abc import ABC, abstractmethod
-import torch
-
-
-class DenoisingLoss(ABC):
-    @abstractmethod
-    def __call__(
-        self, x: torch.Tensor, x_pred: torch.Tensor,
-        noise: torch.Tensor, noise_pred: torch.Tensor,
-        alphas_cumprod: torch.Tensor,
-        timestep: torch.Tensor,
-        **kwargs
-    ) -> torch.Tensor:
-        """
-        Base class for denoising loss.
-        Input:
-            - x: the clean data with shape [B, F, C, H, W]
-            - x_pred: the predicted clean data with shape [B, F, C, H, W]
-            - noise: the noise with shape [B, F, C, H, W]
-            - noise_pred: the predicted noise with shape [B, F, C, H, W]
-            - alphas_cumprod: the cumulative product of alphas (defining the noise schedule) with shape [T]
-            - timestep: the current timestep with shape [B, F]
-        """
-        pass
-
-
-class X0PredLoss(DenoisingLoss):
-    def __call__(
-        self, x: torch.Tensor, x_pred: torch.Tensor,
-        noise: torch.Tensor, noise_pred: torch.Tensor,
-        alphas_cumprod: torch.Tensor,
-        timestep: torch.Tensor,
-        **kwargs
-    ) -> torch.Tensor:
-        return torch.mean((x - x_pred) ** 2)
-
-
-class VPredLoss(DenoisingLoss):
-    def __call__(
-        self, x: torch.Tensor, x_pred: torch.Tensor,
-        noise: torch.Tensor, noise_pred: torch.Tensor,
-        alphas_cumprod: torch.Tensor,
-        timestep: torch.Tensor,
-        **kwargs
-    ) -> torch.Tensor:
-        weights = 1 / (1 - alphas_cumprod[timestep].reshape(*timestep.shape, 1, 1, 1))
-        return torch.mean(weights * (x - x_pred) ** 2)
-
-
-class NoisePredLoss(DenoisingLoss):
-    def __call__(
-        self, x: torch.Tensor, x_pred: torch.Tensor,
-        noise: torch.Tensor, noise_pred: torch.Tensor,
-        alphas_cumprod: torch.Tensor,
-        timestep: torch.Tensor,
-        **kwargs
-    ) -> torch.Tensor:
-        return torch.mean((noise - noise_pred) ** 2)
-
-
-class FlowPredLoss(DenoisingLoss):
-    def __call__(
-        self, x: torch.Tensor, x_pred: torch.Tensor,
-        noise: torch.Tensor, noise_pred: torch.Tensor,
-        alphas_cumprod: torch.Tensor,
-        timestep: torch.Tensor,
-        **kwargs
-    ) -> torch.Tensor:
-        return torch.mean((kwargs["flow_pred"] - (noise - x)) ** 2)
-
-
-NAME_TO_CLASS = {
-    "x0": X0PredLoss,
-    "v": VPredLoss,
-    "noise": NoisePredLoss,
-    "flow": FlowPredLoss
-}
-
-
-def get_denoising_loss(loss_type: str) -> DenoisingLoss:
-    return NAME_TO_CLASS[loss_type]

fastvideo/sf_utils/misc.py

@@ -1,39 +0,0 @@
-import numpy as np
-import random
-import torch
-
-
-def set_seed(seed: int, deterministic: bool = False):
-    """
-    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
-
-    Args:
-        seed (`int`):
-            The seed to set.
-        deterministic (`bool`, *optional*, defaults to `False`):
-            Whether to use deterministic algorithms where available. Can slow down training.
-    """
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-    if deterministic:
-        torch.use_deterministic_algorithms(True)
-
-
-def merge_dict_list(dict_list):
-    if len(dict_list) == 1:
-        return dict_list[0]
-
-    merged_dict = {}
-    for k, v in dict_list[0].items():
-        if isinstance(v, torch.Tensor):
-            if v.ndim == 0:
-                merged_dict[k] = torch.stack([d[k] for d in dict_list], dim=0)
-            else:
-                merged_dict[k] = torch.cat([d[k] for d in dict_list], dim=0)
-        else:
-            # for non-tensor values, we just copy the value from the first item
-            merged_dict[k] = v
-    return merged_dict

fastvideo/sf_utils/scheduler.py

@@ -1,194 +0,0 @@
-from abc import abstractmethod, ABC
-import torch
-
-
-class SchedulerInterface(ABC):
-    """
-    Base class for diffusion noise schedule.
-    """
-    alphas_cumprod: torch.Tensor  # [T], alphas for defining the noise schedule
-
-    @abstractmethod
-    def add_noise(
-        self, clean_latent: torch.Tensor,
-        noise: torch.Tensor, timestep: torch.Tensor
-    ):
-        """
-        Diffusion forward corruption process.
-        Input:
-            - clean_latent: the clean latent with shape [B, C, H, W]
-            - noise: the noise with shape [B, C, H, W]
-            - timestep: the timestep with shape [B]
-        Output: the corrupted latent with shape [B, C, H, W]
-        """
-        pass
-
-    def convert_x0_to_noise(
-        self, x0: torch.Tensor, xt: torch.Tensor,
-        timestep: torch.Tensor
-    ) -> torch.Tensor:
-        """
-        Convert the diffusion network's x0 prediction to noise predidction.
-        x0: the predicted clean data with shape [B, C, H, W]
-        xt: the input noisy data with shape [B, C, H, W]
-        timestep: the timestep with shape [B]
-
-        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t) (eq 11 in https://arxiv.org/abs/2311.18828)
-        """
-        # use higher precision for calculations
-        original_dtype = x0.dtype
-        x0, xt, alphas_cumprod = map(
-            lambda x: x.double().to(x0.device), [x0, xt,
-                                                 self.alphas_cumprod]
-        )
-
-        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
-        beta_prod_t = 1 - alpha_prod_t
-
-        noise_pred = (xt - alpha_prod_t **
-                      (0.5) * x0) / beta_prod_t ** (0.5)
-        return noise_pred.to(original_dtype)
-
-    def convert_noise_to_x0(
-        self, noise: torch.Tensor, xt: torch.Tensor,
-        timestep: torch.Tensor
-    ) -> torch.Tensor:
-        """
-        Convert the diffusion network's noise prediction to x0 predidction.
-        noise: the predicted noise with shape [B, C, H, W]
-        xt: the input noisy data with shape [B, C, H, W]
-        timestep: the timestep with shape [B]
-
-        x0 = (x_t - sqrt(beta_t) * noise) / sqrt(alpha_t) (eq 11 in https://arxiv.org/abs/2311.18828)
-        """
-        # use higher precision for calculations
-        original_dtype = noise.dtype
-        noise, xt, alphas_cumprod = map(
-            lambda x: x.double().to(noise.device), [noise, xt,
-                                                    self.alphas_cumprod]
-        )
-        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
-        beta_prod_t = 1 - alpha_prod_t
-
-        x0_pred = (xt - beta_prod_t **
-                   (0.5) * noise) / alpha_prod_t ** (0.5)
-        return x0_pred.to(original_dtype)
-
-    def convert_velocity_to_x0(
-        self, velocity: torch.Tensor, xt: torch.Tensor,
-        timestep: torch.Tensor
-    ) -> torch.Tensor:
-        """
-        Convert the diffusion network's velocity prediction to x0 predidction.
-        velocity: the predicted noise with shape [B, C, H, W]
-        xt: the input noisy data with shape [B, C, H, W]
-        timestep: the timestep with shape [B]
-
-        v = sqrt(alpha_t) * noise - sqrt(beta_t) x0
-        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t)
-        given v, x_t, we have
-        x0 = sqrt(alpha_t) * x_t - sqrt(beta_t) * v
-        see derivations https://chatgpt.com/share/679fb6c8-3a30-8008-9b0e-d1ae892dac56
-        """
-        # use higher precision for calculations
-        original_dtype = velocity.dtype
-        velocity, xt, alphas_cumprod = map(
-            lambda x: x.double().to(velocity.device), [velocity, xt,
-                                                       self.alphas_cumprod]
-        )
-        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
-        beta_prod_t = 1 - alpha_prod_t
-
-        x0_pred = (alpha_prod_t ** 0.5) * xt - (beta_prod_t ** 0.5) * velocity
-        return x0_pred.to(original_dtype)
-
-
-class FlowMatchScheduler():
-
-    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):
-        self.num_train_timesteps = num_train_timesteps
-        self.shift = shift
-        self.sigma_max = sigma_max
-        self.sigma_min = sigma_min
-        self.inverse_timesteps = inverse_timesteps
-        self.extra_one_step = extra_one_step
-        self.reverse_sigmas = reverse_sigmas
-        self.set_timesteps(num_inference_steps)
-
-    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):
-        sigma_start = self.sigma_min + \
-            (self.sigma_max - self.sigma_min) * denoising_strength
-        if self.extra_one_step:
-            self.sigmas = torch.linspace(
-                sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
-        else:
-            self.sigmas = torch.linspace(
-                sigma_start, self.sigma_min, num_inference_steps)
-        if self.inverse_timesteps:
-            self.sigmas = torch.flip(self.sigmas, dims=[0])
-        self.sigmas = self.shift * self.sigmas / \
-            (1 + (self.shift - 1) * self.sigmas)
-        if self.reverse_sigmas:
-            self.sigmas = 1 - self.sigmas
-        self.timesteps = self.sigmas * self.num_train_timesteps
-        if training:
-            x = self.timesteps
-            y = torch.exp(-2 * ((x - num_inference_steps / 2) /
-                          num_inference_steps) ** 2)
-            y_shifted = y - y.min()
-            bsmntw_weighing = y_shifted * \
-                (num_inference_steps / y_shifted.sum())
-            self.linear_timesteps_weights = bsmntw_weighing
-
-    def step(self, model_output, timestep, sample, to_final=False):
-        if timestep.ndim == 2:
-            timestep = timestep.flatten(0, 1)
-        self.sigmas = self.sigmas.to(model_output.device)
-        self.timesteps = self.timesteps.to(model_output.device)
-        timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
-        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
-            sigma_ = 1 if (
-                self.inverse_timesteps or self.reverse_sigmas) else 0
-        else:
-            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
-        prev_sample = sample + model_output * (sigma_ - sigma)
-        return prev_sample
-
-    def add_noise(self, original_samples, noise, timestep):
-        """
-        Diffusion forward corruption process.
-        Input:
-            - clean_latent: the clean latent with shape [B*T, C, H, W]
-            - noise: the noise with shape [B*T, C, H, W]
-            - timestep: the timestep with shape [B*T]
-        Output: the corrupted latent with shape [B*T, C, H, W]
-        """
-        if timestep.ndim == 2:
-            timestep = timestep.flatten(0, 1)
-        self.sigmas = self.sigmas.to(noise.device)
-        self.timesteps = self.timesteps.to(noise.device)
-        timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
-        sample = (1 - sigma) * original_samples + sigma * noise
-        return sample.type_as(noise)
-
-    def training_target(self, sample, noise, timestep):
-        target = noise - sample
-        return target
-
-    def training_weight(self, timestep):
-        """
-        Input:
-            - timestep: the timestep with shape [B*T]
-        Output: the corresponding weighting [B*T]
-        """
-        if timestep.ndim == 2:
-            timestep = timestep.flatten(0, 1)
-        self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)
-        timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)
-        weights = self.linear_timesteps_weights[timestep_id]
-        return weights

fastvideo/sf_utils/wan_wrapper.py

@@ -1,381 +0,0 @@
-import types
-from typing import List, Optional
-import torch
-from torch import nn
-
-from fastvideo.sf_utils.scheduler import SchedulerInterface, FlowMatchScheduler
-from fastvideo.wan.modules.tokenizers import HuggingfaceTokenizer
-from fastvideo.wan.modules.model import WanModel, RegisterTokens, GanAttentionBlock
-from fastvideo.wan.modules.vae import _video_vae
-from fastvideo.wan.modules.t5 import umt5_xxl
-from fastvideo.wan.modules.causal_model import CausalWanModel
-
-
-class WanTextEncoder(torch.nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-        self.text_encoder = umt5_xxl(
-            encoder_only=True,
-            return_tokenizer=False,
-            dtype=torch.float32,
-            device=torch.device('cpu')
-        ).eval().requires_grad_(False)
-        self.text_encoder.load_state_dict(
-            torch.load("wan_models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",
-                       map_location='cpu', weights_only=False)
-        )
-
-        self.tokenizer = HuggingfaceTokenizer(
-            name="wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/", seq_len=512, clean='whitespace')
-
-    @property
-    def device(self):
-        # Assume we are always on GPU
-        return torch.cuda.current_device()
-
-    def forward(self, text_prompts: List[str]) -> dict:
-        ids, mask = self.tokenizer(
-            text_prompts, return_mask=True, add_special_tokens=True)
-        ids = ids.to(self.device)
-        mask = mask.to(self.device)
-        seq_lens = mask.gt(0).sum(dim=1).long()
-        context = self.text_encoder(ids, mask)
-
-        for u, v in zip(context, seq_lens):
-            u[v:] = 0.0  # set padding to 0.0
-
-        return {
-            "prompt_embeds": context
-        }
-
-
-class WanVAEWrapper(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        mean = [
-            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
-            0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
-        ]
-        std = [
-            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
-            3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
-        ]
-        self.mean = torch.tensor(mean, dtype=torch.float32)
-        self.std = torch.tensor(std, dtype=torch.float32)
-
-        # init model
-        self.model = _video_vae(
-            pretrained_path="wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth",
-            z_dim=16,
-        ).eval().requires_grad_(False)
-
-    def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
-        # pixel: [batch_size, num_channels, num_frames, height, width]
-        device, dtype = pixel.device, pixel.dtype
-        scale = [self.mean.to(device=device, dtype=dtype),
-                 1.0 / self.std.to(device=device, dtype=dtype)]
-
-        output = [
-            self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)
-            for u in pixel
-        ]
-        output = torch.stack(output, dim=0)
-        # from [batch_size, num_channels, num_frames, height, width]
-        # to [batch_size, num_frames, num_channels, height, width]
-        output = output.permute(0, 2, 1, 3, 4)
-        return output
-
-    def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor:
-        # from [batch_size, num_frames, num_channels, height, width]
-        # to [batch_size, num_channels, num_frames, height, width]
-        zs = latent.permute(0, 2, 1, 3, 4)
-        if use_cache:
-            assert latent.shape[0] == 1, "Batch size must be 1 when using cache"
-
-        device, dtype = latent.device, latent.dtype
-        scale = [self.mean.to(device=device, dtype=dtype),
-                 1.0 / self.std.to(device=device, dtype=dtype)]
-
-        if use_cache:
-            decode_function = self.model.cached_decode
-        else:
-            decode_function = self.model.decode
-
-        output = []
-        for u in zs:
-            output.append(decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0))
-        output = torch.stack(output, dim=0)
-        # from [batch_size, num_channels, num_frames, height, width]
-        # to [batch_size, num_frames, num_channels, height, width]
-        output = output.permute(0, 2, 1, 3, 4)
-        return output
-
-
-class WanDiffusionWrapper(torch.nn.Module):
-    def __init__(
-            self,
-            model_name="Wan2.1-T2V-1.3B",
-            timestep_shift=8.0,
-            is_causal=False,
-            local_attn_size=-1,
-            sink_size=0,
-            config=None,
-            **kwargs
-    ):
-        super().__init__()
-
-        assert config is not None, "config is required"
-
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.num_attention_heads = config.num_attention_heads
-        # self.num_single_layers = config.num_single_layers
-        self.num_layers = config.num_layers
-        self.hidden_size = config.hidden_size
-        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_channels_latents = config.num_channels_latents
-        self.patch_size = config.patch_size
-        self.text_len = config.text_len
-        self.local_attn_size = config.local_attn_size
-        self.independent_first_frame = False
-        # self.num_single_layers = config.num_single_layers
-        # self.num_registers = config.num_registers
-        # self.num_frame_per_block = config.num_frame_per_block
-        # self.independent_first_frame = config.independent_first_frame
-        # self.local_attn_size = config.local_attn_size
-
-        if is_causal:
-            self.model = CausalWanModel.from_pretrained(
-                f"wan_models/{model_name}/", local_attn_size=local_attn_size, sink_size=sink_size)
-            from fastvideo.distributed import get_local_torch_device
-            self.model = self.model.to(torch.bfloat16).to(get_local_torch_device())
-        else:
-            self.model = WanModel.from_pretrained(f"wan_models/{model_name}/")
-        self.model.eval()
-
-        # For non-causal diffusion, all frames share the same timestep
-        self.uniform_timestep = not is_causal
-
-        self.scheduler = FlowMatchScheduler(
-            shift=timestep_shift, sigma_min=0.0, extra_one_step=True
-        )
-        self.scheduler.set_timesteps(1000, training=True)
-
-        self.seq_len = 32760  # [1, 21, 16, 60, 104]
-        self.post_init()
-
-    @property
-    def blocks(self):
-        return self.model.blocks
-
-    def enable_gradient_checkpointing(self) -> None:
-        self.model.enable_gradient_checkpointing()
-
-    def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_dim=0) -> None:
-        # NOTE: This is hard coded for WAN2.1-T2V-1.3B for now!!!!!!!!!!!!!!!!!!!!
-        self._cls_pred_branch = nn.Sequential(
-            # Input: [B, 384, 21, 60, 104]
-            nn.LayerNorm(atten_dim * 3 + time_embed_dim),
-            nn.Linear(atten_dim * 3 + time_embed_dim, 1536),
-            nn.SiLU(),
-            nn.Linear(atten_dim, num_class)
-        )
-        self._cls_pred_branch.requires_grad_(True)
-        num_registers = 3
-        self._register_tokens = RegisterTokens(num_registers=num_registers, dim=atten_dim)
-        self._register_tokens.requires_grad_(True)
-
-        gan_ca_blocks = []
-        for _ in range(num_registers):
-            block = GanAttentionBlock()
-            gan_ca_blocks.append(block)
-        self._gan_ca_blocks = nn.ModuleList(gan_ca_blocks)
-        self._gan_ca_blocks.requires_grad_(True)
-        # self.has_cls_branch = True
-
-    def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
-        """
-        Convert flow matching's prediction to x0 prediction.
-        flow_pred: the prediction with shape [B, C, H, W]
-        xt: the input noisy data with shape [B, C, H, W]
-        timestep: the timestep with shape [B]
-
-        pred = noise - x0
-        x_t = (1-sigma_t) * x0 + sigma_t * noise
-        we have x0 = x_t - sigma_t * pred
-        see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e
-        """
-        # use higher precision for calculations
-        original_dtype = flow_pred.dtype
-        flow_pred, xt, sigmas, timesteps = map(
-            lambda x: x.double().to(flow_pred.device), [flow_pred, xt,
-                                                        self.scheduler.sigmas,
-                                                        self.scheduler.timesteps]
-        )
-
-        timestep_id = torch.argmin(
-            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
-        x0_pred = xt - sigma_t * flow_pred
-        return x0_pred.to(original_dtype)
-
-    @staticmethod
-    def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
-        """
-        Convert x0 prediction to flow matching's prediction.
-        x0_pred: the x0 prediction with shape [B, C, H, W]
-        xt: the input noisy data with shape [B, C, H, W]
-        timestep: the timestep with shape [B]
-
-        pred = (x_t - x_0) / sigma_t
-        """
-        # use higher precision for calculations
-        original_dtype = x0_pred.dtype
-        x0_pred, xt, sigmas, timesteps = map(
-            lambda x: x.double().to(x0_pred.device), [x0_pred, xt,
-                                                      scheduler.sigmas,
-                                                      scheduler.timesteps]
-        )
-        timestep_id = torch.argmin(
-            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
-        flow_pred = (xt - x0_pred) / sigma_t
-        return flow_pred.to(original_dtype)
-
-
-    def forward(self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
-        timestep: torch.LongTensor,
-        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
-        | None = None,
-        kv_cache = None,
-        crossattn_cache = None,
-        current_start = None,
-        **kwargs) -> torch.Tensor:
-        assert encoder_hidden_states_image is None, "encoder_hidden_states_image is not supported"
-
-        return self._forward(
-            noisy_image_or_video=hidden_states,
-            conditional_dict={'prompt_embeds': encoder_hidden_states},
-            timestep=timestep,
-            kv_cache=kv_cache,
-            crossattn_cache=crossattn_cache,
-            current_start=current_start,
-        )
-
-
-    def _forward(
-        self,
-        noisy_image_or_video: torch.Tensor, conditional_dict: dict,
-        timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None,
-        crossattn_cache: Optional[List[dict]] = None,
-        current_start: Optional[int] = None,
-        classify_mode: Optional[bool] = False,
-        concat_time_embeddings: Optional[bool] = False,
-        clean_x: Optional[torch.Tensor] = None,
-        aug_t: Optional[torch.Tensor] = None,
-        cache_start: Optional[int] = None
-    ) -> torch.Tensor:
-        noisy_image_or_video = noisy_image_or_video.permute(0, 2, 1, 3, 4)
-        prompt_embeds = conditional_dict["prompt_embeds"]
-
-        # [B, F] -> [B]
-        print(f"timestep: {timestep}")
-        print(f"self.uniform_timestep: {self.uniform_timestep}")
-        print(f"timestep.ndim: {timestep.ndim}")
-        print(f"timestep.shape: {timestep.shape}")
-        print(f"noisy_image_or_video.shape: {noisy_image_or_video.shape}")
-        if self.uniform_timestep:
-            # input_timestep = timestep[:, 0]
-            # if timestep.ndim == 1:
-            #     input_timestep = timestep.unsqueeze(0)
-            # else:
-            input_timestep = timestep
-            pass
-        else:
-            if timestep.ndim == 1:
-                print(f"not uniform timestep, timestep.ndim == 1")
-                input_timestep = timestep.unsqueeze(0)
-            else:
-                input_timestep = timestep
-
-        logits = None
-        # X0 prediction
-        if kv_cache is not None:
-            flow_pred = self.model(
-                noisy_image_or_video.permute(0, 2, 1, 3, 4),
-                t=input_timestep, context=prompt_embeds,
-                seq_len=self.seq_len,
-                kv_cache=kv_cache,
-                crossattn_cache=crossattn_cache,
-                current_start=current_start,
-                cache_start=cache_start
-            ).permute(0, 2, 1, 3, 4)
-        else:
-            if clean_x is not None:
-                # teacher forcing
-                flow_pred = self.model(
-                    noisy_image_or_video.permute(0, 2, 1, 3, 4),
-                    t=input_timestep, context=prompt_embeds,
-                    seq_len=self.seq_len,
-                    clean_x=clean_x.permute(0, 2, 1, 3, 4),
-                    aug_t=aug_t,
-                ).permute(0, 2, 1, 3, 4)
-            else:
-                if classify_mode:
-                    flow_pred, logits = self.model(
-                        noisy_image_or_video.permute(0, 2, 1, 3, 4),
-                        t=input_timestep, context=prompt_embeds,
-                        seq_len=self.seq_len,
-                        classify_mode=True,
-                        register_tokens=self._register_tokens,
-                        cls_pred_branch=self._cls_pred_branch,
-                        gan_ca_blocks=self._gan_ca_blocks,
-                        concat_time_embeddings=concat_time_embeddings
-                    )
-                    flow_pred = flow_pred.permute(0, 2, 1, 3, 4)
-                else:
-                    flow_pred = self.model(
-                        noisy_image_or_video.permute(0, 2, 1, 3, 4),
-                        t=input_timestep, context=prompt_embeds,
-                        seq_len=self.seq_len
-                    ).permute(0, 2, 1, 3, 4)
-
-        # pred_x0 = self._convert_flow_pred_to_x0(
-        #     flow_pred=flow_pred.flatten(0, 1),
-        #     xt=noisy_image_or_video.flatten(0, 1),
-        #     timestep=timestep.flatten(0, 1)
-        # ).unflatten(0, flow_pred.shape[:2])
-
-        if logits is not None:
-            return flow_pred.permute(0, 2, 1, 3, 4), pred_x0.permute(0, 2, 1, 3, 4), logits
-
-        return flow_pred.permute(0, 2, 1, 3, 4)
-        # return flow_pred.permute(0, 2, 1, 3, 4), pred_x0.permute(0, 2, 1, 3, 4)
-
-    def get_scheduler(self) -> SchedulerInterface:
-        """
-        Update the current scheduler with the interface's static method
-        """
-        scheduler = self.scheduler
-        scheduler.convert_x0_to_noise = types.MethodType(
-            SchedulerInterface.convert_x0_to_noise, scheduler)
-        scheduler.convert_noise_to_x0 = types.MethodType(
-            SchedulerInterface.convert_noise_to_x0, scheduler)
-        scheduler.convert_velocity_to_x0 = types.MethodType(
-            SchedulerInterface.convert_velocity_to_x0, scheduler)
-        self.scheduler = scheduler
-        return scheduler
-
-    def post_init(self):
-        """
-        A few custom initialization steps that should be called after the object is created.
-        Currently, the only one we have is to bind a few methods to scheduler.
-        We can gradually add more methods here if needed.
-        """
-        self.get_scheduler()

fastvideo/tests/dataset/test_parquet_io.py

@@ -0,0 +1,108 @@
+import os
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from fastvideo.dataset.dataloader.parquet_io import (
+    ParquetDatasetWriter,
+    records_to_table,
+)
+
+
+def test_records_to_table_types():
+    schema = pa.schema([
+        pa.field("id", pa.string()),
+        pa.field("vae_latent_bytes", pa.binary()),
+        pa.field("vae_latent_shape", pa.list_(pa.int64())),
+        pa.field("duration_sec", pa.float64()),
+        pa.field("width", pa.int64()),
+    ])
+    records = [{
+        "id": "a",
+        "vae_latent_bytes": b"\x00\x01",
+        "vae_latent_shape": [1, 2, 3],
+        "duration_sec": 1.5,
+        "width": 640,
+    }]
+
+    table = records_to_table(records, schema)
+    assert table.schema == schema
+    assert table.num_rows == 1
+    cols = {name: table.column(name).to_pylist()[0] for name in schema.names}
+    assert cols["id"] == "a"
+    assert isinstance(cols["vae_latent_bytes"], (bytes, bytearray))
+    assert cols["vae_latent_shape"] == [1, 2, 3]
+    assert abs(cols["duration_sec"] - 1.5) < 1e-6
+    assert cols["width"] == 640
+
+
+def test_writer_flush_and_remainder(tmp_path: Path):
+    schema = pa.schema([pa.field("id", pa.string())])
+    records = [{"id": str(i)} for i in range(25)]
+    table = records_to_table(records, schema)
+
+    out_dir = tmp_path / "out"
+    writer = ParquetDatasetWriter(str(out_dir), samples_per_file=10)
+    writer.append_table(table)
+    written = writer.flush(num_workers=1)
+    assert written == 20
+
+    files = sorted(out_dir.rglob("*.parquet"))
+    assert len(files) == 2
+    total_rows = sum(pq.read_table(str(f)).num_rows for f in files)
+    assert total_rows == 20
+
+    # Append remainder to complete another chunk
+    extra = records_to_table([{"id": str(i)} for i in range(5)], schema)
+    writer.append_table(extra)
+    written2 = writer.flush(num_workers=1)
+    assert written2 == 10
+    files2 = sorted(out_dir.rglob("*.parquet"))
+    assert len(files2) == 3
+    total_rows2 = sum(pq.read_table(str(f)).num_rows for f in files2)
+    assert total_rows2 == 30
+
+
+def test_writer_flush_write_remainder(tmp_path: Path):
+    schema = pa.schema([pa.field("id", pa.string())])
+    # 25 rows, 10 per file => 2 full files + 1 remainder(5)
+    records = [{"id": str(i)} for i in range(25)]
+    table = records_to_table(records, schema)
+
+    out_dir = tmp_path / "out_last"
+    writer = ParquetDatasetWriter(str(out_dir), samples_per_file=10)
+    writer.append_table(table)
+    # First flush writes 20
+    written1 = writer.flush(num_workers=1)
+    assert written1 == 20
+    # Final flush with remainder
+    written2 = writer.flush(num_workers=1, write_remainder=True)
+    assert written2 == 5
+    files = sorted(out_dir.rglob("*.parquet"))
+    assert len(files) == 3
+    total_rows = sum(pq.read_table(str(f)).num_rows for f in files)
+    assert total_rows == 25
+
+
+def test_writer_parallel_workers(tmp_path: Path):
+    schema = pa.schema([pa.field("id", pa.string())])
+    # 40 rows, 10 per file => 4 files
+    records = [{"id": str(i)} for i in range(40)]
+    table = records_to_table(records, schema)
+
+    out_dir = tmp_path / "out_parallel"
+    writer = ParquetDatasetWriter(str(out_dir), samples_per_file=10)
+    writer.append_table(table)
+    written = writer.flush(num_workers=2)
+    assert written == 40
+
+    # Ensure files exist under worker subdirs
+    worker_dirs = [p for p in out_dir.iterdir() if p.is_dir() and p.name.startswith("worker_")]
+    assert len(worker_dirs) >= 1
+    files = sorted(out_dir.rglob("*.parquet"))
+    assert len(files) == 4
+    total_rows = sum(pq.read_table(str(f)).num_rows for f in files)
+    assert total_rows == 40
+
+

fastvideo/tests/dataset/test_schema_record_creator.py

@@ -0,0 +1,123 @@
+import numpy as np
+
+from fastvideo.dataset.dataloader.record_schema import (
+    basic_t2v_record_creator,
+    i2v_record_creator,
+    ode_text_only_record_creator,
+    text_only_record_creator,
+)
+from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
+
+
+def _mk_basic_batch(N: int) -> PreprocessBatch:
+    batch = PreprocessBatch(data_type="video")
+    batch.video_file_name = [f"vid_{i}" for i in range(N)]
+    batch.prompt = [f"caption_{i}" for i in range(N)]
+    batch.width = [640 for _ in range(N)]
+    batch.height = [360 for _ in range(N)]
+    batch.fps = [4 for _ in range(N)]
+    batch.num_frames = [2 for _ in range(N)]
+    # Latents: shape (N, C, T, H, W); per-record use latents[idx]
+    batch.latents = np.zeros((N, 4, 2, 8, 8), dtype=np.float32)
+    # Prompt embeds: list of per-record arrays [Seq, Dim]
+    batch.prompt_embeds = [np.ones((6, 16), dtype=np.float32) for _ in range(N)]
+    return batch
+
+
+def test_basic_t2v_record_creator_fields():
+    N = 2
+    batch = _mk_basic_batch(N)
+
+    records = basic_t2v_record_creator(batch)
+    assert isinstance(records, list) and len(records) == N
+
+    for i, rec in enumerate(records):
+        assert rec["id"] == batch.video_file_name[i]
+        # Latents bytes/shape/dtype
+        assert isinstance(rec["vae_latent_bytes"], (bytes, bytearray))
+        assert rec["vae_latent_shape"] == list(batch.latents[i].shape)
+        assert rec["vae_latent_dtype"] == str(batch.latents[i].dtype)
+        # Text embedding
+        assert isinstance(rec["text_embedding_bytes"], (bytes, bytearray))
+        assert rec["text_embedding_shape"] == list(batch.prompt_embeds[i].shape)
+        assert rec["text_embedding_dtype"] == str(batch.prompt_embeds[i].dtype)
+        # Meta
+        assert rec["caption"] == batch.prompt[i]
+        assert rec["media_type"] == "video"
+        assert rec["width"] == int(batch.width[i])
+        assert rec["height"] == int(batch.height[i])
+        assert rec["num_frames"] == batch.latents[i].shape[1]
+
+
+def test_i2v_record_creator_additional_fields():
+    N = 3
+    batch = _mk_basic_batch(N)
+    # image_embeds is a list of length 1, with an array of shape [N, D]
+    batch.image_embeds = [np.ones((N, 32), dtype=np.float32)]
+    # first frame latent per record
+    batch.image_latent = np.zeros((N, 4, 1, 8, 8), dtype=np.float32)
+    # pil image per record
+    batch.pil_image = np.zeros((N, 8, 8, 3), dtype=np.uint8)
+
+    records = i2v_record_creator(batch)
+    assert isinstance(records, list) and len(records) == N
+
+    for i, rec in enumerate(records):
+        # clip feature
+        assert isinstance(rec["clip_feature_bytes"], (bytes, bytearray))
+        assert rec["clip_feature_shape"] == list(batch.image_embeds[0][i].shape)
+        assert rec["clip_feature_dtype"] == str(batch.image_embeds[0][i].dtype)
+        # first frame latent
+        assert isinstance(rec["first_frame_latent_bytes"], (bytes, bytearray))
+        assert rec["first_frame_latent_shape"] == list(batch.image_latent[i].shape)
+        assert rec["first_frame_latent_dtype"] == str(batch.image_latent[i].dtype)
+        # pil image
+        assert isinstance(rec["pil_image_bytes"], (bytes, bytearray))
+        assert rec["pil_image_shape"] == list(batch.pil_image[i].shape)
+        assert rec["pil_image_dtype"] == str(batch.pil_image[i].dtype)
+
+
+def test_ode_text_only_record_creator():
+    video_name = "ex"
+    caption = "a prompt"
+    text_embedding = np.ones((6, 16), dtype=np.float32)
+    traj = np.ones((5, 4, 2, 2), dtype=np.float32)
+    tsteps = np.arange(5, dtype=np.float32)
+
+    rec = ode_text_only_record_creator(
+        video_name=video_name,
+        text_embedding=text_embedding,
+        caption=caption,
+        trajectory_latents=traj,
+        trajectory_timesteps=tsteps,
+    )
+    assert rec["id"] == f"text_{video_name}"
+    assert isinstance(rec["text_embedding_bytes"], (bytes, bytearray))
+    assert rec["text_embedding_shape"] == list(text_embedding.shape)
+    assert rec["text_embedding_dtype"] == str(text_embedding.dtype)
+    assert rec["file_name"] == video_name
+    assert rec["caption"] == caption
+    assert rec["media_type"] == "text"
+    # Trajectory fields
+    assert isinstance(rec["trajectory_latents_bytes"], (bytes, bytearray))
+    assert rec["trajectory_latents_shape"] == list(traj.shape)
+    assert rec["trajectory_latents_dtype"] == str(traj.dtype)
+    assert isinstance(rec["trajectory_timesteps_bytes"], (bytes, bytearray))
+    assert rec["trajectory_timesteps_shape"] == list(tsteps.shape)
+    assert rec["trajectory_timesteps_dtype"] == str(tsteps.dtype)
+
+
+def test_text_only_record_creator():
+    text_name = "note1"
+    caption = "a prompt"
+    text_embedding = np.ones((7, 16), dtype=np.float32)
+    rec = text_only_record_creator(
+        text_name=text_name,
+        text_embedding=text_embedding,
+        caption=caption,
+    )
+    assert rec["id"] == f"text_{text_name}"
+    assert isinstance(rec["text_embedding_bytes"], (bytes, bytearray))
+    assert rec["text_embedding_shape"] == list(text_embedding.shape)
+    assert rec["text_embedding_dtype"] == str(text_embedding.dtype)
+    assert rec["caption"] == caption

fastvideo/tests/modal/pr_test.py

@@ -62,7 +62,7 @@ def run_test(pytest_command: str):
     
     sys.exit(result.returncode)
 
-@app.function(gpu="L40S:1", image=image, timeout=900)
+@app.function(gpu="H100:1", image=image, timeout=900)
 def run_encoder_tests():
     run_test("pytest ./fastvideo/tests/encoders -vs")
 
@@ -117,3 +117,11 @@ def run_inference_lora_tests():
 @app.function(gpu="L40S:2", image=image, timeout=900)
 def run_distill_dmd_tests():
     run_test("pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
+
+@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
+def run_self_forcing_tests():
+    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs")
+
+@app.function(gpu="L40S:1", image=image, timeout=900)
+def run_unit_test():
+    run_test("pytest ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ -vs")

fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py

@@ -0,0 +1,184 @@
+import os
+from pathlib import Path
+from huggingface_hub import snapshot_download
+import subprocess
+import sys
+from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+import shutil
+
+# Import the training pipeline
+sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
+
+NUM_NODES = "1"
+MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+
+# preprocessing
+DATA_DIR = "data"
+LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "crush-smol"))
+NUM_GPUS_PER_NODE_PREPROCESSING = "1"
+PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocess.py"
+
+LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "crush-smol_processed_t2v"))
+
+
+# training
+NUM_GPUS_PER_NODE_TRAINING = "4"
+TRAINING_ENTRY_FILE_PATH = "fastvideo/training/wan_distillation_pipeline.py"
+LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_parquet_dataset")
+LOCAL_VALIDATION_DATASET_FILE = "examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
+LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs"))
+
+def download_data():
+    # create the data dir if it doesn't exist
+    data_dir = Path(DATA_DIR)
+
+    print(f"Creating data directory at {data_dir}")
+    os.makedirs(data_dir, exist_ok=True)
+
+    print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
+    try:
+        result = snapshot_download(
+            repo_id="wlsaidhi/crush-smol-merged",
+            local_dir=str(LOCAL_RAW_DATA_DIR),
+            repo_type="dataset",
+            resume_download=True,
+            token=os.environ.get("HF_TOKEN"),  # In case authentication is needed
+        )
+        print(f"Download completed successfully. Files downloaded to: {result}")
+        
+        # Verify the download
+        if not LOCAL_RAW_DATA_DIR.exists():
+            raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
+            
+        # List downloaded files
+        print("Downloaded files:")
+        for file in LOCAL_RAW_DATA_DIR.rglob("*"):
+            if file.is_file():
+                print(f"  - {file.relative_to(LOCAL_RAW_DATA_DIR)}")
+                
+    except Exception as e:
+        print(f"Error during download: {str(e)}")
+        raise
+
+
+def run_preprocessing():
+    # remove the local_preprocessed_data_dir if it exists
+    if LOCAL_PREPROCESSED_DATA_DIR.exists():
+        print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}")
+        shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
+
+    # Run torchrun command
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
+        PREPROCESSING_ENTRY_FILE_PATH,
+        "--model_path", MODEL_PATH,
+        "--seed", "42",
+        "--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge.txt"),
+        "--preprocess_video_batch_size", "1",
+        "--max_height", "480",
+        "--max_width", "832",
+        "--num_frames", "81",
+        "--dataloader_num_workers", "0",
+        "--output_dir", LOCAL_PREPROCESSED_DATA_DIR,
+        "--train_fps", "16",
+        "--samples_per_file", "1",
+        "--flush_frequency", "1",
+        "--video_length_tolerance_range", "5",
+        "--preprocess_task", "t2v",
+    ]
+
+    process = subprocess.run(cmd, check=True)
+
+
+def run_training():
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
+        TRAINING_ENTRY_FILE_PATH,
+        "--model_path", MODEL_PATH,
+        "--inference_mode", "False",
+        "--pretrained_model_name_or_path", MODEL_PATH,
+        "--data_path", LOCAL_TRAINING_DATA_DIR,
+        "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
+        "--train_batch_size", "1",
+        "--num_latent_t", "8",
+        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
+        "--sp_size", "1",
+        "--tp_size", "1",
+        "--hsdp_replicate_dim", "1",
+        "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
+        "--train_sp_batch_size", "1",
+        "--dataloader_num_workers", "10",
+        "--gradient_accumulation_steps", "1",
+        "--max_train_steps", "501",
+        "--learning_rate", "2e-6",
+        "--fake_score_learning_rate", "2e-6",
+        "--mixed_precision", "bf16",
+        "--training_state_checkpointing_steps", "1000",
+        "--weight_only_checkpointing_steps", "1000",
+        "--validation_steps", "50",
+        "--validation_sampling_steps", "3",
+        "--log_validation",
+        "--checkpoints_total_limit", "3",
+        "--ema_start_step", "0",
+        "--training_cfg_rate", "0.0",
+        "--output_dir", LOCAL_OUTPUT_DIR,
+        "--tracker_project_name", "ci_wan_t2v_dmd_overfit",
+        "--num_height", "480",
+        "--num_width", "832",
+        "--num_frames", "81",
+        "--flow_shift", "8",
+        "--validation_guidance_scale", "6.0",
+        "--weight_decay", "0.01",
+        "--generator_update_interval", "5",
+        "--dmd_denoising_steps", "1000,757,522",
+        "--min_timestep_ratio", "0.02",
+        "--max_timestep_ratio", "0.98",
+        "--seed", "1000",
+        "--real_score_guidance_scale", "3.5",
+        "--dit_precision", "fp32",
+        "--max_grad_norm", "1.0",
+        "--enable_gradient_checkpointing_type", "full",
+    ]
+
+    print(f"Running training with command: {cmd}")
+    process = subprocess.run(cmd, check=True)
+
+
+def test_e2e_overfit_single_sample():
+    os.environ["WANDB_MODE"] = "online"
+
+    download_data()
+    run_preprocessing()
+    run_training()
+
+    reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4")
+    print(f"reference_video_file: {reference_video_file}")
+    final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
+    print(f"final_validation_video_file: {final_validation_video_file}")
+
+
+    # Ensure both files exist
+    assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}"
+    assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}"
+
+    # Compute SSIM
+    mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
+        reference_video_file,
+        final_validation_video_file,
+        use_ms_ssim=True  # Using MS-SSIM for better quality assessment
+    )
+
+    print("\n===== SSIM Results for Step 900 Validation =====")
+    print(f"Mean MS-SSIM: {mean_ssim:.4f}")
+    print(f"Min MS-SSIM: {min_ssim:.4f}")
+    print(f"Max MS-SSIM: {max_ssim:.4f}")
+
+    assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}"
+
+
+if __name__ == "__main__":
+    test_e2e_overfit_single_sample()

fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py

@@ -111,7 +111,8 @@ def run_training():
         "--max_train_steps", "901",
         "--learning_rate", "1e-5",
         "--mixed_precision", "bf16",
-        "--checkpointing_steps", "6000",
+        "--weight_only_checkpointing_steps", "6000",
+        "--training_state_checkpointing_steps", "6000",
         "--validation_steps", "100",
         "--validation_sampling_steps", "50",
         "--log_validation",

fastvideo/tests/nightly/test_e2e_overfit_single_sample.py

@@ -62,6 +62,11 @@ def download_data():
 
 
 def run_preprocessing():
+    # remove the local_preprocessed_data_dir if it exists
+    if LOCAL_PREPROCESSED_DATA_DIR.exists():
+        print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}")
+        shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
+
     # Run torchrun command
     cmd = [
         "torchrun",
@@ -111,7 +116,8 @@ def run_training():
         "--max_train_steps", "901",
         "--learning_rate", "1e-5",
         "--mixed_precision", "bf16",
-        "--checkpointing_steps", "6000",
+        "--weight_only_checkpointing_steps", "6000",
+        "--training_state_checkpointing_steps", "6000",
         "--validation_steps", "100",
         "--validation_sampling_steps", "50",
         "--log_validation",