wan2.2 training doesn't crash

.buildkite/pipeline.yml

@@ -198,15 +198,4 @@ steps:
                 env:  
                   - TEST_TYPE=inference_vmoba
                 agents:
-                  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"
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -118,10 +118,6 @@ 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_unit_test:
-        description: "Run unit-test"
+      run_nightly_test:
+        description: "Run nightly-test"
         required: false
         default: false
         type: boolean
@@ -93,7 +93,6 @@ 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
@@ -103,8 +102,6 @@ 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/**'
@@ -158,9 +155,6 @@ jobs:
             precision-test-VSA:
               - *common-paths
               - *vsa-kernel-paths
-            unit-test:
-              - 'fastvideo/**'
-              - *common-paths
 
   encoder-test:
     needs: change-filter
@@ -339,42 +333,23 @@ jobs:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
 
-  unit-test:
-    needs: change-filter
+  nightly-test:
     if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.unit-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_unit_test == 'true')
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
     uses: ./.github/workflows/runpod-test.yml
     with:
-      job_id: "unit-test"
-      gpu_type: "NVIDIA L40S"
-      gpu_count: 1
+      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: "uv pip install -e .[test] && pytest ./fastvideo/dataset/ -vs && pytest ./fastvideo/workflow/ -vs"
+      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 }}
-
-  # 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 }}
+      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
 
   runpod-cleanup:
     # Add other jobs to this list as you create them

.gitignore

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

csrc/attn/vmoba_attn/setup.py

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

csrc/attn/vmoba_attn/vmoba/vmoba.py

@@ -6,16 +6,8 @@ import time
 import os
 import torch
 from typing import Tuple
-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 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
 from functools import lru_cache
 from einops import rearrange
 

examples/dataset/vidprom/README.md

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

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

examples/datasets/crush_smol/download_dataset.sh

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

examples/datasets/crush_smol/validation.json

@@ -1,76 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "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.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
-      "image_path": null,
-      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "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.",
-      "image_path": null,
-      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "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.",
-      "image_path": null,
-      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "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.",
-      "image_path": null,
-      "video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}

examples/datasets/crush_smol/validation_single.json

@@ -1,13 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
-      "image_path": null,
-      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}

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

@@ -16,23 +16,23 @@ 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=29501
+export MASTER_PORT=29503
 export TOKENIZERS_PARALLELISM=false
-export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
 export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=offline
+export WANDB_MODE=online
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 
 # Configs
-NUM_GPUS=4
+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-14B-Diffusers"  # Teacher model
 FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
 
-DATA_DIR="data/mixkit-64_processed/Node_0_GPU_1_File_1/combined_parquet_dataset"
-VALIDATION_DATASET_FILE="data/mixkit-64_processed/validation.json"
+DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
@@ -84,7 +84,7 @@ dataset_args=(
 validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 100
+  --validation_steps 50
   --validation_sampling_steps "4"
   --validation_guidance_scale "6.0" # not used for dmd inference
 )

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

@@ -0,0 +1,151 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:1
+#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
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29503
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
+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.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/"
+VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+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"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 16
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --num_frames 61  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 4  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 16
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS # 64
+  --sp_size 4
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim 8
+)
+
+# 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 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"
+  --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 arguments
+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
+  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
+)
+
+# Self-forcing DMD arguments
+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 specific arguments
+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)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
+)
+
+torchrun \
+--nnodes 1 \
+--master_port $MASTER_PORT \
+--nproc_per_node $NUM_GPUS \
+    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/inference/basic/basic.py

@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator
 
-from fastvideo.configs.sample import SamplingParam
+# from fastvideo.configs.sample import SamplingParam
 
 OUTPUT_PATH = "video_samples"
 def main():
@@ -9,26 +9,23 @@ 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-1.3B-Diffusers",
+        "Wan-AI/Wan2.1-T2V-14B-Diffusers",
         # FastVideo will automatically handle distributed setup
-        num_gpus=1,
+        num_gpus=4,
         use_fsdp_inference=True,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
         pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
-        ti2v_task=True,
         # image_encoder_cpu_offload=False,
     )
 
-    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
     # sampling_param.num_frames = 45
-    sampling_param.image_path = "test.jpg"
+    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
     # Generate videos with the same simple API, regardless of GPU count
     prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+        "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open."
     )
     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/")
@@ -36,13 +33,9 @@ def main():
     # Generate another video with a different prompt, without reloading the
     # model!
     prompt2 = (
-        "A majestic lion strides across the golden savanna, its powerful frame "
-        "glistening under the warm afternoon sun. The tall grass ripples gently in "
-        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
-        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
-        "cinematic.")
+        "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")
     video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
 
 
 if __name__ == "__main__":
-    main()
+    main()

examples/inference/basic/basic_self_forcing_causal.py

@@ -19,15 +19,13 @@ def main():
     )
 
     sampling_param = SamplingParam.from_pretrained(model_name)
-    sampling_param.num_frames = 81
 
-    prompts = [
-        "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.",
-        "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.",
-    ]
-
-    for prompt in prompts:
-        video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
 
 if __name__ == "__main__":
     main()

examples/training/consistency_finetune/causal_ode_init/crush_smol_prompts.txt

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

@@ -1,93 +0,0 @@
-#!/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"
-  --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/preprocess_wan_data_crush_smol.sh

@@ -1,25 +0,0 @@
-#!/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

@@ -1,40 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "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": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}

examples/training/consistency_finetune/ode_init/finetune_ode_init.sh

@@ -1,99 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-# DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
-DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-extended-t2v-1-3b/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=1
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_ode_init"
-  --output_dir "wan_ode_init_70k"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "fixed_wan_ode_init_70k_6e-6"
-  # --resume_from_checkpoint "ode_init_diffusers/"
-  --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"
-  --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
-  --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
-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/ode_init/finetune_ode_init_16k_vidprom.slurm

@@ -1,135 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=1e5B2_16kFV_warp_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_warp/Dode_vidprom8b16k_1e-5.out
-#SBATCH --error=ode_vidprom16k_warp/Dode_vidprom8b16k_1e-5.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-fv2
-
-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 WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# 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"
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-16k-t2v-1-3b-81/"
-VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
-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 "Dwarp_vidprom_8b16k_test_warp_1e-5"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "Dwarp_vidprom_8b16k_wan_ode_init_1e-5"
-  # --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 77
-  --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 "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --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/ode_init/finetune_ode_init_2k_vidprom.slurm

@@ -1,131 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=ode_vidprom2k
-#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_vidprom2k_output/ode_vidprom2k.out
-#SBATCH --error=ode_vidprom2k_output/ode_vidprom2k.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-fv2
-
-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 WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# 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"
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing/"
-VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
-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 "wan_ode_init_vidprom2k"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "vidprom2k_wan_ode_init_5e-6"
-  # --resume_from_checkpoint "ode_init_diffusers/"
-  --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 77
-  # --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 8
-  --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 100
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 5e-6
-  --mixed_precision "bf16"
-  --checkpointing_steps 2000
-  --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/ode_init/finetune_ode_init_crush.slurm

@@ -1,132 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=ode_crush
-#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_crush_output/ode_crush.out
-#SBATCH --error=ode_crush_output/ode_crush.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-fv2
-
-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 WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# 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"
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
-NUM_GPUS=2
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_ode_init"
-  --output_dir "wan_ode_init_warp_2"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "2warp_fixed_wan_ode_init_5e-6"
-  # --resume_from_checkpoint "ode_init_diffusers/"
-  # --warp_denoising_step
-  --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 77
-  # --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 20
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 5e-6
-  --mixed_precision "bf16"
-  --checkpointing_steps 2000
-  --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/ode_init/finetune_ode_init_lmdb.sh

@@ -1,103 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=offline
-export TOKENIZERS_PARALLELISM=false
-export MASTER_PORT=29501
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/mixkit-64_processed/Node_0_GPU_1_File_1/combined_parquet_dataset"
-# DATA_DIR="/mnt/weka/home/hao.zhang/wl/Self-Forcing/ode_single_lmdb_sf/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=1
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_ode_init"
-  --output_dir "Dwarp_vidprom_8b16k_test_warp_1e-5"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "Dwarp_vidprom_8b16k_wan_ode_init_1e-5"
-  # --resume_from_checkpoint "ode_init_diffusers/"
-  --warp_denoising_step
-  --log_visualization
-  --max_train_steps 10
-  --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
-  --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 "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --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
-  --seed 1024
-  # --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
-torchrun \
-  --nnodes 1 \
-  --master_port $MASTER_PORT \
-  --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/ode_init/finetune_ode_init_overfit.sh

@@ -1,98 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="/mnt/weka/home/hao.zhang/wl/FastVideo2/data/crush-smol_processed_t2v_1_3b_ode_init_single"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=1
-# export CUDA_VISIBLE_DEVICES=4,5
-# 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 "overfitwan_ode_init_crush_smol"
-  # --resume_from_checkpoint "ode_init_diffusers/"
-  --max_train_steps 2001
-  # --warp_denoising_step
-  --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
-  # --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 20
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --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
-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/ode_init/finetune_ode_init_single_lmdb.sh

@@ -1,100 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-# DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
-DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-16k-t2v-1-3b/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=1
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_ode_init"
-  --output_dir "debug_ode_init"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "debug_ode_init"
-  # --resume_from_checkpoint "ode_init_diffusers/"
-  --max_train_steps 1000
-  --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
-  --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 10
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --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
-  --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
-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/ode_init/preprocess_wan_data_ode.sh

@@ -1,25 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/vidprom_1.txt"
-OUTPUT_DIR="data/ode_vidprom_1_fv/"
-
-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 1 \
-    --flush_frequency 1 \
-    --video_length_tolerance_range 5 \
-    --preprocess_task "ode_trajectory" 

examples/training/consistency_finetune/ode_init/validation.json

@@ -1,76 +0,0 @@
-{
-  "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/wan_t2v_1.3B/crush_smol/finetune_t2v.sh

@@ -6,8 +6,8 @@ export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v_old"
-VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
+DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
 NUM_GPUS=4
 # export CUDA_VISIBLE_DEVICES=4,5
 
@@ -52,7 +52,7 @@ dataset_args=(
 validation_args=(
   --log_validation 
   --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 50
+  --validation_steps 200
   --validation_sampling_steps "50" 
   --validation_guidance_scale "6.0"
 )

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

@@ -4,7 +4,7 @@ GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v_old/"
+OUTPUT_DIR="data/crush-smol_processed_t2v/"
 
 torchrun --nproc_per_node=$GPU_NUM \
     fastvideo/pipelines/preprocess/v1_preprocess.py \

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

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-GPU_NUM=1 # 2,4,8
+GPU_NUM=2 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATASET_PATH="data/crush-smol/"
 OUTPUT_DIR="data/crush-smol_processed_t2v/"
@@ -14,7 +14,7 @@ torchrun --nproc_per_node=$GPU_NUM \
     --preprocess.dataset_type merged \
     --preprocess.dataset_path $DATASET_PATH \
     --preprocess.dataset_output_dir $OUTPUT_DIR \
-    --preprocess.preprocess_video_batch_size 8 \
+    --preprocess.preprocess_video_batch_size 2 \
     --preprocess.dataloader_num_workers 0 \
     --preprocess.max_height 480 \
     --preprocess.max_width 832 \

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

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

examples/training/t2v_to_i2v/wan_t2v/finetune_t2v_to_i2v.sh

@@ -1,94 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v_old"
-VALIDATION_DATASET_FILE="examples/datasets/crush_smol/validation.json"
-NUM_GPUS=8
-# export CUDA_VISIBLE_DEVICES=4,5
-
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_t2v_i2v_finetune"
-  --output_dir "checkpoints/wan_t2v_i2v_finetune"
-  --max_train_steps 5000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 2
-  --num_latent_t 20
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS 
-  --sp_size 4
-  --tp_size 1
-  --hsdp_replicate_dim 2
-  --hsdp_shard_dim 4
-)
-
-# 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 5e-5
-  --mixed_precision "bf16"
-  --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
-  --enable_gradient_checkpointing_type "full"
-  --t2v_as_i2v_task True
-  # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500"
-)
-
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_t2v_i2v_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"

examples/training/t2v_to_i2v/wan_t2v/preprocess_wan_data_t2v.sh

@@ -1,24 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v_i2v_1_3b/"
-
-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 2 \
-    --seed 42 \
-    --max_height 480 \
-    --max_width 832 \
-    --num_frames 77 \
-    --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 "t2v_ode_trajectory" 

fastvideo/attention/backends/vmoba.py

@@ -5,6 +5,7 @@ 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)
@@ -133,8 +134,6 @@ 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)
@@ -170,6 +169,7 @@ 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 = self.pad_input(hidden_states, indices_q, batch_size,
-                                       sequence_length)
+        hidden_states = 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/attention/layer.py

@@ -3,7 +3,6 @@
 import torch
 import torch.nn as nn
 
-import fastvideo.envs as envs
 from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
 from fastvideo.distributed.communication_op import (
     sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D)
@@ -37,7 +36,7 @@ class DistributedAttention(nn.Module):
         if num_kv_heads is None:
             num_kv_heads = num_heads
 
-        dtype = torch.bfloat16 if envs.FASTVIDEO_FORCE_ATTN_BF16 else get_compute_dtype()
+        dtype = get_compute_dtype()
         attn_backend = get_attn_backend(
             head_size,
             dtype,
@@ -222,7 +221,7 @@ class LocalAttention(nn.Module):
         if num_kv_heads is None:
             num_kv_heads = num_heads
 
-        dtype = torch.bfloat16 if envs.FASTVIDEO_FORCE_ATTN_BF16 else get_compute_dtype()
+        dtype = get_compute_dtype()
         attn_backend = get_attn_backend(
             head_size,
             dtype,

fastvideo/configs/models/dits/base.py

@@ -27,7 +27,6 @@ class DiTArchConfig(ArchConfig):
     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,9 +92,6 @@ 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,13 +4,12 @@ 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 (SelfForcingWanT2V480PConfig,
-                                             WanI2V480PConfig, WanI2V720PConfig,
+from fastvideo.configs.pipelines.wan import (WanI2V480PConfig, WanI2V720PConfig,
                                              WanT2V480PConfig, WanT2V720PConfig)
 
 __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
     "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
     "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig",
-    "SelfForcingWanT2V480PConfig", "get_pipeline_config_cls_from_name"
+    "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/base.py

@@ -45,13 +45,10 @@ class PipelineConfig:
     embedded_cfg_scale: float = 6.0
     flow_shift: float | None = None
     disable_autocast: bool = False
-    ti2v_task: bool = False
-    t2v_as_i2v_task: bool = False
 
     # Model configuration
     dit_config: DiTConfig = field(default_factory=DiTConfig)
     dit_precision: str = "bf16"
-    dit_forward_precision: str = "bf16"
 
     # VAE configuration
     vae_config: VAEConfig = field(default_factory=VAEConfig)
@@ -90,7 +87,6 @@ class PipelineConfig:
 
     # Wan2.2 TI2V parameters
     ti2v_task: bool = False
-    boundary_ratio: float | None = None
 
     # Compilation
     # enable_torch_compile: bool = False
@@ -218,24 +214,6 @@ class PipelineConfig:
             "Comma-separated list of denoising steps (e.g., '1000,757,522')",
         )
 
-        # TI2V task
-        parser.add_argument(
-            f"--{prefix_with_dot}ti2v-task",
-            action=StoreBoolean,
-            dest=f"{prefix_with_dot.replace('-', '_')}ti2v_task",
-            default=PipelineConfig.ti2v_task,
-            help="Enable TI2V",
-        )
-
-        # T2V to I2V task
-        parser.add_argument(
-            f"--{prefix_with_dot}t2v-as-i2v-task",
-            action=StoreBoolean,
-            dest=f"{prefix_with_dot.replace('-', '_')}t2v_as_i2v_task",
-            default=PipelineConfig.t2v_as_i2v_task,
-            help="Enable T2V to I2V task",
-        )
-
         # Add VAE configuration arguments
         from fastvideo.configs.models.vaes.base import VAEConfig
         VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
@@ -267,9 +245,7 @@ class PipelineConfig:
         """
         from fastvideo.configs.pipelines.registry import (
             get_pipeline_config_cls_from_name)
-        logger.info("WTF model_path: %s", model_path)
         pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)
-        logger.info("pipeline_config_cls: %s", pipeline_config_cls)
 
         return cast(PipelineConfig, pipeline_config_cls(model_path=model_path))
 

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,
-    Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
-    WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
-    SelfForcingWanT2V480PConfig)
+    SelfForcingWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config,
+    Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig,
+    WanT2V720PConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -48,9 +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(),
-    "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -62,7 +60,6 @@ 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

@@ -82,7 +82,7 @@ class WanI2V480PConfig(WanT2V480PConfig):
         default_factory=CLIPVisionConfig)
     image_encoder_precision: str = "fp32"
 
-    def __post_init__(self) -> None:
+    def __post_init__(self):
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
 
@@ -108,17 +108,19 @@ 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
@@ -130,21 +132,12 @@ class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config):
 
 @dataclass
 class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
-    flow_shift: float | None = 12.0
-    boundary_ratio: float | None = 0.875
-
-    def __post_init__(self) -> None:
-        self.dit_config.boundary_ratio = self.boundary_ratio
+    pass
 
 
 @dataclass
-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
+class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
+    pass
 
 
 # =============================================

fastvideo/configs/sample/base.py

@@ -40,7 +40,6 @@ 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
@@ -48,8 +47,6 @@ 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"
@@ -170,12 +167,6 @@ 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",
@@ -207,18 +198,6 @@ 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,22 +144,18 @@ 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  # high_noise
-    guidance_scale_2: float = 3.0  # low_noise
+    guidance_scale: float = 4.0
+    guidance_scale_2: float = 3.0
     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  # high_noise
-    guidance_scale_2: float = 3.5  # low_noise
+    guidance_scale: float = 3.5
+    guidance_scale_2: float = 3.5
     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, TextDataset
+from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset
 from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                          TemporalRandomCrop)
 from fastvideo.dataset.validation_dataset import ValidationDataset
@@ -37,15 +37,9 @@ 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", "TextDataset"
+    "VideoCaptionMergedDataset"
 ]

fastvideo/dataset/dataloader/parquet_io.py

@@ -1,264 +0,0 @@
-"""
-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/schema.py

@@ -50,7 +50,6 @@ pyarrow_schema_i2v = pa.schema([
     pa.field("fps", pa.float64()),
 ])
 
-
 pyarrow_schema_t2v = pa.schema([
     pa.field("id", pa.string()),
     # --- Image/Video VAE latents ---
@@ -79,40 +78,3 @@ 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/lmdb_datasets.py

@@ -1,43 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo.dataset.lmdb_utils import get_array_shape_from_lmdb, retrieve_row_from_lmdb
-from torch.utils.data import Dataset
-import numpy as np
-import torch
-import lmdb
-
-# from Self-Forcing: https://github.com/guandeh17/Self-Forcing/blob/main/utils/dataset.py
-class ODERegressionLMDBDataset(Dataset):
-    def __init__(self, data_path: str, max_pair: int = int(1e8)):
-        print(f"data_path: {data_path}")
-        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)
-        }
-

fastvideo/dataset/lmdb_utils.py

@@ -1,75 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# from Self-Forcing: https://github.com/guandeh17/Self-Forcing/blob/main/utils/lmdb.py
-
-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/dataset/preprocessing_datasets.py

@@ -628,134 +628,3 @@ 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

@@ -3,12 +3,9 @@ from typing import Any, cast
 
 import numpy as np
 import torch
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
 
 
-def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]:
+def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
     """
     Pad or crop an embedding [L, D] to exactly padding_length tokens.
     Return:

fastvideo/entrypoints/video_generator.py

@@ -344,9 +344,6 @@ 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

@@ -18,7 +18,6 @@ if TYPE_CHECKING:
     FASTVIDEO_LOGGING_PREFIX: str = ""
     FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None
     FASTVIDEO_TRACE_FUNCTION: int = 0
-    FASTVIDEO_FORCE_ATTN_BF16: bool = False
     FASTVIDEO_ATTENTION_BACKEND: str | None = None
     FASTVIDEO_ATTENTION_CONFIG: str | None = None
     FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork"
@@ -169,10 +168,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "FASTVIDEO_TRACE_FUNCTION":
     lambda: int(os.getenv("FASTVIDEO_TRACE_FUNCTION", "0")),
 
-    # if set, fastvideo will force attention to be computed in bfloat16
-    "FASTVIDEO_FORCE_ATTN_BF16":
-    lambda: bool(int(os.getenv("FASTVIDEO_FORCE_ATTN_BF16", "0"))),
-
     # Backend for attention computation
     # Available options:
     # - "TORCH_SDPA": use torch.nn.MultiheadAttention

fastvideo/fastvideo_args.py

@@ -158,7 +158,6 @@ class FastVideoArgs:
         "transformer": True,
         "vae": True,
     })
-    override_transformer_cls_name: str | None = None
 
     # # DMD parameters
     # dmd_denoising_steps: List[int] | None = field(default=None)
@@ -397,12 +396,6 @@ 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",
-        )
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 
@@ -705,7 +698,6 @@ class TrainingArgs(FastVideoArgs):
     # simulate generator forward to match inference
     simulate_generator_forward: bool = False
     warp_denoising_step: bool = False
-    intermediate_latents_visualization: bool = False
 
     # Self-forcing specific arguments
     num_frame_per_block: int = 3

fastvideo/layers/layernorm.py

@@ -276,4 +276,4 @@ class LayerNormScaleShift(nn.Module):
         if self.compute_dtype == torch.float32:
             output = output.to(x.dtype)
 
-        return output
+        return output

fastvideo/layers/lora/linear.py

@@ -77,11 +77,9 @@ class BaseLayerWithLoRA(nn.Module):
             lora_A = self.lora_A.to_local()
 
         if not self.merged and not self.disable_lora:
-            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
+            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)))
             if self.lora_alpha != self.lora_rank:
                 delta = delta * (
                     self.lora_alpha / self.lora_rank  # type: ignore

fastvideo/models/dits/causal_wanvideo.py

@@ -556,8 +556,6 @@ class CausalWanTransformer3DModel(BaseDiT):
                 start_frame: int = 0,
                 **kwargs) -> torch.Tensor:
 
-        logger.info("timestep dtype: %s, timestep sum: %s", timestep.dtype, timestep.float().sum().item())
-
         orig_dtype = hidden_states.dtype
         if not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
@@ -652,11 +650,10 @@ class CausalWanTransformer3DModel(BaseDiT):
         **kwargs
     ):
         if kwargs.get('kv_cache', None) is not None:
-            noise_pred = self._forward_inference(*args, **kwargs)
+            return self._forward_inference(*args, **kwargs)
         else:
-            noise_pred = self._forward_train(*args, **kwargs)
+            return self._forward_train(*args, **kwargs)
 
-        return noise_pred
 
     def unpatchify(self, x, grid_sizes):
         r"""
@@ -682,4 +679,4 @@ class CausalWanTransformer3DModel(BaseDiT):
             # 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
+        return out

fastvideo/models/dits/wanvideo.py

@@ -172,13 +172,8 @@ class WanT2VCrossAttention(WanSelfAttention):
             k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
             v = self.to_v(context)[0].view(b, -1, n, d)
 
-        if envs.FASTVIDEO_FORCE_ATTN_BF16:
-            out_dtype = v.dtype
-            # compute attention
-            x = self.attn(q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)).to(out_dtype)
-        else:
-            # compute attention
-            x = self.attn(q, k, v)
+        # compute attention
+        x = self.attn(q, k, v)
 
         # output
         x = x.flatten(2)
@@ -364,12 +359,7 @@ class WanTransformerBlock(nn.Module):
                                        is_neox_style=False), _apply_rotary_emb(
                                            key, cos, sin, is_neox_style=False)
 
-        if envs.FASTVIDEO_FORCE_ATTN_BF16:
-            out_dtype = value.dtype
-            attn_output, _ = self.attn1(query.to(torch.bfloat16), key.to(torch.bfloat16), value.to(torch.bfloat16))
-            attn_output = attn_output.to(out_dtype)
-        else:
-            attn_output, _ = self.attn1(query, key, value)
+        attn_output, _ = self.attn1(query, key, value)
         attn_output = attn_output.flatten(2)
         attn_output, _ = self.to_out(attn_output)
         attn_output = attn_output.squeeze(1)

fastvideo/models/loader/component_loader.py

@@ -238,7 +238,6 @@ class TextEncoderLoader(ComponentLoader):
                 1]
 
         target_device = get_local_torch_device()
-        logger.info("Loading text encoder in %s precision", encoder_precision)
         # TODO(will): add support for other dtypes
         return self.load_model(model_path, encoder_config, target_device,
                                fastvideo_args, encoder_precision)
@@ -416,10 +415,6 @@ class TransformerLoader(ComponentLoader):
             raise ValueError(
                 "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
 
@@ -450,8 +445,6 @@ class TransformerLoader(ComponentLoader):
 
         default_dtype = PRECISION_TO_TYPE[
             fastvideo_args.pipeline_config.dit_precision]
-        param_dtype = PRECISION_TO_TYPE[
-            fastvideo_args.pipeline_config.dit_forward_precision]
 
         # Load the model using FSDP loader
         logger.info("Loading model from %s, default_dtype: %s", cls_name,
@@ -471,8 +464,7 @@ 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=param_dtype,
+            param_dtype=torch.bfloat16,
             reduce_dtype=torch.float32,
             output_dtype=None,
             training_mode=fastvideo_args.training_mode)
@@ -481,11 +473,9 @@ class TransformerLoader(ComponentLoader):
         total_params = sum(p.numel() for p in model.parameters())
         logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
 
-        # Need to convert the model to the default_dtype
-        # Otherwise, the model master weights will be in param_dtype, and the gradients will also be in param_dtype
-        # This means the param update will be in lower precision, causing precision loss
-        logger.info("Converting model to dtype: %s", default_dtype)
-        model = model.to(default_dtype)
+        dtypes = set(param.dtype for param in model.parameters())
+        if len(dtypes) > 1:
+            model = model.to(default_dtype)
         model = model.eval()
         return model
 

fastvideo/models/loader/fsdp_load.py

@@ -62,7 +62,6 @@ 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,
@@ -88,7 +87,7 @@ def maybe_load_fsdp_model(
         mp_policy=mp_policy,
     )
 
-    with set_default_dtype(default_dtype), torch.device("meta"):
+    with set_default_dtype(param_dtype), torch.device("meta"):
         model = model_cls(**init_params)
 
     # Check if we should use FSDP
@@ -126,7 +125,7 @@ def maybe_load_fsdp_model(
         model,
         weight_iterator,
         device,
-        default_dtype,
+        param_dtype,
         strict=True,
         cpu_offload=cpu_offload,
         param_names_mapping=param_names_mapping_fn,

fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py

@@ -64,15 +64,8 @@ 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,34 +171,12 @@ 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.double().to(device)
-    noise_input_latent = noise_input_latent.double().to(device)
-    sigmas = scheduler.sigmas.double().to(device)
-    timesteps = scheduler.timesteps.double().to(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)
     timestep_id = torch.argmin(
         (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
     sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
     pred_video = noise_input_latent - sigma_t * pred_noise
     return pred_video.to(dtype)
-
-def pred_video_to_pred_noise(x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> 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)

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, Any] = {}
+    modules: dict[str, torch.nn.Module] = {}
     post_init_called: bool = False
 
     # TODO(will): args should support both inference args and training args
@@ -121,25 +121,14 @@ class ComposedPipelineBase(ABC):
                         model_path: str,
                         device: str | None = None,
                         torch_dtype: torch.dtype | None = None,
-                        pipeline_config: PipelineConfig | None = None,
+                        pipeline_config: str | PipelineConfig | None = None,
                         args: argparse.Namespace | None = None,
                         required_config_modules: list[str] | None = None,
                         loaded_modules: dict[str, torch.nn.Module]
                         | None = None,
                         **kwargs) -> "ComposedPipelineBase":
         """
-        Load a pipeline from a pretrained model. 
-        Few different patterns are supported:
-        - Only provide model_path:
-            - This will load the pipeline in inference mode.
-            - The pipeline will be initialized with the default config.
-            - The pipeline will be initialized with the default modules.
-            - The pipeline will be initialized with the default stages.
-            - The pipeline will be initialized with the default stages.
-        - override the default config using pipeline_config or args or kwargs
-        - override the default modules using loaded_modules
-        - override the pipelineconfig
-
+        Load a pipeline from a pretrained model.
         loaded_modules: Optional[Dict[str, torch.nn.Module]] = None,
         If provided, loaded_modules will be used instead of loading from config/pretrained weights.
         """
@@ -147,18 +136,9 @@ class ComposedPipelineBase(ABC):
 
             kwargs['model_path'] = model_path
             fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
-            if pipeline_config is not None:
-                fastvideo_args.pipeline_config = pipeline_config
-            if fastvideo_args.override_transformer_cls_name is not None:
-                pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers")
-                fastvideo_args.pipeline_config = pipeline_config
         else:
             assert args is not None, "args must be provided for training mode"
             fastvideo_args = TrainingArgs.from_cli_args(args)
-            if fastvideo_args.override_transformer_cls_name is not None:
-                pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers")
-                fastvideo_args.pipeline_config = pipeline_config
-                logger.info("in 2 Overriding transformer cls name to %s", fastvideo_args.override_transformer_cls_name)
             # TODO(will): fix this so that its not so ugly
             fastvideo_args.model_path = model_path
             for key, value in kwargs.items():
@@ -169,8 +149,7 @@ class ComposedPipelineBase(ABC):
             # model is loaded with the correct precision. Subsequently we will
             # use FSDP2's MixedPrecisionPolicy to set the precision for the
             # fwd, bwd, and other operations' precision.
-            fastvideo_args.pipeline_config.dit_precision = 'fp32'
-            # assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
+            assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
 
         logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
 
@@ -258,19 +237,20 @@ 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")
-            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"]
+            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"]
 
         model_index.pop("boundary_ratio", None)
-        # used by Wan2.2 ti2v
         model_index.pop("expand_timesteps", None)
 
         # some sanity checks
@@ -303,8 +283,8 @@ class ComposedPipelineBase(ABC):
                           architecture) in model_index.items():
             if transformers_or_diffusers is None:
                 logger.warning(
-                    "Module %s in model_index.json has null value, removing from required_config_modules",
-                    module_name)
+                    "Module in model_index.json has null value, removing from required_config_modules"
+                )
                 if module_name in self.required_config_modules:
                     self.required_config_modules.remove(module_name)
                 continue

fastvideo/pipelines/pipeline_batch_info.py

@@ -129,7 +129,6 @@ 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
@@ -148,12 +147,7 @@ class ForwardBatch:
     modules: dict[str, Any] = field(default_factory=dict)
 
     # Final output (after pipeline completion)
-    output: torch.Tensor | None = None
-    return_trajectory_latents: bool = False
-    return_trajectory_decoded: bool = False
-    trajectory_timesteps: list[int] | None = None
-    trajectory_latents: torch.Tensor | None = None
-    trajectory_decoded: list[torch.Tensor] | None = None
+    output: Any = None
 
     # Extra parameters that might be needed by specific pipeline implementations
     extra: dict[str, Any] = field(default_factory=dict)
@@ -212,10 +206,6 @@ 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
@@ -246,7 +236,6 @@ class TrainingBatch:
     fake_score_loss: float = 0.0
 
     dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
-    latent_vis_dict: dict[str, torch.Tensor] = field(default_factory=dict)
     fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
 
 

fastvideo/pipelines/preprocess/preprocess_pipeline_base.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
+import multiprocessing
 import os
+from concurrent.futures import ProcessPoolExecutor
 from typing import Any
 
 import numpy as np
@@ -10,8 +12,6 @@ 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,13 +54,9 @@ class BasePreprocessPipeline(ComposedPipelineBase):
         """Get additional features specific to the pipeline type. Override in subclasses."""
         return {}
 
-    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()]
+        """Get the schema fields for the pipeline type. Override in subclasses."""
+        raise NotImplementedError
 
     def create_record_for_schema(self,
                                  preprocess_batch: PreprocessBatch,
@@ -404,22 +400,166 @@ 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")
-                table = records_to_table(batch_data, self.get_pyarrow_schema())
+                # 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())
                 write_pbar.update(1)
                 write_pbar.close()
 
-                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)
+                # Store the table in a list for later processing
+                if not hasattr(self, 'all_tables'):
+                    self.all_tables = []
+                self.all_tables.append(table)
+
                 logger.info("Collected batch with %s samples", len(table))
 
             if num_processed_samples >= args.flush_frequency:
-                written = self.dataset_writer.flush()
-                logger.info("Flushed %s samples to parquet", written)
+                self._flush_tables(num_processed_samples, args,
+                                   combined_parquet_dir)
                 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_pyarrow_schema(self):
-        """Return the PyArrow schema for I2V pipeline."""
-        return pyarrow_schema_i2v
+    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_extra_features(self, valid_data: dict[str, Any],
                            fastvideo_args: FastVideoArgs) -> dict[str, Any]:

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -1,443 +0,0 @@
-# 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 numpy as np
-import pyarrow as pa
-import torch
-from PIL import Image
-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
-
-from fastvideo.forward_context import set_forward_context
-from fastvideo.models.utils import pred_noise_to_pred_video, pred_video_to_pred_noise
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.distributed import get_local_torch_device
-
-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)
-
-        fastvideo_args.model_loaded["transformer"] = False
-        loader = TransformerLoader()
-        fastvideo_args.pipeline_config.dit_precision = "fp32" # Overwrite the precision to fp32 for transformer
-        fastvideo_args.pipeline_config.dit_forward_precision = "fp32"
-        self.transformer = loader.load(
-            fastvideo_args.model_paths["transformer"], fastvideo_args)
-        self.add_module("transformer", self.transformer)
-        fastvideo_args.model_loaded["transformer"] = True
-
-        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):
-            # logger.info("transformer weight sum: %s", sum(p.float().sum().item() for p in self.transformer.parameters()))
-            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"]
-                self.prompt_encoding_stage.text_encoders[0] = self.prompt_encoding_stage.text_encoders[0].to(dtype=torch.bfloat16).to(dtype=torch.float32)
-                # 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]
-                logger.info("prompt_embeds sum: %s, prompt_embeds shape: %s, prompt_embeds dtype: %s", prompt_embeds.float().sum(), prompt_embeds.shape, prompt_embeds.dtype)
-                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)
-
-                negative_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
-
-                # 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(
-                        negative_prompt,
-                        fastvideo_args,
-                        encoder_index=[0],
-                        return_attention_mask=True,
-                    )
-                    negative_prompt_embed = negative_prompt_embeds_list[0]
-                    negative_prompt_attention_mask = negative_prompt_masks_list[
-                        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 = True
-                    batch.height = args.max_height
-                    batch.width = args.max_width
-                    batch.fps = args.train_fps
-                    batch.guidance_scale = 3.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)
-                    noisy_input = []
-                    # latents = result_batch.latents.permute(0, 2, 1, 3, 4)
-                    latents = torch.randn(
-                        [1, 21, 16, 60, 104], dtype=torch.float32, device=get_local_torch_device()
-                    )
-                    # logger.info("transformer weight sum: %s", sum(p.float().sum().item() for p in self.transformer.parameters()))
-                    logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype)
-
-                    logger.info("scheduler timesteps: %s", self.get_module("scheduler").timesteps)
-                    for progress_id, t in enumerate(tqdm(self.get_module("scheduler").timesteps)):
-                        timestep = t * \
-                            torch.ones([1, 21], device=latents.device, dtype=torch.float32)
-
-                        noisy_input.append(latents)
-
-                        with set_forward_context(
-                            current_timestep=0,
-                            attn_metadata=None,
-                            forward_batch=None,
-                        ):
-                            # logger.info("prompt_embed sum: %s, prompt_embed shape: %s, prompt_embed dtype: %s", prompt_embed.float().sum(), prompt_embed.shape, prompt_embed.dtype)
-                            # logger.info("timestep: %s", timestep[:, 0])
-                            # Run transformer
-                            cond_pred_noise_btchw = self.transformer(
-                                hidden_states=latents.permute(0, 2, 1, 3, 4),
-                                encoder_hidden_states=prompt_embed,
-                                timestep=timestep[:, 0]
-                            ).permute(0, 2, 1, 3, 4)
-
-                        # logger.info("cond_pred_noise_btchw sum: %s, cond_pred_noise_btchw shape: %s, cond_pred_noise_btchw dtype: %s", cond_pred_noise_btchw.float().sum(), cond_pred_noise_btchw.shape, cond_pred_noise_btchw.dtype)
-
-                        cond_pred_video_btchw = pred_noise_to_pred_video(
-                            pred_noise=cond_pred_noise_btchw.flatten(0, 1),
-                            noise_input_latent=latents.flatten(0, 1),
-                            timestep=timestep.flatten(0, 1),
-                            scheduler=self.get_module("scheduler")).unflatten(
-                                0, cond_pred_noise_btchw.shape[:2])
-
-                        # logger.info("cond_pred_video_btchw sum: %s, cond_pred_video_btchw shape: %s, cond_pred_video_btchw dtype: %s", cond_pred_video_btchw.float().sum(), cond_pred_video_btchw.shape, cond_pred_video_btchw.dtype)
-
-                        with set_forward_context(
-                            current_timestep=t,
-                            attn_metadata=None,
-                            forward_batch=result_batch,
-                        ):
-                            # logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype)
-                            # Run transformer
-                            uncond_pred_noise_btchw = self.transformer(
-                                latents.permute(0, 2, 1, 3, 4),
-                                negative_prompt_embed,
-                                timestep[:, 0]
-                            ).permute(0, 2, 1, 3, 4)
-
-                        # logger.info("uncond_pred_noise_btchw sum: %s, uncond_pred_noise_btchw shape: %s, uncond_pred_noise_btchw dtype: %s", uncond_pred_noise_btchw.float().sum(), uncond_pred_noise_btchw.shape, uncond_pred_noise_btchw.dtype)
-
-                        uncond_pred_video_btchw = pred_noise_to_pred_video(
-                            pred_noise=uncond_pred_noise_btchw.flatten(0, 1),
-                            noise_input_latent=latents.flatten(0, 1),
-                            timestep=timestep.flatten(0, 1),
-                            scheduler=self.get_module("scheduler")).unflatten(
-                                0, uncond_pred_noise_btchw.shape[:2])
-
-                        pred_video_btchw = uncond_pred_video_btchw + batch.guidance_scale * (
-                            cond_pred_video_btchw - uncond_pred_video_btchw
-                        )
-
-                        # logger.info("pred_video_btchw sum: %s, pred_video_btchw shape: %s, pred_video_btchw dtype: %s", pred_video_btchw.float().sum(), pred_video_btchw.shape, pred_video_btchw.dtype)
-
-                        pred_noise_btchw = pred_video_to_pred_noise(
-                            x0_pred=pred_video_btchw.flatten(0, 1),
-                            xt=latents.flatten(0, 1),
-                            timestep=timestep.flatten(0, 1),
-                            scheduler=self.get_module("scheduler")).unflatten(
-                                0, pred_video_btchw.shape[:2])
-
-                        # logger.info("pred_noise_btchw sum: %s, pred_noise_btchw shape: %s, pred_noise_btchw dtype: %s", pred_noise_btchw.float().sum(), pred_noise_btchw.shape, pred_noise_btchw.dtype)
-
-                        latents = self.get_module("scheduler").step(
-                            pred_noise_btchw.flatten(0, 1),
-                            self.get_module("scheduler").timesteps[progress_id] * torch.ones(
-                                [1, 21], device=latents.device, dtype=torch.long).flatten(0, 1),
-                            latents.flatten(0, 1)
-                        )[0].unflatten(dim=0, sizes=pred_noise_btchw.shape[:2])
-
-                        # logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype)
-
-                    noisy_input.append(latents)
-
-                    noisy_inputs = torch.stack(noisy_input, dim=1)
-
-                    noisy_inputs = noisy_inputs[:, [0, 12, 24, 36, -1]].half()
-
-                    logger.info("noisy inputs sum: %s, noisy inputs shape: %s, noisy inputs dtype: %s", noisy_inputs.float().sum(), noisy_inputs.shape, noisy_inputs.dtype)
-
-                    result_batch.trajectory_latents = noisy_inputs.permute(0, 1, 3, 2, 4, 5)
-                    result_batch.trajectory_timesteps = torch.tensor([self.get_module("scheduler").timesteps[i] for i in [0, 12, 24, 36, -1]])
-                    result_batch.latents = latents.permute(0, 2, 1, 3, 4)
-                    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)
-
-                trajectory_latents = torch.stack(trajectory_latents, dim=0).squeeze(0)
-                # trajecotry_latents = trajectory_latents[:, [0, 12, 24, 36, -1]]
-
-                # 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_pyarrow_schema(self):
-        """Return the PyArrow schema for T2V pipeline."""
-        return pyarrow_schema_t2v
+    def get_schema_fields(self):
+        """Get the schema fields for T2V pipeline."""
+        return [f.name for f in pyarrow_schema_t2v]
 
 
 EntryClass = PreprocessPipeline_T2V

fastvideo/pipelines/preprocess/preprocess_pipeline_text.py

@@ -1,184 +0,0 @@
-# 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,6 +1,5 @@
 import argparse
 import os
-from typing import Any
 
 from fastvideo import PipelineConfig
 from fastvideo.configs.models.vaes import WanVAEConfig
@@ -10,12 +9,8 @@ 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__)
@@ -26,22 +21,12 @@ 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: 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),
-        }
+    kwargs = {
+        "vae_precision": "fp32",
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
+    }
     pipeline_config.update_config_from_dict(kwargs)
-
     fastvideo_args = FastVideoArgs(
         model_path=args.model_path,
         num_gpus=get_world_size(),
@@ -50,23 +35,7 @@ def main(args) -> None:
         text_encoder_cpu_offload=False,
         pipeline_config=pipeline_config,
     )
-    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__)
-
+    PreprocessPipeline = PreprocessPipeline_I2V if args.preprocess_task == "i2v" else PreprocessPipeline_T2V
     pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 
@@ -105,12 +74,7 @@ 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("--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("--preprocess_task", type=str, default="t2v")
     parser.add_argument("--train_fps", type=int, default=30)
     parser.add_argument("--use_image_num", type=int, default=0)
     parser.add_argument("--text_max_length", type=int, default=256)

fastvideo/pipelines/stages/causal_denoising.py

@@ -78,8 +78,6 @@ class CausalDMDDenosingStage(DenoisingStage):
                                              torch.tensor([0],
                                                           dtype=torch.float32)))
             timesteps = scheduler_timesteps[1000 - timesteps]
-        else:
-            assert False, "warp_denoising_step must be true"
         timesteps = timesteps.to(get_local_torch_device())
         logger.info("Using timesteps: %s", timesteps)
 

fastvideo/pipelines/stages/decoding.py

@@ -50,63 +50,6 @@ 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,
@@ -116,28 +59,13 @@ 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 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
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
             
         Returns:
-            Modified batch with:
-                - output: Decoded frames (batch, channels, frames, height, width) as CPU float32
-                - trajectory_decoded (if requested): List of decoded frames per timestep
+            The batch with decoded outputs.
         """
-        # 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()
@@ -147,29 +75,58 @@ class DecodingStage(PipelineStage):
                 pipeline.add_module("vae", self.vae)
             fastvideo_args.model_loaded["vae"] = True
 
-        if fastvideo_args.output_type == "latent":
-            frames = batch.latents
-        else:
-            frames = self.decode(batch.latents, fastvideo_args)
+        self.vae = self.vae.to(get_local_torch_device())
 
-        # 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())
+        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
+        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
+
+            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)
 
         # Convert to CPU float32 for compatibility
-        frames = frames.cpu().float()
+        image = image.cpu().float()
 
         # Update batch with decoded image
-        batch.output = frames
+        batch.output = image
 
         # Offload models if needed
         if hasattr(self, 'maybe_free_model_hooks'):

fastvideo/pipelines/stages/denoising.py

@@ -140,12 +140,11 @@ class DenoisingStage(PipelineStage):
             latents = latents[:, :, rank_in_sp_group, :, :, :]
             batch.latents = latents
             if batch.image_latent is not None:
-                if not fastvideo_args.pipeline_config.ti2v_task and not fastvideo_args.pipeline_config.t2v_as_i2v_task:
-                    image_latent = rearrange(batch.image_latent,
-                                             "b c (n t) h w -> b c n t h w",
-                                             n=sp_world_size).contiguous()
-                    image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
-                    batch.image_latent = image_latent
+                image_latent = rearrange(batch.image_latent,
+                                         "b c (n t) h w -> b c n t h w",
+                                         n=sp_world_size).contiguous()
+                image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
+                batch.image_latent = image_latent
         # Get timesteps and calculate warmup steps
         timesteps = batch.timesteps
         # TODO(will): remove this once we add input/output validation for stages
@@ -205,14 +204,8 @@ 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
-        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
+        if fastvideo_args.boundary_ratio is not None:
+            boundary_timestep = fastvideo_args.boundary_ratio * self.scheduler.num_train_timesteps
         else:
             boundary_timestep = None
         latent_model_input = latents.to(target_dtype)
@@ -254,9 +247,6 @@ class DenoisingStage(PipelineStage):
                                                              patch_size[2])
             seq_len = int(math.ceil(seq_len / sp_world_size)) * sp_world_size
 
-        trajectory_timesteps: list[int] = []
-        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):
@@ -284,27 +274,14 @@ class DenoisingStage(PipelineStage):
 
                 # Expand latents for I2V
                 latent_model_input = latents.to(target_dtype)
-                if batch.image_latent is not None and not fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                if batch.image_latent is not None:
                     assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
                     latent_model_input = torch.cat(
                         [latent_model_input, batch.image_latent],
                         dim=1).to(target_dtype)
-                elif batch.image_latent is not None and fastvideo_args.pipeline_config.t2v_as_i2v_task:
-                    assert batch.image_latent is not None, "image latents should be provided for T2V to I2V task"
-                    if rank_in_sp_group == 0:
-                        logger.info("latent_model_input.shape: %s",
-                                    latent_model_input.shape)
-                        latent_model_input = torch.cat([
-                            batch.image_latent,
-                            latent_model_input[:, :, 1:, :, :],
-                        ],
-                                                       dim=2).to(target_dtype)
-                        logger.info("latent_model_input.shape: %s",
-                                    latent_model_input.shape)
 
                 assert not torch.isnan(
                     latent_model_input).any(), "latent_model_input contains nan"
-
                 if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
                     timestep = torch.stack([t]).to(get_local_torch_device())
                     temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
@@ -319,13 +296,6 @@ class DenoisingStage(PipelineStage):
 
                 latent_model_input = self.scheduler.scale_model_input(
                     latent_model_input, t)
-                if fastvideo_args.pipeline_config.t2v_as_i2v_task:
-                    if rank_in_sp_group == 0:
-                        latent_model_input = torch.cat([
-                            batch.image_latent,
-                            latent_model_input[:, :, 1:, :, :],
-                        ],
-                                                       dim=2).to(target_dtype)
 
                 # Prepare inputs for transformer
                 guidance_expand = (
@@ -457,11 +427,6 @@ 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
@@ -469,35 +434,9 @@ 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:
-            batch.trajectory_timesteps = torch.tensor(trajectory_timesteps).cpu()
-            batch.trajectory_latents = trajectory_tensor.cpu()
-
-        if fastvideo_args.pipeline_config.t2v_as_i2v_task:
-            latents = torch.cat([
-                batch.image_latent,
-                latents[:, :, 1:, :, :],
-            ],
-                                dim=2)
 
         # Update batch with final latents
         batch.latents = latents

fastvideo/pipelines/stages/image_encoding.py

@@ -105,81 +105,6 @@ class ImageVAEEncodingStage(PipelineStage):
     def __init__(self, vae: ParallelTiledVAE) -> None:
         self.vae: ParallelTiledVAE = vae
 
-    def encode_image(self,
-                     image: PIL.Image.Image,
-                     height: int,
-                     width: int,
-                     fastvideo_args: FastVideoArgs,
-                     generator: torch.Generator | None = None) -> torch.Tensor:
-        """
-        Encode image into latent space.
-        """
-        image = self.preprocess(
-            image,
-            vae_scale_factor=self.vae.spatial_compression_ratio,
-            height=height,
-            width=width).to(get_local_torch_device(), dtype=torch.float32)
-
-        # (B, C, H, W) -> (B, C, 1, H, W)
-        print(f"image.shape: {image.shape}")
-        image = image.unsqueeze(2)
-        print(f"after unsqueeze image.shape: {image.shape}")
-        return self.encode_tensor(image, fastvideo_args, generator)
-
-    def encode_tensor(self,
-                      video_condition: torch.Tensor,
-                      fastvideo_args: FastVideoArgs,
-                      generator: torch.Generator | None = None) -> torch.Tensor:
-        """
-        Encode frames into latent space.
-        """
-        self.vae = self.vae.to(get_local_torch_device())
-        video_condition = video_condition.to(device=get_local_torch_device(),
-                                             dtype=torch.float32)
-
-        # 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
-
-        # Encode Image
-        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:
-                video_condition = video_condition.to(vae_dtype)
-            encoder_output = self.vae.encode(video_condition)
-
-        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
-            latent_condition = encoder_output.mean
-        else:
-            generator = generator
-            if generator is None:
-                raise ValueError("Generator must be provided")
-            latent_condition = self.retrieve_latents(encoder_output, generator)
-
-        # 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):
-                latent_condition -= self.vae.shift_factor.to(
-                    latent_condition.device, latent_condition.dtype)
-            else:
-                latent_condition -= self.vae.shift_factor
-
-        if isinstance(self.vae.scaling_factor, torch.Tensor):
-            latent_condition = latent_condition * self.vae.scaling_factor.to(
-                latent_condition.device, latent_condition.dtype)
-        else:
-            latent_condition = latent_condition * self.vae.scaling_factor
-
-        return latent_condition
-
     def forward(
         self,
         batch: ForwardBatch,
@@ -232,28 +157,57 @@ class ImageVAEEncodingStage(PipelineStage):
         # (B, C, H, W) -> (B, C, 1, H, W)
         image = image.unsqueeze(2)
 
-        if fastvideo_args.pipeline_config.t2v_as_i2v_task:
-            # repeat the image self.vae.temporal_compression_ratio times
-            video_condition = image.repeat(1, 1,
-                                           self.vae.temporal_compression_ratio,
-                                           1, 1)
-            # video_condition = image
-            logger.info("video_condition.shape: %s", video_condition.shape)
-        else:
-            video_condition = torch.cat([
-                image,
-                image.new_zeros(image.shape[0], image.shape[1], num_frames - 1,
-                                image.shape[3], image.shape[4])
-            ],
-                                        dim=2)
+        video_condition = torch.cat([
+            image,
+            image.new_zeros(image.shape[0], image.shape[1], num_frames - 1,
+                            image.shape[3], image.shape[4])
+        ],
+                                    dim=2)
+        video_condition = video_condition.to(device=get_local_torch_device(),
+                                             dtype=torch.float32)
 
-        latent_condition = self.encode_tensor(video_condition, fastvideo_args,
-                                              batch.generator)
+        # 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
+
+        # Encode Image
+        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:
+                video_condition = video_condition.to(vae_dtype)
+            encoder_output = self.vae.encode(video_condition)
+
+        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
+            latent_condition = encoder_output.mean
+        else:
+            generator = batch.generator
+            if generator is None:
+                raise ValueError("Generator must be provided")
+            latent_condition = self.retrieve_latents(encoder_output, generator)
+
+        # 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):
+                latent_condition -= self.vae.shift_factor.to(
+                    latent_condition.device, latent_condition.dtype)
+            else:
+                latent_condition -= self.vae.shift_factor
+
+        if isinstance(self.vae.scaling_factor, torch.Tensor):
+            latent_condition = latent_condition * self.vae.scaling_factor.to(
+                latent_condition.device, latent_condition.dtype)
+        else:
+            latent_condition = latent_condition * self.vae.scaling_factor
 
         if fastvideo_args.mode == ExecutionMode.PREPROCESS:
-            batch.image_latent = latent_condition
-        elif fastvideo_args.pipeline_config.t2v_as_i2v_task:
-            logger.info("latent_condition.shape: %s", latent_condition.shape)
             batch.image_latent = latent_condition
         else:
             mask_lat_size = torch.ones(1, 1, num_frames, latent_height,

fastvideo/pipelines/stages/input_validation.py

@@ -35,15 +35,9 @@ class InputValidationStage(PipelineStage):
         """Generate seeds for the inference"""
         seed = batch.seed
         num_videos_per_prompt = batch.num_videos_per_prompt
-        if isinstance(batch.prompt, list):
-            num_prompts = len(batch.prompt)
-        else:
-            num_prompts = 1
-
-        total_num_videos = num_prompts * num_videos_per_prompt
 
         assert seed is not None
-        seeds = [seed + i for i in range(total_num_videos)]
+        seeds = [seed + i for i in range(num_videos_per_prompt)]
         batch.seeds = seeds
         # Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
         batch.generator = [

fastvideo/pipelines/stages/latent_preparation.py

@@ -3,7 +3,6 @@
 Latent preparation stage for diffusion pipelines.
 """
 
-import torch
 from diffusers.utils.torch_utils import randn_tensor
 
 from fastvideo.distributed import get_local_torch_device
@@ -94,7 +93,7 @@ class LatentPreparationStage(PipelineStage):
         # Generate or use provided latents
         if latents is None:
             latents = randn_tensor(shape,
-                                   generator=torch.Generator(device="cuda").manual_seed(1024),
+                                   generator=generator,
                                    device=device,
                                    dtype=dtype)
         else:

fastvideo/platforms/__init__.py

@@ -82,8 +82,8 @@ def rocm_platform_plugin() -> str | None:
                 logger.info("ROCm platform is available")
         finally:
             amdsmi.amdsmi_shut_down()
-    except Exception:
-        pass
+    except Exception as e:
+        logger.info("ROCm platform is unavailable: %s", e)
 
     return "fastvideo.platforms.rocm.RocmPlatform" if is_rocm else None
 

fastvideo/tests/dataset/test_parquet_io.py

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

@@ -1,123 +0,0 @@
-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/encoders/test_t5_encoder.py

@@ -8,9 +8,6 @@ from torch.distributed.tensor import DTensor
 from torch.testing import assert_close
 from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
 
-from fastvideo.wan.modules.tokenizers import HuggingfaceTokenizer
-from fastvideo.wan.modules.t5 import umt5_xxl
-
 from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
@@ -41,22 +38,9 @@ def test_t5_encoder():
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     precision_str = "fp32"
     precision = PRECISION_TO_TYPE[precision_str]
-    # model1 = UMT5EncoderModel.from_pretrained(TEXT_ENCODER_PATH).to(
-    #     precision).to(device).eval()
-    # tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
-    model1 = umt5_xxl(
-        encoder_only=True,
-        return_tokenizer=False,
-        dtype=torch.float32,
-        device=device,
-    ).eval().requires_grad_(False)
-    model1.load_state_dict(
-        torch.load("/mnt/weka/home/hao.zhang/wei/Self-Forcing-clean/wan_models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",
-                    map_location='cpu', weights_only=False)
-    )
-
-    tokenizer1 = HuggingfaceTokenizer(
-        name="/mnt/weka/home/hao.zhang/wei/Self-Forcing-clean/wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/", seq_len=512, clean='whitespace')
+    model1 = UMT5EncoderModel.from_pretrained(TEXT_ENCODER_PATH).to(
+        precision).to(device).eval()
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
 
 
     args = FastVideoArgs(model_path=TEXT_ENCODER_PATH,
@@ -65,9 +49,8 @@ def test_t5_encoder():
                         pin_cpu_memory=False)
     loader = TextEncoderLoader()
     model2 = loader.load(TEXT_ENCODER_PATH, args)
-    model2 = model2.to(dtype=torch.bfloat16).to(precision)
+    model2 = model2.to(precision)
     model2.eval()
-    tokenizer2 = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
 
     # Sanity check weights between the two models
     logger.info("Comparing model weights for sanity check...")
@@ -78,13 +61,8 @@ def test_t5_encoder():
     logger.info("Model1 has %s parameters", len(params1))
     logger.info("Model2 has %s parameters", len(params2))
 
-    model1_weight_sum = sum(p.float().sum().item() for p in model1.parameters())
-    model2_weight_sum = sum(p.float().sum().item() for p in model2.parameters())
-    logger.info("Model1 weight sum: %s", model1_weight_sum)
-    logger.info("Model2 weight sum: %s", model2_weight_sum)
-
-    # weight_diffs = []
-    # # check if embed_tokens are the same
+    weight_diffs = []
+    # check if embed_tokens are the same
     weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.{}.layer.0.SelfAttention.relative_attention_bias.weight", \
                "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\
                 "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \
@@ -92,21 +70,18 @@ def test_t5_encoder():
     
     for idx in range(hf_config.num_hidden_layers):
         for w in weights:
-            # name1 = w.format(idx)
+            name1 = w.format(idx)
             name2 = w.format(idx)
-            # p1 = params1[name1]
+            p1 = params1[name1]
             p2 = params2[name2]
-            p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(device)
-            # assert_close(p1, p2, atol=1e-4, rtol=1e-4)
+            p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1)
+            assert_close(p1, p2, atol=1e-4, rtol=1e-4)
     
 
     # Test with some sample prompts
-    # prompts = [
-    #     "Once upon a time", "The quick brown fox jumps over",
-    #     "In a galaxy far, far away"
-    # ]
     prompts = [
-        "A vibrant scene of Kenyan golfers at a lush green golf course on a sunny day. The golfers, dressed in casual yet stylish attire, are teeing off with animated expressions, showcasing their enthusiasm for the game. Rolling hills and pristine greens stretch out behind them, creating a picturesque backdrop. In the foreground, a golf buggy and a caddy stand ready, adding to the serene atmosphere. The camera captures the action from a mid-shot angle, focusing on the golfers' dynamic motions as they swing their clubs."
+        "Once upon a time", "The quick brown fox jumps over",
+        "In a galaxy far, far away"
     ]
 
     logger.info("Testing T5 encoder with sample prompts")
@@ -116,8 +91,7 @@ def test_t5_encoder():
             logger.info("Testing prompt: %s", prompt)
 
             # Tokenize the prompt
-            tokens1, mask = tokenizer1(prompt, return_mask=True, add_special_tokens=True)
-            tokens2 = tokenizer2(prompt,
+            tokens = tokenizer(prompt,
                                padding="max_length",
                                max_length=512,
                                truncation=True,
@@ -127,23 +101,22 @@ def test_t5_encoder():
             # filter out padding input_ids
             # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1]
             # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1]
-            outputs1 = model1(tokens1.to(device),
-                              mask.to(device))
+            outputs1 = model1(input_ids=tokens.input_ids,
+                              attention_mask=tokens.attention_mask,
+                              output_hidden_states=True).last_hidden_state
             print("--------------------------------")
             logger.info("Testing model2")
 
             # Get outputs from our implementation
             with set_forward_context(current_timestep=0, attn_metadata=None):
                 outputs2 = model2(
-                    input_ids=tokens2.input_ids,
-                    attention_mask=tokens2.attention_mask,
+                    input_ids=tokens.input_ids,
+                    attention_mask=tokens.attention_mask,
                 ).last_hidden_state
 
             # Compare last hidden states
-            last_hidden_state1 = outputs1[mask == 1]
-            last_hidden_state2 = outputs2[tokens2.attention_mask == 1]
-            logger.info("last_hidden_state1 sum: %s", last_hidden_state1.float().sum())
-            logger.info("last_hidden_state2 sum: %s", last_hidden_state2.float().sum())
+            last_hidden_state1 = outputs1[tokens.attention_mask == 1]
+            last_hidden_state2 = outputs2[tokens.attention_mask == 1]
 
             assert last_hidden_state1.shape == last_hidden_state2.shape, \
                 f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}"

fastvideo/tests/modal/pr_test.py

@@ -117,7 +117,3 @@ 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:1", image=image, timeout=900)
-def run_unit_test():
-    run_test("pytest ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ -vs")

fastvideo/tests/transformers/test_ori_wanvideo.py

@@ -22,7 +22,6 @@ logger = init_logger(__name__)
 os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
-os.environ["FASTVIDEO_FORCE_ATTN_BF16"] = "1"
 BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
                                   local_dir=os.path.join(
@@ -33,21 +32,19 @@ TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
 @pytest.mark.usefixtures("distributed_setup")
 def test_ori_wan_transformer():
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.float32
-    precision_str = "fp32"
+    precision = torch.bfloat16
+    precision_str = "bf16"
     args = FastVideoArgs(model_path=TRANSFORMER_PATH,
                          dit_cpu_offload=True,
-                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str, dit_forward_precision=precision_str))
+                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
     args.device = device
 
     loader = TransformerLoader()
-    model2 = loader.load(TRANSFORMER_PATH, args)
+    model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision)
 
     model1 = WanModel.from_pretrained(
-        "/mnt/weka/home/hao.zhang/wei/Self-Forcing-clean/wan_models/Wan2.1-T2V-1.3B")
-    model1.eval()
-    model1 = model1.to(device).to(precision)
-    model1.requires_grad_(False)
+        "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
 
     total_params = sum(p.numel() for p in model1.parameters())
     # Calculate weight sum for model1 (converting to float64 to avoid overflow)
@@ -78,35 +75,29 @@ def test_ori_wan_transformer():
 
     # Create identical inputs for both models
     batch_size = 1
-    text_seq_len = 120
-    seq_len = math.ceil((104 * 60) /
+    text_seq_len = 30
+    seq_len = math.ceil((160 * 90) /
                             (2 * 2) *
                             21)
 
     # Video latents [B, C, T, H, W]
     hidden_states = torch.randn(batch_size,
-                                21,
                                 16,
-                                60,
-                                104,
+                                21,
+                                160,
+                                90,
                                 device=device,
-                                generator=torch.Generator("cuda").manual_seed(1024),
                                 dtype=precision)
 
-    logger.info("Hidden states sum: %s, Hidden states shape: %s, Hidden states dtype: %s", hidden_states.float().sum(), hidden_states.shape, hidden_states.dtype)
-
     # Text embeddings [B, L, D] (including global token)
-    # encoder_hidden_states = torch.randn(batch_size,
-    #                                     text_seq_len + 1,
-    #                                     4096,
-    #                                     device=device,
-    #                                     dtype=precision)
-    encoder_hidden_states = torch.load("../sf_cond_prompt_embeds.pt").to(device, dtype=precision)
-
-    logger.info("Encoder hidden states sum: %s, Encoder hidden states shape: %s, Encoder hidden states dtype: %s", encoder_hidden_states.float().sum(), encoder_hidden_states.shape, encoder_hidden_states.dtype)
+    encoder_hidden_states = torch.randn(batch_size,
+                                        text_seq_len + 1,
+                                        4096,
+                                        device=device,
+                                        dtype=precision)
 
     # Timestep
-    timestep = torch.tensor([995.7627], device=device, dtype=precision)
+    timestep = torch.tensor([500], device=device, dtype=precision)
 
     forward_batch = ForwardBatch(
         data_type="dummy",
@@ -114,19 +105,19 @@ def test_ori_wan_transformer():
 
     # with torch.amp.autocast('cuda', dtype=precision):
     output1 = model1(
-        x=hidden_states.permute(0, 2, 1, 3, 4),
+        x=hidden_states,
         context=encoder_hidden_states,
         t=timestep,
         seq_len=seq_len,
-    ).permute(0, 2, 1, 3, 4)
+    )
     with set_forward_context(
             current_timestep=0,
             attn_metadata=None,
             forward_batch=forward_batch,
     ):
-        output2 = model2(hidden_states=hidden_states.permute(0, 2, 1, 3, 4),
+        output2 = model2(hidden_states=hidden_states,
                             encoder_hidden_states=encoder_hidden_states,
-                            timestep=timestep).permute(0, 2, 1, 3, 4)
+                            timestep=timestep)
 
     # Check if outputs have the same shape
     assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
@@ -137,8 +128,6 @@ def test_ori_wan_transformer():
     mean_diff = torch.mean(torch.abs(output1 - output2))
     logger.info("Max Diff: %s", max_diff.item())
     logger.info("Mean Diff: %s", mean_diff.item())
-    logger.info("Output 1 sum: %s", output1.float().sum())
-    logger.info("Output 2 sum: %s", output2.float().sum())
     assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}"
     # mean diff
     assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}"

fastvideo/tests/workflow/test_components_parquet_saver.py

@@ -1,66 +0,0 @@
-from pathlib import Path
-
-import numpy as np
-import pyarrow as pa
-import pyarrow.parquet as pq
-import torch
-
-from fastvideo.workflow.preprocess.components import ParquetDatasetSaver
-from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
-
-
-def _simple_record_creator(batch: PreprocessBatch) -> list[dict]:
-    # batch.latents will be converted to numpy by the saver before this call
-    assert isinstance(batch.latents, np.ndarray)
-    num = len(batch.video_file_name)
-    records = []
-    for i in range(num):
-        arr = batch.latents[i]
-        records.append({
-            "id": batch.video_file_name[i],
-            "data_bytes": arr.tobytes(),
-            "data_shape": list(arr.shape),
-        })
-    return records
-
-
-def test_parquet_dataset_saver_flush_and_last(tmp_path: Path):
-    # Schema for the simple record creator
-    schema = pa.schema([
-        pa.field("id", pa.string()),
-        pa.field("data_bytes", pa.binary()),
-        pa.field("data_shape", pa.list_(pa.int64())),
-    ])
-
-    B = 5
-    # Build a minimal PreprocessBatch
-    batch = PreprocessBatch(
-        data_type="video",
-        latents=torch.randn(B, 2),
-        prompt_embeds=[torch.randn(B, 1, 1)],
-        # Attention mask should be integer dtype in real pipelines
-        prompt_attention_mask=[torch.ones(B, 1, dtype=torch.int64)],
-    )
-    batch.video_file_name = [f"vid_{i}" for i in range(B)]
-
-    saver = ParquetDatasetSaver(
-        flush_frequency=10,  # higher than B to avoid auto-flush
-        samples_per_file=3,
-        schema=schema,
-        record_creator=_simple_record_creator,
-    )
-
-    out_dir = tmp_path / "saver_out"
-    saver.save_and_write_parquet_batch(batch, str(out_dir))
-    # First flush: should write one full file (3 rows), keep 2 in buffer
-    saver.flush_tables()
-    files = sorted(out_dir.rglob("*.parquet"))
-    assert len(files) == 1
-    assert pq.read_table(str(files[0])).num_rows == 3
-
-    # Final flush: write remainder 2 rows
-    saver.flush_tables(write_remainder=True)
-    files2 = sorted(out_dir.rglob("*.parquet"))
-    assert len(files2) == 2
-    total = sum(pq.read_table(str(f)).num_rows for f in files2)
-    assert total == 5

fastvideo/training/distillation_pipeline.py

@@ -40,7 +40,7 @@ from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
     EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
     get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint,
-    shift_timestep)
+    shift_timestep, compute_density_for_timestep_sampling, get_sigmas)
 from fastvideo.utils import (is_vsa_available, maybe_download_model,
                              set_random_seed, verify_model_config_and_directory)
 
@@ -91,6 +91,9 @@ class DistillationPipeline(TrainingPipeline):
         self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
             shift=self.timestep_shift)
 
+        self.transformer_2 = self.get_module("transformer_2", None)
+        self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
+        
         if training_args.real_score_model_path:
             logger.info(
                 f"Loading real score transformer from: {training_args.real_score_model_path}"
@@ -127,6 +130,39 @@ class DistillationPipeline(TrainingPipeline):
                 self.real_score_transformer,
                 checkpointing_type=training_args.
                 enable_gradient_checkpointing_type)
+            if self.transformer_2 is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2,
+                    checkpointing_type=training_args.
+                    enable_gradient_checkpointing_type)
+
+        if self.transformer_2 is not None:
+            self.transformer_2.train()
+            self.transformer_2.requires_grad_(True)
+            params_to_optimize_2 = self.transformer_2.parameters()
+            params_to_optimize_2 = list(
+                filter(lambda p: p.requires_grad, params_to_optimize_2))
+            
+            betas_str = training_args.betas
+            betas = tuple(float(x.strip()) for x in betas_str.split(","))
+            
+            self.optimizer_2 = torch.optim.AdamW(
+                params_to_optimize_2,
+                lr=training_args.learning_rate,
+                betas=betas,
+                weight_decay=training_args.weight_decay,
+                eps=1e-8,
+            )
+            self.lr_scheduler_2 = get_scheduler(
+                training_args.lr_scheduler,
+                optimizer=self.optimizer_2,
+                num_warmup_steps=training_args.lr_warmup_steps,
+                num_training_steps=training_args.max_train_steps,
+                num_cycles=training_args.lr_num_cycles,
+                power=training_args.lr_power,
+                min_lr_ratio=training_args.min_lr_ratio,
+                last_epoch=self.init_steps - 1,
+            )
 
         # Initialize optimizers
         fake_score_params = list(
@@ -281,6 +317,9 @@ class DistillationPipeline(TrainingPipeline):
         """Prepare training environment for distillation."""
         self.transformer.requires_grad_(True)
         self.transformer.train()
+        if self.transformer_2 is not None:
+            self.transformer_2.requires_grad_(True)
+            self.transformer_2.train()
         self.fake_score_transformer.requires_grad_(True)
         self.fake_score_transformer.train()
 
@@ -436,7 +475,9 @@ class DistillationPipeline(TrainingPipeline):
         training_batch = self._build_distill_input_kwargs(
             noisy_latent, timestep, training_batch.conditional_dict,
             training_batch)
-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+        
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
+        pred_noise = current_model(**training_batch.input_kwargs).permute(
             0, 2, 1, 3, 4)
         pred_video = pred_noise_to_pred_video(
             pred_noise=pred_noise.flatten(0, 1),
@@ -478,6 +519,7 @@ class DistillationPipeline(TrainingPipeline):
         max_target_idx = len(self.denoising_step_list) - 1
         noise_latents = []
         noise_latent_index = target_timestep_idx_int - 1
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
         if max_target_idx > 0:
             # Run student model for all steps before the target timestep
             with torch.no_grad():
@@ -489,7 +531,7 @@ class DistillationPipeline(TrainingPipeline):
                     training_batch_temp = self._build_distill_input_kwargs(
                         current_noise_latents, current_timestep_tensor,
                         training_batch.conditional_dict, training_batch)
-                    pred_flow = self.transformer(
+                    pred_flow = current_model(
                         **training_batch_temp.input_kwargs).permute(
                             0, 2, 1, 3, 4)
                     pred_clean = pred_noise_to_pred_video(
@@ -532,7 +574,7 @@ class DistillationPipeline(TrainingPipeline):
         training_batch = self._build_distill_input_kwargs(
             noisy_input, target_timestep, training_batch.conditional_dict,
             training_batch)
-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+        pred_noise = current_model(**training_batch.input_kwargs).permute(
             0, 2, 1, 3, 4)
         pred_video = pred_noise_to_pred_video(
             pred_noise=pred_noise.flatten(0, 1),
@@ -730,11 +772,12 @@ class DistillationPipeline(TrainingPipeline):
             "encoder_hidden_states": training_batch.encoder_hidden_states,
             "encoder_attention_mask": training_batch.encoder_attention_mask,
         }
-        unconditional_dict = {
-            "encoder_hidden_states": self.negative_prompt_embeds,
-            "encoder_attention_mask": self.negative_prompt_attention_mask,
-        }
-        training_batch.unconditional_dict = unconditional_dict
+        if getattr(self, "negative_prompt_embeds", None) is not None:
+            unconditional_dict = {
+                "encoder_hidden_states": self.negative_prompt_embeds,
+                "encoder_attention_mask": self.negative_prompt_attention_mask,
+            }
+            training_batch.unconditional_dict = unconditional_dict
 
         training_batch.dmd_latent_vis_dict = {}
         training_batch.fake_score_latent_vis_dict = {}
@@ -773,6 +816,9 @@ class DistillationPipeline(TrainingPipeline):
             batches.append(batch)
 
         self.optimizer.zero_grad()
+        # TODO: confirm this
+        if self.transformer_2 is not None:
+            self.optimizer_2.zero_grad()
         total_dmd_loss = 0.0
         dmd_latent_vis_dict = {}
         fake_score_latent_vis_dict = {}
@@ -801,15 +847,31 @@ class DistillationPipeline(TrainingPipeline):
                         attn_metadata=batch_gen.attn_metadata_vsa):
                     (dmd_loss / gradient_accumulation_steps).backward()
                 total_dmd_loss += dmd_loss.detach().item()
-            self._clip_model_grad_norm_(batch_gen, self.transformer)
-            for param in self.transformer.parameters():
-                # check if the gradient is not None and not zero
-                assert param.grad is not None and param.grad.abs().sum() > 0
-            self.optimizer.step()
-            self.optimizer.zero_grad(set_to_none=True)
+            
+            # Only clip gradients for the model that is currently training
+            if self.train_transformer_2 and self.transformer_2 is not None:
+                self._clip_model_grad_norm_(batch_gen, self.transformer_2)
+                for param in self.transformer_2.parameters():
+                    # check if the gradient is not None and not zero
+                    assert param.grad is not None and param.grad.abs().sum() > 0
+                self.optimizer_2.step()
+                self.optimizer_2.zero_grad(set_to_none=True)
+            else:
+                self._clip_model_grad_norm_(batch_gen, self.transformer)
+                for param in self.transformer.parameters():
+                    # check if the gradient is not None and not zero
+                    assert param.grad is not None and param.grad.abs().sum() > 0
+                self.optimizer.step()
+                self.optimizer.zero_grad(set_to_none=True)
 
             if self.generator_ema is not None:
-                self.generator_ema.update(self.transformer)
+                # TODO: support EMA for transformer_2?
+                if self.train_transformer_2 and self.transformer_2 is not None:
+                    # Note: EMA currently only supports the main transformer
+                    # Could be extended to support transformer_2 in the future
+                    pass
+                else:
+                    self.generator_ema.update(self.transformer)
 
             avg_dmd_loss = torch.tensor(total_dmd_loss /
                                         gradient_accumulation_steps,
@@ -839,7 +901,13 @@ class DistillationPipeline(TrainingPipeline):
             assert param.grad is not None and param.grad.abs().sum() > 0
         self.fake_score_optimizer.step()
         self.fake_score_lr_scheduler.step()
-        self.lr_scheduler.step()
+        
+        # Step the appropriate scheduler
+        if self.train_transformer_2 and self.transformer_2 is not None:
+            self.lr_scheduler_2.step()
+        else:
+            self.lr_scheduler.step()
+            
         self.fake_score_optimizer.zero_grad(set_to_none=True)
         avg_fake_score_loss = torch.tensor(total_fake_score_loss /
                                            gradient_accumulation_steps,

fastvideo/training/ode_causal_pipeline.py

@@ -1,492 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import sys
-from copy import deepcopy
-from typing import cast
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-import wandb
-from fastvideo.dataset.dataloader.schema import (
-    pyarrow_schema_ode_trajectory_text_only)
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (
-    SelfForcingFlowMatchScheduler)
-from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import (
-    WanCausalDMDPipeline)
-from fastvideo.pipelines.pipeline_batch_info import TrainingBatch
-from fastvideo.training.training_pipeline import TrainingPipeline
-from fastvideo.training.training_utils import (
-    clip_grad_norm_while_handling_failing_dtensor_cases)
-
-logger = init_logger(__name__)
-
-
-class ODEInitTrainingPipeline(TrainingPipeline):
-    """
-    Training pipeline for ODE-init using precomputed denoising trajectories.
-
-    Supervision: predict the next latent in the stored trajectory by
-    - feeding current latent at timestep t into the transformer to predict noise
-    - stepping the scheduler with the predicted noise
-    - minimizing MSE to the stored next latent at timestep t_next
-    """
-
-    _required_config_modules = ["scheduler", "transformer", "vae"]
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        # Match the preprocess/generation scheduler for consistent stepping
-        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=1000,
-                                                training=True)
-
-    def set_schemas(self):
-        self.train_dataset_schema = pyarrow_schema_ode_trajectory_text_only
-
-    def initialize_training_pipeline(self, training_args: TrainingArgs):
-        super().initialize_training_pipeline(training_args)
-
-        self.noise_scheduler = self.get_module("scheduler")
-        self.vae = self.get_module("vae")
-        self.vae.requires_grad_(False)
-
-        self.timestep_shift = self.training_args.pipeline_config.flow_shift
-        assert self.timestep_shift == 5.0, "flow_shift must be 5.0"
-        self.noise_scheduler = SelfForcingFlowMatchScheduler(
-            shift=self.timestep_shift, sigma_min=0.0, extra_one_step=True)
-        self.noise_scheduler.set_timesteps(num_inference_steps=1000,
-                                           training=True)
-
-        # logger.info(f"ARG dmd_denoising_steps: {training_args.pipeline_config.dmd_denoising_steps}")
-        logger.info(
-            f"ARG dmd_denoising_steps: {self.training_args.pipeline_config.dmd_denoising_steps}"
-        )
-        self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250],
-                                                dtype=torch.long,
-                                                device=get_local_torch_device())
-        # self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], dtype=torch.long, device=get_local_torch_device())
-        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
-            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
-                                   torch.tensor([0],
-                                                dtype=torch.float32))).cuda()
-            logger.info(f"timesteps: {timesteps}")
-            self.dmd_denoising_steps = timesteps[1000 -
-                                                 self.dmd_denoising_steps]
-            logger.info(
-                f"warped self.dmd_denoising_steps: {self.dmd_denoising_steps}")
-            # assert False, "warp_denoising_step must be false"
-        else:
-            assert False, "warp_denoising_step must be true"
-            logger.info("not warped")
-        self.dmd_denoising_steps = self.dmd_denoising_steps.to(
-            get_local_torch_device())
-
-        logger.info(f"denoising_step_list: {self.dmd_denoising_steps}")
-
-        logger.info(
-            "Initialized ODE-init training pipeline with %s denoising steps",
-            len(self.dmd_denoising_steps))
-        # Cache for nearest trajectory index per DMD step (computed lazily on first batch)
-        self._cached_closest_idx_per_dmd = None
-        self.num_train_timestep = self.noise_scheduler.num_train_timesteps
-        # self.min_timestep = int(self.training_args.min_timestep_ratio *
-        #                         self.num_train_timestep)
-        # self.max_timestep = int(self.training_args.max_timestep_ratio *
-        #                         self.num_train_timestep)
-        # self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
-
-    def initialize_validation_pipeline(self, training_args: TrainingArgs):
-        logger.info("Initializing validation pipeline...")
-        args_copy = deepcopy(training_args)
-        args_copy.inference_mode = True
-        # Warm start validation with current transformer
-        self.validation_pipeline = WanCausalDMDPipeline.from_pretrained(
-            # training_args.model_path,
-            "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
-            args=args_copy,  # type: ignore
-            inference_mode=True,
-            loaded_modules={
-                "transformer": self.get_module("transformer"),
-            },
-            tp_size=training_args.tp_size,
-            sp_size=training_args.sp_size,
-            num_gpus=training_args.num_gpus,
-            pin_cpu_memory=training_args.pin_cpu_memory,
-            dit_cpu_offload=True)
-
-    def _get_next_batch(self, training_batch):  # type: ignore[override]
-        # batch = next(self.train_loader_iter, None)  # type: ignore
-        # if batch is None:
-        #     self.current_epoch += 1
-        #     logger.info("Starting epoch %s", self.current_epoch)
-        #     self.train_loader_iter = iter(self.train_dataloader)
-        #     batch = next(self.train_loader_iter)
-
-        # # Required fields from parquet (ODE trajectory schema)
-        # encoder_hidden_states = batch['text_embedding']
-        # encoder_attention_mask = batch['text_attention_mask']
-        # infos = batch['info_list']
-
-        # # Trajectory tensors may include a leading singleton batch dim per row
-        # trajectory_latents = batch['trajectory_latents']
-        # if trajectory_latents.dim() == 7:
-        #     # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W]
-        #     trajectory_latents = trajectory_latents[:, 0]
-        # elif trajectory_latents.dim() == 6:
-        #     # already [B, S, C, T, H, W]
-        #     pass
-        # else:
-        #     raise ValueError(
-        #         f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}"
-        #     )
-
-        # trajectory_timesteps = batch['trajectory_timesteps']
-        # if trajectory_timesteps.dim() == 3:
-        #     # [B, 1, S] -> [B, S]
-        #     trajectory_timesteps = trajectory_timesteps[:, 0]
-        # elif trajectory_timesteps.dim() == 2:
-        #     # [B, S]
-        #     pass
-        # else:
-        #     raise ValueError(
-        #         f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}"
-        #     )
-        # # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
-        # trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)
-
-        # # Move to device
-        device = get_local_torch_device()
-        # training_batch.encoder_hidden_states = encoder_hidden_states.to(
-        #     device, dtype=torch.bfloat16)
-        # training_batch.encoder_attention_mask = encoder_attention_mask.to(
-        #     device, dtype=torch.bfloat16)
-        # training_batch.infos = infos
-
-        # return training_batch, trajectory_latents.to(
-        #     device, dtype=torch.bfloat16), trajectory_timesteps.to(device)
-
-        ## TEMP
-        path = "/mnt/weka/home/hao.zhang/wei/FastVideo/data/ode_vidprom_1_fv/00000.pt"
-        b = torch.load(path)
-        for k, v in b.items():
-            logger.info(f"b[{k}]: {type(v)}")
-            if isinstance(v, torch.Tensor):
-                logger.info(f"b[{k}]: {v.shape}")
-            else:
-                logger.info(f"b[{k}]: {v}")
-        training_batch.encoder_hidden_states = b["text_embedding"][0].unsqueeze(
-            0).to(device, dtype=torch.bfloat16)
-        trajectory_latents = b["ode_latent"].to(device, dtype=torch.bfloat16)
-        logger.info(f"trajectory_latents: {trajectory_latents.shape}")
-        logger.info(
-            f"encoder_hidden_states: {training_batch.encoder_hidden_states.shape}"
-        )
-
-        return training_batch, trajectory_latents.to(device,
-                                                     dtype=torch.bfloat16), None
-
-    def _get_timestep(self,
-                      min_timestep: int,
-                      max_timestep: int,
-                      batch_size: int,
-                      num_frame: int,
-                      num_frame_per_block: int,
-                      uniform_timestep: bool = False) -> torch.Tensor:
-        if uniform_timestep:
-            timestep = torch.randint(min_timestep,
-                                     max_timestep, [batch_size, 1],
-                                     device=self.device,
-                                     dtype=torch.long).repeat(1, num_frame)
-            return timestep
-        else:
-            timestep = torch.randint(min_timestep,
-                                     max_timestep, [batch_size, num_frame],
-                                     device=self.device,
-                                     dtype=torch.long)
-            # logger.info(f"individual timestep: {timestep}")
-            # make the noise level the same within every block
-            timestep = timestep.reshape(timestep.shape[0], -1,
-                                        num_frame_per_block)
-            timestep[:, :, 1:] = timestep[:, :, 0:1]
-            timestep = timestep.reshape(timestep.shape[0], -1)
-            return timestep
-
-    def _step_predict_next_latent(
-        self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_attention_mask: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str,
-                                                              torch.Tensor]]:
-        latent_vis_dict = {}
-        device = get_local_torch_device()
-        target_latent = traj_latents[:, -1]
-
-        # logger.info(f"traj_latents: {traj_latents.shape}")
-        # logger.info(f"traj_timesteps: {traj_timesteps.shape}")
-
-        # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
-        B, S, num_frames, num_channels, height, width = traj_latents.shape
-
-        # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps
-        if self._cached_closest_idx_per_dmd is None:
-            # Use the first sample's trajectory timesteps; assumed identical across batches
-            # s_steps = traj_timesteps[0].to(torch.long)  # [S]
-            # dmd = cast(torch.Tensor, self.dmd_denoising_steps).to(s_steps.device)  # [K]
-            # distances_ks: [K, S] = |s_steps - dmd|
-            # distances_ks = (s_steps.unsqueeze(0) - dmd.unsqueeze(1)).abs()
-            # self._cached_closest_idx_per_dmd = distances_ks.argmin(dim=1).to(torch.long).cpu()  # [K]
-            self._cached_closest_idx_per_dmd = torch.tensor(
-                [0, 12, 24, 36], dtype=torch.long).cpu()
-            logger.info(
-                f"self._cached_closest_idx_per_dmd: {self._cached_closest_idx_per_dmd}"
-            )
-            logger.info(
-                f"corresponding timesteps: {self.noise_scheduler.timesteps[self._cached_closest_idx_per_dmd]}"
-            )
-
-        # logger.info(f"traj_latents: {traj_latents.shape}")
-        # Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd
-        # traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K]
-        # Output: [B, K, C, T, H, W]
-        # relevant_traj_latents = torch.index_select(
-        #     traj_latents,
-        #     dim=1,
-        #     index=self._cached_closest_idx_per_dmd.to(traj_latents.device))
-        relevant_traj_latents = traj_latents
-        logger.info(f"relevant_traj_latents sum: {relevant_traj_latents.float().sum().item()}, relevant_traj_latents: {relevant_traj_latents.shape}")
-        # assert relevant_traj_latents.shape[0] == 1
-
-        indexes = self._get_timestep(  # [B, num_frames]
-            0,
-            len(self.dmd_denoising_steps),
-            B,
-            num_frames,
-            3,
-            uniform_timestep=False)
-        logger.info(f"indexes: {indexes.shape}")
-        logger.info(f"indexes: {indexes}")
-        # noisy_input = relevant_traj_latents[indexes]
-        noisy_input = torch.gather(
-            relevant_traj_latents,
-            dim=1,
-            index=indexes.reshape(B, 1, num_frames, 1, 1,
-                                  1).expand(-1, -1, -1, num_channels, height,
-                                            width).to(self.device)).squeeze(1)
-        # noisy_input = noisy_input.unsqueeze(0)
-
-        # # Sample a single DMD step for the whole batch and fetch its cached nearest S-index
-        # K = len(self.dmd_denoising_steps)
-        # dmd_idx = torch.randint(0, K, (1,), device=device)
-        # logger.info(f"dmd_idx: {dmd_idx}")
-        # assert self._cached_closest_idx_per_dmd is not None
-        # nearest_s_idx = int(self._cached_closest_idx_per_dmd[int(dmd_idx.item())])
-        # nearest_idx = torch.full((B,), nearest_s_idx, device=device, dtype=torch.long)
-
-        # batch_indices = torch.arange(B, device=device)
-        # noisy_input = traj_latents[batch_indices, nearest_idx]  # [B, C, T, H, W]
-        # target_latent = traj_latents[batch_indices, -1]  # [B, C, T, H, W]
-        # t = traj_timesteps[batch_indices, nearest_idx]  # [B]
-
-        # Scale model input as in inference for consistency with stored trajectories
-        # noisy_input = self.modules["scheduler"].scale_model_input(noisy_input, t)
-        # logger.info(f"indexes: {indexes.shape}")
-        # logger.info(f"indexes: {indexes}")
-        timestep = self.dmd_denoising_steps[indexes]
-        # logger.info(f"timestep: {timestep.shape}")
-        # logger.info(f"timestep: {timestep}")
-
-        # Prepare inputs for transformer
-        latent_vis_dict["noisy_input"] = noisy_input.permute(
-            0, 2, 1, 3, 4).detach().clone().cpu()
-        latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3,
-                                                      4).detach().clone().cpu()
-
-        logger.info("noisy input sum: %s, noisy input shape: %s, noisy input dtype: %s", noisy_input.float().sum().item(), noisy_input.shape, noisy_input.dtype)
-        logger.info("encoder hidden states sum: %s, encoder hidden states shape: %s, encoder hidden states dtype: %s", encoder_hidden_states.float().sum().item(), encoder_hidden_states.shape, encoder_hidden_states.dtype)
-        logger.info("timestep sum: %s, timestep shape: %s, timestep dtype: %s", timestep.float().sum().item(), timestep.shape, timestep.dtype)
-        logger.info("model dtype set: %s. model weight sum: %s", set(p.dtype for p in self.transformer.parameters()), sum(p.float().sum().item() for p in self.transformer.parameters()))
-
-        # model_dtype = next(self.transformer.parameters()).dtype
-        model_dtype = torch.bfloat16
-        input_kwargs = {
-            "hidden_states": noisy_input.permute(0, 2, 1, 3, 4),
-            "encoder_hidden_states": encoder_hidden_states,
-            "timestep": timestep,
-            "return_dict": False,
-            "scheduler": self.modules["scheduler"]
-        }
-        # Predict noise and step the scheduler to obtain next latent
-        with set_forward_context(current_timestep=timestep,
-                                 attn_metadata=None,
-                                 forward_batch=None):
-            pred_video = self.transformer(**input_kwargs)
-            # logger.info(f"noise_pred: {noise_pred.shape}")
-        # logger.info("noise pred sum: %s, noise pred shape: %s, noise pred dtype: %s", noise_pred.float().sum().item(), noise_pred.shape, noise_pred.dtype)
-        # if isinstance(noise_pred, (tuple, list)):
-        #     noise_pred = noise_pred[0]
-
-        # from fastvideo.models.utils import pred_noise_to_pred_video
-        # pred_video = pred_noise_to_pred_video(
-        #     pred_noise=noise_pred.flatten(0, 1),
-        #     noise_input_latent=noisy_input.flatten(0, 1),
-        #     timestep=timestep.to(dtype=model_dtype).flatten(0, 1),
-        #     scheduler=self.modules["scheduler"]).unflatten(
-        #         0, noise_pred.shape[:2])
-        logger.info("pred video sum: %s, pred video shape: %s, pred video dtype: %s", pred_video.float().sum().item(), pred_video.shape, pred_video.dtype)
-
-        latent_vis_dict["pred_video"] = pred_video.permute(
-            0, 2, 1, 3, 4).detach().clone().cpu()
-
-        # noisy_input = pred_noise_to_pred_video(noise_pred, noisy_input, t, self.modules["scheduler"])
-        # next_latent_pred = self.modules["scheduler"].step(
-        #     noise_pred, t, current_latents, return_dict=False)[0]
-        return pred_video, target_latent, timestep, latent_vis_dict
-
-    def train_one_step(self, training_batch):  # type: ignore[override]
-        self.transformer.train()
-        self.optimizer.zero_grad()
-        training_batch.total_loss = 0.0
-        args = cast(TrainingArgs, self.training_args)
-
-        # Using cached nearest index per DMD step; computation happens in _step_predict_next_latent
-
-        for _ in range(args.gradient_accumulation_steps):
-            training_batch, traj_latents, traj_timesteps = self._get_next_batch(
-                training_batch)
-            text_embeds = training_batch.encoder_hidden_states
-            text_attention_mask = training_batch.encoder_attention_mask
-            assert traj_latents.shape[0] == 1
-
-            # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
-            B, S = traj_latents.shape[0], traj_latents.shape[1]
-            if S < 2:
-                raise ValueError("Trajectory must contain at least 2 steps")
-
-            # Sample per-sample current step i in [0, S-2]
-
-            # idx = torch.randint(low=0, high=S - 1, size=(B, ),
-            #                     device=traj_latents.device)
-
-            # Gather current latents and next latents
-            # batch_indices = torch.arange(B, device=traj_latents.device)
-            # current_latents = traj_latents[batch_indices, idx]  # [B, C, T,H,W]
-            # current_latent = traj_timesteps[:, -1, :, :, :, :]
-            # target_latents = traj_latents[:, -1, :, :, :, :]
-
-            # Corresponding timesteps t (long) -> cast per sample
-            # t = traj_timesteps[:, -1, :, :, :, :]
-            # if t.dtype != torch.long:
-            #     t = t.long()
-
-            # Forward to predict next latent by stepping scheduler with predicted noise
-            noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent(
-                traj_latents, training_batch.current_timestep, text_embeds, text_attention_mask)
-
-            training_batch.latent_vis_dict.update(latent_vis_dict)
-
-            mask = t != 0
-
-            # Compute loss
-            loss = F.mse_loss(noise_pred[mask],
-                              target_latent[mask],
-                              reduction="mean")
-            loss = loss / args.gradient_accumulation_steps
-
-            with set_forward_context(current_timestep=t,
-                                     attn_metadata=None,
-                                     forward_batch=None):
-                loss.backward()
-
-            logger.info("Loss sum: %s, Loss dtype: %s", loss.float().sum().item(), loss.dtype)
-            assert self.transformer.blocks[0].to_q.weight.dtype == torch.float32
-            logger.info("blocks[0].to_q param sum before backprop: %s", self.transformer.blocks[0].to_q.weight.float().sum().item())
-            logger.info("blocks[0].to_q param grad sum: %s", self.transformer.blocks[0].to_q.weight.grad.float().sum().item())
-            logger.info("Transformer grad dtype: %s", set(p.grad.dtype for p in self.transformer.parameters()))
-            logger.info("Transformer param dtype before backprop: %s", set(p.dtype for p in self.transformer.parameters()))
-            avg_loss = loss.detach().clone()
-            training_batch.total_loss += avg_loss.item()
-
-        # Clip grad and step optimizers
-        # grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
-        #     [p for p in self.transformer.parameters() if p.requires_grad],
-        #     args.max_grad_norm if args.max_grad_norm is not None else 0.0)
-        grad_norm = None
-
-        self.optimizer.step()
-        logger.info("Transformer param dtype after backprop: %s", set(p.dtype for p in self.transformer.parameters()))
-        logger.info("blocks[0].to_q param param after backprop sum: %s", self.transformer.blocks[0].to_q.weight.float().sum().item())
-
-        assert False
-        # self.lr_scheduler.step()
-
-        if grad_norm is None:
-            grad_value = 0.0
-        else:
-            try:
-                if isinstance(grad_norm, torch.Tensor):
-                    grad_value = float(grad_norm.detach().float().item())
-                else:
-                    grad_value = float(grad_norm)
-            except Exception:
-                grad_value = 0.0
-        training_batch.grad_norm = grad_value
-        return training_batch
-
-    def visualize_intermediate_latents(self, training_batch: TrainingBatch,
-                                       training_args: TrainingArgs, step: int):
-        """Add visualization data to wandb logging and save frames to disk."""
-        wandb_loss_dict = {}
-        latents_vis_dict = training_batch.latent_vis_dict
-        latent_log_keys = ['noisy_input', 'x0', 'pred_video']
-        for latent_key in latent_log_keys:
-            assert latent_key in latents_vis_dict and latents_vis_dict[
-                latent_key] is not None
-            latent = latents_vis_dict[latent_key]
-            pixel_latent = self.validation_pipeline.decoding_stage.decode(
-                latent, training_args)
-
-            video = pixel_latent.cpu().float()
-            video = video.permute(0, 2, 1, 3, 4)
-            video = (video * 255).numpy().astype(np.uint8)
-            wandb_loss_dict[latent_key] = wandb.Video(
-                video, fps=16, format="mp4")  # change to 16 for Wan2.1
-            # Clean up references
-            del video, pixel_latent, latent
-
-        # Log to wandb
-        if self.global_rank == 0:
-            wandb.log(wandb_loss_dict, step=step)
-
-        # dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
-        # fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
-        # fake_score_log_keys = ['generator_pred_video']
-        # dmd_log_keys = ['faker_score_pred_video', 'real_score_pred_video']
-
-
-def main(args) -> None:
-    logger.info("Starting ODE-init training pipeline...")
-    logger.info(f"ARG dmd_denoising_steps: {args.dmd_denoising_steps}")
-    pipeline = ODEInitTrainingPipeline.from_pretrained(
-        args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("ODE-init training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/training/training_pipeline.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+import gc
 import dataclasses
 import math
 import os
@@ -69,6 +70,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
     train_dataloader: StatefulDataLoader
     train_loader_iter: Iterator[dict[str, Any]]
     current_epoch: int = 0
+    train_transformer_2: bool = False
 
     def __init__(
             self,
@@ -104,6 +106,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         self.sp_world_size = self.sp_group.world_size
         self.local_rank = world_group.local_rank
         self.transformer = self.get_module("transformer")
+        self.transformer_2 = self.get_module("transformer_2", None)
         self.seed = training_args.seed
         self.set_schemas()
 
@@ -116,6 +119,11 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 self.transformer,
                 checkpointing_type=training_args.
                 enable_gradient_checkpointing_type)
+            if self.transformer_2 is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2,
+                    checkpointing_type=training_args.
+                    enable_gradient_checkpointing_type)
 
         noise_scheduler = self.modules["scheduler"]
         self.set_trainable()
@@ -147,6 +155,30 @@ class TrainingPipeline(LoRAPipeline, ABC):
             min_lr_ratio=training_args.min_lr_ratio,
             last_epoch=self.init_steps - 1,
         )
+        if self.transformer_2 is not None:
+            # Ensure transformer_2 has trainable parameters before creating optimizer
+            self.transformer_2.train()
+            self.transformer_2.requires_grad_(True)
+            params_to_optimize_2 = self.transformer_2.parameters()
+            params_to_optimize_2 = list(
+                filter(lambda p: p.requires_grad, params_to_optimize_2))
+            self.optimizer_2 = torch.optim.AdamW(
+                params_to_optimize_2,
+                lr=training_args.learning_rate,
+                betas=(0.9, 0.999),
+                weight_decay=training_args.weight_decay,
+                eps=1e-8,
+            )
+            self.lr_scheduler_2 = get_scheduler(
+                training_args.lr_scheduler,
+                optimizer=self.optimizer_2,
+                num_warmup_steps=training_args.lr_warmup_steps,
+                num_training_steps=training_args.max_train_steps,
+                num_cycles=training_args.lr_num_cycles,
+                power=training_args.lr_power,
+                min_lr_ratio=training_args.min_lr_ratio,
+                last_epoch=self.init_steps - 1,
+            )
 
         self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
             training_args.data_path,
@@ -161,7 +193,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
             seed=self.seed)
 
         self.noise_scheduler = noise_scheduler
-
+        self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
         self.num_update_steps_per_epoch = math.ceil(
             len(self.train_dataloader) /
             training_args.gradient_accumulation_steps * training_args.sp_size /
@@ -187,9 +219,25 @@ class TrainingPipeline(LoRAPipeline, ABC):
     def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
         self.transformer.train()
         self.optimizer.zero_grad()
+        if self.transformer_2 is not None:
+            self.transformer_2.train()
+            self.optimizer_2.zero_grad()
         training_batch.total_loss = 0.0
         return training_batch
 
+    def _enable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        """Enable training mode and gradients for the specified model."""
+        for param in model.parameters():
+            param.requires_grad = True
+        model.train()
+        optimizer.zero_grad()
+
+    def _disable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        """Disable training mode and gradients for the specified model."""
+        for param in model.parameters():
+            param.requires_grad = False
+        optimizer.zero_grad(set_to_none=True)
+
     def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
         batch = next(self.train_loader_iter, None)  # type: ignore
         if batch is None:
@@ -233,17 +281,17 @@ class TrainingPipeline(LoRAPipeline, ABC):
                             generator=self.noise_gen_cuda,
                             device=latents.device,
                             dtype=latents.dtype)
-        u = compute_density_for_timestep_sampling(
-            weighting_scheme=self.training_args.weighting_scheme,
-            batch_size=batch_size,
-            generator=self.noise_random_generator,
-            logit_mean=self.training_args.logit_mean,
-            logit_std=self.training_args.logit_std,
-            mode_scale=self.training_args.mode_scale,
-        )
-        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
-        timesteps = self.noise_scheduler.timesteps[indices].to(
-            device=latents.device)
+        timesteps = self._sample_timesteps(batch_size, latents.device)
+
+        # Enable training for the model that will be trained next and disable the other
+        if self.train_transformer_2:
+            self._enable_training(self.transformer_2, self.optimizer_2)
+            self._disable_training(self.transformer, self.optimizer)
+        else:
+            self._enable_training(self.transformer, self.optimizer)
+            if self.transformer_2 is not None:
+                self._disable_training(self.transformer_2, self.optimizer_2)
+
         if self.training_args.sp_size > 1:
             # Make sure that the timesteps are the same across all sp processes.
             sp_group = get_sp_group()
@@ -266,6 +314,38 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         return training_batch
 
+    def _sample_timesteps(self, batch_size, device):
+        # Determine which model to train based on the boundary timestep
+        if (self.transformer_2 is not None and self.boundary_timestep is not None and
+                torch.rand(1, generator=self.noise_random_generator).item() <= self.training_args.boundary_ratio):
+            self.train_transformer_2 = True
+        else:
+            self.train_transformer_2 = False
+
+        # Broadcast the decision to all processes
+        decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0, device=self.device)
+        dist.broadcast(decision, src=0)
+        self.train_transformer_2 = decision.item() == 1.0
+
+        # Sample u from the appropriate range
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=self.training_args.weighting_scheme,
+            batch_size=batch_size,
+            generator=self.noise_random_generator,
+            logit_mean=self.training_args.logit_mean,
+            logit_std=self.training_args.logit_std,
+            mode_scale=self.training_args.mode_scale,
+        )
+
+        boundary_ratio = self.training_args.boundary_ratio
+        if self.train_transformer_2:
+            u = (1 - boundary_ratio) + u * boundary_ratio # min: 1 - boundary_ratio, max: 1
+        else:
+            u = u * (1 - boundary_ratio) # min: 0, max: 1 - boundary_ratio
+
+        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        return self.noise_scheduler.timesteps[indices].to(device=device)
+        
     def _build_attention_metadata(
             self, training_batch: TrainingBatch) -> TrainingBatch:
         latents_shape = training_batch.raw_latent_shape
@@ -281,20 +361,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 patch_size=patch_size,
                 VSA_sparsity=current_vsa_sparsity,
                 device=get_local_torch_device())
-        # elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
-        #     moba_params = self.training_args.moba_config.copy()
-        #     moba_params.update({
-        #         "current_timestep":
-        #         training_batch.timesteps,
-        #         "raw_latent_shape":
-        #         training_batch.raw_latent_shape[2:5],
-        #         "patch_size":
-        #         self.training_args.pipeline_config.dit_config.patch_size,
-        #         "device":
-        #         get_local_torch_device(),
-        #     })
-        #     training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder(
-        #     ).build(**moba_params)
         else:
             training_batch.attn_metadata = None
 
@@ -319,7 +385,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
     def _transformer_forward_and_compute_loss(
             self, training_batch: TrainingBatch) -> TrainingBatch:
-        # if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
         if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
             assert training_batch.attn_metadata is not None
         else:
@@ -331,11 +396,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
         #         [1000.0],
         #         device=training_batch.noisy_model_input.device,
         #         dtype=torch.bfloat16)
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
 
         with set_forward_context(
                 current_timestep=training_batch.current_timestep,
                 attn_metadata=training_batch.attn_metadata):
-            model_pred = self.transformer(**input_kwargs)
+            model_pred = current_model(**input_kwargs)
             if self.training_args.precondition_outputs:
                 assert training_batch.sigmas is not None
                 model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
@@ -366,7 +432,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
         # the following:
         # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
         if max_grad_norm is not None:
-            model_parts = [self.transformer]
+            # Only clip gradients for the model that is currently training
+            if self.train_transformer_2 and self.transformer_2 is not None:
+                model_parts = [self.transformer_2]
+            else:
+                model_parts = [self.transformer]
+            
             grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
                 [p for m in model_parts for p in m.parameters()],
                 max_grad_norm,
@@ -411,9 +482,14 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         training_batch = self._clip_grad_norm(training_batch)
 
-        self.optimizer.step()
-        self.lr_scheduler.step()
-
+        # Only step the optimizer and scheduler for the model that is currently training
+        if self.train_transformer_2 and self.transformer_2 is not None:
+            self.optimizer_2.step()
+            self.lr_scheduler_2.step()
+        else:
+            self.optimizer.step()
+            self.lr_scheduler.step()
+            
         training_batch.total_loss = training_batch.total_loss
         training_batch.grad_norm = training_batch.grad_norm
         return training_batch
@@ -445,6 +521,11 @@ class TrainingPipeline(LoRAPipeline, ABC):
         logger.info("Starting training with %s B trainable parameters",
                     round(num_trainable_params / 1e9, 3))
 
+        if getattr(self, "transformer_2", None) is not None:
+            num_trainable_params = _get_trainable_params(self.transformer_2)
+            logger.info("Transformer 2: Starting training with %s B trainable parameters",
+                        round(num_trainable_params / 1e9, 3))
+
         # Set random seeds for deterministic training
         self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
             self.seed)
@@ -465,7 +546,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         self._log_training_info()
 
-        self._log_validation(self.transformer, self.training_args,
+        self._log_validation(self.training_args,
                              self.init_steps)
 
         # Train!
@@ -486,9 +567,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 current_decay_times = min(step // vsa_decay_interval_steps,
                                           vsa_sparsity // vsa_decay_rate)
                 current_vsa_sparsity = current_decay_times * vsa_decay_rate
-            # elif vmoba_available:
-            #     # TODO: add vmoba sparsity scheduling here
-            #     current_vsa_sparsity = 0.0
             else:
                 current_vsa_sparsity = 0.0
 
@@ -530,11 +608,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 self.transformer.train()
                 self.sp_group.barrier()
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                if self.training_args.log_visualization:
-                    self.visualize_intermediate_latents(training_batch,
-                                                        self.training_args,
-                                                        step)
-                self._log_validation(self.transformer, self.training_args, step)
+                self._log_validation(self.training_args, step)
                 gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
                 trainable_params = round(
                     _get_trainable_params(self.transformer) / 1e9, 3)
@@ -615,12 +689,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
         return batch
 
     @torch.no_grad()
-    def _log_validation(self, transformer, training_args, global_step) -> None:
+    def _log_validation(self, training_args, global_step) -> None:
         """
         Generate a validation video and log it to wandb to check the quality during training.
         """
         training_args.inference_mode = True
-        training_args.dit_cpu_offload = True
+        training_args.dit_cpu_offload = False
         if not training_args.log_validation:
             return
         if self.validation_pipeline is None:
@@ -642,7 +716,9 @@ class TrainingPipeline(LoRAPipeline, ABC):
                                            batch_size=None,
                                            num_workers=0)
 
-        transformer.eval()
+        self.transformer.eval()
+        if getattr(self, "transformer_2", None) is not None:
+            self.transformer_2.eval()
 
         validation_steps = training_args.validation_sampling_steps.split(",")
         validation_steps = [int(step) for step in validation_steps]
@@ -734,9 +810,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         # Re-enable gradients for training
         training_args.inference_mode = False
-        transformer.train()
-    
-    def visualize_intermediate_latents(self, training_batch: TrainingBatch,
-                                       training_args: TrainingArgs, step: int):
-        """Add visualization data to wandb logging and save frames to disk."""
-        raise NotImplementedError("Visualize intermediate latents is not implemented for training pipeline")
+        self.transformer.train()
+        if getattr(self, "transformer_2", None) is not None:
+            self.transformer_2.train()

fastvideo/training/training_utils.py

@@ -349,14 +349,10 @@ def load_checkpoint(transformer,
     """
     if not os.path.exists(checkpoint_path):
         logger.warning("Checkpoint path %s does not exist", checkpoint_path)
-        assert False
         return 0
 
     # Extract step number from checkpoint path
-    try:
-        step = int(os.path.basename(checkpoint_path).split('-')[-1])
-    except:
-        step = 1
+    step = int(os.path.basename(checkpoint_path).split('-')[-1])
 
     if rank == 0:
         logger.info("Loading checkpoint from step %s", step)

fastvideo/training/wan_t2v_i2v_training_pipeline.py

@@ -1,211 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import sys
-from copy import deepcopy
-from typing import Any
-
-import torch
-
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.configs.sample import SamplingParam
-from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
-    FlowUniPCMultistepScheduler)
-from fastvideo.pipelines.basic.wan.wan_i2v_pipeline import (
-    WanImageToVideoPipeline)
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch
-from fastvideo.training.training_pipeline import TrainingPipeline
-from fastvideo.utils import is_vsa_available, shallow_asdict
-
-vsa_available = is_vsa_available()
-
-logger = init_logger(__name__)
-
-
-class WanT2VI2VTrainingPipeline(TrainingPipeline):
-    """
-    A training pipeline for Wan.
-    """
-    _required_config_modules = ["scheduler", "transformer", "vae"]
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift)
-
-    def create_training_stages(self, training_args: TrainingArgs):
-        """
-        May be used in future refactors.
-        """
-        pass
-
-    def set_schemas(self):
-        self.train_dataset_schema = pyarrow_schema_t2v
-
-    def initialize_validation_pipeline(self, training_args: TrainingArgs):
-        logger.info("Initializing validation pipeline...")
-        args_copy = deepcopy(training_args)
-
-        args_copy.inference_mode = True
-        args_copy.dit_cpu_offload = True
-        # args_copy.pipeline_config.vae_config.load_encoder = False
-        # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained(
-        pipeline_config = PipelineConfig.from_pretrained(
-            training_args.model_path)
-        pipeline_config.vae_config.load_encoder = True
-        self.validation_pipeline = WanImageToVideoPipeline.from_pretrained(
-            training_args.model_path,
-            args=None,
-            inference_mode=True,
-            pipeline_config=pipeline_config,
-            loaded_modules={
-                "transformer": self.get_module("transformer"),
-            },
-            required_config_modules=[
-                "scheduler", "transformer", "vae", "text_encoder", "tokenizer"
-            ],
-            tp_size=training_args.tp_size,
-            sp_size=training_args.sp_size,
-            num_gpus=training_args.num_gpus,
-            dit_cpu_offload=True,
-        )
-
-    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
-        batch = next(self.train_loader_iter, None)  # type: ignore
-        if batch is None:
-            self.current_epoch += 1
-            logger.info("Starting epoch %s", self.current_epoch)
-            # Reset iterator for next epoch
-            self.train_loader_iter = iter(self.train_dataloader)
-            # Get first batch of new epoch
-            batch = next(self.train_loader_iter)
-
-        latents = batch['vae_latent']
-        latents = latents[:, :, :self.training_args.num_latent_t]
-        encoder_hidden_states = batch['text_embedding']
-        encoder_attention_mask = batch['text_attention_mask']
-        # clip_features = batch['clip_feature']
-        # image_latents = batch['first_frame_latent']
-        # image_latents = image_latents[:, :, :self.training_args.num_latent_t]
-        # pil_image = batch['pil_image']
-        infos = batch['info_list']
-
-        training_batch.latents = latents.to(get_local_torch_device(),
-                                            dtype=torch.bfloat16)
-        training_batch.encoder_hidden_states = encoder_hidden_states.to(
-            get_local_torch_device(), dtype=torch.bfloat16)
-        training_batch.encoder_attention_mask = encoder_attention_mask.to(
-            get_local_torch_device(), dtype=torch.bfloat16)
-        # training_batch.preprocessed_image = pil_image.to(
-        #     get_local_torch_device())
-        # training_batch.image_embeds = clip_features.to(get_local_torch_device())
-        # training_batch.image_latents = image_latents.to(
-        #     get_local_torch_device())
-        training_batch.infos = infos
-
-        return training_batch
-
-    def _prepare_dit_inputs(self,
-                            training_batch: TrainingBatch) -> TrainingBatch:
-        """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning."""
-
-        # First, call parent method to prepare noise, timesteps, etc. for video latents
-        training_batch = super()._prepare_dit_inputs(training_batch)
-
-        latents = training_batch.latents
-        logger.info("latents.shape: %s", latents.shape)
-        first_frame_latent = latents[:, :, 0, :, :]
-
-        logger.info("first_frame_latent.shape: %s", first_frame_latent.shape)
-        logger.info("training_batch.noisy_model_input.shape: %s",
-                    training_batch.noisy_model_input.shape)
-
-        training_batch.noisy_model_input = torch.cat([
-            first_frame_latent.unsqueeze(2),
-            training_batch.noisy_model_input[:, :, 1:, :, :]
-        ],
-                                                     dim=2)
-
-        return training_batch
-
-    def _build_input_kwargs(self,
-                            training_batch: TrainingBatch) -> TrainingBatch:
-
-        # Image Embeds for conditioning
-        # image_embeds = training_batch.image_embeds
-        # assert torch.isnan(image_embeds).sum() == 0
-        # image_embeds = image_embeds.to(get_local_torch_device(),
-        #                                dtype=torch.bfloat16)
-        # encoder_hidden_states_image = image_embeds
-
-        # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs
-        training_batch.input_kwargs = {
-            "hidden_states":
-            training_batch.noisy_model_input,
-            "encoder_hidden_states":
-            training_batch.encoder_hidden_states,
-            "timestep":
-            training_batch.timesteps.to(get_local_torch_device(),
-                                        dtype=torch.bfloat16),
-            "encoder_attention_mask":
-            training_batch.encoder_attention_mask,
-            # "encoder_hidden_states_image":
-            # encoder_hidden_states_image,
-            "return_dict":
-            False,
-        }
-        return training_batch
-
-    def _prepare_validation_batch(self, sampling_param: SamplingParam,
-                                  training_args: TrainingArgs,
-                                  validation_batch: dict[str, Any],
-                                  num_inference_steps: int) -> ForwardBatch:
-        sampling_param.prompt = validation_batch['prompt']
-        sampling_param.height = training_args.num_height
-        sampling_param.width = training_args.num_width
-        sampling_param.image_path = validation_batch['video_path']
-        sampling_param.num_inference_steps = num_inference_steps
-        sampling_param.data_type = "video"
-        assert self.seed is not None
-        sampling_param.seed = self.seed
-
-        latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
-                        sampling_param.height // 8, sampling_param.width // 8]
-        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
-        temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
-        num_frames = (training_args.num_latent_t -
-                      1) * temporal_compression_factor + 1
-        sampling_param.num_frames = num_frames
-        batch = ForwardBatch(
-            **shallow_asdict(sampling_param),
-            latents=None,
-            generator=torch.Generator(device="cpu").manual_seed(self.seed),
-            n_tokens=n_tokens,
-            eta=0.0,
-            VSA_sparsity=training_args.VSA_sparsity,
-        )
-
-        return batch
-
-
-def main(args) -> None:
-    logger.info("Starting training pipeline...")
-
-    pipeline = WanT2VI2VTrainingPipeline.from_pretrained(
-        args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/utils.py

@@ -2,10 +2,6 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py
 
 import argparse
-from einops import rearrange
-import torchvision
-import numpy as np
-import imageio
 import ctypes
 import hashlib
 import importlib
@@ -27,14 +23,10 @@ from typing import Any, TypeVar, cast
 
 import cloudpickle
 import filelock
-import imageio
-import numpy as np
 import torch
-import torchvision
 import yaml
 from diffusers.loaders.lora_base import (
     _best_guess_weight_name)  # watch out for potetential removal from diffusers
-from einops import rearrange
 from huggingface_hub import snapshot_download
 from remote_pdb import RemotePdb
 from torch.distributed.fsdp import MixedPrecisionPolicy
@@ -894,17 +886,3 @@ def best_output_size(w, h, dw, dh, expected_area):
         return ow1, oh1
     else:
         return ow2, oh2
-
-
-def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor],
-                                  output_path: str, fps: int):
-    # Process outputs
-    videos = rearrange(decoded_latents, "b c t h w -> t b c h w")
-    frames = []
-    for x in videos:
-        x = torchvision.utils.make_grid(x, nrow=6)
-        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-        frames.append((x * 255).numpy().astype(np.uint8))
-
-    os.makedirs(os.path.dirname(output_path), exist_ok=True)
-    imageio.mimsave(output_path, frames, fps=fps, format="mp4")

fastvideo/wan/README.md

@@ -1,2 +0,0 @@
-Code in this folder is modified from https://github.com/Wan-Video/Wan2.1
-Apache-2.0 License 

fastvideo/wan/__init__.py

@@ -1,3 +0,0 @@
-from . import configs, distributed, modules
-from .image2video import WanI2V
-from .text2video import WanT2V