update

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

.buildkite/pipeline.yml

@@ -198,4 +198,15 @@ steps:
                 env:  
                   - TEST_TYPE=inference_vmoba
                 agents:
-                  queue: "default"
+                  queue: "default"
+            - path:
+                - "fastvideo/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Unit Tests"
+                env:
+                  - TEST_TYPE=unit_test
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -118,6 +118,10 @@ case "$TEST_TYPE" in
         log "Running V-MoBA precision tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_vmoba"
         ;;
+    "unit_test")
+        log "Running unit tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_unit_test"
+        ;;
     *)
         log "Error: Unknown test type: $TEST_TYPE"
         exit 1

.github/workflows/pr-test.yml

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

.gitignore

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

README.md

@@ -155,8 +155,8 @@ If you find FastVideo useful, please considering citing our work:
 }
 
 @article{zhang2025vsa,
-  title={VSA: Faster Video Diffusion with Trainable Sparse Attention},
-  author={Zhang, Peiyuan and Huang, Haofeng and Chen, Yongqi and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
+  title={Vsa: Faster video diffusion with trainable sparse attention},
+  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
   journal={arXiv preprint arXiv:2505.13389},
   year={2025}
 }

csrc/attn/vmoba_attn/setup.py

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

csrc/attn/vmoba_attn/vmoba/vmoba.py

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

examples/dataset/vidprom/README.md

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

examples/dataset/vidprom/download_dataset.sh

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

examples/training/consistency_finetune/causal_ode_init/preprocess_wan_data_crush_smol.sh

@@ -15,6 +15,7 @@ torchrun --nproc_per_node=$GPU_NUM \
     --max_height 480 \
     --max_width 832 \
     --num_frames 81 \
+    --flow_shift 5.0 \
     --dataloader_num_workers 0 \
     --output_dir=$OUTPUT_DIR \
     --train_fps 16 \

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

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

fastvideo/attention/backends/vmoba.py

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

fastvideo/configs/models/dits/base.py

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

fastvideo/configs/pipelines/base.py

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

fastvideo/configs/pipelines/wan.py

@@ -82,7 +82,7 @@ class WanI2V480PConfig(WanT2V480PConfig):
         default_factory=CLIPVisionConfig)
     image_encoder_precision: str = "fp32"
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
 
@@ -108,19 +108,17 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 757, 522])
 
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True
-
 
 @dataclass
 class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
     flow_shift: float | None = 5.0
     ti2v_task: bool = True
+    expand_timesteps: bool = True
 
     def __post_init__(self) -> None:
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
+        self.dit_config.expand_timesteps = self.expand_timesteps
 
 
 @dataclass
@@ -132,12 +130,21 @@ class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config):
 
 @dataclass
 class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
-    pass
+    flow_shift: float | None = 12.0
+    boundary_ratio: float | None = 0.875
+
+    def __post_init__(self) -> None:
+        self.dit_config.boundary_ratio = self.boundary_ratio
 
 
 @dataclass
-class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
-    pass
+class Wan2_2_I2V_A14B_Config(WanI2V480PConfig):
+    flow_shift: float | None = 5.0
+    boundary_ratio: float | None = 0.900
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.dit_config.boundary_ratio = self.boundary_ratio
 
 
 # =============================================

fastvideo/configs/sample/base.py

@@ -40,6 +40,7 @@ class SamplingParam:
     num_inference_steps: int = 50
     guidance_scale: float = 1.0
     guidance_rescale: float = 0.0
+    boundary_ratio: float | None = None
 
     # TeaCache parameters
     enable_teacache: bool = False
@@ -169,6 +170,12 @@ class SamplingParam:
             default=SamplingParam.guidance_rescale,
             help="Guidance rescale factor",
         )
+        parser.add_argument(
+            "--boundary-ratio",
+            type=float,
+            default=SamplingParam.boundary_ratio,
+            help="Boundary timestep ratio",
+        )
         parser.add_argument(
             "--save-video",
             action="store_true",

fastvideo/configs/sample/wan.py

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

fastvideo/dataset/dataloader/record_schema.py

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

fastvideo/dataset/dataloader/schema.py

@@ -102,3 +102,17 @@ pyarrow_schema_ode_trajectory_text_only = pa.schema([
     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/layers/lora/linear.py

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

fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py

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

fastvideo/pipelines/composed_pipeline_base.py

@@ -237,20 +237,19 @@ class ComposedPipelineBase(ABC):
         # remove keys that are not pipeline modules
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
-        # @TODO(Wei): Temporary hack
         if "boundary_ratio" in model_index and model_index[
                 "boundary_ratio"] is not None:
             logger.info(
                 "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..."
             )
             self.required_config_modules.append("transformer_2")
-            if fastvideo_args.boundary_ratio is None:
-                logger.info(
-                    "MoE pipeline detected. Setting boundary ratio to %s",
-                    model_index["boundary_ratio"])
-                fastvideo_args.boundary_ratio = model_index["boundary_ratio"]
+            logger.info("MoE pipeline detected. Setting boundary ratio to %s",
+                        model_index["boundary_ratio"])
+            fastvideo_args.pipeline_config.dit_config.boundary_ratio = model_index[
+                "boundary_ratio"]
 
         model_index.pop("boundary_ratio", None)
+        # used by Wan2.2 ti2v
         model_index.pop("expand_timesteps", None)
 
         # some sanity checks
@@ -283,8 +282,8 @@ class ComposedPipelineBase(ABC):
                           architecture) in model_index.items():
             if transformers_or_diffusers is None:
                 logger.warning(
-                    "Module in model_index.json has null value, removing from required_config_modules"
-                )
+                    "Module %s in model_index.json has null value, removing from required_config_modules",
+                    module_name)
                 if module_name in self.required_config_modules:
                     self.required_config_modules.remove(module_name)
                 continue

fastvideo/pipelines/pipeline_batch_info.py

@@ -129,6 +129,7 @@ class ForwardBatch:
     timesteps: torch.Tensor | None = None
     timestep: torch.Tensor | float | int | None = None
     step_index: int | None = None
+    boundary_ratio: float | None = None
 
     # Scheduler parameters
     num_inference_steps: int = 50

fastvideo/pipelines/preprocess/preprocess_pipeline_base.py

@@ -10,6 +10,8 @@ from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from fastvideo.dataset import getdataset
+from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter,
+                                                     records_to_table)
 from fastvideo.dataset.preprocessing_datasets import PreprocessBatch
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -17,8 +19,6 @@ from fastvideo.logger import init_logger
 from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages import TextEncodingStage
-from fastvideo.workflow.preprocess.parquet_io import (ParquetDatasetWriter,
-                                                      records_to_table)
 
 logger = init_logger(__name__)
 
@@ -423,7 +423,3 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                 written = self.dataset_writer.flush()
                 logger.info("Flushed %s samples to parquet", written)
                 num_processed_samples = 0
-
-    def _final_flush_if_any(self):
-        if hasattr(self, 'dataset_writer'):
-            self.dataset_writer.flush()

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -12,7 +12,6 @@ import os
 from collections.abc import Iterator
 from typing import Any
 
-import numpy as np
 import pyarrow as pa
 import torch
 from torch.utils.data import DataLoader
@@ -21,6 +20,10 @@ 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
@@ -36,8 +39,6 @@ from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage,
                                         TextEncodingStage,
                                         TimestepPreparationStage)
 from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict
-from fastvideo.workflow.preprocess.parquet_io import (ParquetDatasetWriter,
-                                                      records_to_table)
 
 logger = init_logger(__name__)
 
@@ -59,26 +60,13 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
 
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         """Set up pipeline stages with proper dependency injection."""
-        logger.info('WTF flow_shift: %s',
-                    fastvideo_args.pipeline_config.flow_shift)
         assert fastvideo_args.pipeline_config.flow_shift == 5
-        # self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
-        # shift=fastvideo_args.pipeline_config.flow_shift)
         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)
-        # logger.info('WTF scheduler timesteps: %s',
-        #             self.modules["scheduler"].timesteps)
-        # scheduler = FlowMatchScheduler(
-        #     shift=8.0, sigma_min=0.0, extra_one_step=True)
-        # device = get_local_torch_device()
-        # # scheduler.num_train_timesteps = 100
-        # scheduler.set_timesteps(num_inference_steps=50, denoising_strength=1.0)
-        # scheduler.sigmas = scheduler.sigmas.to(device)
-        # self.modules["scheduler"] = scheduler
 
         self.add_stage(stage_name="input_validation_stage",
                        stage=InputValidationStage())
@@ -186,6 +174,8 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                     ]
                     batch.num_inference_steps = 48
                     batch.return_trajectory_latents = True
+                    # Enabling this will save the decoded trajectory videos.
+                    # Used for debugging.
                     batch.return_trajectory_decoded = False
                     batch.height = args.max_height
                     batch.width = args.max_width
@@ -225,7 +215,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                                 args.train_fps)
 
                 # Prepare batch data for Parquet dataset
-                batch_data = []
+                batch_data: list[dict[str, Any]] = []
 
                 # Add progress bar for saving outputs
                 save_pbar = tqdm(enumerate(valid_data["path"]),
@@ -254,14 +244,16 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                                 else:
                                     sample_extra_features[key] = value[idx]
 
-                    # Create record for Parquet dataset (without VAE latents for text-only)
-                    record = self.create_text_only_record(
-                        args,
+                    # 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,
-                        valid_data=valid_data,
-                        idx=idx,
-                        extra_features=sample_extra_features)
+                        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:
@@ -289,78 +281,10 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
 
         # Final flush for any remaining samples
         if hasattr(self, 'dataset_writer'):
-            written = self.dataset_writer.flush()
+            written = self.dataset_writer.flush(write_remainder=True)
             if written:
                 logger.info("Final flush wrote %s samples", written)
 
-    def create_text_only_record(
-            self,
-            args,
-            video_name: str,
-            text_embedding: np.ndarray,
-            valid_data: dict[str, Any],
-            idx: int,
-            extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
-        """Create a record for text-only preprocessing using text-only schema."""
-
-        # Create base record using only fields from text-only schema
-        record = {
-            "id": f"text_{video_name}_{idx}",
-            "text_embedding_bytes": text_embedding.tobytes(),
-            "text_embedding_shape": list(text_embedding.shape),
-            "text_embedding_dtype": str(text_embedding.dtype),
-            "file_name": video_name,
-            "caption": valid_data["text"][idx],
-            "media_type": "text",
-        }
-
-        assert extra_features is not None, "extra_features is required"
-        assert "trajectory_latents" in extra_features, "trajectory_latents is required"
-        assert "trajectory_timesteps" in extra_features, "trajectory_timesteps is required"
-
-        # Add trajectory data if available
-        if extra_features and "trajectory_latents" in extra_features:
-            trajectory_latents = extra_features[
-                "trajectory_latents"][idx] if isinstance(
-                    extra_features["trajectory_latents"],
-                    list) else extra_features["trajectory_latents"]
-            record.update({
-                "trajectory_latents_bytes":
-                trajectory_latents.tobytes(),
-                "trajectory_latents_shape":
-                list(trajectory_latents.shape),
-                "trajectory_latents_dtype":
-                str(trajectory_latents.dtype),
-            })
-        else:
-            record.update({
-                "trajectory_latents_bytes": b"",
-                "trajectory_latents_shape": [],
-                "trajectory_latents_dtype": "",
-            })
-
-        if extra_features and "trajectory_timesteps" in extra_features:
-            trajectory_timesteps = extra_features[
-                "trajectory_timesteps"][idx] if isinstance(
-                    extra_features["trajectory_timesteps"],
-                    list) else extra_features["trajectory_timesteps"]
-            record.update({
-                "trajectory_timesteps_bytes":
-                trajectory_timesteps.tobytes(),
-                "trajectory_timesteps_shape":
-                list(trajectory_timesteps.shape),
-                "trajectory_timesteps_dtype":
-                str(trajectory_timesteps.dtype),
-            })
-        else:
-            record.update({
-                "trajectory_timesteps_bytes": b"",
-                "trajectory_timesteps_shape": [],
-                "trajectory_timesteps_dtype": "",
-            })
-
-        return record
-
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
         if not self.post_init_called:
             self.post_init()

fastvideo/pipelines/preprocess/preprocess_pipeline_text.py

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

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+from typing import Any
 
 from fastvideo import PipelineConfig
 from fastvideo.configs.models.vaes import WanVAEConfig
@@ -13,6 +14,8 @@ 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__)
@@ -23,13 +26,22 @@ def main(args) -> None:
     maybe_init_distributed_environment_and_model_parallel(1, 1)
     num_gpus = int(os.environ["WORLD_SIZE"])
     assert num_gpus == 1, "Only support 1 GPU"
+
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
-    kwargs = {
-        "vae_precision": "fp32",
-        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
-        "flow_shift": 5,
-    }
+
+    kwargs: dict[str, Any] = {}
+    if args.preprocess_task == "text_only":
+        kwargs = {
+            "text_encoder_cpu_offload": False,
+        }
+    else:
+        # Full config for video/image processing
+        kwargs = {
+            "vae_precision": "fp32",
+            "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
+        }
     pipeline_config.update_config_from_dict(kwargs)
+
     fastvideo_args = FastVideoArgs(
         model_path=args.model_path,
         num_gpus=get_world_size(),
@@ -42,13 +54,19 @@ def main(args) -> None:
         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}")
+        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. "
+                         f"Valid options: t2v, i2v, ode_trajectory, text_only")
 
     logger.info("Preprocess task: %s using %s", args.preprocess_task,
                 PreprocessPipeline.__name__)
+
     pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 
@@ -87,7 +105,12 @@ if __name__ == "__main__":
     parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
     parser.add_argument("--group_frame", action="store_true")  # TODO
     parser.add_argument("--group_resolution", action="store_true")  # TODO
-    parser.add_argument("--preprocess_task", type=str, default="t2v")
+    parser.add_argument("--flow_shift", type=float, default=None)
+    parser.add_argument("--preprocess_task",
+                        type=str,
+                        default="t2v",
+                        choices=["t2v", "i2v", "text_only", "ode_trajectory"],
+                        help="Type of preprocessing task to run")
     parser.add_argument("--train_fps", type=int, default=30)
     parser.add_argument("--use_image_num", type=int, default=0)
     parser.add_argument("--text_max_length", type=int, default=256)

fastvideo/pipelines/stages/denoising.py

@@ -204,8 +204,14 @@ class DenoisingStage(PipelineStage):
                 neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
 
         # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
-        if fastvideo_args.boundary_ratio is not None:
-            boundary_timestep = fastvideo_args.boundary_ratio * self.scheduler.num_train_timesteps
+        boundary_ratio = fastvideo_args.pipeline_config.dit_config.boundary_ratio
+        if batch.boundary_ratio is not None:
+            logger.info("Overriding boundary ratio from %s to %s",
+                        boundary_ratio, batch.boundary_ratio)
+            boundary_ratio = batch.boundary_ratio
+
+        if boundary_ratio is not None:
+            boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
         else:
             boundary_timestep = None
         latent_model_input = latents.to(target_dtype)

fastvideo/tests/dataset/test_schema_record_creator.py

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

fastvideo/tests/modal/pr_test.py

@@ -117,3 +117,7 @@ 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/workflow/test_components_parquet_saver.py

@@ -38,7 +38,8 @@ def test_parquet_dataset_saver_flush_and_last(tmp_path: Path):
         data_type="video",
         latents=torch.randn(B, 2),
         prompt_embeds=[torch.randn(B, 1, 1)],
-        prompt_attention_mask=[torch.ones(B, 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)]
 
@@ -52,16 +53,14 @@ def test_parquet_dataset_saver_flush_and_last(tmp_path: Path):
     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(str(out_dir))
+    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_last(str(out_dir))
+    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/training_pipeline.py

@@ -478,7 +478,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 current_vsa_sparsity = current_decay_times * vsa_decay_rate
             elif vmoba_available:
                 # TODO: add vmoba sparsity scheduling here
-                pass
+                current_vsa_sparsity = 0.0
             else:
                 current_vsa_sparsity = 0.0
 

fastvideo/workflow/preprocess/components.py

@@ -222,7 +222,7 @@ class ParquetDatasetSaver:
 
         # If flush is needed
         if self.num_processed_samples >= self.flush_frequency:
-            self.flush_tables(output_dir)
+            self.flush_tables()
 
     def _process_non_padded_embeddings(
             self, prompt_embeds: torch.Tensor,
@@ -245,12 +245,7 @@ class ParquetDatasetSaver:
 
         return non_padded_embeds
 
-    def _convert_batch_to_pyarrow_table(self,
-                                        batch_data: list[dict]) -> pa.Table:
-        # Deprecated path, kept for backward compatibility if needed.
-        return records_to_table(batch_data, self.schema)
-
-    def flush_tables(self, output_dir: str, write_remainder: bool = False):
+    def flush_tables(self, write_remainder: bool = False):
         """Flush buffered records to disk.
 
         Args:
@@ -266,16 +261,16 @@ class ParquetDatasetSaver:
         remainder = self.num_processed_samples % self.samples_per_file
         self.num_processed_samples = 0 if write_remainder else remainder
 
-    def flush_last(self, output_dir: str):
-        """Flush and write any remaining rows (final flush)."""
-        self.flush_tables(output_dir, write_remainder=True)
-
     def clean_up(self) -> None:
         """Clean up all tables"""
+        self.flush_tables(write_remainder=True)
         self._writer = None
         self.num_processed_samples = 0
         gc.collect()
 
+    def __del__(self):
+        self.clean_up()
+
 
 def build_dataset(preprocess_config: PreprocessConfig, split: str,
                   validator: Callable[[dict[str, Any]], bool]) -> Dataset:

fastvideo/workflow/preprocess/preprocess_workflow.py

@@ -4,6 +4,8 @@ from typing import cast
 from torch.utils.data import DataLoader
 
 from fastvideo.configs.configs import PreprocessConfig
+from fastvideo.dataset.dataloader.record_schema import (
+    basic_t2v_record_creator, i2v_record_creator)
 from fastvideo.dataset.dataloader.schema import (pyarrow_schema_i2v,
                                                  pyarrow_schema_t2v)
 from fastvideo.distributed.parallel_state import get_world_rank
@@ -13,8 +15,6 @@ from fastvideo.pipelines.pipeline_registry import PipelineType
 from fastvideo.workflow.preprocess.components import (
     ParquetDatasetSaver, PreprocessingDataValidator, VideoForwardBatchBuilder,
     build_dataset)
-from fastvideo.workflow.preprocess.record_schema import (
-    basic_t2v_record_creator, i2v_record_creator)
 from fastvideo.workflow.workflow_base import WorkflowBase
 
 logger = init_logger(__name__)

fastvideo/workflow/preprocess/preprocess_workflow_i2v.py

@@ -35,8 +35,7 @@ class PreprocessWorkflowI2V(PreprocessWorkflow):
             self.processed_dataset_saver.save_and_write_parquet_batch(
                 forward_batch, self.training_dataset_output_dir)
 
-        self.processed_dataset_saver.flush_tables(
-            self.training_dataset_output_dir)
+        self.processed_dataset_saver.flush_tables()
         self.processed_dataset_saver.clean_up()
 
         # Validation dataset preprocessing
@@ -51,6 +50,5 @@ class PreprocessWorkflowI2V(PreprocessWorkflow):
 
                 self.processed_dataset_saver.save_and_write_parquet_batch(
                     forward_batch, self.validation_dataset_output_dir)
-            self.processed_dataset_saver.flush_tables(
-                self.validation_dataset_output_dir)
+            self.processed_dataset_saver.flush_tables()
             self.processed_dataset_saver.clean_up()

fastvideo/workflow/preprocess/preprocess_workflow_t2v.py

@@ -35,8 +35,7 @@ class PreprocessWorkflowT2V(PreprocessWorkflow):
             self.processed_dataset_saver.save_and_write_parquet_batch(
                 forward_batch, self.training_dataset_output_dir)
 
-        self.processed_dataset_saver.flush_tables(
-            self.training_dataset_output_dir)
+        self.processed_dataset_saver.flush_tables()
         self.processed_dataset_saver.clean_up()
 
         # Validation dataset preprocessing
@@ -51,6 +50,5 @@ class PreprocessWorkflowT2V(PreprocessWorkflow):
 
                 self.processed_dataset_saver.save_and_write_parquet_batch(
                     forward_batch, self.validation_dataset_output_dir)
-            self.processed_dataset_saver.flush_tables(
-                self.validation_dataset_output_dir)
+            self.processed_dataset_saver.flush_tables()
             self.processed_dataset_saver.clean_up()

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
 
     # Miscellaneous Utilities
     "tqdm", "pytest", "PyYAML==6.0.1", "protobuf>=5.28.3",
-    "gradio>=5.22.0", "moviepy>=2.0.0", "flask",
+    "gradio==5.41.0", "moviepy>=2.0.0", "flask",
     "flask_restful", "aiohttp", "huggingface_hub", "cloudpickle",
     # System & Monitoring Tools
     "gpustat", "watch", "remote-pdb",

scripts/preprocess/launch_slurm_text.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# Create output directory if it doesn't exist
+mkdir -p preprocess_output_text
+
+# Launch 8 jobs, one for each node
+# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node)
+for node_id in {0..7}; do
+    # Calculate the starting file number for this node
+    start_file=$((node_id * 8 + 1))
+    
+    echo "Launching text-only node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt"
+    echo "sbatch --job-name=text-${node_id} --output=preprocess_output_text/preprocess-text-node-${node_id}.out --error=preprocess_output_text/preprocess-text-node-${node_id}.err scripts/preprocess/syn_text.slurm $start_file $node_id"
+    
+    sbatch --job-name=text-${node_id} \
+           --output=preprocess_output_text/preprocess-text-node-${node_id}.out \
+           --error=preprocess_output_text/preprocess-text-node-${node_id}.err \
+           scripts/preprocess/syn_text.slurm $start_file $node_id
+done
+
+echo "All 8 text-only nodes launched successfully!" 

scripts/preprocess/syn_text.slurm

@@ -0,0 +1,88 @@
+#!/bin/bash
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=16
+#SBATCH --mem=960G
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+# conda init
+# source ~/conda/miniconda/bin/activate
+# PYTHON_VIRTUAL_ENVIRONMENT=fastvideo-train-yq
+# conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+nvidia-smi
+nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices for TEXT-ONLY preprocessing"
+echo ""
+echo " Run started at:- "
+date
+
+# Accept parameters from launch script
+START_FILE=${1:-1}  # Starting file number for this node
+NODE_ID=${2:-0}  # Node identifier (0-7)
+
+num_gpus=1
+export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+# Start port number - we'll increment for each job
+base_port=$((29603 + NODE_ID * 100))  # Different port range per node
+
+# Create an array of CUDA device IDs
+gpu_ids=(0 1 2 3 4 5 6 7)
+
+GPU_NUM=1
+MODEL_TYPE="wan"
+
+echo "NODE_ID: $NODE_ID"
+echo "START_FILE: $START_FILE"
+echo "Base port for this node: $base_port"
+echo "Processing TEXT-ONLY data"
+
+# Run 8 parallel preprocessing jobs on this node
+for i in {1..8}; do
+    # Calculate port for this job
+    port=$((base_port + i))
+    
+    # Get GPU ID using modulo to cycle through available GPUs
+    gpu=${gpu_ids[((i-1))]}
+    
+    # Calculate which file this GPU should process
+    file_num=$((START_FILE + i - 1))
+    DATA_MERGE_PATH="prompts/v2m_${file_num}.txt"
+    
+    # Create unique output directory based on node and GPU for text-only processing
+    OUTPUT_DIR="data/test-text-preprocessing/Node_${NODE_ID}_GPU_${i}_File_${file_num}"
+    
+    start_cpu=$(( (i-1)*2 ))  # Reduced CPU allocation for 8 nodes
+    end_cpu=$(( start_cpu+1 ))
+    
+    echo "Starting GPU $gpu processing text-only file v2m_${file_num}.txt on port $port, output: $OUTPUT_DIR"
+    
+    # Run the text-only preprocessing command in background
+    CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=$GPU_NUM --master_port $port \
+        fastvideo/pipelines/preprocess/v1_preprocess.py \
+        --model_path $MODEL_BASE \
+        --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 "text_only" &
+done
+
+# Wait for all jobs on this node to complete
+wait
+
+echo "All text-only processing blocks completed!" 

scripts/preprocess/v1_preprocess_wan_data_i2v.sh

@@ -21,4 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
     --validation_dataset_file $VALIDATION_PATH \
     --samples_per_file 8 \
     --flush_frequency 8 \
-    --preprocess_task "i2v" 
+    --preprocess_task "i2v" 

scripts/preprocess/v1_preprocess_wan_data_t2v.sh

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