update

.github/workflows/pr-test.yml

@@ -14,13 +14,9 @@ on:
       - ".github/workflows/pr-test.yml"
       - "pyproject.toml"
       - "docker/Dockerfile.python3.12"
+      - "csrc/**"
   workflow_dispatch:
     inputs:
-      custom_image:
-        description: "Custom image from this repository (default: fastvideo-dev:py3.12-latest)"
-        required: false
-        default: "fastvideo-dev:py3.12-latest"
-        type: string
       run_encoder_test:
         description: "Run encoder-test"
         required: false
@@ -56,6 +52,16 @@ on:
         required: false
         default: false
         type: boolean
+      run_precision_test_STA:
+        description: "Run precision-test-STA"
+        required: false
+        default: false
+        type: boolean
+      run_precision_test_VSA:
+        description: "Run precision-test-VSA"
+        required: false
+        default: false
+        type: boolean
       run_nightly_test:
         description: "Run nightly-test"
         required: false
@@ -65,6 +71,7 @@ on:
 env:
   PYTHONUNBUFFERED: "1"
 
+
 concurrency:
   group: pr-test-${{ github.ref }}
   cancel-in-progress: true
@@ -84,44 +91,69 @@ jobs:
       training-test: ${{ steps.filter.outputs.training-test }}
       training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }}
       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 }}
     steps:
       - uses: actions/checkout@v4
       - uses: dorny/paths-filter@v3
         id: filter
         with:
           filters: |
+            # Define reusable path patterns
+            common-paths: &common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
+            sta-kernel-paths: &sta-kernel-paths
+              - 'csrc/attn/st_attn/**'
+              - 'csrc/attn/setup_sta.py'
+              - 'csrc/attn/config_sta.py'
+              - 'csrc/attn/st_attn.cpp'
+            vsa-kernel-paths: &vsa-kernel-paths
+              - 'csrc/attn/vsa/**'
+              - 'csrc/attn/tk/**'
+              - 'csrc/attn/setup_vsa.py'
+              - 'csrc/attn/config_vsa.py'
+              - 'csrc/attn/vsa.cpp'
+            vsa-paths: &vsa-paths
+              - 'fastvideo/v1/**'
+              - *common-paths
+              - *vsa-kernel-paths
+            
+            # Actual tests
             encoder-test:
               - 'fastvideo/v1/models/encoders/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/encoders/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
+              - *common-paths
             vae-test:
               - 'fastvideo/v1/models/vaes/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/vaes/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
+              - *common-paths
             transformer-test:
               - 'fastvideo/v1/models/dits/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/transformers/**'
               - 'fastvideo/v1/layers/**'
               - 'fastvideo/v1/attention/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
+              - *common-paths
             training-test:
               - 'fastvideo/v1/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
+              - *common-paths
             training-test-VSA:
               - 'fastvideo/v1/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
+              - *common-paths
+              - *vsa-kernel-paths
             inference-test-STA:
               - 'fastvideo/v1/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
+              - *common-paths
+              - *sta-kernel-paths
+            precision-test-STA:
+              - *common-paths
+              - *sta-kernel-paths
+            precision-test-VSA:
+              - *common-paths
+              - *vsa-kernel-paths
 
   encoder-test:
     needs: change-filter
@@ -134,7 +166,7 @@ jobs:
       gpu_type: "NVIDIA A40"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
       test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
       timeout_minutes: 30
     secrets:
@@ -152,7 +184,7 @@ jobs:
       gpu_type: "NVIDIA A40"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
       test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
       timeout_minutes: 30
     secrets:
@@ -170,7 +202,7 @@ jobs:
       gpu_type: "NVIDIA L40S"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
       test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
       timeout_minutes: 30
     secrets:
@@ -216,7 +248,7 @@ jobs:
       gpu_count: 4
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
       test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/Vanilla -srP"
       timeout_minutes: 30
     secrets:
@@ -236,7 +268,7 @@ jobs:
       gpu_count: 1
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
       test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/VSA -srP"
       timeout_minutes: 30
     secrets:
@@ -256,13 +288,51 @@ jobs:
       gpu_count: 1
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
       test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/inference/STA -srP"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
 
+  precision-test-STA:
+    needs: change-filter
+    if: >-
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_STA == 'true')
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "precision-test-STA"
+      gpu_type: "NVIDIA H100 NVL"
+      gpu_count: 1
+      volume_size: 100
+      disk_size: 100
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_sta.py"
+      timeout_minutes: 30
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
+
+  precision-test-VSA:
+    needs: change-filter
+    if: >-
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_VSA == 'true')
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "precision-test-VSA"
+      gpu_type: "NVIDIA H100 NVL"
+      gpu_count: 1
+      volume_size: 100
+      disk_size: 100
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_block_sparse.py"
+      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')
@@ -273,7 +343,7 @@ jobs:
       gpu_count: 4
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
       test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs"
       timeout_minutes: 30
     secrets:
@@ -282,7 +352,8 @@ jobs:
       WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
 
   runpod-cleanup:
-    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
+    # Add other jobs to this list as you create them
+    needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, inference-test-STA, precision-test-STA, precision-test-VSA] 
     if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
     runs-on: ubuntu-latest
     steps:
@@ -299,7 +370,7 @@ jobs:
 
       - name: Cleanup all RunPod instances
         env:
-          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]'
+          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
         run: python .github/scripts/runpod_cleanup.py

csrc/attn/README.md

@@ -4,7 +4,7 @@
 
 
 ## Installation
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
+We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only support H100/H200, because ThunderKittens uses TMA but doesn't support Blackwell yet.
 First, install C++20 for ThunderKittens:
 ```bash
 sudo apt update
@@ -53,8 +53,14 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
 
 ## Test
 ```bash
-python test/test_sta.py
+python tests/test_sta.py # test STA
+python tests/test_block_sparse.py # test VSA
 ```
+## Benchmark
+```bash
+python benchmarks/bench_sta.py
+```
+
 
 ## How Does STA Work?
 We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 

csrc/attn/benchmarks/bench_sta.py

@@ -5,6 +5,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from st_attn import sliding_tile_attention
+from triton.testing import do_bench
 
 
 def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
@@ -13,16 +14,16 @@ def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
     return f if mode == "fwd" else (2.5 * f if mode == "bwd" else 3.5 * f)
 
 
-def efficiency(flop, time):
-    flop = flop / 1e12
-    time = time / 1e6
-    return flop / time
+def compute_TFLOPS(flops, ms):
+    flops = flops / 1e12
+    ms = ms / 1e3
+    return flops / ms
 
 
 def benchmark_attention(configurations):
     results = {'fwd': defaultdict(list), 'bwd': defaultdict(list)}
 
-    for B, H, N, D, causal in configurations:
+    for B, H, N, D, causal, dit_seq_shape, window_size in configurations:
         print("=" * 60)
         print(f"Timing forward and backward pass for B={B}, H={H}, N={N}, D={D}, causal={causal}")
 
@@ -30,38 +31,31 @@ def benchmark_attention(configurations):
         k = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
         v = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
 
-        grad_output = torch.randn_like(q, requires_grad=False).contiguous()
+        # grad_output = torch.randn_like(q, requires_grad=False).contiguous()
+        # qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
+        # kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
+        # vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
 
-        qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
-        kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
-        vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
 
-        # Prepare for timing forward pass
-        start_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
-        end_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+        # # Warmup for forward pass
+        # for _ in range(10):
+        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
 
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+        # # Time the forward pass
+        # for i in range(10):
+        #     start_events_fwd[i].record()
+        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
+        #     end_events_fwd[i].record()
+        ms = do_bench(lambda: sliding_tile_attention(q, k, v, [window_size] * 24, 0, False, dit_seq_shape))
 
-        # Warmup for forward pass
-        for _ in range(10):
-            o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80')
+        # times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
+        # time_us_fwd = np.mean(times_fwd) * 1000
 
-        # Time the forward pass
-        for i in range(10):
-            start_events_fwd[i].record()
-            o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80')
-            end_events_fwd[i].record()
-
-        torch.cuda.synchronize()
-        times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
-        time_us_fwd = np.mean(times_fwd) * 1000
-
-        tflops_fwd = efficiency(flops(B, N, H, D, causal, 'fwd'), time_us_fwd)
+        tflops_fwd = compute_TFLOPS(flops(B, N, H, D, causal, 'fwd'), ms)
         results['fwd'][(D, causal)].append((N, tflops_fwd))
 
-        print(f"Average time for forward pass in us: {time_us_fwd:.2f}")
-        print(f"Average efficiency for forward pass in TFLOPS: {tflops_fwd}")
+        print(f"Average time for forward pass (ms): {ms:.2f}")
+        print(f"Average TFLOPS: {tflops_fwd}")
         print("-" * 60)
 
         # torch.cuda.empty_cache()
@@ -85,15 +79,14 @@ def benchmark_attention(configurations):
         # times_bwd = [s.elapsed_time(e) for s, e in zip(start_events_bwd, end_events_bwd)]
         # time_us_bwd = np.mean(times_bwd) * 1000
 
-        # tflops_bwd = efficiency(flops(B, N, H, D, causal, 'bwd'), time_us_bwd)
+        # tflops_bwd = compute_TFLOPS(flops(B, N, H, D, causal, 'bwd'), ms)
         # results['bwd'][(D, causal)].append((N, tflops_bwd))
 
-        # print(f"Average time for backward pass in us: {time_us_bwd:.2f}")
-        # print(f"Average efficiency for backward pass in TFLOPS: {tflops_bwd}")
-        print("=" * 60)
+        # print(f"Average time for backward pass(ms): {ms:.2f}")
+        # print(f"Average TFLOPS: {tflops_bwd}")
+        # print("=" * 60)
 
         torch.cuda.empty_cache()
-        torch.cuda.synchronize()
 
     return results
 
@@ -124,7 +117,10 @@ def plot_results(results):
 
 # Example list of configurations to test
 configurations = [
-    (2, 24, 69120, 128, False),
+    (2, 24, 69120, 128, False, '18x48x80', [3, 6, 10]),
+    (2, 24, 69120, 128, True, '18x48x80', [3, 6, 10]),
+    (2, 24, 82944, 128, False, '36x48x48', [3, 3, 6]), # Stepvideo
+    (2, 24, 82944, 128, True, '36x48x48', [3, 3, 6]),
     # (16, 16, 768*16,  128, False),
     # (16, 16, 768*2,  128, False),
     # (16, 16, 768*4,  128, False),

csrc/attn/st_attn/st_attn_h100.cu

@@ -4,9 +4,17 @@
 #include <cooperative_groups.h>
 #include <iostream>
 #include <stdio.h>
+#include <c10/cuda/CUDAGuard.h>
+
+// #define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
+__device__ __forceinline__ int clamp_int(int value, int min, int max) {
+    return (value < min) ? min : ((value > max) ? max : value);
+}
+// #define ABS(x) ((x) < 0 ? -(x) : (x))
+__device__ __forceinline__ int abs_int(int value) {
+    return (value < 0) ? -value : value;
+}
 
-#define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
-#define ABS(x) ((x) < 0 ? -(x) : (x))
 
 constexpr int CONSUMER_WARPGROUPS = (3); 
 constexpr int PRODUCER_WARPGROUPS = (1); 
@@ -117,16 +125,16 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
             int qt = seq_idx / 6 / (CH * CW);
             int qh = (seq_idx / 6) % (CH * CW) / CW;
             int qw = (seq_idx / 6) % CW;
-            qt = CLAMP(qt, DT, CT-DT-1);
-            qh = CLAMP(qh, DH, CH-DH-1);
-            qw = CLAMP(qw, DW, CW-DW-1);
+            qt = clamp_int(qt, DT, CT-DT-1);
+            qh = clamp_int(qh, DH, CH-DH-1);
+            qw = clamp_int(qw, DW, CW-DW-1);
             int count = 0;
             int j = 0;
             while (count < K::stages - 1) {
                 int kt = j / 3 / (CH * CW);
                 int kh = (j / 3) % (CH * CW) / CW;
                 int kw = (j / 3) % CW;
-                bool mask = (ABS(qt - kt) <= DT) && (ABS(qh - kh) <= DH) && (ABS(qw - kw) <= DW);
+                bool mask = (abs_int(qt - kt) <= DT) && (abs_int(qh - kh) <= DH) && (abs_int(qw - kw) <= DW);
                 if (mask){
                     coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0};
                     tma::expect_bytes(k_smem_arrived[count], sizeof(k_tile));
@@ -167,15 +175,15 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
                 int qt = seq_idx / 6 / (CH * CW);
                 int qh = (seq_idx / 6) % (CH * CW) / CW;
                 int qw = (seq_idx / 6) % CW;
-                qt = CLAMP(qt, DT, CT-DT-1);
-                qh = CLAMP(qh, DH, CH-DH-1);
-                qw = CLAMP(qw, DW, CW-DW-1);
-                int k_t_min = CLAMP(qt-DT, 0, CT-1);
-                int k_t_max = CLAMP(qt+DT, 0, CT-1);
-                int k_h_min = CLAMP(qh-DH, 0, CH-1);
-                int k_h_max = CLAMP(qh+DH, 0, CH-1);
-                int k_w_min = CLAMP(qw-DW, 0, CW-1);
-                int k_w_max = CLAMP(qw+DW, 0, CW-1);
+                qt = clamp_int(qt, DT, CT-DT-1);
+                qh = clamp_int(qh, DH, CH-DH-1);
+                qw = clamp_int(qw, DW, CW-DW-1);
+                int k_t_min = clamp_int(qt-DT, 0, CT-1);
+                int k_t_max = clamp_int(qt+DT, 0, CT-1);
+                int k_h_min = clamp_int(qh-DH, 0, CH-1);
+                int k_h_max = clamp_int(qh+DH, 0, CH-1);
+                int k_w_min = clamp_int(qw-DW, 0, CW-1);
+                int k_w_max = clamp_int(qw+DW, 0, CW-1);
                 int count = 0;
                 for (int kt = k_t_min; kt <= k_t_max; kt++) {
                     for (int kh = k_h_min; kh <= k_h_max; kh++) {
@@ -234,7 +242,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
             // the last three kv blocks are for text, we process them separately
             kv_iters = img_kv_blocks - 1;
         } else {
-            kv_iters = CLAMP(DT*2+1, 1, CT) * CLAMP(DH*2+1, 1, CH) * CLAMP(DW*2+1, 1, CW) * 3 - 1 ; 
+            kv_iters = clamp_int(DT*2+1, 1, CT) * clamp_int(DH*2+1, 1, CH) * clamp_int(DW*2+1, 1, CW) * 3 - 1 ; 
         }
 
         kittens::wait(qsmem_semaphore, 0);
@@ -415,8 +423,9 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
     float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
     float* d_l   = reinterpret_cast<float*>(l_ptr);
 
-    cudaDeviceSynchronize();
-    auto stream = at::cuda::getCurrentCUDAStream().stream(); 
+    //cudadevicesynchronize();
+    const c10::cuda::OptionalCUDAGuard device_guard(q.device());
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
 
 
     if (head_dim == 128) {
@@ -442,8 +451,8 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
 
         globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len),  static_cast<int>(text_length), static_cast<int>(hr)};
 
-        auto mem_size = kittens::MAX_SHARED_MEMORY;
-        auto threads  = NUM_WORKERS * kittens::WARP_THREADS;
+        constexpr int mem_size = kittens::MAX_SHARED_MEMORY;
+        int threads = NUM_WORKERS * kittens::WARP_THREADS;
         if (has_text) {
             // TORCH_CHECK(seq_len % (CONSUMER_WARPGROUPS*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 192");
             dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4)-2, qo_heads, batch);
@@ -823,10 +832,10 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
 
         }
         CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
     }
 
     return o;
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
 }
 

csrc/attn/tests/test_block_sparse.py

@@ -7,6 +7,7 @@ from vsa import BLOCK_M, BLOCK_N
 
 import numpy as np
 import random
+import gc
 
 def set_seed(seed: int = 42):
     # Python random module
@@ -20,15 +21,6 @@ def set_seed(seed: int = 42):
     torch.cuda.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)  # if using multi-GPU
 
-def parse_arguments():
-    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
-    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
-    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
-    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
-    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
-    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
-    parser.add_argument('--num_iterations', type=int, default=100, help='Number of test iterations to run')
-    return parser.parse_args()
 
 @torch.no_grad
 def precision_metric(quant_o, fa2_o): 
@@ -135,9 +127,7 @@ def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device=
                 
     return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
 
-def main():
-    args = parse_arguments()
-
+def main(args):
     set_seed(42)
     
     # Extract parameters
@@ -191,23 +181,36 @@ def main():
             block_mask_expanded = block_sparse_mask.unsqueeze(-1).unsqueeze(-2)  # [b, h, num_q_blocks, num_kv_blocks, 1, 1]
             block_mask_expanded = block_mask_expanded.expand(-1, -1, -1, -1, BLOCK_M, BLOCK_N)  # [b, h, num_q_blocks, num_kv_blocks, BLOCK_M, BLOCK_N]
             full_mask = block_mask_expanded.permute(0, 1, 2, 4, 3, 5).reshape(batch, head, seq_len, seq_len)
-            
-            q_sdpa = q.clone()
-            k_sdpa = k.clone()
-            v_sdpa = v.clone()
 
             q.requires_grad = True
             k.requires_grad = True
             v.requires_grad = True
-            q_sdpa.requires_grad = True
-            k_sdpa.requires_grad = True
-            v_sdpa.requires_grad = True
+
 
             # testing forward
             o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
+            del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded
+            grad_o = torch.randn_like(o)
+            o.backward(grad_o)
+            # clear memory
+            q_sdpa = q.detach().clone()
+            k_sdpa = k.detach().clone()
+            v_sdpa = v.detach().clone()
+            q_sdpa.requires_grad = True
+            k_sdpa.requires_grad = True
+            v_sdpa.requires_grad = True
+            q.data = torch.empty(0, device=q.device)
+            k.data = torch.empty(0, device=k.device)
+            v.data = torch.empty(0, device=v.device)
+            torch.cuda.empty_cache()
+
             o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask)
 
+
             sim, l1, rmse = precision_metric(o, o_sdpa)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 8e-5, f"l1 too large: {l1}"
+            assert rmse < 2e-5, f"RMSE too large: {rmse}"
             forward_metrics['sim'].append(sim)
             forward_metrics['l1'].append(l1)
             forward_metrics['rmse'].append(rmse)
@@ -215,52 +218,72 @@ def main():
             print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
 
             # test backward
-            grad_o = torch.randn_like(o)
-            o.backward(grad_o)
             o_sdpa.backward(grad_o)
 
             sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad)
+            # Error bounds collected on H100
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 3e-4, f"RMSE too large: {rmse}"
             grad_q_metrics['sim'].append(sim)
             grad_q_metrics['l1'].append(l1)
             grad_q_metrics['rmse'].append(rmse)
             print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
             
             sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 2e-4, f"RMSE too large: {rmse}"
             grad_k_metrics['sim'].append(sim)
             grad_k_metrics['l1'].append(l1)
             grad_k_metrics['rmse'].append(rmse)
             print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
             
             sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 1e-4, f"l1 too large: {l1}"
+            assert rmse < 2e-5, f"RMSE too large: {rmse}"
             grad_v_metrics['sim'].append(sim)
             grad_v_metrics['l1'].append(l1)
             grad_v_metrics['rmse'].append(rmse)
             print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
-        
+            
+            del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa
+            gc.collect()
+            torch.cuda.empty_cache()
+
         # Print summary statistics if multiple iterations were run
         if num_iterations > 1:
             print("\n" + "="*50)
             print(f"Summary Statistics (over {num_iterations} iterations):")
             
             print("\nForward metrics:")
-            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}")
             
             print("\nGradient Q metrics:")
-            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}")
             
             print("\nGradient K metrics:")
-            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}")
             
             print("\nGradient V metrics:")
-            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}")
 
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
+    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
+    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
+    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
+    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
+    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
+    parser.add_argument('--num_iterations', type=int, default=50, help='Number of test iterations to run')
+    args = parser.parse_args()
+    main(args)

csrc/attn/tests/test_sta.py

@@ -81,5 +81,7 @@ std = 10
 
 # Run correctness check directly
 results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output')
+assert results['TK vs FLEX']['avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large"
+assert results['TK vs FLEX']['max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large"
 print(f"Average difference: {results['TK vs FLEX']['avg_diff']}")
 print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}")

csrc/attn/vsa/block_sparse_h100.cu

@@ -3,6 +3,8 @@
 #include "kittens.cuh"
 #include <cooperative_groups.h>
 #include <iostream>
+#include <c10/cuda/CUDAGuard.h>
+
 
 using namespace kittens;
 namespace cg = cooperative_groups;
@@ -940,8 +942,9 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
     float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
     float* d_l   = reinterpret_cast<float*>(l_ptr);
 
-    cudaDeviceSynchronize();
-    auto stream = at::cuda::getCurrentCUDAStream().stream(); 
+    //cudadevicesynchronize();
+    const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
 
     if (head_dim == 64) {
         using q_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
@@ -966,7 +969,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
 
         globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr())};
 
-        auto mem_size = 54000;
+        constexpr int mem_size = 54000;
 
         dim3 grid(seq_len/(64), qo_heads, batch);
 
@@ -979,7 +982,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
         fwd_attend_ker<64><<<grid, (128), mem_size, stream>>>(g);
 
         CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
     }
 
     if (head_dim == 128) {
@@ -1005,7 +1008,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
 
         globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr())};
 
-        auto mem_size = 54000;
+        constexpr int mem_size = 54000;
 
         dim3 grid(seq_len/(64), qo_heads, batch);
 
@@ -1018,11 +1021,11 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
         fwd_attend_ker<128><<<grid, (128), mem_size, stream>>>(g);
 
         CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
     }
 
     return {o, l_vec};
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
 }
 
 std::vector<torch::Tensor> 
@@ -1132,13 +1135,14 @@ block_sparse_attention_backward(torch::Tensor q,
     float* d_kg = reinterpret_cast<float*>(kg_ptr);
     float* d_vg = reinterpret_cast<float*>(vg_ptr);
 
-    auto mem_size = kittens::MAX_SHARED_MEMORY; 
-    auto threads  = 4 * kittens::WARP_THREADS;
+    constexpr int mem_size = kittens::MAX_SHARED_MEMORY; 
+    int threads  = 4 * kittens::WARP_THREADS;
 
-    cudaDeviceSynchronize();
-    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    //cudadevicesynchronize();
+    const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
 
-    cudaStreamSynchronize(stream);
+    //  cudaStreamSynchronize(stream);
 
     // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256");
     dim3 grid_bwd(seq_len/(4*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
@@ -1222,7 +1226,7 @@ block_sparse_attention_backward(torch::Tensor q,
         dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
         threads = 128;
 
-        cudaDeviceSynchronize();
+        //cudadevicesynchronize();
 
         {
             cudaFuncSetAttribute(
@@ -1240,8 +1244,8 @@ block_sparse_attention_backward(torch::Tensor q,
         }
 
         // CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
-        cudaDeviceSynchronize();
+        // cudaStreamSynchronize(stream);
+        //cudadevicesynchronize();
         // const auto kernel_end = std::chrono::high_resolution_clock::now();
         // std::cout << "Kernel Time: " << std::chrono::duration_cast<std::chrono::microseconds>(kernel_end - start).count() << "us" << std::endl;
         // std::cout << "---" << std::endl;
@@ -1326,7 +1330,7 @@ block_sparse_attention_backward(torch::Tensor q,
         dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
         threads = 128;
 
-        cudaDeviceSynchronize();
+        //cudadevicesynchronize();
 
         {
             cudaFuncSetAttribute(
@@ -1338,10 +1342,10 @@ block_sparse_attention_backward(torch::Tensor q,
             bwd_attend_ker<128><<<grid_bwd_2, threads, 113000, stream>>>(bwd_global); 
         }
 
-        cudaStreamSynchronize(stream);
-        cudaDeviceSynchronize();
+        // cudaStreamSynchronize(stream);
+        //cudadevicesynchronize();
     }
 
     return {qg, kg, vg};
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
 }

examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json

@@ -0,0 +1,31 @@
+{
+  "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": "examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4",
+      "num_inference_steps": 50,
+      "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": "examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4",
+      "num_inference_steps": 50,
+      "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": "examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4",
+      "num_inference_steps": 50,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}

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

@@ -12,6 +12,7 @@ class DiTArchConfig(ArchConfig):
     _fsdp_shard_conditions: list = field(default_factory=list)
     _compile_conditions: list = field(default_factory=list)
     _param_names_mapping: dict = field(default_factory=dict)
+    _reverse_param_names_mapping: dict = field(default_factory=dict)
     _lora_param_names_mapping: dict = field(default_factory=dict)
     _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
         AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,

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

@@ -147,6 +147,9 @@ class HunyuanVideoArchConfig(DiTArchConfig):
             r"final_layer.linear.\1",
         })
 
+    # Reverse mapping for saving checkpoints: training -> diffusers
+    _reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
     patch_size: int = 2
     patch_size_t: int = 1
     in_channels: int = 16

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

@@ -49,9 +49,13 @@ class WanVideoArchConfig(DiTArchConfig):
             r"blocks.\1.ffn.fc_in.\2",
             r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
             r"blocks.\1.ffn.fc_out.\2",
-            r"blocks\.(\d+)\.norm2\.(.*)$":
+            r"^blocks\.(\d+)\.norm2\.(.*)$":
             r"blocks.\1.self_attn_residual_norm.norm.\2",
         })
+
+    # Reverse mapping for saving checkpoints: training -> diffusers
+    _reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
     # Some LoRA adapters use the original official layer names instead of hf layer names,
     # so apply this before the param_names_mapping
     _lora_param_names_mapping: dict = field(

fastvideo/v1/dataset/__init__.py

@@ -1,19 +1,17 @@
-import os
-
+# SPDX-License-Identifier: Apache-2.0
 from torchvision import transforms
 from torchvision.transforms import Lambda
-from transformers import AutoTokenizer
 
-from fastvideo.v1.dataset.t2v_datasets import T2V_dataset
+from fastvideo.v1.dataset.parquet_dataset_map_style import (
+    build_parquet_map_style_dataloader)
+from fastvideo.v1.dataset.preprocessing_datasets import (
+    VideoCaptionMergedDataset)
 from fastvideo.v1.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                             TemporalRandomCrop)
-
-from .parquet_dataset_map_style import build_parquet_map_style_dataloader
-
-__all__ = ["build_parquet_map_style_dataloader"]
+from fastvideo.v1.dataset.validation_dataset import ValidationDataset
 
 
-def getdataset(args, start_idx=0) -> T2V_dataset:
+def getdataset(args) -> VideoCaptionMergedDataset:
     temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
     norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
     resize_topcrop = [
@@ -31,15 +29,17 @@ def getdataset(args, start_idx=0) -> T2V_dataset:
         *resize_topcrop,
         norm_fun,
     ])
-    tokenizer_path = os.path.join(args.model_path, "tokenizer")
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
-                                              cache_dir=args.cache_dir)
     if args.dataset == "t2v":
-        return T2V_dataset(args,
-                           transform=transform,
-                           temporal_sample=temporal_sample,
-                           tokenizer=tokenizer,
-                           transform_topcrop=transform_topcrop,
-                           start_idx=start_idx)
+        return VideoCaptionMergedDataset(data_merge_path=args.data_merge_path,
+                                         args=args,
+                                         transform=transform,
+                                         temporal_sample=temporal_sample,
+                                         transform_topcrop=transform_topcrop)
 
     raise NotImplementedError(args.dataset)
+
+
+__all__ = [
+    "build_parquet_map_style_dataloader", "ValidationDataset",
+    "VideoCaptionMergedDataset"
+]

fastvideo/v1/dataset/dataloader/schema.py

@@ -48,6 +48,48 @@ pyarrow_schema_i2v = pa.schema([
     pa.field("fps", pa.float64()),
 ])
 
+pyarrow_schema_i2v_validation = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Image/Video VAE latents ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("vae_latent_bytes", pa.binary()),
+    # e.g., [C, T, H, W] or [C, H, W]
+    pa.field("vae_latent_shape", pa.list_(pa.int64())),
+    # e.g., 'float32'
+    pa.field("vae_latent_dtype", pa.string()),
+    # --- Text encoder output tensor ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("text_embedding_bytes", pa.binary()),
+    # e.g., [SeqLen, Dim]
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    # e.g., 'bfloat16' or 'float32'
+    pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_attention_mask_bytes", pa.binary()),
+    # e.g., [SeqLen]
+    pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+    # e.g., 'bool' or 'int8'
+    pa.field("text_attention_mask_dtype", pa.string()),
+    #I2V
+    pa.field("clip_feature_bytes", pa.binary()),
+    pa.field("clip_feature_shape", pa.list_(pa.int64())),
+    pa.field("clip_feature_dtype", pa.string()),
+    # I2V Validation
+    pa.field("pil_image_bytes", pa.binary()),
+    pa.field("pil_image_shape", pa.list_(pa.int64())),
+    pa.field("pil_image_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # 'image' or 'video'
+    pa.field("width", pa.int64()),
+    pa.field("height", pa.int64()),
+    # -- Video-specific (can be null/default for images) ---
+    # Number of frames processed (e.g., 1 for image, N for video)
+    pa.field("num_frames", pa.int64()),
+    pa.field("duration_sec", pa.float64()),
+    pa.field("fps", pa.float64()),
+])
+
 pyarrow_schema_t2v = pa.schema([
     pa.field("id", pa.string()),
     # --- Image/Video VAE latents ---
@@ -80,4 +122,38 @@ pyarrow_schema_t2v = pa.schema([
     pa.field("num_frames", pa.int64()),
     pa.field("duration_sec", pa.float64()),
     pa.field("fps", pa.float64()),
-])
+])
+
+pyarrow_schema_t2v_validation = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Image/Video VAE latents ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("vae_latent_bytes", pa.binary()),
+    # e.g., [C, T, H, W] or [C, H, W]
+    pa.field("vae_latent_shape", pa.list_(pa.int64())),
+    # e.g., 'float32'
+    pa.field("vae_latent_dtype", pa.string()),
+    # --- Text encoder output tensor ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("text_embedding_bytes", pa.binary()),
+    # e.g., [SeqLen, Dim]
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    # e.g., 'bfloat16' or 'float32'
+    pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_attention_mask_bytes", pa.binary()),
+    # e.g., [SeqLen]
+    pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+    # e.g., 'bool' or 'int8'
+    pa.field("text_attention_mask_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # 'image' or 'video'
+    pa.field("width", pa.int64()),
+    pa.field("height", pa.int64()),
+    # -- Video-specific (can be null/default for images) ---
+    # Number of frames processed (e.g., 1 for image, N for video)
+    pa.field("num_frames", pa.int64()),
+    pa.field("duration_sec", pa.float64()),
+    pa.field("fps", pa.float64()),
+])

fastvideo/v1/dataset/preprocessing_datasets.py

@@ -0,0 +1,592 @@
+# SPDX-License-Identifier: Apache-2.0
+import json
+import math
+import os
+import random
+from abc import ABC, abstractmethod
+from collections import Counter
+from dataclasses import dataclass
+from os.path import join as opj
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+from PIL import Image
+from transformers import AutoTokenizer
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class DatasetBatch:
+    """
+    Batch information for dataset processing stages.
+    
+    This class holds all the information about a video-caption or image-caption pair
+    as it moves through the processing pipeline. Fields are populated by different stages.
+    """
+    # Raw metadata
+    path: str
+    cap: Union[str, List[str]]
+    resolution: Optional[Dict] = None
+    fps: Optional[float] = None
+    duration: Optional[float] = None
+
+    # Processed metadata
+    num_frames: Optional[int] = None
+    sample_frame_index: Optional[List[int]] = None
+    sample_num_frames: Optional[int] = None
+
+    # Processed data
+    pixel_values: Optional[torch.Tensor] = None
+    text: Optional[str] = None
+    input_ids: Optional[torch.Tensor] = None
+    cond_mask: Optional[torch.Tensor] = None
+
+    @property
+    def is_video(self) -> bool:
+        """Check if this is a video item."""
+        return self.path.endswith(".mp4")
+
+    @property
+    def is_image(self) -> bool:
+        """Check if this is an image item."""
+        return self.path.endswith(".jpg")
+
+
+class DatasetStage(ABC):
+    """
+    Abstract base class for dataset processing stages.
+    
+    Similar to PipelineStage but designed for dataset preprocessing operations.
+    """
+
+    @abstractmethod
+    def process(self, batch: DatasetBatch, **kwargs) -> DatasetBatch:
+        """
+        Process the dataset batch.
+        
+        Args:
+            batch: Dataset batch to process
+            **kwargs: Additional processing parameters
+            
+        Returns:
+            Processed batch
+        """
+        raise NotImplementedError
+
+
+class DatasetFilterStage(ABC):
+    """
+    Abstract base class for dataset filtering stages.
+    
+    These stages can filter out items during metadata processing.
+    """
+
+    @abstractmethod
+    def should_keep(self, batch: DatasetBatch, **kwargs) -> bool:
+        """
+        Check if batch should be kept.
+        
+        Args:
+            batch: Dataset batch to check
+            **kwargs: Additional parameters
+            
+        Returns:
+            True if batch should be kept, False otherwise
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def process(self, batch: DatasetBatch, **kwargs) -> DatasetBatch:
+        """
+        Process the dataset batch (for non-filtering operations).
+        
+        Args:
+            batch: Dataset batch to process
+            **kwargs: Additional processing parameters
+            
+        Returns:
+            Processed batch
+        """
+        raise NotImplementedError
+
+
+class DataValidationStage(DatasetFilterStage):
+    """Stage for validating data items."""
+
+    def should_keep(self, batch: DatasetBatch, **kwargs) -> bool:
+        """
+        Validate data item.
+        
+        Args:
+            batch: Dataset batch to validate
+            
+        Returns:
+            True if valid, False if invalid
+        """
+        # Check for caption
+        if batch.cap is None:
+            return False
+
+        if batch.is_video:
+            # Validate video-specific fields
+            if batch.duration is None or batch.fps is None:
+                return False
+        elif not batch.is_image:
+            return False
+
+        return True
+
+    def process(self, batch: DatasetBatch, **kwargs) -> DatasetBatch:
+        """Process does nothing for validation - filtering is handled by should_keep."""
+        return batch
+
+
+class ResolutionFilterStage(DatasetFilterStage):
+    """Stage for filtering data items based on resolution constraints."""
+
+    def __init__(self,
+                 max_h_div_w_ratio: float = 17 / 16,
+                 min_h_div_w_ratio: float = 8 / 16,
+                 max_height: int = 1024,
+                 max_width: int = 1024):
+        self.max_h_div_w_ratio = max_h_div_w_ratio
+        self.min_h_div_w_ratio = min_h_div_w_ratio
+        self.max_height = max_height
+        self.max_width = max_width
+
+    def should_keep(self, batch: DatasetBatch, **kwargs) -> bool:
+        """
+        Check if data item passes resolution filtering.
+        
+        Args:
+            batch: Dataset batch with resolution information
+            
+        Returns:
+            True if passes filter, False otherwise
+        """
+        # Only apply to videos
+        if not batch.is_video:
+            return True
+
+        if batch.resolution is None:
+            return False
+
+        height = batch.resolution.get("height", None)
+        width = batch.resolution.get("width", None)
+        if height is None or width is None:
+            return False
+
+        # Check aspect ratio
+        aspect = self.max_height / self.max_width
+        hw_aspect_thr = 1.5
+
+        return self.filter_resolution(
+            height,
+            width,
+            max_h_div_w_ratio=hw_aspect_thr * aspect,
+            min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
+        )
+
+    def process(self, batch: DatasetBatch, **kwargs) -> DatasetBatch:
+        """Process does nothing for resolution filtering - filtering is handled by should_keep."""
+        return batch
+
+    def filter_resolution(self, h: int, w: int, max_h_div_w_ratio: float,
+                          min_h_div_w_ratio: float) -> bool:
+        """Filter based on height/width ratio."""
+        return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
+
+
+class FrameSamplingStage(DatasetFilterStage):
+    """Stage for temporal frame sampling and indexing."""
+
+    def __init__(self,
+                 num_frames: int,
+                 train_fps: int,
+                 speed_factor: int = 1,
+                 video_length_tolerance_range: float = 5.0,
+                 drop_short_ratio: float = 0.0):
+        self.num_frames = num_frames
+        self.train_fps = train_fps
+        self.speed_factor = speed_factor
+        self.video_length_tolerance_range = video_length_tolerance_range
+        self.drop_short_ratio = drop_short_ratio
+
+    def should_keep(self, batch: DatasetBatch, **kwargs) -> bool:
+        """
+        Check if video should be kept based on length constraints.
+        
+        Args:
+            batch: Dataset batch
+            
+        Returns:
+            True if should be kept, False otherwise
+        """
+        if batch.is_image:
+            return True
+
+        if batch.duration is None or batch.fps is None:
+            return False
+
+        num_frames = math.ceil(batch.fps * batch.duration)
+
+        # Check if video is too long
+        if (num_frames / batch.fps > self.video_length_tolerance_range *
+            (self.num_frames / self.train_fps * self.speed_factor)):
+            return False
+
+        # Resample frame indices to check length
+        frame_interval = batch.fps / self.train_fps
+        start_frame_idx = 0
+        frame_indices = np.arange(start_frame_idx, num_frames,
+                                  frame_interval).astype(int)
+
+        # Filter short videos
+        return not (len(frame_indices) < self.num_frames
+                    and random.random() < self.drop_short_ratio)
+
+    def process(self,
+                batch: DatasetBatch,
+                temporal_sample_fn=None,
+                **kwargs) -> DatasetBatch:
+        """
+        Process frame sampling for video data items.
+        
+        Args:
+            batch: Dataset batch
+            temporal_sample_fn: Function for temporal sampling
+            
+        Returns:
+            Updated batch with frame sampling info
+        """
+        if batch.is_image:
+            # For images, just add sample info
+            batch.sample_frame_index = [0]
+            batch.sample_num_frames = 1
+            return batch
+
+        assert batch.duration is not None and batch.fps is not None
+        batch.num_frames = math.ceil(batch.fps * batch.duration)
+
+        # Resample frame indices
+        frame_interval = batch.fps / self.train_fps
+        start_frame_idx = 0
+        frame_indices = np.arange(start_frame_idx, batch.num_frames,
+                                  frame_interval).astype(int)
+
+        # Temporal crop if too long
+        if len(frame_indices
+               ) > self.num_frames and temporal_sample_fn is not None:
+            begin_index, end_index = temporal_sample_fn(len(frame_indices))
+            frame_indices = frame_indices[begin_index:end_index]
+
+        batch.sample_frame_index = frame_indices.tolist()
+        batch.sample_num_frames = len(frame_indices)
+
+        return batch
+
+
+class VideoTransformStage(DatasetStage):
+    """Stage for video data transformation."""
+
+    def __init__(self, transform) -> None:
+        self.transform = transform
+
+    def process(self, batch: DatasetBatch, **kwargs) -> DatasetBatch:
+        """
+        Transform video data.
+        
+        Args:
+            batch: Dataset batch with video information
+            
+        Returns:
+            Batch with transformed video tensor
+        """
+        if not batch.is_video:
+            return batch
+
+        assert os.path.exists(batch.path), f"file {batch.path} do not exist!"
+        assert batch.sample_frame_index is not None, "Frame indices must be set before transformation"
+
+        torchvision_video, _, metadata = torchvision.io.read_video(
+            batch.path, output_format="TCHW")
+        video = torchvision_video[batch.sample_frame_index]
+        if self.transform is not None:
+            video = self.transform(video)
+        video = rearrange(video, "t c h w -> c t h w")
+        video = video.to(torch.uint8)
+
+        h, w = video.shape[-2:]
+        assert (
+            h / w <= 17 / 16 and h / w >= 8 / 16
+        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({batch.path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
+
+        video = video.float() / 127.5 - 1.0
+        batch.pixel_values = video
+        return batch
+
+
+class ImageTransformStage(DatasetStage):
+    """Stage for image data transformation."""
+
+    def __init__(self, transform, transform_topcrop) -> None:
+        self.transform = transform
+        self.transform_topcrop = transform_topcrop
+
+    def process(self, batch: DatasetBatch, **kwargs) -> DatasetBatch:
+        """
+        Transform image data.
+        
+        Args:
+            batch: Dataset batch with image information
+            
+        Returns:
+            Batch with transformed image tensor
+        """
+        if not batch.is_image:
+            return batch
+
+        image = Image.open(batch.path).convert("RGB")
+        image = torch.from_numpy(np.array(image))
+        image = rearrange(image, "h w c -> c h w").unsqueeze(0)
+
+        if self.transform_topcrop is not None:
+            image = self.transform_topcrop(image)
+        elif self.transform is not None:
+            image = self.transform(image)
+
+        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
+        image = image.float() / 127.5 - 1.0
+        batch.pixel_values = image
+        return batch
+
+
+class TextEncodingStage(DatasetStage):
+    """Stage for text tokenization and encoding."""
+
+    def __init__(self, tokenizer, text_max_length: int, cfg_rate: float = 0.0):
+        self.tokenizer = tokenizer
+        self.text_max_length = text_max_length
+        self.cfg_rate = cfg_rate
+
+    def process(self, batch: DatasetBatch, **kwargs) -> DatasetBatch:
+        """
+        Process text data.
+        
+        Args:
+            batch: Dataset batch with caption information
+            
+        Returns:
+            Batch with encoded text information
+        """
+        text = batch.cap
+        if not isinstance(text, list):
+            text = [text]
+        text = [random.choice(text)]
+
+        text = text[0] if random.random() > self.cfg_rate else ""
+        text_tokens_and_mask = self.tokenizer(
+            text,
+            max_length=self.text_max_length,
+            padding="max_length",
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+
+        batch.text = text
+        batch.input_ids = text_tokens_and_mask["input_ids"]
+        batch.cond_mask = text_tokens_and_mask["attention_mask"]
+        return batch
+
+
+class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
+                                torch.distributed.checkpoint.stateful.Stateful):
+    """
+    Merged dataset for video and caption data with stage-based processing.
+    
+    This dataset processes video and image data through a series of stages:
+    - Data validation
+    - Resolution filtering  
+    - Frame sampling
+    - Transformation
+    - Text encoding
+    """
+
+    def __init__(self,
+                 data_merge_path: str,
+                 args,
+                 transform,
+                 temporal_sample,
+                 transform_topcrop,
+                 start_idx: int = 0):
+        self.data_merge_path = data_merge_path
+        self.start_idx = start_idx
+        self.args = args
+        self.temporal_sample = temporal_sample
+
+        # Initialize tokenizer
+        tokenizer_path = os.path.join(args.model_path, "tokenizer")
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
+                                                  cache_dir=args.cache_dir)
+
+        # Initialize processing stages
+        self._init_stages(args, transform, transform_topcrop, tokenizer)
+
+        # Process metadata
+        self.processed_batches = self._process_metadata()
+
+    def _init_stages(self, args, transform, transform_topcrop,
+                     tokenizer) -> None:
+        """Initialize all processing stages."""
+        self.validation_stage = DataValidationStage()
+        self.resolution_filter_stage = ResolutionFilterStage(
+            max_height=args.max_height, max_width=args.max_width)
+        self.frame_sampling_stage = FrameSamplingStage(
+            num_frames=args.num_frames,
+            train_fps=args.train_fps,
+            speed_factor=args.speed_factor,
+            video_length_tolerance_range=args.video_length_tolerance_range,
+            drop_short_ratio=args.drop_short_ratio)
+        self.video_transform_stage = VideoTransformStage(transform)
+        self.image_transform_stage = ImageTransformStage(
+            transform, transform_topcrop)
+        self.text_encoding_stage = TextEncodingStage(
+            tokenizer=tokenizer,
+            text_max_length=args.text_max_length,
+            cfg_rate=args.cfg)
+
+    def _load_raw_data(self) -> List[Dict]:
+        """Load raw data from JSON files."""
+        all_data = []
+
+        # Read folder-annotation pairs
+        with open(self.data_merge_path) as f:
+            folder_anno_pairs = [
+                line.strip().split(",") for line in f if line.strip()
+            ]
+
+        # Process each folder-annotation pair
+        for folder, annotation_file in folder_anno_pairs:
+            with open(annotation_file) as f:
+                data_items = json.load(f)
+
+            # Update paths with folder prefix
+            for item in data_items:
+                item["path"] = opj(folder, item["path"])
+
+            all_data.extend(data_items)
+
+        return all_data[self.start_idx:]
+
+    def _process_metadata(self) -> List[DatasetBatch]:
+        """Process the raw metadata through all filtering stages."""
+        raw_data = self._load_raw_data()
+        processed_batches = []
+
+        # Initialize counters
+        filter_counts = {
+            "validation_failed": 0,
+            "resolution_failed": 0,
+            "frame_sampling_failed": 0
+        }
+        sample_num_frames: List[int] = []
+
+        for item in raw_data:
+            batch = DatasetBatch(path=item["path"],
+                                 cap=item["cap"],
+                                 resolution=item.get("resolution"),
+                                 fps=item.get("fps"),
+                                 duration=item.get("duration"))
+
+            # Apply filtering stages
+            if not self._apply_filter_stages(batch, filter_counts):
+                continue
+
+            # Apply frame sampling processing
+            batch = self.frame_sampling_stage.process(
+                batch, temporal_sample_fn=self.temporal_sample)
+
+            processed_batches.append(batch)
+            assert batch.sample_num_frames is not None
+            sample_num_frames.append(batch.sample_num_frames)
+
+        self._log_filtering_stats(filter_counts, sample_num_frames,
+                                  len(raw_data), len(processed_batches))
+        return processed_batches
+
+    def _apply_filter_stages(self, batch: DatasetBatch,
+                             filter_counts: Dict[str, int]) -> bool:
+        """Apply all filter stages and update counters. Returns True if batch should be kept."""
+        if not self.validation_stage.should_keep(batch):
+            filter_counts["validation_failed"] += 1
+            return False
+
+        if not self.resolution_filter_stage.should_keep(batch):
+            filter_counts["resolution_failed"] += 1
+            return False
+
+        if not self.frame_sampling_stage.should_keep(batch):
+            filter_counts["frame_sampling_failed"] += 1
+            return False
+
+        return True
+
+    def _log_filtering_stats(self, filter_counts: Dict[str, int],
+                             sample_num_frames: List[int], before_count: int,
+                             after_count: int):
+        """Log filtering statistics."""
+        logger.info(
+            "validation_failed: %d, resolution_failed: %d, frame_sampling_failed: %d, "
+            "Counter(sample_num_frames): %s, before filter: %d, after filter: %d",
+            filter_counts['validation_failed'],
+            filter_counts['resolution_failed'],
+            filter_counts['frame_sampling_failed'], Counter(sample_num_frames),
+            before_count, after_count)
+
+    def __iter__(self):
+        """Iterate through processed data items."""
+        for idx in range(len(self.processed_batches)):
+            yield self._get_item(idx)
+
+    def __len__(self):
+        return len(self.processed_batches)
+
+    def _get_item(self, idx: int) -> Dict:
+        """Get a single processed data item."""
+        batch = self.processed_batches[idx]
+
+        # Apply transformation stages
+        batch = self.video_transform_stage.process(batch)
+        batch = self.image_transform_stage.process(batch)
+        batch = self.text_encoding_stage.process(batch)
+
+        # Build result dictionary
+        result = {
+            "pixel_values": batch.pixel_values,
+            "text": batch.text,
+            "input_ids": batch.input_ids,
+            "cond_mask": batch.cond_mask,
+            "path": batch.path,
+        }
+
+        # Add video-specific fields
+        if batch.is_video:
+            result.update({"fps": batch.fps, "duration": batch.duration})
+
+        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/v1/dataset/t2v_datasets.py

@@ -1,352 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import json
-import math
-import os
-import random
-from collections import Counter
-from os.path import join as opj
-
-import numpy as np
-import torch
-import torchvision
-from einops import rearrange
-from PIL import Image
-from torch.utils.data import Dataset
-
-from fastvideo.utils.dataset_utils import DecordInit
-from fastvideo.utils.logging_ import main_print
-
-
-class SingletonMeta(type):
-    _instances: dict[type, 'SingletonMeta'] = {}
-
-    def __call__(cls, *args, **kwargs):
-        if cls not in cls._instances:
-            instance = super().__call__(*args, **kwargs)
-            cls._instances[cls] = instance
-        return cls._instances[cls]
-
-
-class DataSetProg(metaclass=SingletonMeta):
-
-    def __init__(self) -> None:
-        self.cap_list: list[dict] = []
-        self.elements: list[int] = []
-        self.num_workers = 1
-        self.n_elements = 0
-        self.worker_elements: dict[int, list[int]] = {}
-        self.n_used_elements: dict[int, int] = {}
-
-    def set_cap_list(self, num_workers, cap_list, n_elements) -> None:
-        self.num_workers = num_workers
-        self.cap_list = cap_list
-        self.n_elements = n_elements
-        self.elements = list(range(n_elements))
-        random.shuffle(self.elements)
-        print(f"n_elements: {len(self.elements)}", flush=True)
-
-        for i in range(self.num_workers):
-            self.n_used_elements[i] = 0
-            per_worker = int(
-                math.ceil(len(self.elements) / float(self.num_workers)))
-            start = i * per_worker
-            end = min(start + per_worker, len(self.elements))
-            self.worker_elements[i] = self.elements[start:end]
-
-    def get_item(self, work_info) -> int:
-        worker_id = 0 if work_info is None else work_info.id
-
-        idx = self.worker_elements[worker_id][
-            self.n_used_elements[worker_id] %
-            len(self.worker_elements[worker_id])]
-        self.n_used_elements[worker_id] += 1
-        return idx
-
-
-dataset_prog = DataSetProg()
-
-
-def filter_resolution(h: int,
-                      w: int,
-                      max_h_div_w_ratio: float = 17 / 16,
-                      min_h_div_w_ratio: float = 8 / 16) -> bool:
-    return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
-
-
-class T2V_dataset(Dataset):
-
-    def __init__(self,
-                 args,
-                 transform,
-                 temporal_sample,
-                 tokenizer,
-                 transform_topcrop,
-                 start_idx=0) -> None:
-        self.start_idx = start_idx
-        self.data = args.data_merge_path
-        self.num_frames = args.num_frames
-        self.train_fps = args.train_fps
-        self.use_image_num = args.use_image_num
-        self.transform = transform
-        self.transform_topcrop = transform_topcrop
-        self.temporal_sample = temporal_sample
-        self.tokenizer = tokenizer
-        self.text_max_length = args.text_max_length
-        self.cfg = args.cfg
-        self.speed_factor = args.speed_factor
-        self.max_height = args.max_height
-        self.max_width = args.max_width
-        self.drop_short_ratio = args.drop_short_ratio
-        assert self.speed_factor >= 1
-        self.v_decoder = DecordInit()
-        self.video_length_tolerance_range = args.video_length_tolerance_range
-        self.support_Chinese = True
-        if "mt5" not in args.text_encoder_name:
-            self.support_Chinese = False
-
-        cap_list = self.get_cap_list()
-
-        assert len(cap_list) > 0
-        cap_list, self.sample_num_frames = self.define_frame_index(cap_list)
-        self.lengths = self.sample_num_frames
-
-        n_elements = len(cap_list)
-        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list,
-                                  n_elements)
-
-        print(f"video length: {len(dataset_prog.cap_list)}", flush=True)
-
-    def set_checkpoint(self, n_used_elements):
-        for i in range(len(dataset_prog.n_used_elements)):
-            dataset_prog.n_used_elements[i] = n_used_elements
-
-    def __len__(self):
-        return dataset_prog.n_elements
-
-    def __getitem__(self, idx):
-
-        data = self.get_data(idx)
-        return data
-
-    def get_data(self, idx) -> dict:
-        path = dataset_prog.cap_list[idx]["path"]
-        if path.endswith(".mp4"):
-            return self.get_video(idx)
-        else:
-            return self.get_image(idx)
-
-    def get_video(self, idx) -> dict:
-        video_path = dataset_prog.cap_list[idx]["path"]
-        assert os.path.exists(video_path), f"file {video_path} do not exist!"
-        frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
-
-        torchvision_video, _, metadata = torchvision.io.read_video(
-            video_path, output_format="TCHW")
-        video = torchvision_video[frame_indices]
-        video = self.transform(video)
-        video = rearrange(video, "t c h w -> c t h w")
-        video = video.to(torch.uint8)
-        assert video.dtype == torch.uint8
-
-        h, w = video.shape[-2:]
-        assert (
-            h / w <= 17 / 16 and h / w >= 8 / 16
-        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({video_path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
-
-        video = video.float() / 127.5 - 1.0
-
-        text = dataset_prog.cap_list[idx]["cap"]
-        if not isinstance(text, list):
-            text = [text]
-        text = [random.choice(text)]
-
-        text = text[0] if random.random() > self.cfg else ""
-        text_tokens_and_mask = self.tokenizer(
-            text,
-            max_length=self.text_max_length,
-            padding="max_length",
-            truncation=True,
-            return_attention_mask=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        input_ids = text_tokens_and_mask["input_ids"]
-        cond_mask = text_tokens_and_mask["attention_mask"]
-        return dict(pixel_values=video,
-                    text=text,
-                    input_ids=input_ids,
-                    cond_mask=cond_mask,
-                    path=video_path,
-                    fps=dataset_prog.cap_list[idx]["fps"],
-                    duration=dataset_prog.cap_list[idx]["duration"])
-
-    def get_image(self, idx) -> dict:
-        image_data = dataset_prog.cap_list[
-            idx]  # [{'path': path, 'cap': cap}, ...]
-
-        image = Image.open(image_data["path"]).convert("RGB")  # [h, w, c]
-        image = torch.from_numpy(np.array(image))  # [h, w, c]
-        image = rearrange(image, "h w c -> c h w").unsqueeze(0)  #  [1 c h w]
-        # for i in image:
-        #     h, w = i.shape[-2:]
-        #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
-
-        image = (self.transform_topcrop(image) if "human_images"
-                 in image_data["path"] else self.transform(image)
-                 )  #  [1 C H W] -> num_img [1 C H W]
-        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
-
-        image = image.float() / 127.5 - 1.0
-
-        caps: list[str] = (image_data["cap"] if isinstance(
-            image_data["cap"], list) else [image_data["cap"]])
-        caps = [random.choice(caps)]
-        text = caps
-        input_ids, cond_mask = [], []
-        single_text = text[0] if random.random() > self.cfg else ""
-        text_tokens_and_mask = self.tokenizer(
-            single_text,
-            max_length=self.text_max_length,
-            padding="max_length",
-            truncation=True,
-            return_attention_mask=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        input_ids = text_tokens_and_mask["input_ids"]  # 1, l
-        cond_mask = text_tokens_and_mask["attention_mask"]  # 1, l
-        return dict(
-            pixel_values=image,
-            text=text,
-            input_ids=input_ids,
-            cond_mask=cond_mask,
-            path=image_data["path"],
-        )
-
-    def define_frame_index(self, cap_list) -> tuple[list[dict], list[int]]:
-        new_cap_list = []
-        sample_num_frames = []
-        cnt_too_long = 0
-        cnt_too_short = 0
-        cnt_no_cap = 0
-        cnt_no_resolution = 0
-        cnt_resolution_mismatch = 0
-        cnt_movie = 0
-        cnt_img = 0
-        for i in cap_list:
-            path = i["path"]
-            cap = i.get("cap", None)
-            # ======no caption=====
-            if cap is None:
-                cnt_no_cap += 1
-                continue
-            if path.endswith(".mp4"):
-                # ======no fps and duration=====
-                duration = i.get("duration", None)
-                fps = i.get("fps", None)
-                if fps is None or duration is None:
-                    continue
-
-                # ======resolution mismatch=====
-                resolution = i.get("resolution", None)
-                if resolution is None:
-                    cnt_no_resolution += 1
-                    continue
-                else:
-                    if (resolution.get("height", None) is None
-                            or resolution.get("width", None) is None):
-                        cnt_no_resolution += 1
-                        continue
-                    height, width = i["resolution"]["height"], i["resolution"][
-                        "width"]
-                    aspect = self.max_height / self.max_width
-                    hw_aspect_thr = 1.5
-                    is_pick = filter_resolution(
-                        height,
-                        width,
-                        max_h_div_w_ratio=hw_aspect_thr * aspect,
-                        min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
-                    )
-                    if not is_pick:
-                        print("resolution mismatch")
-                        cnt_resolution_mismatch += 1
-                        continue
-
-                # if path == 'finetrainers/3dgs-dissolve/videos/1.mp4':
-                #     from IPython import embed; embed()
-                i["num_frames"] = math.ceil(fps * duration)
-                # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
-                if i["num_frames"] / fps > self.video_length_tolerance_range * (
-                        self.num_frames / self.train_fps * self.speed_factor
-                ):  # too long video is not suitable for this training stage (self.num_frames)
-                    cnt_too_long += 1
-                    continue
-
-                # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
-                frame_interval = fps / self.train_fps
-                start_frame_idx = 0
-                frame_indices = np.arange(start_frame_idx, i["num_frames"],
-                                          frame_interval).astype(int)
-
-                # comment out it to enable dynamic frames training
-                if (len(frame_indices) < self.num_frames
-                        and random.random() < self.drop_short_ratio):
-                    cnt_too_short += 1
-                    continue
-
-                #  too long video will be temporal-crop randomly
-                if len(frame_indices) > self.num_frames:
-                    begin_index, end_index = self.temporal_sample(
-                        len(frame_indices))
-                    frame_indices = frame_indices[begin_index:end_index]
-                    # frame_indices = frame_indices[:self.num_frames]  # head crop
-                i["sample_frame_index"] = frame_indices.tolist()
-                new_cap_list.append(i)
-                i["sample_num_frames"] = len(
-                    i["sample_frame_index"]
-                )  # will use in dataloader(group sampler)
-                sample_num_frames.append(i["sample_num_frames"])
-            elif path.endswith(".jpg"):  # image
-                cnt_img += 1
-                new_cap_list.append(i)
-                i["sample_num_frames"] = 1
-                sample_num_frames.append(i["sample_num_frames"])
-            else:
-                raise NameError(
-                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image"
-                )
-        # import ipdb;ipdb.set_trace()
-        main_print(
-            f"no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, "
-            f"no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, "
-            f"Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, "
-            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}"
-        )
-        return new_cap_list, sample_num_frames
-
-    def decord_read(self, path, frame_indices) -> torch.Tensor:
-        decord_vr = self.v_decoder(path)
-        video_data = decord_vr.get_batch(frame_indices).asnumpy()
-        video_data = torch.from_numpy(video_data)
-        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
-        return video_data
-
-    def read_jsons(self, data) -> list[dict]:
-        cap_lists = []
-        with open(data) as f:
-            folder_anno = [
-                i.strip().split(",") for i in f.readlines()
-                if len(i.strip()) > 0
-            ]
-        print(folder_anno)
-        for folder, anno in folder_anno:
-            with open(anno) as f:
-                sub_list = json.load(f)
-            for i in range(len(sub_list)):
-                sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
-            cap_lists += sub_list
-        return cap_lists
-
-    def get_cap_list(self) -> list:
-        cap_lists = self.read_jsons(self.data)[self.start_idx:]
-        return cap_lists

fastvideo/v1/dataset/validation_dataset.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# adapted from: https://github.com/a-r-r-o-w/finetrainers/blob/main/finetrainers/data/dataset.py
+import pathlib
+
+import datasets
+import torch
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.vision_utils import load_image, load_video
+
+logger = init_logger(__name__)
+
+
+class ValidationDataset(torch.utils.data.IterableDataset):
+
+    def __init__(self, filename: str):
+        super().__init__()
+
+        self.filename = pathlib.Path(filename)
+
+        if not self.filename.exists():
+            raise FileNotFoundError(
+                f"File {self.filename.as_posix()} does not exist")
+
+        if self.filename.suffix == ".csv":
+            data = datasets.load_dataset("csv",
+                                         data_files=self.filename.as_posix(),
+                                         split="train")
+        elif self.filename.suffix == ".json":
+            data = datasets.load_dataset("json",
+                                         data_files=self.filename.as_posix(),
+                                         split="train",
+                                         field="data")
+        elif self.filename.suffix == ".parquet":
+            data = datasets.load_dataset("parquet",
+                                         data_files=self.filename.as_posix(),
+                                         split="train")
+        elif self.filename.suffix == ".arrow":
+            data = datasets.load_dataset("arrow",
+                                         data_files=self.filename.as_posix(),
+                                         split="train")
+        else:
+            _SUPPORTED_FILE_FORMATS = [".csv", ".json", ".parquet", ".arrow"]
+            raise ValueError(
+                f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}"
+            )
+
+        self._data = data.to_iterable_dataset()
+
+    def __iter__(self):
+        for sample in self._data:
+            # For consistency reasons, we mandate that "caption" is always present in the validation dataset.
+            # However, since the model specifications use "prompt", we create an alias here.
+            sample["prompt"] = sample["caption"]
+
+            # Load image or video if the path is provided
+            # TODO(aryan): need to handle custom columns here for control conditions
+            sample["image"] = None
+            sample["video"] = None
+
+            if sample.get("image_path", None) is not None:
+                image_path = sample["image_path"]
+                if not pathlib.Path(image_path).is_file(
+                ) and not image_path.startswith("http"):
+                    logger.warning("Image file %s does not exist.",
+                                   image_path.as_posix())
+                else:
+                    sample["image"] = load_image(sample["image_path"])
+
+            if sample.get("video_path", None) is not None:
+                video_path = sample["video_path"]
+                if not pathlib.Path(video_path).is_file(
+                ) and not video_path.startswith("http"):
+                    logger.warning("Video file %s does not exist.",
+                                   video_path.as_posix())
+                else:
+                    sample["video"] = load_video(sample["video_path"])
+
+            if sample.get("control_image_path", None) is not None:
+                control_image_path = sample["control_image_path"]
+                if not pathlib.Path(control_image_path).is_file(
+                ) and not control_image_path.startswith("http"):
+                    logger.warning("Control Image file %s does not exist.",
+                                   control_image_path.as_posix())
+                else:
+                    sample["control_image"] = load_image(
+                        sample["control_image_path"])
+
+            if sample.get("control_video_path", None) is not None:
+                control_video_path = sample["control_video_path"]
+                if not pathlib.Path(control_video_path).is_file(
+                ) and not control_video_path.startswith("http"):
+                    logger.warning("Control Video file %s does not exist.",
+                                   control_video_path)
+                else:
+                    sample["control_video"] = load_video(
+                        sample["control_video_path"])
+
+            sample = {k: v for k, v in sample.items() if v is not None}
+            yield sample

fastvideo/v1/fastvideo_args.py

@@ -388,7 +388,8 @@ class TrainingArgs(FastVideoArgs):
     precondition_outputs: bool = False
 
     # validation & logs
-    validation_prompt_dir: str = ""
+    validation_dataset_file: str = ""
+    validation_preprocessed_path: str = ""
     validation_sampling_steps: str = ""
     validation_guidance_scale: str = ""
     validation_steps: float = 0.0
@@ -536,9 +537,12 @@ class TrainingArgs(FastVideoArgs):
             help="Whether to precondition the outputs of the model")
 
         # Validation and logging
-        parser.add_argument("--validation-prompt-dir",
+        parser.add_argument("--validation-dataset-file",
                             type=str,
-                            help="Directory containing validation prompts")
+                            help="Path to unprocessed validation dataset")
+        parser.add_argument("--validation-preprocessed-path",
+                            type=str,
+                            help="Path to processed validation dataset")
         parser.add_argument("--validation-sampling-steps",
                             type=str,
                             help="Validation sampling steps")

fastvideo/v1/models/dits/base.py

@@ -14,6 +14,7 @@ class BaseDiT(nn.Module, ABC):
     _fsdp_shard_conditions: list = []
     _compile_conditions: list = []
     _param_names_mapping: dict
+    _reverse_param_names_mapping: dict
     hidden_size: int
     num_attention_heads: int
     num_channels_latents: int
@@ -78,6 +79,7 @@ class CachableDiT(BaseDiT):
     # These are required class attributes that should be overridden by concrete implementations
     _fsdp_shard_conditions = []
     _param_names_mapping = {}
+    _reverse_param_names_mapping = {}
     _lora_param_names_mapping: dict = {}
     # Ensure these instance attributes are properly defined in subclasses
     hidden_size: int

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -442,6 +442,8 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
     _supported_attention_backends = HunyuanVideoConfig(
     )._supported_attention_backends
     _param_names_mapping = HunyuanVideoConfig()._param_names_mapping
+    _reverse_param_names_mapping = HunyuanVideoConfig(
+    )._reverse_param_names_mapping
     _lora_param_names_mapping = HunyuanVideoConfig()._lora_param_names_mapping
 
     def __init__(self, config: HunyuanVideoConfig, hf_config: dict[str, Any]):

fastvideo/v1/models/dits/stepvideo.py

@@ -460,6 +460,8 @@ class StepVideoModel(BaseDiT):
         # lambda n, m: "pos_embed" in n  # If needed for the patch embedding.
     ]
     _param_names_mapping = StepVideoConfig()._param_names_mapping
+    _reverse_param_names_mapping = StepVideoConfig(
+    )._reverse_param_names_mapping
     _lora_param_names_mapping = StepVideoConfig()._lora_param_names_mapping
     _supported_attention_backends = StepVideoConfig(
     )._supported_attention_backends

fastvideo/v1/models/dits/wanvideo.py

@@ -518,6 +518,7 @@ class WanTransformer3DModel(CachableDiT):
     _supported_attention_backends = WanVideoConfig(
     )._supported_attention_backends
     _param_names_mapping = WanVideoConfig()._param_names_mapping
+    _reverse_param_names_mapping = WanVideoConfig()._reverse_param_names_mapping
     _lora_param_names_mapping = WanVideoConfig()._lora_param_names_mapping
 
     def __init__(self, config: WanVideoConfig, hf_config: dict[str,

fastvideo/v1/models/loader/fsdp_load.py

@@ -222,10 +222,14 @@ def load_model_from_full_model_state_dict(
     used_keys = set()
     sharded_sd = {}
     to_merge_params: DefaultDict[str, Dict[Any, Any]] = defaultdict(dict)
+    reverse_param_names_mapping = {}
+    assert param_names_mapping is not None
     for source_param_name, full_tensor in full_sd_iterator:
-        assert param_names_mapping is not None
         target_param_name, merge_index, num_params_to_merge = param_names_mapping(
             source_param_name)
+        reverse_param_names_mapping[target_param_name] = (source_param_name,
+                                                          merge_index,
+                                                          num_params_to_merge)
         used_keys.add(target_param_name)
         if merge_index is not None:
             to_merge_params[target_param_name][merge_index] = full_tensor
@@ -260,6 +264,7 @@ def load_model_from_full_model_state_dict(
                 sharded_tensor = sharded_tensor.cpu()
         sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)
 
+    model._reverse_param_names_mapping = reverse_param_names_mapping
     unused_keys = set(meta_sd.keys()) - used_keys
     if unused_keys:
         logger.warning("Found new parameters in meta state dict: %s",

fastvideo/v1/models/vision_utils.py

@@ -1,8 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
+import tempfile
 from typing import Callable, List, Optional, Tuple, Union
+from urllib.parse import unquote, urlparse
 
+import imageio
 import numpy as np
 import PIL.Image
 import PIL.ImageOps
@@ -86,6 +89,7 @@ def normalize(
     return 2.0 * images - 1.0
 
 
+# adapted from diffusers.utils import load_image
 def load_image(
     image: Union[str, PIL.Image.Image],
     convert_method: Optional[Callable[[PIL.Image.Image],
@@ -131,6 +135,85 @@ def load_image(
     return image
 
 
+# adapted from diffusers.utils import load_video
+def load_video(
+    video: str,
+    convert_method: Optional[Callable[[List[PIL.Image.Image]],
+                                      List[PIL.Image.Image]]] = None,
+) -> List[PIL.Image.Image]:
+    """
+    Loads `video` to a list of PIL Image.
+    Args:
+        video (`str`):
+            A URL or Path to a video to convert to a list of PIL Image format.
+        convert_method (Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]], *optional*):
+            A conversion method to apply to the video after loading it. When set to `None` the images will be converted
+            to "RGB".
+    Returns:
+        `List[PIL.Image.Image]`:
+            The video as a list of PIL images.
+    """
+    is_url = video.startswith("http://") or video.startswith("https://")
+    is_file = os.path.isfile(video)
+    was_tempfile_created = False
+
+    if not (is_url or is_file):
+        raise ValueError(
+            f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path."
+        )
+
+    if is_url:
+        response = requests.get(video, stream=True)
+        if response.status_code != 200:
+            raise ValueError(
+                f"Failed to download video. Status code: {response.status_code}"
+            )
+
+        parsed_url = urlparse(video)
+        file_name = os.path.basename(unquote(parsed_url.path))
+
+        suffix = os.path.splitext(file_name)[1] or ".mp4"
+        with tempfile.NamedTemporaryFile(suffix=suffix,
+                                         delete=False) as temp_file:
+            video_path = temp_file.name
+            video_data = response.iter_content(chunk_size=8192)
+            for chunk in video_data:
+                temp_file.write(chunk)
+
+        video = video_path
+
+    pil_images = []
+    if video.endswith(".gif"):
+        gif = PIL.Image.open(video)
+        try:
+            while True:
+                pil_images.append(gif.copy())
+                gif.seek(gif.tell() + 1)
+        except EOFError:
+            pass
+
+    else:
+        try:
+            imageio.plugins.ffmpeg.get_exe()
+        except AttributeError:
+            raise AttributeError(
+                "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg"
+            ) from None
+
+        with imageio.get_reader(video) as reader:
+            # Read all frames
+            for frame in reader:
+                pil_images.append(PIL.Image.fromarray(frame))
+
+    if was_tempfile_created:
+        os.remove(video_path)
+
+    if convert_method is not None:
+        pil_images = convert_method(pil_images)
+
+    return pil_images
+
+
 def get_default_height_width(
     image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
     vae_scale_factor: int,

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -11,6 +11,7 @@ import pprint
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List, Optional, Union
 
+import PIL.Image
 import torch
 
 from fastvideo.v1.attention import AttentionMetadata
@@ -37,6 +38,7 @@ class ForwardBatch:
     # Image inputs
     image_path: Optional[str] = None
     image_embeds: List[torch.Tensor] = field(default_factory=list)
+    pil_image: Optional[PIL.Image.Image] = None
 
     # Text inputs
     prompt: Optional[Union[str, List[str]]] = None

fastvideo/v1/pipelines/preprocess/preprocess_pipeline_base.py

@@ -3,6 +3,7 @@ import gc
 import multiprocessing
 import os
 from concurrent.futures import ProcessPoolExecutor
+from itertools import chain
 from typing import Any, Dict, List, Optional
 
 import numpy as np
@@ -10,17 +11,16 @@ import pyarrow as pa
 import pyarrow.parquet as pq
 import torch
 from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
 from tqdm import tqdm
 
 from fastvideo.v1.configs.sample import SamplingParam
-from fastvideo.v1.dataset import getdataset
+from fastvideo.v1.dataset import ValidationDataset, getdataset
 from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.pipelines.stages import TextEncodingStage
+from fastvideo.v1.pipelines.stages import EncodingStage, TextEncodingStage
 
 logger = init_logger(__name__)
 
@@ -36,6 +36,9 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                            tokenizers=[self.get_module("tokenizer")],
                        ))
 
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=EncodingStage(vae=self.get_module("vae"), ))
+
     @torch.no_grad()
     def forward(
         self,
@@ -46,7 +49,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
         # 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_validation_text(fastvideo_args, args)
+        self.preprocess_validation(fastvideo_args, args)
         self.preprocess_video_and_text(fastvideo_args, args)
 
     def get_extra_features(self, valid_data: Dict[str, Any],
@@ -103,7 +106,6 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                                             "combined_parquet_dataset")
         os.makedirs(combined_parquet_dir, exist_ok=True)
         local_rank = int(os.getenv("RANK", 0))
-        world_size = int(os.getenv("WORLD_SIZE", 1))
 
         # Get how many samples have already been processed
         start_idx = 0
@@ -114,14 +116,10 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                     start_idx += table.num_rows
 
         # Loading dataset
-        train_dataset = getdataset(args, start_idx=start_idx)
-        sampler = DistributedSampler(train_dataset,
-                                     rank=local_rank,
-                                     num_replicas=world_size,
-                                     shuffle=False)
+        train_dataset = getdataset(args)
+
         train_dataloader = DataLoader(
             train_dataset,
-            sampler=sampler,
             batch_size=args.preprocess_video_batch_size,
             num_workers=args.dataloader_num_workers,
         )
@@ -285,7 +283,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                 num_processed_samples = 0
                 self.all_tables = []
 
-    def preprocess_validation_text(self, fastvideo_args: FastVideoArgs, args):
+    def preprocess_validation(self, fastvideo_args: FastVideoArgs, args):
         """Process validation text prompts and save them to parquet files.
         
         This base implementation handles the common validation text processing logic.
@@ -296,22 +294,32 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                                               "validation_parquet_dataset")
         os.makedirs(validation_parquet_dir, exist_ok=True)
 
-        with open(args.validation_prompt_txt, encoding="utf-8") as file:
-            lines = file.readlines()
-        prompts = [line.strip() for line in lines]
+        validation_dataset = ValidationDataset(args.validation_dataset_file)
 
         # Prepare batch data for Parquet dataset
         batch_data = []
         sampling_param = SamplingParam.from_pretrained(
             fastvideo_args.model_path)
         if sampling_param.negative_prompt:
-            prompts = [sampling_param.negative_prompt] + prompts
+            negative_prompt = {
+                'caption': sampling_param.negative_prompt,
+                'image_path': None,
+                'video_path': None,
+            }
+            validation_iterable = chain([negative_prompt], validation_dataset)
+        else:
+            negative_prompt = None
+            validation_iterable = validation_dataset
+
         # Add progress bar for validation text preprocessing
-        pbar = tqdm(enumerate(prompts),
+        pbar = tqdm(enumerate(validation_iterable),
                     desc="Processing validation prompts",
                     unit="prompt")
-        for prompt_idx, prompt in pbar:
+        for idx, sample in pbar:
             with torch.inference_mode():
+                prompt = sample["caption"]
+                # is_negative_prompt = idx == 0
+
                 # Text Encoder
                 batch = ForwardBatch(
                     data_type="video",

fastvideo/v1/pipelines/preprocess/v1_preprocess.py

@@ -44,7 +44,7 @@ if __name__ == "__main__":
     parser.add_argument("--model_path", type=str, default="data/mochi")
     parser.add_argument("--model_type", type=str, default="mochi")
     parser.add_argument("--data_merge_path", type=str, required=True)
-    parser.add_argument("--validation_prompt_txt", type=str)
+    parser.add_argument("--validation_dataset_file", type=str)
     parser.add_argument("--num_frames", type=int, default=163)
     parser.add_argument(
         "--dataloader_num_workers",

fastvideo/v1/pipelines/stages/encoding.py

@@ -12,8 +12,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 from fastvideo.v1.models.vision_utils import (get_default_height_width,
-                                              load_image, normalize,
-                                              numpy_to_pt, pil_to_numpy, resize)
+                                              normalize, numpy_to_pt,
+                                              pil_to_numpy, resize)
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.pipelines.stages.validators import V  # Import validators
@@ -60,7 +60,7 @@ class EncodingStage(PipelineStage):
         latent_height = batch.height // self.vae.spatial_compression_ratio
         latent_width = batch.width // self.vae.spatial_compression_ratio
 
-        image = load_image(image_path)
+        image = batch.pil_image
         image = self.preprocess(
             image,
             vae_scale_factor=self.vae.spatial_compression_ratio,
@@ -97,7 +97,7 @@ class EncodingStage(PipelineStage):
         generator = batch.generator
         if generator is None:
             raise ValueError("Generator must be provided")
-        latent_condition = self.retrieve_latents(encoder_output, generator[0])
+        latent_condition = self.retrieve_latents(encoder_output, generator)
 
         # Apply shifting if needed
         if (hasattr(self.vae, "shift_factor")
@@ -181,7 +181,7 @@ class EncodingStage(PipelineStage):
                      fastvideo_args: FastVideoArgs) -> VerificationResult:
         """Verify encoding stage inputs."""
         result = VerificationResult()
-        result.add_check("image_path", batch.image_path, V.string_not_empty)
+        result.add_check("pil_image", batch.pil_image, V.not_none)
         result.add_check("height", batch.height, V.positive_int)
         result.add_check("width", batch.width, V.positive_int)
         result.add_check("generator", batch.generator,

fastvideo/v1/pipelines/stages/image_encoding.py

@@ -11,7 +11,6 @@ from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.vision_utils import load_image
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.pipelines.stages.validators import StageValidators as V
@@ -58,7 +57,7 @@ class ImageEncodingStage(PipelineStage):
         if fastvideo_args.use_cpu_offload:
             self.image_encoder = self.image_encoder.to(get_torch_device())
 
-        image = load_image(batch.image_path)
+        image = batch.pil_image
 
         image_inputs = self.image_processor(
             images=image, return_tensors="pt").to(get_torch_device())
@@ -78,7 +77,7 @@ class ImageEncodingStage(PipelineStage):
                      fastvideo_args: FastVideoArgs) -> VerificationResult:
         """Verify image encoding stage inputs."""
         result = VerificationResult()
-        result.add_check("image_path", batch.image_path, V.string_not_empty)
+        result.add_check("pil_image", batch.pil_image, V.not_none)
         result.add_check("image_embeds", batch.image_embeds, V.is_list)
         return result
 

fastvideo/v1/pipelines/stages/input_validation.py

@@ -7,6 +7,7 @@ import torch
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.vision_utils import load_image
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.pipelines.stages.validators import (StageValidators,
@@ -91,6 +92,11 @@ class InputValidationStage(PipelineStage):
                 f"Guidance scale must be positive, but got {batch.guidance_scale}"
             )
 
+        # for i2v, get image from image_path
+        if batch.image_path is not None:
+            image = load_image(batch.image_path)
+            batch.pil_image = image
+
         return batch
 
     def verify_input(self, batch: ForwardBatch,

fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py

@@ -76,4 +76,36 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
                        stage=DecodingStage(vae=self.get_module("vae")))
 
 
+class WanImageToVideoValidationPipeline(ComposedPipelineBase):
+    """
+    I2V Validation pipeline for Wan2.1, assumes that the input are preprocess latents.
+    """
+    _required_config_modules = ["vae", "scheduler", "transformer"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="image_latent_preparation_stage",
+                       stage=EncodingStage(vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
 EntryClass = WanImageToVideoPipeline

fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py

@@ -80,7 +80,7 @@ def run_preprocessing():
         "--dataloader_num_workers", "0",
         "--output_dir", LOCAL_PREPROCESSED_DATA_DIR,
         "--train_fps", "16",
-        "--validation_prompt_txt", os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt_1_sample.txt"),
+        "--validation_dataset_file", os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt_1_sample.json"),
         "--samples_per_file", "1",
         "--flush_frequency", "1",
         "--video_length_tolerance_range", "5",
@@ -100,7 +100,7 @@ def run_training():
         "--inference_mode", "False",
         "--pretrained_model_name_or_path", MODEL_PATH,
         "--data_path", LOCAL_TRAINING_DATA_DIR,
-        "--validation_prompt_dir", LOCAL_VALIDATION_DATA_DIR,
+        "--validation_preprocessed_path", LOCAL_VALIDATION_DATA_DIR,
         "--train_batch_size", "1",
         "--num_latent_t", "8",
         "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,

fastvideo/v1/tests/slurm/training/wan_t2v_1.3b/finetune.sh

@@ -0,0 +1,51 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+DATA_DIR="data/crush-smol_processed_main_t2v/latents/combined_parquet_dataset"
+VALIDATION_DIR="data/crush-smol_processed_main_t2v/latents/validation_parquet_dataset"
+NUM_GPUS=4
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
+    fastvideo/v1/training/wan_training_pipeline.py\
+    --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --inference_mode False\
+    --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --data_path "$DATA_DIR"\
+    --validation_preprocessed_path "$VALIDATION_DIR"\
+    --train_batch_size=1 \
+    --num_latent_t 8 \
+    --sp_size 4 \
+    --tp_size 4 \
+    --hsdp_replicate_dim 1 \
+    --hsdp_shard_dim 4 \
+    --num_gpus $NUM_GPUS \
+    --train_sp_batch_size 1\
+    --dataloader_num_workers 1\
+    --gradient_accumulation_steps=8 \
+    --max_train_steps=5000 \
+    --learning_rate=1e-5\
+    --mixed_precision="bf16"\
+    --checkpointing_steps=6000 \
+    --validation_steps 50\
+    --validation_sampling_steps "50" \
+    --log_validation \
+    --checkpoints_total_limit 3\
+    --allow_tf32\
+    --ema_start_step 0\
+    --cfg 0.0\
+    --output_dir="$DATA_DIR/outputs/wan_finetune"\
+    --tracker_project_name wan_finetune \
+    --num_height 480 \
+    --num_width 832 \
+    --num_frames  77 \
+    --validation_guidance_scale "1.0" \
+    --num_euler_timesteps 50 \
+    --multi_phased_distill_schedule "4000-1" \
+    --weight_decay 0.01 \
+    --not_apply_cfg_solver \
+    --dit_precision "fp32" \
+    --max_grad_norm 1.0

fastvideo/v1/tests/slurm/training/wan_t2v_1.3b/preprocess.sh

@@ -0,0 +1,24 @@
+# export WANDB_MODE="offline"
+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_main_t2v/latents"
+VALIDATION_PATH="examples/training/finetune/wan_t2v_1.3b/crush_smol/validation.json"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 8 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 77 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --model_type $MODEL_TYPE \
+    --train_fps 16 \
+    --validation_dataset_file $VALIDATION_PATH \
+    --samples_per_file 8 \
+    --flush_frequency 8 \
+    --preprocess_task "t2v"

fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py

@@ -33,7 +33,7 @@ def run_worker():
         "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         "--cache_dir", "/home/.cache",
         "--data_path", "data/mini_dataset_i2v_VSA/combined_parquet_dataset",
-        "--validation_prompt_dir", "data/mini_dataset_i2v_VSA/validation_parquet_dataset",
+        "--validation_preprocessed_path", "data/mini_dataset_i2v_VSA/validation_parquet_dataset",
         "--train_batch_size", "1",
         "--num_latent_t", "4",
         "--num_gpus", "1",

fastvideo/v1/tests/training/Vanilla/test_training_loss.py

@@ -38,7 +38,7 @@ def run_worker():
         "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         "--cache_dir", "/home/.cache",
         "--data_path", "data/crush-smol_parq/combined_parquet_dataset",
-        "--validation_prompt_dir", "data/crush-smol_parq/validation_parquet_dataset",
+        "--validation_preprocessed_path", "data/crush-smol_parq/validation_parquet_dataset",
         "--train_batch_size", "2",
         "--num_latent_t", "4",
         "--num_gpus", "4",
@@ -116,7 +116,7 @@ def test_distributed_training():
         'avg_step_time': 1.0,
         'grad_norm': 0.2,
         'step_time': 0.5,
-        'train_loss': 0.001
+        'train_loss': 0.0025
     }
 
     failures = []

fastvideo/v1/training/training_pipeline.py

@@ -528,10 +528,10 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         logger.info("Using validation seed: %s", validation_seed)
 
         # Prepare validation prompts
-        logger.info('fastvideo_args.validation_prompt_dir: %s',
-                    training_args.validation_prompt_dir)
+        logger.info('fastvideo_args.validation_preprocessed_path: %s',
+                    training_args.validation_preprocessed_path)
         validation_dataset, validation_dataloader = build_parquet_map_style_dataloader(
-            training_args.validation_prompt_dir,
+            training_args.validation_preprocessed_path,
             batch_size=1,
             num_data_workers=0,
             drop_last=False,

fastvideo/v1/training/training_utils.py

@@ -8,7 +8,6 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dcp
-import torch.distributed.checkpoint.stateful
 from einops import rearrange
 from safetensors.torch import save_file
 
@@ -154,13 +153,20 @@ def save_checkpoint(transformer,
 
     if rank == 0:
         # Save model weights (consolidated)
-        weight_path = os.path.join(save_dir,
+        transformer_save_dir = os.path.join(save_dir, "transformer")
+        os.makedirs(transformer_save_dir, exist_ok=True)
+        weight_path = os.path.join(transformer_save_dir,
                                    "diffusion_pytorch_model.safetensors")
         logger.info("rank: %s, saving consolidated checkpoint to %s",
                     rank,
                     weight_path,
                     local_main_process_only=False)
-        save_file(cpu_state, weight_path)
+
+        # Convert training format to diffusers format and save
+        diffusers_state_dict = convert_training_to_diffusers_format(
+            cpu_state, transformer)
+        save_file(diffusers_state_dict, weight_path)
+
         logger.info("rank: %s, consolidated checkpoint saved to %s",
                     rank,
                     weight_path,
@@ -170,7 +176,7 @@ def save_checkpoint(transformer,
         config_dict = transformer.hf_config
         if "dtype" in config_dict:
             del config_dict["dtype"]  # TODO
-        config_path = os.path.join(save_dir, "config.json")
+        config_path = os.path.join(transformer_save_dir, "config.json")
         # save dict as json
         with open(config_path, "w") as f:
             json.dump(config_dict, f, indent=4)
@@ -479,3 +485,68 @@ def _has_foreach_support(tensors: List[torch.Tensor],
                          device: torch.device) -> bool:
     return _device_has_foreach_support(device) and all(
         t is None or type(t) in [torch.Tensor] for t in tensors)
+
+
+def convert_training_to_diffusers_format(state_dict: Dict[str, Any],
+                                         transformer) -> Dict[str, Any]:
+    """
+    Convert training format state dict to diffusers format using reverse_param_names_mapping.
+    
+    Args:
+        state_dict: State dict in training format
+        transformer: Transformer model object with _reverse_param_names_mapping
+        
+    Returns:
+        State dict in diffusers format
+    """
+    new_state_dict = {}
+
+    # Get the reverse mapping from the transformer
+    reverse_param_names_mapping = transformer._reverse_param_names_mapping
+    assert reverse_param_names_mapping != {}, "reverse_param_names_mapping is empty"
+
+    # Group parameters that need to be split (merged parameters)
+    merge_groups: Dict[str, List[Tuple[str, int, int]]] = {}
+
+    # First pass: collect all merge groups
+    for training_key, (
+            diffusers_key, merge_index,
+            num_params_to_merge) in reverse_param_names_mapping.items():
+        if merge_index is not None:
+            # This is a merged parameter that needs to be split
+            if training_key not in merge_groups:
+                merge_groups[training_key] = []
+            merge_groups[training_key].append(
+                (diffusers_key, merge_index, num_params_to_merge))
+
+    # Second pass: handle merged parameters by splitting them
+    used_keys = set()
+    for training_key, splits in merge_groups.items():
+        if training_key in state_dict:
+            v = state_dict[training_key]
+            # Sort by merge_index to ensure correct order
+            splits.sort(key=lambda x: x[1])
+            total = splits[0][2]
+            split_size = v.shape[0] // total
+            split_tensors = torch.split(v, split_size, dim=0)
+
+            for diffusers_key, split_index, _ in splits:
+                new_state_dict[diffusers_key] = split_tensors[split_index]
+            used_keys.add(training_key)
+
+    # Third pass: handle regular parameters (direct mappings)
+    for training_key, v in state_dict.items():
+        if training_key in used_keys:
+            continue
+
+        if training_key in reverse_param_names_mapping:
+            diffusers_key, merge_index, _ = reverse_param_names_mapping[
+                training_key]
+            if merge_index is None:
+                # Direct mapping
+                new_state_dict[diffusers_key] = v
+        else:
+            # No mapping found, keep as is
+            new_state_dict[training_key] = v
+
+    return new_state_dict

pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
 
     # Machine Learning & Transformers
     "transformers>=4.46.1", "tokenizers>=0.20.1", "sentencepiece==0.2.0",
-    "timm==1.0.11", "peft==0.13.2", "diffusers>=0.33.1", "bitsandbytes",
+    "timm==1.0.11", "peft==0.15.0", "diffusers>=0.33.1", "bitsandbytes",
     "torch==2.7.1", "torchvision",
 
     # Acceleration & Optimization

scripts/finetune/finetune_v1.sh

@@ -15,7 +15,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --inference_mode False\
     --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
     --data_path "$DATA_DIR"\
-    --validation_prompt_dir "$VALIDATION_DIR"\
+    --validation_preprocessed_path "$VALIDATION_DIR"\
     --train_batch_size=4 \
     --num_latent_t 20 \
     --sp_size 4 \

scripts/finetune/finetune_v1_VSA.sh

@@ -21,7 +21,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
     --cache_dir "/home/ray/.cache" \
     --data_path "$DATA_DIR" \
-    --validation_prompt_dir "$VALIDATION_DIR" \
+    --validation_preprocessed_path "$VALIDATION_DIR" \
     --train_batch_size 1 \
     --num_latent_t 16 \
     --sp_size 1 \

scripts/preprocess/v1_preprocess_wan_data_i2v.sh

@@ -18,7 +18,7 @@ torchrun --nproc_per_node=$GPU_NUM \
     --output_dir=$OUTPUT_DIR \
     --model_type $MODEL_TYPE \
     --train_fps 16 \
-    --validation_prompt_txt $VALIDATION_PATH \
+    --validation_dataset_file $VALIDATION_PATH \
     --samples_per_file 8 \
     --flush_frequency 8 \
     --preprocess_task "i2v" 

scripts/preprocess/v1_preprocess_wan_data_t2v.sh

@@ -18,7 +18,7 @@ torchrun --nproc_per_node=$GPU_NUM \
     --output_dir=$OUTPUT_DIR \
     --model_type $MODEL_TYPE \
     --train_fps 16 \
-    --validation_prompt_txt $VALIDATION_PATH \
+    --validation_dataset_file $VALIDATION_PATH \
     --samples_per_file 1 \
     --flush_frequency 1 \
     --video_length_tolerance_range 5 \