i2v validation

Co-authored-by: JerryZhou54 <zhouw.jerry2017@outlook.com>

.github/workflows/sta-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "csrc/sliding_tile_attention/setup.py"
+      - "csrc/attn/setup_sta.py"
   workflow_dispatch:
 
 jobs:
@@ -23,13 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/sliding_tile_attention
+          cd csrc/attn
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_sta.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup_sta.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -136,19 +136,21 @@ jobs:
 
       - name: Build wheel
         run: |
+          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
+          
           # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
           # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
           # However this still fails so I'm using a newer version of setuptools
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/sliding_tile_attention  # Move into the correct folder
-          git submodule update --init --recursive tk # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
+          cd csrc/attn  # Move into the correct folder
+          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
+          python setup_sta.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/sliding_tile_attention
+          cd csrc/attn
 
           CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
           TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2)
@@ -163,7 +165,7 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.wheel_name }}
-          path: csrc/sliding_tile_attention/dist/*.whl
+          path: csrc/attn/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -229,17 +231,19 @@ jobs:
 
       - name: Build source distribution
         run: |
+          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
+          
           # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
           # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
           # However this still fails so I'm using a newer version of setuptools
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/sliding_tile_attention  # Move into the correct folder
-          git submodule update --init --recursive tk # Ensure ThunderKittens submodule is initialized
-          python setup.py sdist --dist-dir=dist
+          cd csrc/attn  # Move into the correct folder
+          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
+          python setup_sta.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/sliding_tile_attention/dist/
+          packages-dir: csrc/attn/dist/

.github/workflows/test.yml

@@ -28,4 +28,4 @@ jobs:
 
       - name: Run Pytest
         run: |
-          pytest --ignore csrc/sliding_tile_attention/test
+          pytest --ignore csrc/attn/test

.gitignore

@@ -27,7 +27,6 @@ env
 **/build/
 **.pyc
 **.txt
-csrc/attn/tk/
 
 # Distribution / packaging
 build/

.gitmodules

@@ -1,3 +1,3 @@
-[submodule "csrc/sliding_tile_attention/tk"]
-	path = csrc/sliding_tile_attention/tk
+[submodule "csrc/attn/tk"]
+	path = csrc/attn/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git

csrc/attn/README.md

@@ -7,26 +7,26 @@
 We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
 First, install C++20 for ThunderKittens:
 
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-sudo apt install clang-11
-```
-Install STA:
+## Environment Setup
+First, set up your CUDA environment:
 ```bash
 export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 git submodule update --init --recursive
-python setup.py install
+```
+
+## Install Sliding Tile Attention (STA)
+```bash
+python setup_sta.py install
+```
+
+## Install Video Sparse Attention (VSA)
+```bash
+python setup_vsa.py install
 ```
 
 ## Usage
-
 ```python
 from st_attn import sliding_tile_attention
 # assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 

csrc/attn/test/test_block_sparse.py

@@ -2,8 +2,8 @@ import torch
 import argparse
 from flash_attn.utils.benchmark import benchmark_forward
 from flash_attn import flash_attn_func
-from st_attn import block_sparse_attention_fwd, block_sparse_attention_backward, BlockSparseAttentionFunction
-from st_attn import BLOCK_M, BLOCK_N
+from vsa import block_sparse_attention_fwd, block_sparse_attention_backward, BlockSparseAttentionFunction
+from vsa import BLOCK_M, BLOCK_N
 
 import numpy as np
 import random

csrc/attn/test/test_sta.py

@@ -9,14 +9,14 @@ flex_attention = torch.compile(flex_attention, dynamic=False)
 
 
 def flex_test(Q, K, V, kernel_size):
-    mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (36, 48, 48), 39, 'cuda', 0)
+    mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (18, 48, 80), 0, 'cuda', 0)
     output = flex_attention(Q, K, V, block_mask=mask)
 
     return output
 
 
 def h100_fwd_kernel_test(Q, K, V, kernel_size):
-    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 39, False)
+    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 0, False, '18x48x80')
     return o
 
 
@@ -37,7 +37,7 @@ def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mo
             'max_diff': 0
         },
     }
-    kernel_size_ls = [(6, 1, 6), (6, 6, 1)]
+    kernel_size_ls = [(3, 3, 5), (3, 1, 10)]
     from tqdm import tqdm
     for kernel_size in tqdm(kernel_size_ls):
         for _ in range(num_iterations):
@@ -74,7 +74,7 @@ def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mo
 
 # Example usage
 b, h, d = 2, 24, 128
-n = 82944  # Sequence length
+n = 69120  # Sequence length
 causal = False
 mean = 1e-1
 std = 10

csrc/attn/vsa/__init__.py

@@ -2,7 +2,7 @@ import math
 
 import torch
 from torch.utils.checkpoint import detach_variable
-
+from typing import Tuple
 try:
     from vsa_cuda import block_sparse_fwd, block_sparse_bwd
 except ImportError:
@@ -12,33 +12,116 @@ except ImportError:
 
 BLOCK_M = 64
 BLOCK_N = 64
-def block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num):
+
+def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None):
     """
-    block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]. 
-        [*, *, i, j] = 1 means the i-th q block should attend to the j-th kv block.
+    q: [batch_size, num_heads, seq_len, head_dim]
+    k: [batch_size, num_heads, seq_len, head_dim]
+    v: [batch_size, num_heads, seq_len, head_dim]
+    topk: int
+    block_size: int or tuple of 3 ints
+    video_shape: tuple of (T, H, W)
+    compress_attn_weight: [batch_size, num_heads, seq_len, head_dim]
+    select_attn_weight: [batch_size, num_heads, seq_len, head_dim]
+    
+    V1 of sparse attention. Include compress attn and sparse attn branch, use average pooling to compress. 
+    Assume q, k, v is flattened in this way: [batch_size, num_heads, T//block_size[0], H//block_size[1], W//block_size[2], block_size[0], block_size[1], block_size[2]]
     """
-    # assert all elements in q2k_block_sparse_num can be devisible by 2
-    o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-    return o, lse
 
-def block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num):
-    grad_q, grad_k, grad_v = block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
-    return grad_q, grad_k, grad_v
+    if isinstance(block_size, int):
+        block_size = (block_size, block_size, block_size)
 
-class BlockSparseAttentionFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
-        o, lse = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-        ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num)
-        return o
+    block_elements = block_size[0] * block_size[1] * block_size[2]
+    assert block_elements % 64 == 0 and block_elements >= 64
+    assert q.shape[2] % block_elements == 0
+    batch_size, num_heads, seq_len, head_dim = q.shape
+    # compress attn
+    q_compress = q.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
+    k_compress = k.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
+    v_compress = v.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
 
-    @staticmethod
-    def backward(ctx, grad_output):
-        q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors
-        grad_q, grad_k, grad_v = block_sparse_attention_backward(
-            q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num
-        )
-        return grad_q, grad_k, grad_v, None, None, None, None
+    output_compress, block_attn_score = torch_attention(q_compress, k_compress,
+                                                        v_compress)
+
+    output_compress = output_compress.view(batch_size, num_heads,
+                                           seq_len // block_elements, 1,
+                                           head_dim)
+    output_compress = output_compress.repeat(1, 1, 1, block_elements,
+                                             1).view(batch_size, num_heads,
+                                                     seq_len, head_dim)
+
+    q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num = generate_topk_block_sparse_pattern(
+        block_attn_score, topk)
+
+    output_select = block_sparse_attn(q, k, v, q2k_block_sparse_index,
+                                      q2k_block_sparse_num,
+                                      k2q_block_sparse_index,
+                                      k2q_block_sparse_num)
+
+    if compress_attn_weight is not None:
+        final_output = output_compress * compress_attn_weight + output_select
+    else:
+        final_output = output_compress + output_select
+    return final_output
+
+def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
+    QK = torch.matmul(q, k.transpose(-2, -1))
+    QK /= (q.size(-1)**0.5)
+
+    # Causal mask removed since causal is always false
+
+    QK = torch.nn.functional.softmax(QK, dim=-1)
+    output = torch.matmul(QK, v)
+    return output, QK
+
+def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,
+                                       topk: int):
+    """
+    Generate a block sparse pattern where each q block attends to exactly topk kv blocks,
+    based on the provided attention scores.
+    
+    Args:
+        block_attn_score: [bs, h, num_q_blocks, num_kv_blocks]
+            Attention scores between query and key blocks
+        topk: int
+            Number of kv blocks each q block attends to
+        
+    Returns:
+        q2k_block_sparse_index: [bs, h, num_q_blocks, topk]
+            Contains the indices of kv blocks that each q block attends to.
+        q2k_block_sparse_num: [bs, h, num_q_blocks]
+            Contains the number of kv blocks that each q block attends to (all equal to topk).
+        k2q_block_sparse_index: [bs, h, num_kv_blocks, max_q_per_kv]
+            Contains the indices of q blocks that attend to each kv block.
+        k2q_block_sparse_num: [bs, h, num_kv_blocks]
+            Contains the number of q blocks that attend to each kv block.
+    """
+    device = block_attn_score.device
+    # Extract dimensions from block_attn_score
+    bs, h, num_q_blocks, num_kv_blocks = block_attn_score.shape
+
+    sorted_result = torch.sort(block_attn_score, dim=-1, descending=True)
+
+    sorted_indice = sorted_result.indices
+
+    q2k_block_sparse_index, _ = torch.sort(sorted_indice[:, :, :, :topk],
+                                           dim=-1)
+    q2k_block_sparse_index = q2k_block_sparse_index.to(dtype=torch.int32)
+    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks),
+                                      topk,
+                                      device=device,
+                                      dtype=torch.int32)
+
+    block_map = topk_index_to_map(q2k_block_sparse_index,
+                                  num_kv_blocks,
+                                  transpose_map=True)
+    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(
+        block_map.transpose(2, 3))
+
+    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
 
 @torch._dynamo.disable
 def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
@@ -61,6 +144,19 @@ def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q
         q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
     )
 
+def block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num):
+    """
+    block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]. 
+        [*, *, i, j] = 1 means the i-th q block should attend to the j-th kv block.
+    """
+    # assert all elements in q2k_block_sparse_num can be devisible by 2
+    o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
+    return o, lse
+
+def block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num):
+    grad_q, grad_k, grad_v = block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
+    return grad_q, grad_k, grad_v
+
 ## pytorch sdpa version of block sparse ##
 import triton
 import triton.language as tl
@@ -128,6 +224,169 @@ def index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, BLOCK_Q, BLOCK_K
 
     return mask
 
+@triton.jit
+def topk_index_to_map_kernel(
+    map_ptr,
+    index_ptr,
+    map_bs_stride,
+    map_h_stride,
+    map_q_stride,
+    map_kv_stride,
+    index_bs_stride,
+    index_h_stride,
+    index_q_stride,
+    index_kv_stride,
+    topk: tl.constexpr,
+):
+    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
+    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
+
+    for i in tl.static_range(topk):
+        index = tl.load(index_ptr_base + i * index_kv_stride)
+        tl.store(map_ptr_base + index * map_kv_stride, 1.0)
+
+@triton.jit
+def map_to_index_kernel(
+    map_ptr,
+    index_ptr,
+    index_num_ptr,
+    map_bs_stride,
+    map_h_stride,
+    map_q_stride,
+    map_kv_stride,
+    index_bs_stride,
+    index_h_stride,
+    index_q_stride,
+    index_kv_stride,
+    index_num_bs_stride,
+    index_num_h_stride,
+    index_num_q_stride,
+    num_kv_blocks: tl.constexpr,
+):
+    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
+    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
+
+    num = 0
+    for i in tl.static_range(num_kv_blocks):
+        map_entry = tl.load(map_ptr_base + i * map_kv_stride)
+        if map_entry:
+            tl.store(index_ptr_base + num * index_kv_stride, i)
+            num += 1
+
+    tl.store(
+        index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
+        q * index_num_q_stride, num)
+
+def topk_index_to_map(index: torch.Tensor,
+                      num_kv_blocks: int,
+                      transpose_map: bool = False):
+    """
+    Convert topk indices to a map.
+    
+    Args:
+        index: [bs, h, num_q_blocks, topk]
+            The topk indices tensor.
+        num_kv_blocks: int
+            The number of key-value blocks in the block_map returned
+        transpose_map: bool
+            If True, the block_map will be transposed on the final two dimensions.
+    
+    Returns:
+        block_map: [bs, h, num_q_blocks, num_kv_blocks]
+            A binary map where 1 indicates that the q block attends to the kv block.
+    """
+    bs, h, num_q_blocks, topk = index.shape
+
+    if transpose_map is False:
+        block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks),
+                                dtype=torch.bool,
+                                device=index.device)
+    else:
+        block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks),
+                                dtype=torch.bool,
+                                device=index.device)
+        block_map = block_map.transpose(2, 3)
+
+    grid = (bs, h, num_q_blocks)
+    topk_index_to_map_kernel[grid](
+        block_map,
+        index,
+        block_map.stride(0),
+        block_map.stride(1),
+        block_map.stride(2),
+        block_map.stride(3),
+        index.stride(0),
+        index.stride(1),
+        index.stride(2),
+        index.stride(3),
+        topk=topk,
+    )
+
+    return block_map
+
+def map_to_index(block_map: torch.Tensor):
+    """
+    Convert a block map to indices and counts.
+    
+    Args:
+        block_map: [bs, h, num_q_blocks, num_kv_blocks]
+            The block map tensor.
+    
+    Returns:
+        index: [bs, h, num_q_blocks, num_kv_blocks]
+            The indices of the blocks.
+        index_num: [bs, h, num_q_blocks]
+            The number of blocks for each q block.
+    """
+    bs, h, num_q_blocks, num_kv_blocks = block_map.shape
+
+    index = torch.full((block_map.shape),
+                       -1,
+                       dtype=torch.int32,
+                       device=block_map.device)
+    index_num = torch.empty((bs, h, num_q_blocks),
+                            dtype=torch.int32,
+                            device=block_map.device)
+
+    grid = (bs, h, num_q_blocks)
+    map_to_index_kernel[grid](
+        block_map,
+        index,
+        index_num,
+        block_map.stride(0),
+        block_map.stride(1),
+        block_map.stride(2),
+        block_map.stride(3),
+        index.stride(0),
+        index.stride(1),
+        index.stride(2),
+        index.stride(3),
+        index_num.stride(0),
+        index_num.stride(1),
+        index_num.stride(2),
+        num_kv_blocks=num_kv_blocks,
+    )
+
+    return index, index_num
+
+class BlockSparseAttentionFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
+        o, lse = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
+        ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num)
+        return o
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors
+        grad_q, grad_k, grad_v = block_sparse_attention_backward(
+            q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num
+        )
+        return grad_q, grad_k, grad_v, None, None, None, None
+
+
 class DummyOperator(torch.autograd.Function):
     @staticmethod
     def forward(ctx, x):
@@ -207,4 +466,4 @@ class BlockSparseAttnTorch:
 
         o = DummyOperator.apply(output)
         o.register_hook(self.recompute)
-        return o
+        return o

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/data_preprocess/preprocess_vae_latents.py

@@ -104,13 +104,7 @@ if __name__ == "__main__":
         default=None,
         help="The output directory where the model predictions and checkpoints will be written.",
     )
-    parser.add_argument(
-        "--logging_dir",
-        type=str,
-        default="logs",
-        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
-    )
+
 
     args = parser.parse_args()
     main(args)

fastvideo/data_preprocess/v1_preprocess.py

@@ -18,7 +18,7 @@ logger = init_logger(__name__)
 
 def main(args):
     args.model_path = maybe_download_model(args.model_path)
-    maybe_init_distributed_environment_and_model_parallel(args.tp_size, args.sp_size)
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
 
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
     kwargs = {
@@ -43,7 +43,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",
@@ -97,13 +97,6 @@ if __name__ == "__main__":
         default=None,
         help="The output directory where the model predictions and checkpoints will be written.",
     )
-    parser.add_argument(
-        "--logging_dir",
-        type=str,
-        default="logs",
-        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
-    )
 
     args = parser.parse_args()
     main(args)

fastvideo/distill.py

@@ -671,13 +671,6 @@ if __name__ == "__main__":
         help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
               ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
-    parser.add_argument(
-        "--logging_dir",
-        type=str,
-        default="logs",
-        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
-    )
 
     # optimizer & scheduler & Training
     parser.add_argument("--num_train_epochs", type=int, default=100)

fastvideo/distill_adv.py

@@ -693,13 +693,6 @@ if __name__ == "__main__":
         help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
               ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
-    parser.add_argument(
-        "--logging_dir",
-        type=str,
-        default="logs",
-        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
-    )
 
     # optimizer & scheduler & Training
     parser.add_argument("--num_train_epochs", type=int, default=100)

fastvideo/train.py

@@ -520,13 +520,7 @@ if __name__ == "__main__":
         help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
               ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
-    parser.add_argument(
-        "--logging_dir",
-        type=str,
-        default="logs",
-        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
-    )
+
 
     # optimizer & scheduler & Training
     parser.add_argument("--num_train_epochs", type=int, default=100)

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -287,16 +287,10 @@ class SlidingTileAttentionImpl(AttentionImpl):
                     forward_batch.mask_search_final_result_pos[timestep].append(
                         layer_loss_save)
         else:
-            # windows = [
-            #     self.mask_strategy[timestep][layer_idx][head_idx + start_head]
-            #     for head_idx in range(head_num)
-            # ]
             windows = [
                 STA_param[head_idx + start_head] for head_idx in range(head_num)
             ]
-            # if has_text is False:
-            #     from IPython import embed
-            #     embed()
+
             hidden_states = sliding_tile_attention(
                 query, key, value, windows, text_length, has_text,
                 self.dit_seq_shape_str).transpose(1, 2)

fastvideo/v1/attention/backends/video_sparse_attn.py

@@ -1,18 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 import math
 from dataclasses import dataclass
-from typing import Any, List, Optional, Type, cast
+from typing import List, Optional, Type
 
 import torch
-import triton
-import triton.language as tl
 from einops import rearrange
-
-try:
-    from vsa import block_sparse_attn
-except ImportError:  # noqa: E722
-    block_sparse_attn = None
-from typing import Tuple
+from vsa import video_sparse_attn
 
 from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
                                                       AttentionImpl,
@@ -181,8 +174,7 @@ class VideoSparseAttentionImpl(AttentionImpl):
             (1 - attn_metadata.VSA_sparsity) *
             (self.img_seq_length / math.prod(self.VSA_base_tile_size)))
 
-        # Cast to Any to bypass type checking for untyped function
-        hidden_states = cast(Any, sparse_attn_c_s_p)(
+        hidden_states = video_sparse_attn(
             query,
             key,
             value,
@@ -191,267 +183,3 @@ class VideoSparseAttentionImpl(AttentionImpl):
             compress_attn_weight=gate_compress).transpose(1, 2)
 
         return hidden_states
-
-
-def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
-    QK = torch.matmul(q, k.transpose(-2, -1))
-    QK /= (q.size(-1)**0.5)
-
-    # Causal mask removed since causal is always false
-
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v)
-    return output, QK
-
-
-def sparse_attn_c_s_p(q, k, v, topk, block_size, compress_attn_weight=None):
-    """
-    q: [batch_size, num_heads, seq_len, head_dim]
-    k: [batch_size, num_heads, seq_len, head_dim]
-    v: [batch_size, num_heads, seq_len, head_dim]
-    topk: int
-    block_size: int or tuple of 3 ints
-    video_shape: tuple of (T, H, W)
-    compress_attn_weight: [batch_size, num_heads, seq_len, head_dim]
-    select_attn_weight: [batch_size, num_heads, seq_len, head_dim]
-    
-    V1 of sparse attention. Include compress attn and sparse attn branch, use average pooling to compress. 
-    Assume q, k, v is flattened in this way: [batch_size, num_heads, T//block_size[0], H//block_size[1], W//block_size[2], block_size[0], block_size[1], block_size[2]]
-    """
-
-    if isinstance(block_size, int):
-        block_size = (block_size, block_size, block_size)
-
-    block_elements = block_size[0] * block_size[1] * block_size[2]
-    assert block_elements % 64 == 0 and block_elements >= 64
-    assert q.shape[2] % block_elements == 0
-    batch_size, num_heads, seq_len, head_dim = q.shape
-    # compress attn
-    q_compress = q.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).mean(dim=3)
-    k_compress = k.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).mean(dim=3)
-    v_compress = v.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).mean(dim=3)
-
-    output_compress, block_attn_score = torch_attention(q_compress, k_compress,
-                                                        v_compress)
-
-    output_compress = output_compress.view(batch_size, num_heads,
-                                           seq_len // block_elements, 1,
-                                           head_dim)
-    output_compress = output_compress.repeat(1, 1, 1, block_elements,
-                                             1).view(batch_size, num_heads,
-                                                     seq_len, head_dim)
-
-    q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num = generate_topk_block_sparse_pattern(
-        block_attn_score, topk)
-
-    output_select = block_sparse_attn(q, k, v, q2k_block_sparse_index,
-                                      q2k_block_sparse_num,
-                                      k2q_block_sparse_index,
-                                      k2q_block_sparse_num)
-
-    if compress_attn_weight is not None:
-        final_output = output_compress * compress_attn_weight + output_select
-    else:
-        final_output = output_compress + output_select
-    return final_output
-
-
-@triton.jit
-def topk_index_to_map_kernel(
-    map_ptr,
-    index_ptr,
-    map_bs_stride,
-    map_h_stride,
-    map_q_stride,
-    map_kv_stride,
-    index_bs_stride,
-    index_h_stride,
-    index_q_stride,
-    index_kv_stride,
-    topk: tl.constexpr,
-):
-    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
-    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
-
-    for i in tl.static_range(topk):
-        index = tl.load(index_ptr_base + i * index_kv_stride)
-        tl.store(map_ptr_base + index * map_kv_stride, 1.0)
-
-
-@triton.jit
-def map_to_index_kernel(
-    map_ptr,
-    index_ptr,
-    index_num_ptr,
-    map_bs_stride,
-    map_h_stride,
-    map_q_stride,
-    map_kv_stride,
-    index_bs_stride,
-    index_h_stride,
-    index_q_stride,
-    index_kv_stride,
-    index_num_bs_stride,
-    index_num_h_stride,
-    index_num_q_stride,
-    num_kv_blocks: tl.constexpr,
-):
-    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
-    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
-
-    num = 0
-    for i in tl.static_range(num_kv_blocks):
-        map_entry = tl.load(map_ptr_base + i * map_kv_stride)
-        if map_entry:
-            tl.store(index_ptr_base + num * index_kv_stride, i)
-            num += 1
-
-    tl.store(
-        index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
-        q * index_num_q_stride, num)
-
-
-def topk_index_to_map(index: torch.Tensor,
-                      num_kv_blocks: int,
-                      transpose_map: bool = False):
-    """
-    Convert topk indices to a map.
-    
-    Args:
-        index: [bs, h, num_q_blocks, topk]
-            The topk indices tensor.
-        num_kv_blocks: int
-            The number of key-value blocks in the block_map returned
-        transpose_map: bool
-            If True, the block_map will be transposed on the final two dimensions.
-    
-    Returns:
-        block_map: [bs, h, num_q_blocks, num_kv_blocks]
-            A binary map where 1 indicates that the q block attends to the kv block.
-    """
-    bs, h, num_q_blocks, topk = index.shape
-
-    if transpose_map is False:
-        block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks),
-                                dtype=torch.bool,
-                                device=index.device)
-    else:
-        block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks),
-                                dtype=torch.bool,
-                                device=index.device)
-        block_map = block_map.transpose(2, 3)
-
-    grid = (bs, h, num_q_blocks)
-    topk_index_to_map_kernel[grid](
-        block_map,
-        index,
-        block_map.stride(0),
-        block_map.stride(1),
-        block_map.stride(2),
-        block_map.stride(3),
-        index.stride(0),
-        index.stride(1),
-        index.stride(2),
-        index.stride(3),
-        topk=topk,
-    )
-
-    return block_map
-
-
-def map_to_index(block_map: torch.Tensor):
-    """
-    Convert a block map to indices and counts.
-    
-    Args:
-        block_map: [bs, h, num_q_blocks, num_kv_blocks]
-            The block map tensor.
-    
-    Returns:
-        index: [bs, h, num_q_blocks, num_kv_blocks]
-            The indices of the blocks.
-        index_num: [bs, h, num_q_blocks]
-            The number of blocks for each q block.
-    """
-    bs, h, num_q_blocks, num_kv_blocks = block_map.shape
-
-    index = torch.full((block_map.shape),
-                       -1,
-                       dtype=torch.int32,
-                       device=block_map.device)
-    index_num = torch.empty((bs, h, num_q_blocks),
-                            dtype=torch.int32,
-                            device=block_map.device)
-
-    grid = (bs, h, num_q_blocks)
-    map_to_index_kernel[grid](
-        block_map,
-        index,
-        index_num,
-        block_map.stride(0),
-        block_map.stride(1),
-        block_map.stride(2),
-        block_map.stride(3),
-        index.stride(0),
-        index.stride(1),
-        index.stride(2),
-        index.stride(3),
-        index_num.stride(0),
-        index_num.stride(1),
-        index_num.stride(2),
-        num_kv_blocks=num_kv_blocks,
-    )
-
-    return index, index_num
-
-
-def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,
-                                       topk: int):
-    """
-    Generate a block sparse pattern where each q block attends to exactly topk kv blocks,
-    based on the provided attention scores.
-    
-    Args:
-        block_attn_score: [bs, h, num_q_blocks, num_kv_blocks]
-            Attention scores between query and key blocks
-        topk: int
-            Number of kv blocks each q block attends to
-        
-    Returns:
-        q2k_block_sparse_index: [bs, h, num_q_blocks, topk]
-            Contains the indices of kv blocks that each q block attends to.
-        q2k_block_sparse_num: [bs, h, num_q_blocks]
-            Contains the number of kv blocks that each q block attends to (all equal to topk).
-        k2q_block_sparse_index: [bs, h, num_kv_blocks, max_q_per_kv]
-            Contains the indices of q blocks that attend to each kv block.
-        k2q_block_sparse_num: [bs, h, num_kv_blocks]
-            Contains the number of q blocks that attend to each kv block.
-    """
-    device = block_attn_score.device
-    # Extract dimensions from block_attn_score
-    bs, h, num_q_blocks, num_kv_blocks = block_attn_score.shape
-
-    sorted_result = torch.sort(block_attn_score, dim=-1, descending=True)
-
-    sorted_indice = sorted_result.indices
-
-    q2k_block_sparse_index, _ = torch.sort(sorted_indice[:, :, :, :topk],
-                                           dim=-1)
-    q2k_block_sparse_index = q2k_block_sparse_index.to(dtype=torch.int32)
-    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks),
-                                      topk,
-                                      device=device,
-                                      dtype=torch.int32)
-
-    block_map = topk_index_to_map(q2k_block_sparse_index,
-                                  num_kv_blocks,
-                                  transpose_map=True)
-    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(
-        block_map.transpose(2, 3))
-
-    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num

fastvideo/v1/dataset/dataloader/schema.py

@@ -35,6 +35,9 @@ pyarrow_schema_i2v = pa.schema([
     pa.field("clip_feature_bytes", pa.binary()),
     pa.field("clip_feature_shape", pa.list_(pa.int64())),
     pa.field("clip_feature_dtype", pa.string()),
+    pa.field("encoded_first_frame_bytes", pa.binary()),
+    pa.field("encoded_first_frame_shape", pa.list_(pa.int64())),
+    pa.field("encoded_first_frame_dtype", pa.string()),
     # --- Metadata ---
     pa.field("file_name", pa.string()),
     pa.field("caption", pa.string()),

fastvideo/v1/dataset/parquet_dataset_map_style.py

@@ -268,6 +268,8 @@ class LatentsParquetMapStyleDataset(Dataset):
         for i, row in enumerate(rows):
             # Get tensors from row
             data = self._get_torch_tensors_from_row_dict(row)
+            print(data)
+            import pdb; pdb.set_trace()
             latents, emb = data["vae_latent"], data["text_embedding"]
 
             padded_emb, mask = self._pad(emb, self.text_padding_length)

fastvideo/v1/fastvideo_args.py

@@ -44,8 +44,8 @@ class FastVideoArgs:
     num_gpus: int = 1
     tp_size: int = -1
     sp_size: int = -1
-    dp_size: int = 1
-    dp_shards: int = -1
+    hsdp_replicate_dim: int = 1
+    hsdp_shard_dim: int = -1
     dist_timeout: Optional[int] = None  # timeout for torch.distributed
 
     # Video generation parameters
@@ -189,17 +189,15 @@ class FastVideoArgs:
             help="The sequence parallelism size.",
         )
         parser.add_argument(
-            "--data-parallel-size",
-            "--dp-size",
+            "--hsdp-replicate-dim",
             type=int,
-            default=FastVideoArgs.dp_size,
+            default=FastVideoArgs.hsdp_replicate_dim,
             help="The data parallelism size.",
         )
         parser.add_argument(
-            "--data-parallel-shards",
-            "--dp-shards",
+            "--hsdp-shard-dim",
             type=int,
-            default=FastVideoArgs.dp_shards,
+            default=FastVideoArgs.hsdp_shard_dim,
             help="The data parallelism shards.",
         )
         parser.add_argument(
@@ -379,10 +377,6 @@ class FastVideoArgs:
                 kwargs[attr] = args.tensor_parallel_size
             elif attr == 'sp_size' and hasattr(args, 'sequence_parallel_size'):
                 kwargs[attr] = args.sequence_parallel_size
-            elif attr == 'dp_size' and hasattr(args, 'data_parallel_size'):
-                kwargs[attr] = args.data_parallel_size
-            elif attr == 'dp_shards' and hasattr(args, 'data_parallel_shards'):
-                kwargs[attr] = args.data_parallel_shards
             elif attr == 'flow_shift' and hasattr(args, 'shift'):
                 kwargs[attr] = args.shift
             # Use getattr with default value from the dataclass for potentially missing attributes
@@ -397,19 +391,20 @@ class FastVideoArgs:
     def check_fastvideo_args(self) -> None:
         """Validate inference arguments for consistency"""
         if not self.inference_mode:
-            assert self.dp_size is not -1, "dp_size must be set for training"
-            assert self.dp_shards is not -1, "dp_shards must be set for training"
-            assert self.sp_size is not -1, "sp_size must be set for training"
+            assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training"
+            assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training"
+            assert self.sp_size != -1, "sp_size must be set for training"
 
-        if self.tp_size is -1:
+        if self.tp_size == -1:
             self.tp_size = self.num_gpus
-        if self.sp_size is -1:
+        if self.sp_size == -1:
             self.sp_size = self.num_gpus
-        if self.dp_shards is -1:
-            self.dp_shards = self.num_gpus
+        if self.hsdp_shard_dim == -1:
+            self.hsdp_shard_dim = self.num_gpus
+
         assert self.sp_size <= self.num_gpus and self.num_gpus % self.sp_size == 0, "num_gpus must >= and be divisible by sp_size"
-        assert self.dp_size <= self.num_gpus and self.num_gpus % self.dp_size == 0, "num_gpus must >= and be divisible by dp_size"
-        assert self.dp_shards <= self.num_gpus and self.num_gpus % self.dp_shards == 0, "num_gpus must >= and be divisible by dp_shards"
+        assert self.hsdp_replicate_dim <= self.num_gpus and self.num_gpus % self.hsdp_replicate_dim == 0, "num_gpus must >= and be divisible by hsdp_replicate_dim"
+        assert self.hsdp_shard_dim <= self.num_gpus and self.num_gpus % self.hsdp_shard_dim == 0, "num_gpus must >= and be divisible by hsdp_shard_dim"
 
         if self.num_gpus < max(self.tp_size, self.sp_size):
             self.num_gpus = max(self.tp_size, self.sp_size)
@@ -528,7 +523,8 @@ class TrainingArgs(FastVideoArgs):
     precondition_outputs: bool = False
 
     # validation & logs
-    validation_prompt_dir: str = ""
+    validation_dataset_file: str = ""
+    validation_path: str = ""
     validation_sampling_steps: str = ""
     validation_guidance_scale: str = ""
     validation_steps: float = 0.0
@@ -541,7 +537,6 @@ class TrainingArgs(FastVideoArgs):
     checkpoints_total_limit: int = 0
     checkpointing_steps: int = 0
     resume_from_checkpoint: bool = False
-    logging_dir: str = ""
 
     # optimizer & scheduler
     num_train_epochs: int = 0
@@ -582,9 +577,6 @@ class TrainingArgs(FastVideoArgs):
     # master_weight_type
     master_weight_type: str = ""
 
-    # For fast checking in LoRA pipeline
-    training_mode: bool = True
-
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
         # Get all fields from the dataclass
@@ -600,16 +592,10 @@ class TrainingArgs(FastVideoArgs):
                 kwargs[attr] = args.sequence_parallel_size
             elif attr == 'flow_shift' and hasattr(args, 'shift'):
                 kwargs[attr] = args.shift
-            elif attr == 'dp_size' and hasattr(args, 'data_parallel_size'):
-                kwargs[attr] = args.data_parallel_size
-            elif attr == 'dp_shards' and hasattr(args, 'data_parallel_shards'):
-                kwargs[attr] = args.data_parallel_shards
-            # Use getattr with default value from the dataclass for potentially missing attributes
             else:
                 default_value = getattr(cls, attr, None)
-                value = getattr(args, attr, default_value)
-                if value is not None:
-                    kwargs[attr] = value
+                if getattr(args, attr, default_value) is not None:
+                    kwargs[attr] = getattr(args, attr, default_value)
 
         return cls(**kwargs)
 
@@ -683,9 +669,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="File containing validation dataset")
+        parser.add_argument("--validation-path",
+                            type=str,
+                            help="Path to validation dataset")
         parser.add_argument("--validation-sampling-steps",
                             type=str,
                             help="Validation sampling steps")

fastvideo/v1/models/dits/wanvideo.py

@@ -618,6 +618,8 @@ class WanTransformer3DModel(CachableDiT):
 
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
             timestep, encoder_hidden_states, encoder_hidden_states_image)
+        if encoder_hidden_states.dim() == 2:
+            encoder_hidden_states = encoder_hidden_states.unsqueeze(0)
         timestep_proj = timestep_proj.unflatten(1, (6, -1))
 
         if encoder_hidden_states_image is not None:

fastvideo/v1/models/loader/component_loader.py

@@ -403,7 +403,7 @@ class TransformerLoader(ComponentLoader):
         # Load the model using FSDP loader
         logger.info("Loading model from %s, default_dtype: %s", cls_name,
                     default_dtype)
-        assert fastvideo_args.dp_shards is not None
+        assert fastvideo_args.hsdp_shard_dim is not None
         model = maybe_load_fsdp_model(
             model_cls=model_cls,
             init_params={
@@ -412,8 +412,8 @@ class TransformerLoader(ComponentLoader):
             },
             weight_dir_list=safetensors_list,
             device=get_torch_device(),
-            data_parallel_size=fastvideo_args.dp_size,
-            data_parallel_shards=fastvideo_args.dp_shards,
+            hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
+            hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
             cpu_offload=fastvideo_args.use_cpu_offload,
             fsdp_inference=fastvideo_args.use_fsdp_inference,
             default_dtype=default_dtype,

fastvideo/v1/models/loader/fsdp_load.py

@@ -60,8 +60,8 @@ def maybe_load_fsdp_model(
     init_params: Dict[str, Any],
     weight_dir_list: List[str],
     device: torch.device,
-    data_parallel_size: int,
-    data_parallel_shards: int,
+    hsdp_replicate_dim: int,
+    hsdp_shard_dim: int,
     default_dtype: torch.dtype,
     param_dtype: torch.dtype,
     reduce_dtype: torch.dtype,
@@ -87,13 +87,15 @@ def maybe_load_fsdp_model(
 
     with set_default_dtype(default_dtype), torch.device("meta"):
         model = model_cls(**init_params)
-
-    dp_size = data_parallel_size if fsdp_inference or training_mode else 1
+    world_size = hsdp_replicate_dim * hsdp_shard_dim
+    if not training_mode and not fsdp_inference:
+        hsdp_replicate_dim = world_size
+        hsdp_shard_dim = 1
     device_mesh = init_device_mesh(
         "cuda",
         # (Replicate(), Shard(dim=0))
-        mesh_shape=(dp_size, data_parallel_shards),
-        mesh_dim_names=("dp_replicate", "dp_shards"),
+        mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
+        mesh_dim_names=("replicate", "shard"),
     )
     shard_model(model,
                 cpu_offload=cpu_offload,

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -12,6 +12,7 @@ from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List, Optional, Union
 
 import torch
+import PIL.Image
 
 from fastvideo.v1.configs.sample.teacache import (TeaCacheParams,
                                                   WanTeaCacheParams)
@@ -36,6 +37,8 @@ class ForwardBatch:
     # Image inputs
     image_path: Optional[str] = None
     image_embeds: List[torch.Tensor] = field(default_factory=list)
+    preprocessed_image: Optional[torch.Tensor] = None
+    pil_image: Optional[PIL.Image.Image] = None
 
     # Text inputs
     prompt: Optional[Union[str, List[str]]] = None

fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py

@@ -9,7 +9,15 @@ from typing import Any, Dict, List, Optional
 
 import numpy as np
 import torch
+import PIL
 from PIL import Image
+import os
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor
+import gc
+import pyarrow as pa
+from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 
 from fastvideo.v1.dataset.dataloader.schema import pyarrow_schema_i2v
 from fastvideo.v1.distributed import get_torch_device
@@ -17,6 +25,13 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.pipelines.preprocess_pipeline_base import (
     BasePreprocessPipeline)
+from fastvideo.v1.models.vision_utils import numpy_to_pt, pil_to_numpy, normalize
+from fastvideo.v1.dataset.validation_dataset import ValidationDataset
+from fastvideo.v1.pipelines.stages import TextEncodingStage, ImageEncodingStage
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
 
 
 class PreprocessPipeline_I2V(BasePreprocessPipeline):
@@ -26,18 +41,235 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
         "text_encoder", "tokenizer", "vae", "image_encoder", "image_processor"
     ]
 
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=ImageEncodingStage(
+                           image_encoder=self.get_module("image_encoder"),
+                           image_processor=self.get_module("image_processor"),
+                       ))
+
+
+    def preprocess_validation_text(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.
+        Subclasses can override this method to add pipeline-specific features.
+        """
+        # Create Parquet dataset directory for validation
+        validation_parquet_dir = os.path.join(args.output_dir,
+                                              "validation_parquet_dataset")
+        os.makedirs(validation_parquet_dir, exist_ok=True)
+
+        validation_dataset = ValidationDataset(args.validation_dataset_file)
+
+        from itertools import chain
+
+
+        # for data in validation_dataset:
+            # print(data)
+
+        # assert False
+
+        # with open(args.validation_prompt_txt, encoding="utf-8") as file:
+        #     lines = file.readlines()
+        # prompts = [line.strip() for line in lines]
+
+        # Prepare batch data for Parquet dataset
+        batch_data = []
+        sampling_param = SamplingParam.from_pretrained(
+            fastvideo_args.model_path)
+        if sampling_param.negative_prompt:
+            negative_prompt = {
+                'caption': sampling_param.negative_prompt,
+                'image_path': None,
+                'video_path': None,
+            }
+        else:
+            negative_prompt = None
+
+        validation_dataset = chain([negative_prompt], validation_dataset)
+
+        
+
+
+
+        # Add progress bar for validation text preprocessing
+        pbar = tqdm(enumerate(validation_dataset),
+                    desc="Processing validation dataset",
+                    unit="sample")
+        for idx, row in pbar:
+            print(idx)
+            # print(type(row))
+            # print(row)
+            prompt = row['caption']
+            print(row)
+            print(prompt)
+
+
+            with torch.inference_mode():
+                # Text Encoder
+                batch = ForwardBatch(
+                    data_type="video",
+                    prompt=prompt,
+                    prompt_embeds=[],
+                    prompt_attention_mask=[],
+                )
+                assert hasattr(self, "prompt_encoding_stage")
+                result_batch = self.prompt_encoding_stage(batch, fastvideo_args)
+
+                if idx != 0:
+                    image = None
+                    if 'image' in row:
+                        image = row['image']
+                    else:
+                        assert 'video' in row
+                        image = row['video'][0]
+                    
+                    assert image is not None
+                    result_batch.pil_image = image
+
+                    # assert image is not None
+                    assert hasattr(self, "image_encoding_stage")
+                    result_batch = self.image_encoding_stage(result_batch, fastvideo_args)
+
+            prompt_embeds = result_batch.prompt_embeds[0]
+            prompt_attention_mask = result_batch.prompt_attention_mask[0]
+
+            file_name = prompt.split(".")[0]
+
+            # Get the sequence length from attention mask (number of 1s)
+            seq_len = prompt_attention_mask.sum().item()
+
+            text_embedding = prompt_embeds[0, :seq_len].cpu().numpy()
+            text_attention_mask = prompt_attention_mask[
+                0, :seq_len].cpu().numpy().astype(np.uint8)
+
+            if idx != 0:
+                image_embeds = result_batch.image_embeds[0]
+                image_embeds = image_embeds.cpu().numpy()
+            else:
+                image_embeds = None
+
+            # Log the shapes after removing padding
+            logger.info(
+                "Shape after removing padding - Embeddings: %s, Mask: %s",
+                text_embedding.shape, text_attention_mask.shape)
+
+
+            # image embedding
+            if idx != 0:
+                clip_feature = {
+                    'clip_feature': image_embeds,
+                }
+            else:
+                clip_feature = None
+
+            # Create record for Parquet dataset
+            record = self.create_record(video_name=file_name,
+                                        vae_latent=np.array([],
+                                                            dtype=np.float32),
+                                        text_embedding=text_embedding,
+                                        text_attention_mask=text_attention_mask,
+                                        valid_data=None,
+                                        idx=0,
+                                        extra_features=clip_feature)
+            batch_data.append(record)
+
+            logger.info("Saved validation sample: %s", file_name)
+
+        if batch_data:
+            # Add progress bar for writing to Parquet dataset
+            write_pbar = tqdm(total=1,
+                              desc="Writing to Parquet dataset",
+                              unit="batch")
+            # Convert batch data to PyArrow arrays
+            arrays = []
+            for field in self.get_schema_fields():
+                if field.endswith('_bytes'):
+                    arrays.append(
+                        pa.array([record[field] for record in batch_data],
+                                 type=pa.binary()))
+                elif field.endswith('_shape'):
+                    arrays.append(
+                        pa.array([record[field] for record in batch_data],
+                                 type=pa.list_(pa.int32())))
+                elif field in ['width', 'height', 'num_frames']:
+                    arrays.append(
+                        pa.array([record[field] for record in batch_data],
+                                 type=pa.int32()))
+                elif field in ['duration_sec', 'fps']:
+                    arrays.append(
+                        pa.array([record[field] for record in batch_data],
+                                 type=pa.float32()))
+                else:
+                    arrays.append(
+                        pa.array([record[field] for record in batch_data]))
+
+            table = pa.Table.from_arrays(arrays, names=self.get_schema_fields())
+            write_pbar.update(1)
+            write_pbar.close()
+
+            logger.info("Total validation samples: %s", len(table))
+
+            work_range = (0, 1, table, 0, validation_parquet_dir, len(table))
+
+            total_written = 0
+            failed_ranges = []
+            with ProcessPoolExecutor(max_workers=1) as executor:
+                futures = {
+                    executor.submit(self.process_chunk_range, work_range):
+                    work_range
+                }
+                for future in tqdm(futures, desc="Processing chunks"):
+                    try:
+                        total_written += future.result()
+                    except Exception as e:
+                        work_range = futures[future]
+                        failed_ranges.append(work_range)
+                        logger.error("Failed to process range %s-%s: %s",
+                                     work_range[0], work_range[1], str(e))
+
+            if failed_ranges:
+                logger.warning("Retrying %s failed ranges sequentially",
+                               len(failed_ranges))
+                for work_range in failed_ranges:
+                    try:
+                        total_written += self.process_chunk_range(work_range)
+                    except Exception as e:
+                        logger.error(
+                            "Failed to process range %s-%s after retry: %s",
+                            work_range[0], work_range[1], str(e))
+
+            logger.info("Total validation samples written: %s", total_written)
+
+            # Clear memory
+            del table
+            gc.collect()  # Force garbage collection
+
     def get_schema_fields(self) -> List[str]:
         """Get the schema fields for I2V pipeline."""
         return [f.name for f in pyarrow_schema_i2v]
 
     def get_extra_features(self, valid_data: Dict[str, Any],
                            fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
+        features = {}
         """Get CLIP features from the first frame of each video."""
         first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(
             0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
+        batch_size, _, num_frames, height, width = valid_data["pixel_values"].shape
+        latent_height = height // self.get_module("vae").spatial_compression_ratio
+        latent_width = width // self.get_module("vae").spatial_compression_ratio
 
         processed_images = []
+        # Frame has values between -1 and 1
         for frame in first_frame:
+            frame = (frame + 1) * 127.5
             frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
             processed_img = self.get_module("image_processor")(
                 images=frame_pil, return_tensors="pt")
@@ -52,8 +284,69 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
             with set_forward_context(current_timestep=0, attn_metadata=None):
                 clip_features = self.get_module("image_encoder")(**image_inputs)
             clip_features = clip_features.last_hidden_state
+        
+        features["clip_feature"] = clip_features
 
-        return {"clip_feature": clip_features}
+        """Get VAE features from the first frame of each video"""
+        video_conditions = []
+        for frame in first_frame:
+            processed_img = frame.to(device="cpu", dtype=torch.float32)
+            processed_img = processed_img.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(2)
+            # (B, H, W, C) -> (B, C, 1, H, W)
+            video_condition = torch.cat([
+                processed_img,
+                processed_img.new_zeros(processed_img.shape[0], processed_img.shape[1],
+                                num_frames - 1, height, width)
+            ],
+                                        dim=2)
+            video_condition = video_condition.to(device=get_torch_device(),
+                                                dtype=torch.float32)
+            video_conditions.append(video_condition)
+        
+        video_conditions = torch.cat(video_conditions, dim=0)
+
+        with torch.autocast(device_type="cuda",
+                            dtype=torch.float32,
+                            enabled=True):
+            encoder_outputs = self.get_module("vae").encode(video_conditions)
+
+        latent_condition = encoder_outputs.mean
+        if (hasattr(self.get_module("vae"), "shift_factor")
+                and self.get_module("vae").shift_factor is not None):
+            if isinstance(self.get_module("vae").shift_factor, torch.Tensor):
+                latent_condition -= self.get_module("vae").shift_factor.to(
+                    latent_condition.device, latent_condition.dtype)
+            else:
+                latent_condition -= self.get_module("vae").shift_factor
+
+        if isinstance(self.get_module("vae").scaling_factor, torch.Tensor):
+            latent_condition = latent_condition * self.get_module("vae").scaling_factor.to(
+                latent_condition.device, latent_condition.dtype)
+        else:
+            latent_condition = latent_condition * self.get_module("vae").scaling_factor
+
+        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
+                                   latent_width)
+        mask_lat_size[:, :, list(range(1, num_frames))] = 0
+        first_frame_mask = mask_lat_size[:, :, 0:1]
+        first_frame_mask = torch.repeat_interleave(
+            first_frame_mask,
+            dim=2,
+            repeats=self.get_module("vae").temporal_compression_ratio)
+        mask_lat_size = torch.concat(
+            [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
+        mask_lat_size = mask_lat_size.view(batch_size, -1,
+                                           self.get_module("vae").temporal_compression_ratio,
+                                           latent_height, latent_width)
+        mask_lat_size = mask_lat_size.transpose(1, 2)
+        mask_lat_size = mask_lat_size.to(latent_condition.device)
+
+        image_latent = torch.concat([mask_lat_size, latent_condition],
+                                          dim=1)
+
+        features["encoded_first_frame"] = image_latent
+
+        return features
 
     def create_record(
             self,
@@ -87,6 +380,20 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
                 "clip_feature_dtype": "",
             })
 
+        if extra_features and "encoded_first_frame" in extra_features:
+            encoded_first_frame = extra_features["encoded_first_frame"]
+            record.update({
+                "encoded_first_frame_bytes": encoded_first_frame.tobytes(),
+                "encoded_first_frame_shape": list(encoded_first_frame.shape),
+                "encoded_first_frame_dtype": str(encoded_first_frame.dtype),
+            })
+        else:
+            record.update({
+                "encoded_first_frame_bytes": b"",
+                "encoded_first_frame_shape": [],
+                "encoded_first_frame_dtype": "",
+            })
+
         return record
 
 

fastvideo/v1/pipelines/preprocess_pipeline_base.py

@@ -419,7 +419,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
             # Clear memory
             del table
             gc.collect()  # Force garbage collection
-
+    
     def _flush_tables(self, num_processed_samples: int, args,
                       combined_parquet_dir: str):
         """Flush collected tables to disk."""

fastvideo/v1/pipelines/stages/denoising.py

@@ -101,7 +101,7 @@ class DenoisingStage(PipelineStage):
                                 n=sp_world_size).contiguous()
             latents = latents[:, :, rank_in_sp_group, :, :, :]
             batch.latents = latents
-            if batch.image_latent is not None:
+            if batch.image_latent is not None and latents.shape[2] != batch.image_latent.shape[2]:
                 image_latent = rearrange(batch.image_latent,
                                          "b t (n s) h w -> b t n s h w",
                                          n=sp_world_size).contiguous()

fastvideo/v1/pipelines/stages/encoding.py

@@ -49,22 +49,30 @@ class EncodingStage(PipelineStage):
         """
         self.vae = self.vae.to(get_torch_device())
 
-        image_path = batch.image_path
+        assert batch.pil_image is not None, "Image must be provided"
+        image = batch.pil_image
+
+        # image_path = batch.image_path
         # TODO(will): remove this once we add input/output validation for stages
-        if image_path is None:
-            raise ValueError("Image Path must be provided")
+        # if image_path is None:
+        #     raise ValueError("Image Path must be provided")
         assert batch.height is not None
         assert batch.width is not None
         latent_height = batch.height // self.vae.spatial_compression_ratio
         latent_width = batch.width // self.vae.spatial_compression_ratio
 
-        image = load_image(image_path)
-        image = self.preprocess(
+        # image = load_image(image_path)
+        image = self.preprocess_image(
             image,
             vae_scale_factor=self.vae.spatial_compression_ratio,
             height=batch.height,
             width=batch.width).to(get_torch_device(), dtype=torch.float32)
         image = image.unsqueeze(2)
+
+        # if batch.preprocessed_image is not None:
+        #     image = batch.preprocessed_image
+
+
         video_condition = torch.cat([
             image,
             image.new_zeros(image.shape[0], image.shape[1],
@@ -148,7 +156,7 @@ class EncodingStage(PipelineStage):
             raise AttributeError(
                 "Could not access latents of provided encoder_output")
 
-    def preprocess(
+    def preprocess_image(
             self,
             image: PIL.Image.Image,
             vae_scale_factor: int,
@@ -172,4 +180,4 @@ class EncodingStage(PipelineStage):
         if do_normalize:
             image = normalize(image)
 
-        return image
+        return image

fastvideo/v1/pipelines/stages/image_encoding.py

@@ -56,7 +56,9 @@ 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 = load_image(batch.image_path)
+        assert batch.pil_image is not None, "Image must be provided"
+        image = batch.pil_image
 
         image_inputs = self.image_processor(
             images=image, return_tensors="pt").to(get_torch_device())

fastvideo/v1/pipelines/stages/input_validation.py

@@ -3,12 +3,21 @@
 Input validation stage for diffusion pipelines.
 """
 
-import torch
+from typing import Optional
 
+import torch
+from fastvideo.v1.models.vision_utils import (load_image, pil_to_numpy,
+                                              numpy_to_pt, normalize,
+                                              get_default_height_width,
+                                              resize)
+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.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
+import PIL.Image
+
+
 
 logger = init_logger(__name__)
 
@@ -85,5 +94,44 @@ class InputValidationStage(PipelineStage):
             raise ValueError(
                 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
+
+            # image = self.preprocess_image(
+            #     image,
+            #     vae_scale_factor=self.vae.spatial_compression_ratio,
+            #     height=batch.height,
+            #     width=batch.width).to(get_torch_device(), dtype=torch.float32)
+            # image = image.unsqueeze(2)
+            # batch.preprocessed_image = image
 
         return batch
+
+    def preprocess_image(
+            self,
+            image: PIL.Image.Image,
+            vae_scale_factor: int,
+            height: Optional[int] = None,
+            width: Optional[int] = None,
+            resize_mode: str = "default",  # "default", "fill", "crop"
+    ) -> torch.Tensor:
+        image = [image]
+
+        height, width = get_default_height_width(image[0], vae_scale_factor,
+                                                 height, width)
+        image = [
+            resize(i, height, width, resize_mode=resize_mode) for i in image
+        ]
+        image = pil_to_numpy(image)  # to np
+        image = numpy_to_pt(image)  # to pt
+
+        do_normalize = True
+        if image.min() < 0:
+            do_normalize = False
+        if do_normalize:
+            image = normalize(image)
+
+        return image

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.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/training/wan_i2v_training_pipeline.py

@@ -0,0 +1,588 @@
+import os
+import random
+import sys
+import time
+from collections import deque
+from copy import deepcopy
+import gc
+
+import torchvision
+import imageio
+from einops import rearrange
+import numpy as np
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from tqdm.auto import tqdm
+
+from fastvideo import SamplingParam
+from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_sp_group,
+                                      get_world_group, get_torch_device)
+from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import (
+    FlowUniPCMultistepScheduler)
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines.wan.wan_i2v_pipeline import WanImageToVideoValidationPipeline
+from fastvideo.v1.training.training_pipeline import TrainingPipeline
+from fastvideo.v1.training.training_utils import (
+    clip_grad_norm_while_handling_failing_dtensor_cases,
+    compute_density_for_timestep_sampling, get_sigmas, load_checkpoint,
+    normalize_dit_input, save_checkpoint)
+# from fastvideo.v1.dataset.parquet_datasets import ParquetVideoTextDataset
+from fastvideo.v1.dataset import build_parquet_map_style_dataloader
+
+import wandb  # isort: skip
+
+logger = init_logger(__name__)
+
+# Manual gradient checking flag - set to True to enable gradient verification
+ENABLE_GRADIENT_CHECK = False
+
+
+class WanI2VTrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for Wan.
+    """
+    _required_config_modules = ["scheduler", "transformer"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
+            shift=fastvideo_args.flow_shift)
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        """
+        May be used in future refactors.
+        """
+        pass
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+
+        args_copy.inference_mode = True
+        args_copy.vae_config.load_encoder = False
+        args_copy.log_validation = False
+        validation_pipeline = WanImageToVideoValidationPipeline(
+            training_args.model_path, 
+            fastvideo_args=args_copy, 
+            # inference_mode=True,
+            loaded_modules={"transformer": self.get_module("transformer")},
+        )
+
+        self.validation_pipeline = validation_pipeline
+        self.latents = None
+
+    def train_one_step(
+        self,
+        transformer,
+        model_type,
+        optimizer,
+        lr_scheduler,
+        loader_iter,
+        noise_scheduler,
+        noise_random_generator,
+        gradient_accumulation_steps,
+        sp_size,
+        precondition_outputs,
+        max_grad_norm,
+        weighting_scheme,
+        logit_mean,
+        logit_std,
+        mode_scale,
+    ) -> tuple[float, float]:
+        assert self.training_args is not None
+        self.modules["transformer"].requires_grad_(True)
+        self.modules["transformer"].train()
+
+        total_loss = 0.0
+        optimizer.zero_grad()
+
+        for _ in range(gradient_accumulation_steps):
+            (
+                self.latents,
+                self.encoder_hidden_states,
+                self.encoder_attention_mask,
+                self.infos,
+                self.extra_latents
+            ) = next(self.train_loader_iter, None)
+            if self.latents is None:
+                self.current_epoch += 1
+                logger.info("Starting epoch %s", self.current_epoch)
+                # Reset iterator for next epoch
+                self.train_loader_iter = iter(self.train_dataloader)
+                # Get first batch of new epoch
+                (
+                    self.latents,
+                    self.encoder_hidden_states,
+                    self.encoder_attention_mask,
+                    self.infos,
+                    self.extra_latents
+                ) = next(self.train_loader_iter)
+
+            latents = self.latents
+            encoder_hidden_states = self.encoder_hidden_states
+            encoder_attention_mask = self.encoder_attention_mask
+            infos = self.infos
+            extra_latents = self.extra_latents
+            
+
+            # logger.info("rank: %s, caption: %s",
+            #             self.rank,
+            #             infos['caption'],
+            #             local_main_process_only=False)
+            # TODO(will): don't hardcode bfloat16
+            latents = latents.to(get_torch_device(),
+                                 dtype=torch.bfloat16)
+            encoder_hidden_states = encoder_hidden_states.to(
+                get_torch_device(), dtype=torch.bfloat16)
+            latents = normalize_dit_input(model_type, latents)
+            batch_size = latents.shape[0]
+            noise = torch.randn_like(latents)
+            u = compute_density_for_timestep_sampling(
+                weighting_scheme=weighting_scheme,
+                batch_size=batch_size,
+                generator=noise_random_generator,
+                logit_mean=logit_mean,
+                logit_std=logit_std,
+                mode_scale=mode_scale,
+            )
+            indices = (u * noise_scheduler.config.num_train_timesteps).long()
+            timesteps = noise_scheduler.timesteps[indices].to(
+                device=latents.device)
+            if sp_size > 1:
+                # Make sure that the timesteps are the same across all sp processes.
+                sp_group = get_sp_group()
+                sp_group.broadcast(timesteps, src=0)
+            sigmas = get_sigmas(
+                noise_scheduler,
+                latents.device,
+                timesteps,
+                n_dim=latents.ndim,
+                dtype=latents.dtype,
+            )
+            noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
+
+            input_kwargs = {}
+
+            # I2V
+            if extra_latents:
+                image_embeds, image_latents = extra_latents["img_embed"], extra_latents["img_lat"]
+                # Image Embeds
+                assert torch.isnan(image_embeds).sum() == 0
+                image_embeds = image_embeds.to(get_torch_device(),
+                                               dtype=torch.bfloat16)
+                input_kwargs["encoder_hidden_states_image"] = image_embeds
+
+                # Image Latents
+                assert torch.isnan(image_latents).sum() == 0
+                image_latents = image_latents.to(get_torch_device(),
+                                                            dtype=torch.bfloat16)
+                noisy_model_input = torch.cat(
+                    [noisy_model_input, image_latents],
+                    dim=1)
+
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                input_kwargs.update({
+                    "hidden_states": noisy_model_input,
+                    "encoder_hidden_states": encoder_hidden_states,
+                    "timestep": timesteps,
+                    "encoder_attention_mask": encoder_attention_mask,  # B, L
+                    "return_dict": False,
+                })
+                if 'hunyuan' in model_type:
+                    input_kwargs["guidance"] = torch.tensor(
+                        [1000.0],
+                        device=noisy_model_input.device,
+                        dtype=torch.bfloat16)
+                with set_forward_context(current_timestep=timesteps,
+                                         attn_metadata=None):
+                    model_pred = transformer(**input_kwargs)
+
+                if precondition_outputs:
+                    model_pred = noisy_model_input - model_pred * sigmas
+                target = latents if precondition_outputs else noise - latents
+
+                loss = (torch.mean((model_pred.float() - target.float())**2) /
+                        gradient_accumulation_steps)
+
+            loss.backward()
+
+            avg_loss = loss.detach().clone()
+            # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}",
+            #             local_main_process_only=False)
+            world_group = get_world_group()
+            world_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
+            total_loss += avg_loss.item()
+
+        # TODO(will): perhaps move this into transformer api so that we can do
+        # the following:
+        # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+        if max_grad_norm is not None:
+            model_parts = [self.transformer]
+            grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
+                [p for m in model_parts for p in m.parameters()],
+                max_grad_norm,
+                foreach=None,
+            )
+            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
+        else:
+            grad_norm = 0.0
+
+        optimizer.step()
+        lr_scheduler.step()
+        return total_loss, grad_norm
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ):
+        assert self.training_args is not None
+
+        # Set random seeds for deterministic training
+        seed = self.training_args.seed if self.training_args.seed is not None else 42
+
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        noise_random_generator = torch.Generator(device="cpu").manual_seed(seed)
+
+        logger.info("Initialized random seeds with seed: %s", seed)
+
+        noise_scheduler = FlowMatchEulerDiscreteScheduler()
+
+        # Train!
+        assert self.training_args.sp_size is not None
+        assert self.training_args.gradient_accumulation_steps is not None
+        total_batch_size = (self.world_size *
+                            self.training_args.gradient_accumulation_steps /
+                            self.training_args.sp_size *
+                            self.training_args.train_sp_batch_size)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %s", len(self.train_dataset))
+        logger.info("  Dataloader size = %s", len(self.train_dataloader))
+        logger.info("  Num Epochs = %s", self.num_train_epochs)
+        logger.info("  Resume training from step %s",
+                    self.init_steps)  # type: ignore
+        logger.info("  Instantaneous batch size per device = %s",
+                    self.training_args.train_batch_size)
+        logger.info(
+            "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
+            total_batch_size)
+        logger.info("  Gradient Accumulation steps = %s",
+                    self.training_args.gradient_accumulation_steps)
+        logger.info("  Total optimization steps = %s",
+                    self.training_args.max_train_steps)
+        logger.info(
+            "  Total training parameters per FSDP shard = %s B",
+            sum(p.numel()
+                for p in self.transformer.parameters() if p.requires_grad) /
+            1e9)
+        # print dtype
+        logger.info("  Master weight dtype: %s",
+                    self.transformer.parameters().__next__().dtype)
+
+        if self.training_args.resume_from_checkpoint:
+            logger.info("Loading checkpoint from %s",
+                        self.training_args.resume_from_checkpoint)
+            resumed_step = load_checkpoint(
+                self.transformer, self.global_rank,
+                self.training_args.resume_from_checkpoint, self.optimizer,
+                self.train_dataloader, self.lr_scheduler,
+                noise_random_generator)
+            if resumed_step > 0:
+                self.init_steps = resumed_step
+                logger.info("Successfully resumed from step %s", resumed_step)
+            else:
+                logger.warning(
+                    "Failed to load checkpoint, starting from step 0")
+                self.init_steps = 0
+
+        progress_bar = tqdm(
+            range(0, self.training_args.max_train_steps),
+            initial=self.init_steps,
+            desc="Steps",
+            # Only show the progress bar once on each machine.
+            disable=self.local_rank > 0,
+        )
+
+        self.train_loader_iter = iter(self.train_dataloader)
+
+        step_times: deque[float] = deque(maxlen=100)
+
+        # TODO(will): fix this
+        # for i in range(self.init_steps):
+        #     next(loader_iter)
+        # get gpu memory usage
+        gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+        logger.info("GPU memory usage before train_one_step: %s MB",
+                    gpu_memory_usage)
+        self._log_validation(self.transformer, self.training_args, 1)
+        for step in range(self.init_steps + 1,
+                          self.training_args.max_train_steps + 1):
+            start_time = time.perf_counter()
+
+            loss, grad_norm = self.train_one_step(
+                self.transformer,
+                # args.model_type,
+                "wan",
+                self.optimizer,
+                self.lr_scheduler,
+                self.train_loader_iter,
+                noise_scheduler,
+                noise_random_generator,
+                self.training_args.gradient_accumulation_steps,
+                self.training_args.sp_size,
+                self.training_args.precondition_outputs,
+                self.training_args.max_grad_norm,
+                self.training_args.weighting_scheme,
+                self.training_args.logit_mean,
+                self.training_args.logit_std,
+                self.training_args.mode_scale,
+            )
+            gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+            logger.info("GPU memory usage after train_one_step: %s MB",
+                        gpu_memory_usage)
+
+            step_time = time.perf_counter() - start_time
+            step_times.append(step_time)
+            avg_step_time = sum(step_times) / len(step_times)
+
+            # Manual gradient checking - only at first step
+            if step == 1 and ENABLE_GRADIENT_CHECK:
+                logger.info("Performing gradient check at step %s", step)
+                self.setup_gradient_check(args, self.train_loader_iter,
+                                          noise_scheduler,
+                                          noise_random_generator)
+
+            progress_bar.set_postfix({
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+            })
+            progress_bar.update(1)
+            if self.global_rank == 0:
+                wandb.log(
+                    {
+                        "train_loss": loss,
+                        "learning_rate": self.lr_scheduler.get_last_lr()[0],
+                        "step_time": step_time,
+                        "avg_step_time": avg_step_time,
+                        "grad_norm": grad_norm,
+                    },
+                    step=step,
+                )
+            if step % self.training_args.checkpointing_steps == 0:
+                save_checkpoint(self.transformer, self.global_rank,
+                                self.training_args.output_dir, step,
+                                self.optimizer, self.train_dataloader,
+                                self.lr_scheduler, noise_random_generator)
+                self.transformer.train()
+                self.sp_group.barrier()
+            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
+                self._log_validation(self.transformer, self.training_args, step)
+
+        save_checkpoint(self.transformer, self.global_rank,
+                        self.training_args.output_dir,
+                        self.training_args.max_train_steps, self.optimizer,
+                        self.train_dataloader, self.lr_scheduler,
+                        noise_random_generator)
+
+        if get_sp_group():
+            cleanup_dist_env_and_memory()
+
+    @torch.no_grad()
+    def _log_validation(self, transformer, training_args, global_step) -> None:
+        assert training_args is not None
+        training_args.inference_mode = True
+        training_args.use_cpu_offload = False
+        if not training_args.log_validation:
+            return
+        if self.validation_pipeline is None:
+            raise ValueError("Validation pipeline is not set")
+
+        logger.info("Starting validation")
+
+        # Create sampling parameters if not provided
+        sampling_param = SamplingParam.from_pretrained(training_args.model_path)
+
+        # Set deterministic seed for validation
+        validation_seed = training_args.seed if training_args.seed is not None else 42
+        torch.manual_seed(validation_seed)
+        torch.cuda.manual_seed_all(validation_seed)
+
+        logger.info("Using validation seed: %s", validation_seed)
+
+        # Prepare validation prompts
+        logger.info('fastvideo_args.validation_path: %s',
+                    training_args.validation_path)
+
+        # validation_dataset, validation_dataloader = build_parquet_map_style_dataloader(
+        #     training_args.validation_path,
+        #     batch_size=1,
+        #     num_data_workers=0,
+        #     drop_last=False,
+        #     cfg_rate=training_args.cfg)
+
+        if sampling_param.negative_prompt:
+            _, negative_prompt_embeds, negative_prompt_attention_mask, _ = validation_dataset.get_validation_negative_prompt(
+            )
+        # validation_dataset = ParquetVideoTextDataset(
+        #     training_args.validation_prompt_dir,
+        #     batch_size=1,
+        #     cfg_rate=training_args.cfg,
+        #     num_latent_t=training_args.num_latent_t,
+        #     validation=True)
+        # if sampling_param.negative_prompt:
+        #     _, negative_prompt_embeds, negative_prompt_attention_mask, _ = validation_dataset.get_validation_negative_prompt(
+        #     )
+
+
+        validation_loader_iter = iter(validation_dataloader)
+        transformer.eval()
+
+        # Process each validation prompt
+        videos = []
+        captions = []
+        for batch in validation_loader_iter:
+            logger.info(batch)
+            (
+                latents,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                infos,
+                extra_latents
+            ) = batch
+            # caption = infos['caption']
+            # captions.extend(caption)
+            prompt_embeds = encoder_hidden_states.to(get_torch_device())
+            prompt_attention_mask = encoder_attention_mask.to(get_torch_device())
+            image_embeds = extra_latents["img_embed"].to(get_torch_device())
+            image_latent = extra_latents["img_lat"].to(get_torch_device())
+
+            # Calculate sizes
+            latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
+                            sampling_param.height // 8,
+                            sampling_param.width // 8]
+            n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+
+            temporal_compression_factor = training_args.vae_config.arch_config.temporal_compression_ratio
+            num_frames = (training_args.num_latent_t -
+                            1) * temporal_compression_factor + 1
+
+            # Prepare batch for validation
+            batch = ForwardBatch(
+                data_type="video",
+                latents=None,
+                seed=validation_seed,  # Use deterministic seed
+                generator=torch.Generator(
+                    device="cpu").manual_seed(validation_seed),
+                prompt_embeds=[prompt_embeds],
+                prompt_attention_mask=[prompt_attention_mask],
+                negative_prompt_embeds=[negative_prompt_embeds],
+                negative_attention_mask=[negative_prompt_attention_mask],
+                image_embeds=[image_embeds],
+                # image_latent=image_latent,
+                # make sure we use the same height, width, and num_frames as the training pipeline
+                height=training_args.num_height,
+                width=training_args.num_width,
+                num_frames=num_frames,
+                # TODO(will): validation_sampling_steps and
+                # validation_guidance_scale are actually passed in as a list of
+                # values, like "10,20,30". The validation should be run for each
+                # combination of values.
+                # num_inference_steps=fastvideo_args.validation_sampling_steps,
+                num_inference_steps=sampling_param.num_inference_steps,
+                # guidance_scale=fastvideo_args.validation_guidance_scale,
+                guidance_scale=sampling_param.guidance_scale,
+                n_tokens=n_tokens,
+                eta=0.0,
+            )
+
+            # Run validation inference
+            with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+                output_batch = self.validation_pipeline.forward(
+                    batch, training_args)
+                samples = output_batch.output
+
+            # Re-enable gradients for training
+            transformer.requires_grad_(True)
+            transformer.train()
+
+            # Process outputs
+            video = rearrange(samples, "b c t h w -> t b c h w")
+            frames = []
+            for x in video:
+                x = torchvision.utils.make_grid(x, nrow=6)
+                x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                frames.append((x * 255).numpy().astype(np.uint8))
+            videos.append(frames)
+
+            # Log validation results
+            world_group = get_world_group()
+            num_sp_groups = world_group.world_size // self.sp_group.world_size
+
+            # Only sp_group leaders (rank_in_sp_group == 0) need to send their
+            # results to global rank 0
+            if self.rank_in_sp_group == 0:
+                if self.global_rank == 0:
+                    # Global rank 0 collects results from all sp_group leaders
+                    all_videos = videos  # Start with own results
+                    all_captions = captions
+
+                    # Receive from other sp_group leaders
+                    for sp_group_idx in range(1, num_sp_groups):
+                        src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
+                        recv_videos = world_group.recv_object(src=src_rank)
+                        recv_captions = world_group.recv_object(src=src_rank)
+                        all_videos.extend(recv_videos)
+                        all_captions.extend(recv_captions)
+
+                    video_filenames = []
+                    for i, (video,
+                            caption) in enumerate(zip(all_videos, all_captions)):
+                        os.makedirs(training_args.output_dir, exist_ok=True)
+                        filename = os.path.join(
+                            training_args.output_dir,
+                            f"validation_step_{global_step}_video_{i}.mp4")
+                        imageio.mimsave(filename, video, fps=sampling_param.fps)
+                        video_filenames.append(filename)
+
+                    logs = {
+                        "validation_videos": [
+                            wandb.Video(filename, caption=caption) for filename,
+                            caption in zip(video_filenames, all_captions)
+                        ]
+                    }
+                    wandb.log(logs, step=global_step)
+                else:
+                    # Other sp_group leaders send their results to global rank 0
+                    world_group.send_object(videos, dst=0)
+                    world_group.send_object(captions, dst=0)
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+def main(args) -> None:
+    logger.info("Starting training pipeline...")
+
+    pipeline = WanI2VTrainingPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.forward(None, args)
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.v1.fastvideo_args import TrainingArgs
+    from fastvideo.v1.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.use_cpu_offload = False
+    main(args)

scripts/finetune/finetune_v1.sh

@@ -21,8 +21,8 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --num_latent_t 20 \
     --sp_size 4 \
     --tp_size 4 \
-    --dp_size  1 \
-    --dp_shards 4 \
+    --hsdp_replicate_dim 1 \
+    --hsdp_shards 4 \
     --num_gpus $NUM_GPUS \
     --train_sp_batch_size 1\
     --dataloader_num_workers 10\

scripts/finetune/finetune_v1_i2v.sh

@@ -0,0 +1,54 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export HOME="/mnt/weka/home/hao.zhang/wei"
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+DATA_DIR=$HOME/FastVideo/data/wei-i2v-dataset/crush-smol_preprocessed/combined_parquet_dataset
+VALIDATION_DIR=$HOME/FastVideo/data/wei-i2v-dataset/crush-smol_preprocessed/validation_parquet_dataset
+NUM_GPUS=4
+# IP=[MASTER NODE IP]
+
+CHECKPOINT_PATH="$DATA_DIR/outputs/wan_i2v_finetune/checkpoint-50"
+
+# 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_i2v_training_pipeline.py\
+    --model_path Wan-AI/Wan2.1-I2V-14B-480P-Diffusers \
+    --inference_mode False\
+    --pretrained_model_name_or_path Wan-AI/Wan2.1-I2V-14B-480P-Diffusers \
+    --cache_dir "$HOME/ray/.cache"\
+    --data_path "$DATA_DIR"\
+    --validation_prompt_dir "$VALIDATION_DIR"\
+    --train_batch_size=1\
+    --num_latent_t 8 \
+    --sp_size $NUM_GPUS \
+    --tp_size $NUM_GPUS \
+    --dp_shards $NUM_GPUS \
+    --train_sp_batch_size 1\
+    --dataloader_num_workers 1\
+    --gradient_accumulation_steps=8 \
+    --max_train_steps=5000 \
+    --learning_rate=1e-6\
+    --mixed_precision="bf16"\
+    --checkpointing_steps=500000 \
+    --validation_steps 20\
+    --validation_sampling_steps "2,4,8" \
+    --log_validation \
+    --checkpoints_total_limit 3\
+    --allow_tf32\
+    --ema_start_step 0\
+    --cfg 0.0\
+    --output_dir="$DATA_DIR/outputs/wan_i2v_finetune"\
+    --tracker_project_name wan_i2v_finetune \
+    --num_height 480 \
+    --num_width 832 \
+    --num_frames  77 \
+    --shift 3 \
+    --validation_guidance_scale "1.0" \
+    --num_euler_timesteps 50 \
+    --multi_phased_distill_schedule "4000-1" \
+    --weight_decay 0.01 \
+    --not_apply_cfg_solver \
+    --master_weight_type "fp32" \
+    --max_grad_norm 1.0 \
+    # --resume_from_checkpoint "$CHECKPOINT_PATH"

scripts/preprocess/v1_preprocess_wan_data_i2v.sh

@@ -1,16 +1,18 @@
-# export WANDB_MODE="offline"
+export WANDB_MODE="offline"
+export HOME="/mnt/user_storage/src/"
 GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 MODEL_TYPE="wan"
-DATA_MERGE_PATH="finetrainers/crush-smol/merge.txt"
-OUTPUT_DIR="crush-smol_preprocess"
-VALIDATION_PATH="assets/prompt.txt"
+DATA_MERGE_PATH="$HOME/FastVideo/data/crush-smol/merge.txt"
+OUTPUT_DIR="$HOME/FastVideo/data/crush-smol_parq_i2v"
+# VALIDATION_PATH="assets/prompt.txt"
+VALIDATION_PATH="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
 
 torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/data_preprocess/preprocess.py \
+    fastvideo/data_preprocess/v1_preprocess.py \
     --model_path $MODEL_PATH \
     --data_merge_path $DATA_MERGE_PATH \
-    --preprocess_video_batch_size 8 \
+    --preprocess_video_batch_size 4 \
     --max_height 480 \
     --max_width 832 \
     --num_frames 77 \
@@ -18,7 +20,7 @@ torchrun --nproc_per_node=$GPU_NUM \
     --output_dir=$OUTPUT_DIR \
     --model_type $MODEL_TYPE \
     --train_fps 16 \
-    --validation_prompt_txt $VALIDATION_PATH \
-    --samples_per_file 16 \
+    --validation_dataset_file $VALIDATION_PATH \
+    --samples_per_file 32 \
     --flush_frequency 32 \
     --preprocess_task "i2v"