wan2.2 training doesn't crash

Co-authored-by: RandNMR73 <notomatthew31@gmail.com>

.buildkite/pipeline.yml

@@ -176,3 +176,26 @@ steps:
                   - TEST_TYPE=precision_vsa
                 agents:
                   queue: "default"
+            - path:
+                - "csrc/attn/vmoba_attn/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests VMoBA"
+                env:
+                  - TEST_TYPE=precision_vmoba
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/vmoba_attn/vmoba/**"
+                - "fastvideo/attention/backends/vmoba.py"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Inference Tests VMoBA"
+                env:  
+                  - TEST_TYPE=inference_vmoba
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -109,6 +109,15 @@ case "$TEST_TYPE" in
         log "Running distillation DMD tests..."
         MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_distill_dmd_tests"
         ;;
+        # run_inference_tests_vmoba
+    "inference_vmoba")
+        log "Running V-MoBA inference tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_vmoba"
+        ;;
+    "precision_vmoba")
+        log "Running V-MoBA precision tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_vmoba"
+        ;;
     *)
         log "Error: Unknown test type: $TEST_TYPE"
         exit 1

.github/workflows/pr-test.yml

@@ -235,7 +235,7 @@ jobs:
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-  
+
   training-test:
     needs: change-filter
     if: >-
@@ -373,4 +373,4 @@ jobs:
           JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
-        run: python .github/scripts/runpod_cleanup.py
+        run: python .github/scripts/runpod_cleanup.py

.gitignore

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

README.md

@@ -7,7 +7,7 @@
 FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
 
 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/rG0QpZdw" target="_blank"> <b> WeChat </b> </a> |
+    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/S7HLCSTh" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">

csrc/attn/vmoba_attn/README.md

@@ -0,0 +1,32 @@
+# Attention Kernel Used in FastVideo
+
+## VMoBA: Mixture-of-Block Attention for Video Diffusion Models (VMoBA)
+
+### Installation
+Please ensure that you have installed FlashAttention version **2.7.1 or higher**, as some interfaces have changed in recent releases.
+
+### Usage
+
+You can use `moba_attn_varlen` in the following ways:
+
+**Install from source:**
+```bash
+python setup.py install
+```
+
+**Import after installation:**
+```python
+from vmoba import moba_attn_varlen
+```
+
+**Or import directly from the project root:**
+```python
+from csrc.attn.vmoba_attn.vmoba import moba_attn_varlen
+```
+
+### Verify if you have successfully installed
+
+```bash
+python csrc/attn/vmoba_attn/vmoba/vmoba.py
+```
+

csrc/attn/vmoba_attn/setup.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from setuptools import find_packages, setup
+
+PACKAGE_NAME = "vmoba"
+VERSION = "0.0.0"
+AUTHOR = "JianzongWu"
+DESCRIPTION = "VMoBA: Mixture-of-Block Attention for Video Diffusion Models"
+URL = "https://github.com/KwaiVGI/VMoBA"
+
+setup(
+    name=PACKAGE_NAME,
+    version=VERSION,
+    author=AUTHOR,
+    description=DESCRIPTION,
+    url=URL,
+    packages=find_packages(),
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+    ],
+    python_requires='>=3.12',
+    install_requires=[]
+)

csrc/attn/vmoba_attn/tests/test_vmoba_attn.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import pytest
+import random
+from csrc.attn.vmoba_attn.vmoba  import moba_attn_varlen
+
+def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"):
+    """
+    Generates random data for testing the variable-length attention function.
+    """
+    torch.manual_seed(42)
+    random.seed(42)
+    torch.cuda.manual_seed_all(42)
+
+    # Generate sequence lengths for each item in the batch
+    if batch_size > 1:
+        # Ensure sequence lengths are reasonably distributed
+        avg_seqlen = total_seqlen // batch_size
+        seqlens = [random.randint(avg_seqlen // 2, avg_seqlen + avg_seqlen // 2) for _ in range(batch_size - 1)]
+        remaining_len = total_seqlen - sum(seqlens)
+        if remaining_len > 0:
+            seqlens.append(remaining_len)
+        else: # Adjust if sum exceeds total_seqlen
+            seqlens.append(avg_seqlen)
+            current_sum = sum(seqlens)
+            seqlens[-1] -= (current_sum - total_seqlen)
+        # Ensure all lengths are positive
+        seqlens = [max(1, s) for s in seqlens]
+        # Final adjustment to match total_seqlen
+        seqlens[-1] += total_seqlen - sum(seqlens)
+
+    else:
+        seqlens = [total_seqlen]
+
+    cu_seqlens = torch.tensor([0] + list(torch.cumsum(torch.tensor(seqlens), 0)), device=device, dtype=torch.int32)
+    max_seqlen = max(seqlens) if seqlens else 0
+
+    q = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+    k = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+    v = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+
+    return q, k, v, cu_seqlens, max_seqlen
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("total_seqlen", [512, 1024])
+@pytest.mark.parametrize("num_heads", [8])
+@pytest.mark.parametrize("head_dim", [64])
+@pytest.mark.parametrize("moba_chunk_size", [64])
+@pytest.mark.parametrize("moba_topk", [2, 4])
+@pytest.mark.parametrize("select_mode", ["topk", "threshold"])
+@pytest.mark.parametrize("threshold_type", ["query_head", "head_global", "overall"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_moba_attn_varlen_forward(
+    batch_size, total_seqlen, num_heads, head_dim, moba_chunk_size, moba_topk, select_mode, threshold_type, dtype
+):
+    """
+    Tests the forward pass of moba_attn_varlen for basic correctness.
+    It checks output shape, dtype, and for the presence of NaNs/Infs.
+    """
+    if dtype == torch.float32:
+        pytest.skip("float32 is not supported in flash attention")
+
+    q, k, v, cu_seqlens, max_seqlen = generate_test_data(
+        batch_size, total_seqlen, num_heads, head_dim, dtype
+    )
+
+    # Ensure chunk size is not larger than the smallest sequence length
+    min_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).min().item()
+    if moba_chunk_size > min_seqlen:
+        pytest.skip("moba_chunk_size is larger than the minimum sequence length in the batch")
+
+    try:
+        output = moba_attn_varlen(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            moba_chunk_size=moba_chunk_size,
+            moba_topk=moba_topk,
+            select_mode=select_mode,
+            threshold_type=threshold_type,
+            simsum_threshold=0.5, # A reasonable default for threshold mode
+        )
+    except Exception as e:
+        pytest.fail(f"moba_attn_varlen forward pass failed with exception: {e}")
+
+    # 1. Check output shape
+    assert output.shape == q.shape, f"Expected output shape {q.shape}, but got {output.shape}"
+
+    # 2. Check output dtype
+    assert output.dtype == q.dtype, f"Expected output dtype {q.dtype}, but got {output.dtype}"
+
+    # 3. Check for NaNs or Infs in the output
+    assert torch.all(torch.isfinite(output)), "Output contains NaN or Inf values"

csrc/attn/vmoba_attn/vmoba/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+from .vmoba import moba_attn_varlen, process_moba_input, process_moba_output

csrc/attn/vmoba_attn/vmoba/vmoba.py

@@ -0,0 +1,860 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapt from https://github.com/KwaiVGI/VMoBA/blob/main/src/vmoba.py
+
+import random
+import time
+import os
+import torch
+from typing import Tuple
+from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
+from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+from functools import lru_cache
+from einops import rearrange
+
+
+@lru_cache(maxsize=16)
+def calc_chunks(cu_seqlen, moba_chunk_size):
+    """
+    Calculate chunk boundaries.
+
+    For vision tasks we include all chunks (even the last one which might be shorter)
+    so that every chunk can be selected.
+    """
+    batch_sizes = cu_seqlen[1:] - cu_seqlen[:-1]
+    batch_num_chunk = (batch_sizes + (moba_chunk_size - 1)) // moba_chunk_size
+    cu_num_chunk = torch.ones(
+        batch_num_chunk.numel() + 1,
+        device=cu_seqlen.device,
+        dtype=batch_num_chunk.dtype,
+    )
+    cu_num_chunk[1:] = batch_num_chunk.cumsum(dim=0)
+    num_chunk = cu_num_chunk[-1]
+    chunk_sizes = torch.full(
+        (num_chunk + 1,), moba_chunk_size, dtype=torch.int32, device=cu_seqlen.device
+    )
+    chunk_sizes[0] = 0
+    batch_last_chunk_size = batch_sizes - (batch_num_chunk - 1) * moba_chunk_size
+    chunk_sizes[cu_num_chunk[1:]] = batch_last_chunk_size
+    cu_chunk = chunk_sizes.cumsum(dim=-1, dtype=torch.int32)
+    chunk_to_batch = torch.zeros(
+        (num_chunk,), dtype=torch.int32, device=cu_seqlen.device
+    )
+    chunk_to_batch[cu_num_chunk[1:-1]] = 1
+    chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32)
+
+    # Do not filter out any chunk
+    filtered_chunk_indices = torch.arange(
+        num_chunk, device=cu_seqlen.device, dtype=torch.int32
+    )
+    num_filtered_chunk = num_chunk
+
+    return cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch
+
+
+# --- Threshold Selection Helper Functions ---
+
+def _select_threshold_query_head(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects chunks for each <query, head> pair based on threshold.
+    Normalization and sorting happen along the chunk dimension (dim=0).
+    """
+    C, H, S = gate.shape
+    eps = 1e-6
+
+    # LSE‐style normalization per <head, query> (across chunks)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    row_min = gate_min_val.amin(dim=0)        # (H, S)
+    row_max = gate_masked.amax(dim=0)         # (H, S)
+    denom = row_max - row_min
+    denom = torch.where(denom <= eps, torch.ones_like(denom), denom) # avoid divide‑by‑zero
+
+    gate_norm = (gate - row_min.unsqueeze(0)) / denom.unsqueeze(0)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) pull out the self‐chunk’s normalized weight for each <head,seq>
+    self_norm = (gate_norm * gate_self_chunk_mask).sum(dim=0)  # (H, S)
+
+    # 2) compute how much more normalized weight we need beyond self
+    total_norm_sum = gate_norm.sum(dim=0) # (H, S)
+    remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps)  # (H, S)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # if already ≥ thresh, no extra needed
+
+    # 3) zero out the self‐chunk in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0
+
+    # 4) sort the other chunks by descending norm, per <head,seq>
+    sorted_norm, sorted_idx = torch.sort(others_norm, descending=True, dim=0) # (C, H, S)
+
+    # 5) cumulative‑sum the sorted norms per <head,seq>
+    cumsum_others = sorted_norm.cumsum(dim=0) # (C, H, S)
+
+    # 6) for each <head,seq>, find the smallest k where cumsum_ratio ≥ remain_ratio
+    ratio = cumsum_others / (total_norm_sum.unsqueeze(0) + eps) # (C, H, S)
+    cond = ratio >= remain_ratio.unsqueeze(0) # (C, H, S) boolean mask
+    any_cond = cond.any(dim=0) # (H, S)
+    # Find the index of the first True value along dim 0. If none, use C-1.
+    cutoff = torch.where(any_cond, cond.float().argmax(dim=0), torch.full_like(any_cond, fill_value=C - 1)) # (H, S)
+
+    # 7) build a mask in sorted order up to that cutoff
+    idx_range = torch.arange(C, device=gate.device).view(-1, 1, 1) # (C, 1, 1)
+    sorted_mask = idx_range <= cutoff.unsqueeze(0) # (C, H, S)
+
+    # 8) scatter it back to original chunk order
+    others_mask = torch.zeros_like(gate, dtype=torch.bool)
+    others_mask.scatter_(0, sorted_idx, sorted_mask)
+
+    # 9) finally, include every self‐chunk plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_block(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <query, head> pairs for each block based on threshold.
+    Normalization and sorting happen across the head and sequence dimensions (dim=1, 2).
+    """
+    C, H, S = gate.shape
+    HS = H * S
+    eps = 1e-6
+
+    # LSE‐style normalization per block (across heads and queries)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1)
+    block_min = gate_min_val.amin(dim=(1, 2), keepdim=True) # (C, 1, 1)
+    block_denom = block_max - block_min
+    block_denom = torch.where(block_denom <= eps, torch.ones_like(block_denom), block_denom) # (C, 1, 1)
+
+    gate_norm = (gate - block_min) / block_denom # (C, H, S)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) identify normalized weights of entries that *are* self-chunks (from query perspective)
+    self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S)
+    # Sum these weights *per block*
+    self_norm_sum_per_block = self_norm_entries.sum(dim=(1, 2)) # (C,)
+
+    # 2) compute how much more normalized weight each block needs beyond its self-chunk contributions
+    total_norm_sum_per_block = gate_norm.sum(dim=(1, 2)) # (C,)
+    remain_ratio = simsum_threshold - self_norm_sum_per_block / (total_norm_sum_per_block + eps) # (C,)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # (C,)
+
+    # 3) zero out the self‐chunk entries in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries
+
+    # 4) sort the other <head, seq> pairs by descending norm, per block
+    others_flat = others_norm.contiguous().view(C, HS) # (C, H*S)
+    sorted_others_flat, sorted_indices_flat = torch.sort(others_flat, dim=1, descending=True) # (C, H*S)
+
+    # 5) cumulative‑sum the sorted norms per block
+    cumsum_others_flat = sorted_others_flat.cumsum(dim=1) # (C, H*S)
+
+    # 6) for each block, find the smallest k where cumsum_ratio ≥ remain_ratio
+    ratio_flat = cumsum_others_flat / (total_norm_sum_per_block.unsqueeze(1) + eps) # (C, H*S)
+    cond_flat = ratio_flat >= remain_ratio.unsqueeze(1) # (C, H*S) boolean mask
+    any_cond = cond_flat.any(dim=1) # (C,)
+    # Find the index of the first True value along dim 1. If none, use HS-1.
+    cutoff_flat = torch.where(any_cond, cond_flat.float().argmax(dim=1), torch.full_like(any_cond, fill_value=HS - 1)) # (C,)
+
+    # 7) build a mask in sorted order up to that cutoff per block
+    idx_range_flat = torch.arange(HS, device=gate.device).unsqueeze(0) # (1, H*S)
+    sorted_mask_flat = idx_range_flat <= cutoff_flat.unsqueeze(1) # (C, H*S)
+
+    # 8) scatter it back to original <head, seq> order per block
+    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C, H*S)
+    others_mask_flat.scatter_(1, sorted_indices_flat, sorted_mask_flat)
+    others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
+
+    # 9) finally, include every self‐chunk entry plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_overall(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <chunk, query, head> triplets globally based on threshold.
+    Normalization and sorting happen across all valid entries.
+    """
+    C, H, S = gate.shape
+    CHS = C * H * S
+    eps = 1e-6
+
+    # LSE‐style normalization globally across all valid entries
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    overall_max = gate_masked.max() # scalar
+    overall_min = gate_min_val.min() # scalar
+    overall_denom = overall_max - overall_min
+    overall_denom = torch.where(overall_denom <= eps, torch.tensor(1.0, device=gate.device, dtype=gate.dtype), overall_denom)
+
+    gate_norm = (gate - overall_min) / overall_denom # (C, H, S)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) identify normalized weights of entries that *are* self-chunks
+    self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S)
+    # Sum these weights globally
+    self_norm_sum_overall = self_norm_entries.sum() # scalar
+
+    # 2) compute how much more normalized weight is needed globally beyond self-chunk contributions
+    total_norm_sum_overall = gate_norm.sum() # scalar
+    remain_ratio = simsum_threshold - self_norm_sum_overall / (total_norm_sum_overall + eps) # scalar
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # scalar
+
+    # 3) zero out the self‐chunk entries in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries
+
+    # 4) sort all other entries by descending norm, globally
+    others_flat = others_norm.flatten() # (C*H*S,)
+    valid_others_mask_flat = valid_gate_mask.flatten() & ~gate_self_chunk_mask.flatten() # Mask for valid, non-self entries
+    
+    # Only sort the valid 'other' entries
+    valid_others_indices = torch.where(valid_others_mask_flat)[0]
+    valid_others_values = others_flat[valid_others_indices]
+
+    sorted_others_values, sort_perm = torch.sort(valid_others_values, descending=True) # (N_valid_others,)
+    sorted_original_indices = valid_others_indices[sort_perm] # Original indices in C*H*S space, sorted by value
+
+    # 5) cumulative‑sum the sorted valid 'other' norms globally
+    cumsum_others_values = sorted_others_values.cumsum(dim=0) # (N_valid_others,)
+
+    # 6) find the smallest k where cumsum_ratio ≥ remain_ratio globally
+    ratio_values = cumsum_others_values / (total_norm_sum_overall + eps) # (N_valid_others,)
+    cond_values = ratio_values >= remain_ratio # (N_valid_others,) boolean mask
+    any_cond = cond_values.any() # scalar
+    
+    # Find the index of the first True value in the *sorted* list. If none, use all valid others.
+    cutoff_idx_in_sorted = torch.where(
+        any_cond, 
+        cond_values.float().argmax(dim=0), 
+        torch.tensor(len(sorted_others_values) - 1, device=gate.device, dtype=torch.long)
+    )
+
+    # 7) build a mask selecting the top-k others based on the cutoff
+    # Select the original indices corresponding to the top entries in the sorted list
+    selected_other_indices = sorted_original_indices[:cutoff_idx_in_sorted + 1]
+
+    # 8) create the mask in the original flat shape
+    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,)
+    if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected
+        others_mask_flat[selected_other_indices] = True
+    others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
+
+    # 9) finally, include every self‐chunk entry plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_head_global(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <chunk, query> globally for each head based on threshold.
+    """
+    C, H, S = gate.shape
+    eps = 1e-6
+
+    # 1) LSE‐style normalization per head (across chunks and sequence dims)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)
+
+    max_per_head = gate_masked.amax(dim=(0, 2), keepdim=True)  # (1, H, 1)
+    min_per_head = gate_min_val.amin(dim=(0, 2), keepdim=True)  # (1, H, 1)
+    denom = max_per_head - min_per_head
+    denom = torch.where(denom <= eps, torch.ones_like(denom), denom)
+
+    gate_norm = (gate - min_per_head) / denom
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0)  # (C, H, S)
+
+    # 2) sum normalized self‐chunk contributions per head
+    self_norm_sum = (gate_norm * gate_self_chunk_mask).sum(dim=(0, 2))  # (H,)
+
+    # 3) total normalized sum per head
+    total_norm_sum = gate_norm.sum(dim=(0, 2))  # (H,)
+
+    # 4) how much more normalized weight needed per head
+    remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps)  # (H,)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0)
+
+    # 5) zero out self‐chunk entries to focus on "others"
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0  # (C, H, S)
+
+    # 6) flatten chunk and sequence dims, per head
+    CS = C * S
+    others_flat = others_norm.permute(1, 0, 2).reshape(H, CS)  # (H, C*S)
+    valid_flat  = (valid_gate_mask & ~gate_self_chunk_mask) \
+                   .permute(1, 0, 2).reshape(H, CS)           # (H, C*S)
+
+    # 7) vectorized selection of “others” per head
+    masked_flat = torch.where(valid_flat, others_flat, torch.zeros_like(others_flat))
+    sorted_vals, sorted_idx = torch.sort(masked_flat, dim=1, descending=True)  # (H, C*S)
+
+    cumsum_vals = sorted_vals.cumsum(dim=1)                                    # (H, C*S)
+    ratio_vals  = cumsum_vals / (total_norm_sum.unsqueeze(1) + eps)            # (H, C*S)
+    cond        = ratio_vals >= remain_ratio.unsqueeze(1)                      # (H, C*S)
+
+    has_cutoff = cond.any(dim=1)                                               # (H,)
+    default    = torch.full((H,), CS - 1, device=gate.device, dtype=torch.long)
+    cutoff     = torch.where(has_cutoff, cond.float().argmax(dim=1), default)  # (H,)
+
+    idx_range   = torch.arange(CS, device=gate.device).unsqueeze(0)            # (1, C*S)
+    sorted_mask = idx_range <= cutoff.unsqueeze(1)                             # (H, C*S)
+
+    selected_flat = torch.zeros_like(valid_flat)                               # (H, C*S)
+    selected_flat.scatter_(1, sorted_idx, sorted_mask)                         # (H, C*S)
+
+    # 8) reshape selection mask back to (C, H, S)
+    others_mask = selected_flat.reshape(H, C, S).permute(1, 0, 2)               # (C, H, S)
+
+    # 9) include self‐chunks plus selected others, and obey valid mask
+    final_gate_mask = valid_gate_mask & (gate_self_chunk_mask | others_mask)
+
+    return final_gate_mask
+
+
+class MixedAttention(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        self_attn_cu_seqlen,
+        moba_q,
+        moba_kv,
+        moba_cu_seqlen_q,
+        moba_cu_seqlen_kv,
+        max_seqlen,
+        moba_chunk_size,
+        moba_q_sh_indices,
+    ):
+        ctx.max_seqlen = max_seqlen
+        ctx.moba_chunk_size = moba_chunk_size
+        ctx.softmax_scale = softmax_scale = q.shape[-1] ** (-0.5)
+
+        # Non-causal self-attention branch
+        # return out, softmax_lse, S_dmask, rng_state
+        self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=self_attn_cu_seqlen,
+            cu_seqlens_k=self_attn_cu_seqlen,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+        )
+        # MOBA attention branch (non-causal)
+        moba_attn_out, moba_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
+            q=moba_q,
+            k=moba_kv[:, 0],
+            v=moba_kv[:, 1],
+            cu_seqlens_q=moba_cu_seqlen_q,
+            cu_seqlens_k=moba_cu_seqlen_kv,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=moba_chunk_size,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+        )
+
+        self_attn_lse_sh = self_attn_lse_hs.t().contiguous()
+        moba_attn_lse = moba_attn_lse_hs.t().contiguous()
+
+        output = torch.zeros((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+        output_2d = output.view(-1, q.shape[2])
+
+        max_lse_1d = self_attn_lse_sh.view(-1)
+        max_lse_1d = max_lse_1d.index_reduce(
+            0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax"
+        )
+        self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh)
+        moba_attn_lse = (
+            moba_attn_lse.view(-1)
+            .sub(max_lse_1d.index_select(0, moba_q_sh_indices))
+            .reshape_as(moba_attn_lse)
+        )
+
+        mixed_attn_se_sh = self_attn_lse_sh.exp()
+        moba_attn_se = moba_attn_lse.exp()
+
+        mixed_attn_se_sh.view(-1).index_add_(
+            0, moba_q_sh_indices, moba_attn_se.view(-1)
+        )
+        mixed_attn_lse_sh = mixed_attn_se_sh.log()
+
+        # Combine self-attention output
+        factor = (self_attn_lse_sh - mixed_attn_lse_sh).exp()  # [S, H]
+        self_attn_out_sh = self_attn_out_sh * factor.unsqueeze(-1)
+        output_2d += self_attn_out_sh.reshape_as(output_2d)
+
+        # Combine MOBA attention output
+        mixed_attn_lse = (
+            mixed_attn_lse_sh.view(-1)
+            .index_select(0, moba_q_sh_indices)
+            .view_as(moba_attn_lse)
+        )
+        factor = (moba_attn_lse - mixed_attn_lse).exp()  # [S, H]
+        moba_attn_out = moba_attn_out * factor.unsqueeze(-1)
+        raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1])
+        output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out)
+        output = output.to(q.dtype)
+        mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh)
+        ctx.save_for_backward(
+            output,
+            mixed_attn_lse_sh,
+            q,
+            k,
+            v,
+            self_attn_cu_seqlen,
+            moba_q,
+            moba_kv,
+            moba_cu_seqlen_q,
+            moba_cu_seqlen_kv,
+            moba_q_sh_indices,
+        )
+
+        return output
+
+    @staticmethod
+    def backward(ctx, d_output):
+
+        max_seqlen = ctx.max_seqlen
+        moba_chunk_size = ctx.moba_chunk_size
+        softmax_scale = ctx.softmax_scale
+
+        (
+            output,
+            mixed_attn_vlse_sh,
+            q,
+            k,
+            v,
+            self_attn_cu_seqlen,
+            moba_q,
+            moba_kv,
+            moba_cu_seqlen_q,
+            moba_cu_seqlen_kv,
+            moba_q_sh_indices,
+        ) = ctx.saved_tensors
+
+        d_output = d_output.contiguous()
+
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        _ = _flash_attn_varlen_backward(
+            dout=d_output,
+            q=q,
+            k=k,
+            v=v,
+            out=output,
+            softmax_lse=mixed_attn_vlse_sh.t().contiguous(),
+            dq=dq,
+            dk=dk,
+            dv=dv,
+            cu_seqlens_q=self_attn_cu_seqlen,
+            cu_seqlens_k=self_attn_cu_seqlen,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=True,
+            window_size_left=-1,
+            window_size_right=-1
+        )
+
+        headdim = q.shape[-1]
+        d_moba_output = (
+            d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)
+        )
+        moba_output = (
+            output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)
+        )
+
+        mixed_attn_vlse = (
+            mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1)
+        )
+
+        dmq = torch.empty_like(moba_q)
+        dmkv = torch.empty_like(moba_kv)
+        _ = _flash_attn_varlen_backward(
+            dout=d_moba_output,
+            q=moba_q,
+            k=moba_kv[:, 0],
+            v=moba_kv[:, 1],
+            out=moba_output,
+            softmax_lse=mixed_attn_vlse,
+            dq=dmq,
+            dk=dmkv[:,0],
+            dv=dmkv[:,1],
+            cu_seqlens_q=moba_cu_seqlen_q,
+            cu_seqlens_k=moba_cu_seqlen_kv,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=moba_chunk_size,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=True,
+            window_size_left=-1,
+            window_size_right=-1
+        )
+
+        return dq, dk, dv, None, dmq, dmkv, None, None, None, None, None
+
+
+def moba_attn_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    max_seqlen: int,
+    moba_chunk_size: int,
+    moba_topk: int,
+    select_mode: str = 'threshold', # "topk" or "threshold"
+    simsum_threshold: float = 0.25,
+    threshold_type: str = 'query_head',
+) -> torch.Tensor:
+    """
+    Accelerated MOBA attention for vision tasks with proper LSE normalization.
+
+    This version:
+      - Splits KV into chunks.
+      - For each query head, selects the top-k relevant KV chunks (including the self chunk)
+        by amplifying the diagonal (self-chunk) logits.
+      - Aggregates the attention outputs from the selected chunks using a log-sum-exp
+        reduction so that attending to each query over the selected chunks is equivalent
+        to the original algorithm.
+    """
+    # Stack keys and values.
+    kv = torch.stack((k, v), dim=1)
+    seqlen, num_head, head_dim = q.shape
+
+    # Compute chunk boundaries.
+    cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(
+        cu_seqlens, moba_chunk_size
+    )
+
+    self_attn_cu_seqlen = cu_chunk
+
+    # Update top-k selection to include the self chunk.
+    moba_topk = min(moba_topk, num_filtered_chunk)
+
+    # --- Build filtered KV from chunks ---
+    chunk_starts = cu_chunk[filtered_chunk_indices]          # [num_filtered_chunk]
+    chunk_ends = cu_chunk[filtered_chunk_indices + 1]          # [num_filtered_chunk]
+    chunk_lengths = chunk_ends - chunk_starts                  # [num_filtered_chunk]
+    max_chunk_len = int(chunk_lengths.max().item())
+
+    range_tensor = torch.arange(max_chunk_len, device=kv.device, dtype=chunk_starts.dtype).unsqueeze(0)
+    indices = chunk_starts.unsqueeze(1) + range_tensor
+    indices = torch.clamp(indices, max=kv.shape[0] - 1)
+    valid_mask = range_tensor < chunk_lengths.unsqueeze(1)
+    gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:])
+    gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered)
+
+    # Compute key_gate_weight over valid tokens.
+    key_values = gathered[:, :, 0].float()  # [num_filtered_chunk, max_chunk_len, num_head, head_dim]
+    valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1)
+    key_sum = (key_values * valid_mask_exp).sum(dim=1)
+    divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1)
+    key_gate_weight = key_sum / divisor  # [num_filtered_chunk, num_head, head_dim]
+
+    # Compute gate logits between key_gate_weight and queries.
+    q_float = q.float()
+    # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float)  # [num_filtered_chunk, num_head, seqlen]
+    gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2)
+
+    # Amplify the diagonal (self chunk) contributions.
+    gate_seq_idx = torch.arange(seqlen, device=q.device, dtype=torch.int32).unsqueeze(0).expand(num_filtered_chunk, seqlen)
+    chunk_start = cu_chunk[filtered_chunk_indices]  # [num_filtered_chunk]
+    chunk_end = cu_chunk[filtered_chunk_indices + 1]  # [num_filtered_chunk]
+    gate_self_chunk_mask = ((gate_seq_idx >= chunk_start.unsqueeze(1)) &
+                            (gate_seq_idx < chunk_end.unsqueeze(1))).unsqueeze(1).expand(-1, num_head, -1)
+    amplification_factor = 1e9  # Example factor; adjust as needed.
+    origin_gate = gate.clone()
+    gate = gate.clone()
+    if select_mode == "topk":
+        gate[gate_self_chunk_mask] += amplification_factor
+
+    # Exclude positions that are outside the valid batch boundaries.
+    batch_starts = cu_seqlens[chunk_to_batch[filtered_chunk_indices]]
+    batch_ends = cu_seqlens[chunk_to_batch[filtered_chunk_indices] + 1]
+    gate_batch_start_mask = gate_seq_idx < batch_starts.unsqueeze(1)
+    gate_batch_end_mask = gate_seq_idx >= batch_ends.unsqueeze(1)
+    gate_inf_mask = gate_batch_start_mask | gate_batch_end_mask
+    gate.masked_fill_(gate_inf_mask.unsqueeze(1), -float("inf"))
+
+    if select_mode == 'topk':
+        # We amplify self‐chunk in gate already, so self entries will rank highest.
+        valid_gate_mask = gate != -float("inf")
+        if threshold_type == 'query_head':
+            # === per‐<head,seq> top-k across chunks (original behavior) ===
+            # gate: (C, H, S)
+            _, gate_topk_idx = torch.topk(gate, k=moba_topk, dim=0, largest=True, sorted=False)
+            gate_idx_mask = torch.zeros_like(gate, dtype=torch.bool)
+            gate_idx_mask.scatter_(0, gate_topk_idx, True)
+            gate_mask = valid_gate_mask & gate_idx_mask
+        elif threshold_type == 'overall':
+            # === global top-k across all (chunk, head, seq) entries ===
+            C, H, S = gate.shape
+            flat_gate = gate.flatten()
+            flat_mask = valid_gate_mask.flatten()
+            flat_gate_masked = torch.where(flat_mask, flat_gate, -float("inf"))
+            # pick topk global entries
+            vals, idx = torch.topk(flat_gate_masked, k=moba_topk * H * S, largest=True, sorted=False)
+            others_mask_flat = torch.zeros_like(flat_mask, dtype=torch.bool)
+            others_mask_flat[idx] = True
+            gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(gate.shape)
+        elif threshold_type == 'head_global':
+            # per-head top-k across all chunks and sequence positions
+            C, H, S = gate.shape
+            CS = C * S
+            flat_gate  = gate.permute(1, 0, 2).reshape(H, CS)
+            flat_valid = valid_gate_mask.permute(1, 0, 2).reshape(H, CS)
+            flat_gate_masked = torch.where(flat_valid, flat_gate, torch.full_like(flat_gate, -float('inf')))
+            # pick top-k indices per head
+            _, topk_idx = torch.topk(flat_gate_masked, k=moba_topk * S, dim=1, largest=True, sorted=False)
+            gate_idx_flat = torch.zeros_like(flat_valid, dtype=torch.bool)
+            gate_idx_flat.scatter_(1, topk_idx, True)
+            gate_mask = gate_idx_flat.reshape(H, C, S).permute(1, 0, 2)
+        else:
+            raise ValueError(
+                f"Invalid threshold_type for topk: {threshold_type}. "
+                "Choose 'query_head', 'block', or 'overall'."
+            )
+    elif select_mode == 'threshold':
+        # Delegate to the specific thresholding function
+        valid_gate_mask = gate != -float("inf")  # (num_chunk, num_head, seqlen)
+        if threshold_type == 'query_head':
+            gate_mask = _select_threshold_query_head(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'block':
+            gate_mask = _select_threshold_block(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'overall':
+            gate_mask = _select_threshold_overall(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'head_global':
+            gate_mask = _select_threshold_head_global(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        else:
+            raise ValueError(f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'.")
+    else:
+        raise ValueError(f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'.")
+    
+    # eliminate self_chunk in MoBA branch
+    gate_mask = gate_mask & ~gate_self_chunk_mask
+    # if gate_mask is all false, perform flash_attn instead
+    if gate_mask.sum() == 0:
+        return flash_attn_varlen_func(
+            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=False
+        )
+
+    # Determine which query positions are selected.
+    # nonzero_indices has shape [N, 3] where each row is [chunk_index, head_index, seq_index].
+    moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [(h s k)]
+    moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen)
+    moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1)
+
+    # Build cumulative sequence lengths for the selected queries.
+    moba_seqlen_q = gate_mask.sum(dim=-1).flatten()
+    q_zero_mask = moba_seqlen_q == 0
+    valid_expert_mask = ~q_zero_mask
+    if q_zero_mask.sum() > 0:
+        moba_seqlen_q = moba_seqlen_q[valid_expert_mask]
+    moba_cu_seqlen_q = torch.cat(
+        (
+            torch.tensor([0], device=q.device, dtype=moba_seqlen_q.dtype),
+            moba_seqlen_q.cumsum(dim=0),
+        ),
+        dim=0,
+    ).to(torch.int32)
+
+    # Rearrange gathered KV for the MOBA branch.
+    experts_tensor = rearrange(gathered, "nc cl two h d -> (nc h) cl two d")
+    valid_expert_lengths = chunk_lengths.unsqueeze(1).expand(num_filtered_chunk, num_head).reshape(-1).to(torch.int32)
+    if q_zero_mask.sum() > 0:
+        experts_tensor = experts_tensor[valid_expert_mask]
+        valid_expert_lengths = valid_expert_lengths[valid_expert_mask]
+
+    seq_range = torch.arange(experts_tensor.shape[1], device=experts_tensor.device).unsqueeze(0)
+    mask = seq_range < valid_expert_lengths.unsqueeze(1)
+    moba_kv = experts_tensor[mask]  # Shape: ((nc h cl_valid) two d)
+    moba_kv = moba_kv.unsqueeze(2)   # Shape: ((nc h cl_valid) two 1 d)
+
+    moba_cu_seqlen_kv = torch.cat(
+        [torch.zeros(1, device=experts_tensor.device, dtype=torch.int32),
+         valid_expert_lengths.cumsum(dim=0)],
+        dim=0,
+    ).to(torch.int32)
+
+    assert (
+        moba_cu_seqlen_kv.shape == moba_cu_seqlen_q.shape
+    ), f"Mismatch between moba_cu_seqlen_kv.shape and moba_cu_seqlen_q.shape: {moba_cu_seqlen_kv.shape} vs {moba_cu_seqlen_q.shape}"
+
+    return MixedAttention.apply(
+        q,
+        k,
+        v,
+        self_attn_cu_seqlen,
+        moba_q,
+        moba_kv,
+        moba_cu_seqlen_q,
+        moba_cu_seqlen_kv,
+        max_seqlen,
+        moba_chunk_size,
+        moba_q_sh_indices,
+    )
+
+
+def process_moba_input(
+    x,
+    patch_resolution,
+    chunk_size,
+):
+    """
+    Process inputs for the attention function.
+
+    Args:
+        x (torch.Tensor): Input tensor with shape [batch_size, num_patches, num_heads, head_dim].
+        patch_resolution (tuple): Tuple containing the patch resolution (t, h, w).
+        chunk_size (int): Size of the chunk. (maybe tuple or int, according to chunk type)
+
+    Returns:
+        torch.Tensor: Processed input tensor.
+    """
+    if isinstance(chunk_size, float) or isinstance(chunk_size, int):
+        moba_chunk_size = int(chunk_size * patch_resolution[1] * patch_resolution[2])
+    else:
+        assert isinstance(chunk_size, (Tuple, list)), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}"
+        if len(chunk_size) == 2:
+            assert patch_resolution[1] % chunk_size[0] == 0 and patch_resolution[2] % chunk_size[1] == 0, f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}"
+            nch, ncw = patch_resolution[1] // chunk_size[0], patch_resolution[2] // chunk_size[1]
+            x = rearrange(x, "b (t nch ch ncw cw) n d -> b (nch ncw t ch cw) n d", t=patch_resolution[0], nch=nch, ncw=ncw, ch=chunk_size[0], cw=chunk_size[1])
+            moba_chunk_size = patch_resolution[0] * chunk_size[0] * chunk_size[1]
+        elif len(chunk_size) == 3:
+            assert patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[1] % chunk_size[1] == 0 and patch_resolution[2] % chunk_size[2] == 0, f"patch_resolution {patch_resolution} should be divisible by 3d chunk_size {chunk_size}"
+            nct, nch, ncw = patch_resolution[0] // chunk_size[0], patch_resolution[1] // chunk_size[1], patch_resolution[2] // chunk_size[2]
+            x = rearrange(x, "b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d", nct=nct, nch=nch, ncw=ncw, ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2])
+            moba_chunk_size = chunk_size[0] * chunk_size[1] * chunk_size[2]
+        else:
+            raise ValueError(f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}")
+        
+    return x, moba_chunk_size
+
+
+def process_moba_output(
+    x,
+    patch_resolution,
+    chunk_size,
+):
+    if isinstance(chunk_size, float) or isinstance(chunk_size, int):
+        pass
+    elif len(chunk_size) == 2:
+        x = rearrange(x, "b (nch ncw t ch cw) n d -> b (t nch ch ncw cw) n d", nch=patch_resolution[1] // chunk_size[0], ncw=patch_resolution[2] // chunk_size[1], t=patch_resolution[0], ch=chunk_size[0], cw=chunk_size[1])
+    elif len(chunk_size) == 3:
+        x = rearrange(x, "b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d", nct=patch_resolution[0] // chunk_size[0], nch=patch_resolution[1] // chunk_size[1], ncw=patch_resolution[2] // chunk_size[2], ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2])
+
+    return x
+
+
+# TEST
+def generate_data(batch_size, seqlen, num_head, head_dim, dtype):
+    random.seed(0)
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+    device = torch.cuda.current_device()
+
+    q = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    k = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    v = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}")
+    cu_seqlens = torch.arange(0, q.shape[0] * q.shape[1] + 1, q.shape[1], dtype=torch.int32, device='cuda')
+    max_seqlen = q.shape[1]
+    q = rearrange(q, "b s ... -> (b s) ...")
+    k = rearrange(k, "b s ... -> (b s) ...")
+    v = rearrange(v, "b s ... -> (b s) ...")
+
+    return q, k, v, cu_seqlens, max_seqlen
+
+
+def test_attn_varlen_moba_speed(batch, head, seqlen, head_dim, moba_chunk_size, moba_topk, dtype=torch.bfloat16, select_mode='threshold', simsum_threshold=0.25, threshold_type='query_head'):
+    """Speed test comparing flash_attn vs moba_attention"""
+    # Get data
+    q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype)
+    print(f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}")
+    vo_grad = torch.randn_like(q)
+    
+    # Warmup
+    warmup_iters = 3
+    perf_test_iters = 10
+
+    # Warmup
+    for _ in range(warmup_iters):
+        o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False)
+        torch.autograd.backward(o, vo_grad)
+    
+    torch.cuda.synchronize()
+    start_flash = time.perf_counter()
+    for _ in range(perf_test_iters):
+        o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False)
+        torch.autograd.backward(o, vo_grad)
+        
+    torch.cuda.synchronize()
+    time_flash = (time.perf_counter() - start_flash) / perf_test_iters * 1000
+
+    # Warmup
+    for _ in range(warmup_iters):
+        om = moba_attn_varlen(q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type)
+        torch.autograd.backward(om, vo_grad)
+
+    torch.cuda.synchronize()
+    start_moba = time.perf_counter()
+    for _ in range(perf_test_iters):
+        om = moba_attn_varlen(q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type)
+        torch.autograd.backward(om, vo_grad)
+    
+    torch.cuda.synchronize()
+    time_moba = (time.perf_counter() - start_moba) / perf_test_iters * 1000
+    
+    print(f"Flash: {time_flash:.2f}ms, MoBA: {time_moba:.2f}ms")
+    print(f"Speedup:  {time_flash / time_moba:.2f}x")
+
+
+
+if __name__ == "__main__":
+    """
+    CUDA_VISIBLE_DEVICES=1 \
+    python -u csrc/attn/vmoba_attn/vmoba/vmoba.py
+    """
+    test_attn_varlen_moba_speed(batch=1, head=12, seqlen=32760, head_dim=128, moba_chunk_size=32760 // 3 // 6 // 4, moba_topk=3, select_mode='threshold', simsum_threshold=0.3, threshold_type='query_head')

examples/datasets/crush_smol/validation.json

@@ -1,76 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
-      "image_path": null,
-      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object.",
-      "image_path": null,
-      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign.",
-      "image_path": null,
-      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.",
-      "image_path": null,
-      "video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}

examples/datasets/crush_smol/validation_single.json

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

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

@@ -0,0 +1,151 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_t2v_output/t2v_%j.out
+#SBATCH --error=dmd_t2v_output/t2v_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29503
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+# Configs
+NUM_GPUS=1
+
+# Model paths for Self-Forcing DMD distillation:
+GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+
+DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
+  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan_t2v_finetune"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 100
+  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
+)
+
+# Self-forcing DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 3.0
+  --fake_score_learning_rate 8e-6
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+# Self-forcing specific arguments
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
+)
+
+torchrun \
+--nnodes 1 \
+--master_port $MASTER_PORT \
+--nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"

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

@@ -4,21 +4,21 @@ GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v_i2v_1_3b/"
+OUTPUT_DIR="data/crush-smol_processed_t2v/"
 
 torchrun --nproc_per_node=$GPU_NUM \
     fastvideo/pipelines/preprocess/v1_preprocess.py \
     --model_path $MODEL_PATH \
     --data_merge_path $DATA_MERGE_PATH \
-    --preprocess_video_batch_size 2 \
+    --preprocess_video_batch_size 8 \
     --seed 42 \
     --max_height 480 \
     --max_width 832 \
-    --num_frames 77 \
+    --num_frames 81 \
     --dataloader_num_workers 0 \
     --output_dir=$OUTPUT_DIR \
     --train_fps 16 \
     --samples_per_file 8 \
     --flush_frequency 8 \
     --video_length_tolerance_range 5 \
-    --preprocess_task "t2v_ode_trajectory" 
+    --preprocess_task "t2v" 

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

@@ -0,0 +1,151 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_t2v_output/t2v_%j.out
+#SBATCH --error=dmd_t2v_output/t2v_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29503
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+# Configs
+NUM_GPUS=8
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+
+DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
+  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan2.2_t2v_finetune"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 16
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --num_frames 61  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 4  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 16
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS # 64
+  --sp_size 4
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim 8
+)
+
+# Model arguments
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 100
+  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
+)
+
+# Self-forcing DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 3.0
+  --fake_score_learning_rate 8e-6
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+# Self-forcing specific arguments
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
+)
+
+torchrun \
+--nnodes 1 \
+--master_port $MASTER_PORT \
+--nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"

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

@@ -98,6 +98,7 @@ dmd_args=(
 torchrun \
 --nnodes 1 \
 --nproc_per_node $NUM_GPUS \
+--master_port $MASTER_PORT \
     fastvideo/training/wan_distillation_pipeline.py \
     "${parallel_args[@]}" \
     "${model_args[@]}" \

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

@@ -0,0 +1,112 @@
+#!/bin/bash
+
+# Basic Info
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export MASTER_PORT=29501
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+# Configs
+NUM_GPUS=1
+MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name wan_t2v_distill_dmd_VSA
+  --output_dir="checkpoints/wan_t2v_finetune"
+  --max_train_steps=4000
+  --train_batch_size=1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps=1
+  --num_latent_t 31
+  --num_height 704
+  --num_width 1280
+  --num_frames 121
+  --enable_gradient_checkpointing_type "full"
+  --training_state_checkpointing_steps=500
+  --weight_only_checkpointing_steps=500
+  --lora_rank 32
+  --lora_training True
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus 1
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 200
+  --validation_sampling_steps "3"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate=1e-4
+  --mixed_precision="bf16"
+  --weight_decay 0.01
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --ema_start_step 0
+  --flow_shift 8
+  --seed 1000
+)
+
+# DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,757,522'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --generator_update_interval 5
+  --real_score_guidance_scale 3.5
+  --VSA_sparsity 0.8
+)
+
+torchrun \
+--nnodes 1 \
+--nproc_per_node $NUM_GPUS \
+--master_port $MASTER_PORT \
+    fastvideo/training/wan_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}"

examples/inference/basic/basic.py

@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator
 
-from fastvideo.configs.sample import SamplingParam
+# from fastvideo.configs.sample import SamplingParam
 
 OUTPUT_PATH = "video_samples"
 def main():
@@ -9,37 +9,32 @@ def main():
     # If a local path is provided, FastVideo will make a best effort
     # attempt to identify the optimal arguments.
     generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        "Wan-AI/Wan2.1-T2V-14B-Diffusers",
         # FastVideo will automatically handle distributed setup
-        num_gpus=1,
+        num_gpus=4,
         use_fsdp_inference=True,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
         pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
-        ti2v_task=True,
         # image_encoder_cpu_offload=False,
     )
 
-    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
     # sampling_param.num_frames = 45
-    sampling_param.image_path = "test.jpg"
+    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
     # Generate videos with the same simple API, regardless of GPU count
     prompt = (
-        "A girl is packing a suitcase when stuff suddently starts flying around the room."
+        "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open."
     )
     video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
     # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
 
     # Generate another video with a different prompt, without reloading the
     # model!
-    # prompt2 = (
-    #     "A majestic lion strides across the golden savanna, its powerful frame "
-    #     "glistening under the warm afternoon sun. The tall grass ripples gently in "
-    #     "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
-    #     "embodying the raw energy of the wild. Low angle, steady tracking shot, "
-    #     "cinematic.")
-    # video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
+    prompt2 = (
+        "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object")
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
 
 
 if __name__ == "__main__":

examples/inference/basic/basic_self_forcing_causal.py

@@ -19,15 +19,13 @@ def main():
     )
 
     sampling_param = SamplingParam.from_pretrained(model_name)
-    sampling_param.num_frames = 81
 
-    prompts = [
-        "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
-        "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.",
-    ]
-
-    for prompt in prompts:
-        video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
 
 if __name__ == "__main__":
     main()

examples/training/consistency_finetune/ode_init/finetune_ode_init.sh

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

examples/training/consistency_finetune/ode_init/finetune_ode_init_16k_vidprom.slurm

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

examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm

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

examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm

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

examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh

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

examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh

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

examples/training/consistency_finetune/ode_init/validation.json

@@ -1,40 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}

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

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

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

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

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

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

examples/training/t2v_to_i2v/wan_t2v/finetune_t2v_to_i2v.sh

@@ -1,94 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v_old"
-VALIDATION_DATASET_FILE="examples/datasets/crush_smol/validation.json"
-NUM_GPUS=8
-# export CUDA_VISIBLE_DEVICES=4,5
-
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_t2v_i2v_finetune"
-  --output_dir "checkpoints/wan_t2v_i2v_finetune"
-  --max_train_steps 5000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 2
-  --num_latent_t 20
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS 
-  --sp_size 4
-  --tp_size 1
-  --hsdp_replicate_dim 2
-  --hsdp_shard_dim 4
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path $DATA_DIR
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation 
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 50
-  --validation_sampling_steps "50" 
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 5e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-  --enable_gradient_checkpointing_type "full"
-  --t2v_as_i2v_task True
-  # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500"
-)
-
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_t2v_i2v_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"

fastvideo/attention/backends/vmoba.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from flash_attn.bert_padding import pad_input
+
+from csrc.attn.vmoba_attn.vmoba import (moba_attn_varlen, process_moba_input,
+                                        process_moba_output)
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class VMOBAAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_name() -> str:
+        return "VMOBA_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["VMOBAAttentionImpl"]:
+        return VMOBAAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["VideoMobaAttentionMetadata"]:
+        return VideoMobaAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["VideoMobaAttentionMetadataBuilder"]:
+        return VideoMobaAttentionMetadataBuilder
+
+
+@dataclass
+class VideoMobaAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+
+    temporal_chunk_size: int
+    temporal_topk: int
+    spatial_chunk_size: tuple[int, int]
+    spatial_topk: int
+    st_chunk_size: tuple[int, int, int]
+    st_topk: int
+
+    moba_select_mode: str
+    moba_threshold: float
+    moba_threshold_type: str
+    patch_resolution: list[int]
+
+    first_full_step: int = 12
+    first_full_layer: int = 0
+    # temporal_layer -> spatial_layer -> st_layer
+    temporal_layer: int = 1
+    spatial_layer: int = 1
+    st_layer: int = 1
+
+
+class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        temporal_chunk_size: int,
+        temporal_topk: int,
+        spatial_chunk_size: tuple[int, int],
+        spatial_topk: int,
+        st_chunk_size: tuple[int, int, int],
+        st_topk: int,
+        moba_select_mode: str = 'threshold',
+        moba_threshold: float = 0.25,
+        moba_threshold_type: str = 'query_head',
+        device: torch.device = None,
+        first_full_layer: int = 0,
+        first_full_step: int = 12,
+        temporal_layer: int = 1,
+        spatial_layer: int = 1,
+        st_layer: int = 1,
+        **kwargs,
+    ) -> VideoMobaAttentionMetadata:
+        if device is None:
+            device = torch.device("cpu")
+        assert raw_latent_shape[0] % patch_size[0] == 0 and raw_latent_shape[
+            1] % patch_size[1] == 0 and raw_latent_shape[2] % patch_size[
+                2] == 0, f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}"
+        patch_resolution = [
+            t // pt for t, pt in zip(raw_latent_shape, patch_size, strict=False)
+        ]
+
+        return VideoMobaAttentionMetadata(
+            current_timestep=current_timestep,
+            temporal_chunk_size=temporal_chunk_size,
+            temporal_topk=temporal_topk,
+            spatial_chunk_size=spatial_chunk_size,
+            spatial_topk=spatial_topk,
+            st_chunk_size=st_chunk_size,
+            st_topk=st_topk,
+            moba_select_mode=moba_select_mode,
+            moba_threshold=moba_threshold,
+            moba_threshold_type=moba_threshold_type,
+            patch_resolution=patch_resolution,
+            first_full_layer=first_full_layer,
+            first_full_step=first_full_step,
+            temporal_layer=temporal_layer,
+            spatial_layer=spatial_layer,
+            st_layer=st_layer,
+        )
+
+
+class VMOBAAttentionImpl(AttentionImpl):
+
+    def __init__(self,
+                 num_heads,
+                 head_size,
+                 softmax_scale,
+                 causal=False,
+                 num_kv_heads=None,
+                 prefix="",
+                 **extra_impl_args) -> None:
+        self.prefix = prefix
+        self.layer_idx = self._get_layer_idx(prefix)
+
+    def _get_layer_idx(self, prefix: str) -> int | None:
+        match = re.search(r"blocks\.(\d+)", prefix)
+        if not match:
+            raise ValueError(f"Invalid prefix: {prefix}")
+        return int(match.group(1))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """
+        query: [B, L, H, D]
+        key:   [B, L, H, D]
+        value: [B, L, H, D]
+        attn_metadata: AttentionMetadata
+        """
+        batch_size, sequence_length, num_heads, head_dim = query.shape
+
+        # select chunk type according to layer idx:
+        loop_layer_num = attn_metadata.temporal_layer + attn_metadata.spatial_layer + attn_metadata.st_layer
+        moba_layer = self.layer_idx - attn_metadata.first_full_layer
+        if moba_layer % loop_layer_num < attn_metadata.temporal_layer:
+            moba_chunk_size = attn_metadata.temporal_chunk_size
+            moba_topk = attn_metadata.temporal_topk
+        elif moba_layer % loop_layer_num < attn_metadata.temporal_layer + attn_metadata.spatial_layer:
+            moba_chunk_size = attn_metadata.spatial_chunk_size
+            moba_topk = attn_metadata.spatial_topk
+        elif moba_layer % loop_layer_num < attn_metadata.temporal_layer + attn_metadata.spatial_layer + attn_metadata.st_layer:
+            moba_chunk_size = attn_metadata.st_chunk_size
+            moba_topk = attn_metadata.st_topk
+
+        # torch.distributed.breakpoint()
+        query, chunk_size = process_moba_input(query,
+                                               attn_metadata.patch_resolution,
+                                               moba_chunk_size)
+        key, chunk_size = process_moba_input(key,
+                                             attn_metadata.patch_resolution,
+                                             moba_chunk_size)
+        value, chunk_size = process_moba_input(value,
+                                               attn_metadata.patch_resolution,
+                                               moba_chunk_size)
+        max_seqlen = query.shape[1]
+        indices_q = torch.arange(0,
+                                 query.shape[0] * query.shape[1],
+                                 device=query.device)
+        cu_seqlens = torch.arange(0,
+                                  query.shape[0] * query.shape[1] + 1,
+                                  query.shape[1],
+                                  dtype=torch.int32,
+                                  device=query.device)
+        query = rearrange(query, "b s ... -> (b s) ...")
+        key = rearrange(key, "b s ... -> (b s) ...")
+        value = rearrange(value, "b s ... -> (b s) ...")
+
+        # current_timestep=attn_metadata.current_timestep
+        hidden_states = moba_attn_varlen(
+            query,
+            key,
+            value,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            moba_chunk_size=chunk_size,
+            moba_topk=moba_topk,
+            select_mode=attn_metadata.moba_select_mode,
+            simsum_threshold=attn_metadata.moba_threshold,
+            threshold_type=attn_metadata.moba_threshold_type,
+        )
+        hidden_states = pad_input(hidden_states, indices_q, batch_size,
+                                  sequence_length)
+        hidden_states = process_moba_output(hidden_states,
+                                            attn_metadata.patch_resolution,
+                                            moba_chunk_size)
+
+        return hidden_states

fastvideo/configs/backend/vmoba/wan_1.3B_77_448_832.json

@@ -0,0 +1,16 @@
+{
+    "temporal_chunk_size": 2,
+    "temporal_topk": 2,
+    "spatial_chunk_size": [4, 13],
+    "spatial_topk": 6,
+    "st_chunk_size": [4, 4, 13],
+    "st_topk": 18,
+    "moba_select_mode": "topk",
+    "moba_threshold": 0.25,
+    "moba_threshold_type": "query_head",
+    "first_full_layer": 0,
+    "first_full_step": 12,
+    "temporal_layer": 1,
+    "spatial_layer": 1,
+    "st_layer": 1
+}

fastvideo/configs/backend/vmoba/wan_1.3B_77_480_832.json

@@ -0,0 +1,16 @@
+{
+    "temporal_chunk_size": 2,
+    "temporal_topk": 3,
+    "spatial_chunk_size": [3, 4],
+    "spatial_topk": 20,
+    "st_chunk_size": [4, 6, 4],
+    "st_topk": 15,
+    "moba_select_mode": "threshold",
+    "moba_threshold": 0.25,
+    "moba_threshold_type": "query_head",
+    "first_full_layer": 0,
+    "first_full_step": 12,
+    "temporal_layer": 1,
+    "spatial_layer": 1,
+    "st_layer": 1
+}

fastvideo/configs/configs.py

@@ -32,6 +32,29 @@ class DatasetType(str, Enum):
         return [dataset_type.value for dataset_type in cls]
 
 
+class VideoLoaderType(str, Enum):
+    """
+    Enumeration for different video loaders.
+    """
+    TORCHCODEC = "torchcodec"
+    TORCHVISION = "torchvision"
+
+    @classmethod
+    def from_string(cls, value: str) -> "VideoLoaderType":
+        """Convert string to VideoLoader enum."""
+        try:
+            return cls(value.lower())
+        except ValueError:
+            raise ValueError(
+                f"Invalid video loader: {value}. Must be one of: {', '.join([m.value for m in cls])}"
+            ) from None
+
+    @classmethod
+    def choices(cls) -> list[str]:
+        """Get all available choices as strings for argparse."""
+        return [video_loader.value for video_loader in cls]
+
+
 @dataclasses.dataclass
 class PreprocessConfig:
     """Configuration for preprocessing operations."""
@@ -51,6 +74,7 @@ class PreprocessConfig:
     flush_frequency: int = 256
 
     # Video processing parameters
+    video_loader_type: VideoLoaderType = VideoLoaderType.TORCHCODEC
     max_height: int = 480
     max_width: int = 848
     num_frames: int = 163
@@ -120,6 +144,12 @@ class PreprocessConfig:
                                      help="How often to save to parquet files")
 
         # Video processing parameters
+        preprocess_args.add_argument(
+            f"--{prefix_with_dot}video-loader-type",
+            type=str,
+            choices=VideoLoaderType.choices(),
+            default=PreprocessConfig.video_loader_type.value,
+            help="Type of the video loader")
         preprocess_args.add_argument(f"--{prefix_with_dot}max-height",
                                      type=int,
                                      default=PreprocessConfig.max_height,
@@ -174,6 +204,10 @@ class PreprocessConfig:
         if 'dataset_type' in kwargs and isinstance(kwargs['dataset_type'], str):
             kwargs['dataset_type'] = DatasetType.from_string(
                 kwargs['dataset_type'])
+        if 'video_loader_type' in kwargs and isinstance(
+                kwargs['video_loader_type'], str):
+            kwargs['video_loader_type'] = VideoLoaderType.from_string(
+                kwargs['video_loader_type'])
 
         preprocess_config = cls()
         if not update_config_from_args(

fastvideo/configs/models/dits/base.py

@@ -15,9 +15,13 @@ class DiTArchConfig(ArchConfig):
     reverse_param_names_mapping: dict = field(default_factory=dict)
     lora_param_names_mapping: dict = field(default_factory=dict)
     _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
-        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
-        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
-        AttentionBackendEnum.VIDEO_SPARSE_ATTN)
+        AttentionBackendEnum.SLIDING_TILE_ATTN,
+        AttentionBackendEnum.SAGE_ATTN,
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+        AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+        AttentionBackendEnum.VMOBA_ATTN,
+    )
 
     hidden_size: int = 0
     num_attention_heads: int = 0

fastvideo/configs/pipelines/base.py

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

fastvideo/configs/pipelines/registry.py

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

fastvideo/configs/pipelines/wan.py

@@ -12,13 +12,13 @@ from fastvideo.configs.models.vaes import WanVAEConfig
 from fastvideo.configs.pipelines.base import PipelineConfig
 
 
-def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
-    mask: torch.tensor = outputs.attention_mask
-    hidden_state: torch.tensor = outputs.last_hidden_state
+def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    mask: torch.Tensor = outputs.attention_mask
+    hidden_state: torch.Tensor = outputs.last_hidden_state
     seq_lens = mask.gt(0).sum(dim=1).long()
     assert torch.isnan(hidden_state).sum() == 0
     prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
-    prompt_embeds_tensor: torch.tensor = torch.stack([
+    prompt_embeds_tensor: torch.Tensor = torch.stack([
         torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
         for u in prompt_embeds
     ],
@@ -39,12 +39,12 @@ class WanT2V480PConfig(PipelineConfig):
     vae_sp: bool = False
 
     # Denoising stage
-    flow_shift: float | None = 8.0
+    flow_shift: float | None = 3.0
 
     # Text encoding stage
     text_encoder_configs: tuple[EncoderConfig, ...] = field(
         default_factory=lambda: (T5Config(), ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor],
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
                                   ...] = field(default_factory=lambda:
                                                (t5_postprocess_text, ))
 
@@ -68,7 +68,7 @@ class WanT2V720PConfig(WanT2V480PConfig):
     # WanConfig-specific parameters with defaults
 
     # Denoising stage
-    flow_shift: int = 5
+    flow_shift: float | None = 5.0
 
 
 @dataclass
@@ -94,7 +94,7 @@ class WanI2V720PConfig(WanI2V480PConfig):
     # WanConfig-specific parameters with defaults
 
     # Denoising stage
-    flow_shift: int = 5
+    flow_shift: float | None = 5.0
 
 
 @dataclass
@@ -104,7 +104,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
     # WanConfig-specific parameters with defaults
 
     # Denoising stage
-    flow_shift: int = 8
+    flow_shift: float | None = 8.0
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 757, 522])
 
@@ -115,7 +115,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
 
 @dataclass
 class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
-    flow_shift: int = 5
+    flow_shift: float | None = 5.0
     ti2v_task: bool = True
 
     def __post_init__(self) -> None:
@@ -125,7 +125,7 @@ class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
 
 @dataclass
 class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config):
-    flow_shift: int = 5
+    flow_shift: float | None = 5.0
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 757, 522])
 
@@ -149,3 +149,4 @@ class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
     flow_shift: float | None = 5.0
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 750, 500, 250])
+    warp_denoising_step: bool = True

fastvideo/configs/sample/base.py

@@ -47,8 +47,6 @@ class SamplingParam:
     # Misc
     save_video: bool = True
     return_frames: bool = False
-    return_trajectory_latents: bool = False  # returns all latents for each timestep
-    return_trajectory_decoded: bool = False  # returns decoded latents for each timestep
 
     def __post_init__(self) -> None:
         self.data_type = "video" if self.num_frames > 1 else "image"
@@ -194,16 +192,11 @@ class SamplingParam:
             help="Path to input image for image-to-video generation",
         )
         parser.add_argument(
-            "--return-trajectory-latents",
-            action="store_true",
-            default=SamplingParam.return_trajectory_latents,
-            help="Whether to return the trajectory",
-        )
-        parser.add_argument(
-            "--return-trajectory-decoded",
-            action="store_true",
-            default=SamplingParam.return_trajectory_decoded,
-            help="Whether to return the decoded trajectory",
+            "--moba-config-path",
+            type=str,
+            default=None,
+            help=
+            "Path to a JSON file containing V-MoBA specific configurations.",
         )
         return parser
 

fastvideo/dataset/dataloader/schema.py

@@ -78,43 +78,3 @@ pyarrow_schema_t2v = pa.schema([
     pa.field("duration_sec", pa.float64()),
     pa.field("fps", pa.float64()),
 ])
-
-pyarrow_schema_ode_trajectory = pa.schema([
-    pa.field("id", pa.string()),
-    # --- Image/Video VAE latents ---
-    # Tensors are stored as raw bytes with shape and dtype info for loading
-    pa.field("vae_latent_bytes", pa.binary()),
-    # e.g., [C, T, H, W] or [C, H, W]
-    pa.field("vae_latent_shape", pa.list_(pa.int64())),
-    # e.g., 'float32'
-    pa.field("vae_latent_dtype", pa.string()),
-    # --- Text encoder output tensor ---
-    # Tensors are stored as raw bytes with shape and dtype info for loading
-    pa.field("text_embedding_bytes", pa.binary()),
-    # e.g., [SeqLen, Dim]
-    pa.field("text_embedding_shape", pa.list_(pa.int64())),
-    # e.g., 'bfloat16' or 'float32'
-    pa.field("text_embedding_dtype", pa.string()),
-    # I2V
-    pa.field("image_condition_latents_bytes", pa.binary()),
-    pa.field("image_condition_latents_shape", pa.list_(pa.int64())),
-    pa.field("image_condition_latents_dtype", pa.string()),
-    # --- ODE Trajectory ---
-    pa.field("trajectory_latents_bytes", pa.binary()),
-    pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
-    pa.field("trajectory_latents_dtype", pa.string()),
-    pa.field("trajectory_timesteps_bytes", pa.binary()),
-    pa.field("trajectory_timesteps_shape", pa.list_(pa.int64())),
-    pa.field("trajectory_timesteps_dtype", pa.string()),
-    # --- Metadata ---
-    pa.field("file_name", pa.string()),
-    pa.field("caption", pa.string()),
-    pa.field("media_type", pa.string()),  # 'image' or 'video'
-    pa.field("width", pa.int64()),
-    pa.field("height", pa.int64()),
-    # -- Video-specific (can be null/default for images) ---
-    # Number of frames processed (e.g., 1 for image, N for video)
-    pa.field("num_frames", pa.int64()),
-    pa.field("duration_sec", pa.float64()),
-    pa.field("fps", pa.float64()),
-])

fastvideo/dataset/utils.py

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

fastvideo/entrypoints/video_generator.py

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

fastvideo/fastvideo_args.py

@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py
 """The arguments of FastVideo Inference."""
-
 import argparse
 import dataclasses
+import json
 from contextlib import contextmanager
 from dataclasses import field
 from enum import Enum
@@ -139,6 +139,10 @@ class FastVideoArgs:
     # VSA parameters
     VSA_sparsity: float = 0.0  # inference/validation sparsity
 
+    # V-MoBA parameters
+    moba_config_path: str | None = None
+    moba_config: dict[str, Any] = field(default_factory=dict)
+
     # Master port for distributed training/inference
     master_port: int | None = None
 
@@ -154,7 +158,6 @@ class FastVideoArgs:
         "transformer": True,
         "vae": True,
     })
-    override_transformer_cls_name: str | None = None
 
     # # DMD parameters
     # dmd_denoising_steps: List[int] | None = field(default=None)
@@ -167,6 +170,16 @@ class FastVideoArgs:
         return not self.inference_mode
 
     def __post_init__(self):
+        if self.moba_config_path:
+            try:
+                with open(self.moba_config_path) as f:
+                    self.moba_config = json.load(f)
+                logger.info("Loaded V-MoBA config from %s",
+                            self.moba_config_path)
+            except (FileNotFoundError, json.JSONDecodeError) as e:
+                logger.error("Failed to load V-MoBA config from %s: %s",
+                             self.moba_config_path, e)
+                raise
         self.check_fastvideo_args()
 
     @staticmethod
@@ -383,12 +396,6 @@ class FastVideoArgs:
             default=FastVideoArgs.enable_stage_verification,
             help="Enable input/output verification for pipeline stages",
         )
-        parser.add_argument(
-            "--override-transformer-cls-name",
-            type=str,
-            default=FastVideoArgs.override_transformer_cls_name,
-            help="Override transformer cls name",
-        )
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 
@@ -598,6 +605,11 @@ class TrainingArgs(FastVideoArgs):
     pretrained_model_name_or_path: str = ""
     dit_model_name_or_path: str = ""
 
+    # DMD model paths - separate paths for each network
+    generator_model_path: str = ""  # path for generator (student) model
+    real_score_model_path: str = ""  # path for real score (teacher) model
+    fake_score_model_path: str = ""  # path for fake score (critic) model
+
     # diffusion setting
     ema_decay: float = 0.0
     ema_start_step: int = 0
@@ -620,6 +632,7 @@ class TrainingArgs(FastVideoArgs):
     checkpoints_total_limit: int = 0
     checkpointing_steps: int = 0
     resume_from_checkpoint: str = ""  # specify the checkpoint folder to resume from
+    init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading
 
     # optimizer & scheduler
     num_train_epochs: int = 0
@@ -651,6 +664,7 @@ class TrainingArgs(FastVideoArgs):
     linear_quadratic_threshold: float = 0.0
     linear_range: float = 0.0
     weight_decay: float = 0.0
+    betas: str = "0.9,0.999"  # betas for optimizer, format: "beta1,beta2"
     use_ema: bool = False
     multi_phased_distill_schedule: str = ""
     pred_decay_weight: float = 0.0
@@ -671,11 +685,13 @@ class TrainingArgs(FastVideoArgs):
 
     # distillation args
     generator_update_interval: int = 5
+    dfake_gen_update_ratio: int = 5  # self-forcing: how often to train generator vs critic
     min_timestep_ratio: float = 0.2
     max_timestep_ratio: float = 0.98
     real_score_guidance_scale: float = 3.5
     fake_score_learning_rate: float = 0.0  # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
     fake_score_lr_scheduler: str = "constant"  # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
+    fake_score_betas: str = "0.9,0.999"  # betas for fake score optimizer, format: "beta1,beta2"
     training_state_checkpointing_steps: int = 0  # for resuming training
     weight_only_checkpointing_steps: int = 0  # for inference
     log_visualization: bool = False
@@ -683,6 +699,16 @@ class TrainingArgs(FastVideoArgs):
     simulate_generator_forward: bool = False
     warp_denoising_step: bool = False
 
+    # Self-forcing specific arguments
+    num_frame_per_block: int = 3
+    independent_first_frame: bool = False
+    enable_gradient_masking: bool = True
+    gradient_mask_last_n_frames: int = 21
+    validate_cache_structure: bool = False  # Debug flag for cache validation
+    same_step_across_blocks: bool = False  # Use same exit timestep for all blocks
+    last_step_only: bool = False  # Only use the last timestep for training
+    context_noise: int = 0  # Context noise level for cache updates
+
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
         provided_args = clean_cli_args(args)
@@ -783,6 +809,20 @@ class TrainingArgs(FastVideoArgs):
                             type=str,
                             help="Directory to cache models")
 
+        # DMD model paths - separate paths for each network
+        parser.add_argument(
+            "--generator-model-path",
+            type=str,
+            help="Path to generator (student) model for DMD distillation")
+        parser.add_argument(
+            "--real-score-model-path",
+            type=str,
+            help="Path to real score (teacher) model for DMD distillation")
+        parser.add_argument(
+            "--fake-score-model-path",
+            type=str,
+            help="Path to fake score (critic) model for DMD distillation")
+
         # Diffusion settings
         parser.add_argument("--ema-decay",
                             type=float,
@@ -853,6 +893,10 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--resume-from-checkpoint",
                             type=str,
                             help="Path to checkpoint to resume from")
+        parser.add_argument(
+            "--init-weights-from-safetensors",
+            type=str,
+            help="Path to safetensors file for initial weight loading")
         parser.add_argument("--logging-dir",
                             type=str,
                             help="Directory for logging")
@@ -957,6 +1001,10 @@ class TrainingArgs(FastVideoArgs):
                             help="Linear quadratic threshold")
         parser.add_argument("--linear-range", type=float, help="Linear range")
         parser.add_argument("--weight-decay", type=float, help="Weight decay")
+        parser.add_argument("--betas",
+                            type=str,
+                            default=TrainingArgs.betas,
+                            help="Betas for optimizer (format: 'beta1,beta2')")
         parser.add_argument("--use-ema",
                             action=StoreBoolean,
                             help="Whether to use EMA")
@@ -993,11 +1041,27 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--lora-rank", type=int, help="LoRA rank")
         parser.add_argument("--lora-alpha", type=int, help="LoRA alpha")
 
+        # V-MoBA parameters
+        parser.add_argument(
+            "--moba-config-path",
+            type=str,
+            default=None,
+            help=
+            "Path to a JSON file containing V-MoBA specific configurations.",
+        )
+
         # Distillation arguments
         parser.add_argument("--generator-update-interval",
                             type=int,
                             default=TrainingArgs.generator_update_interval,
                             help="Ratio of student updates to critic updates.")
+        parser.add_argument(
+            "--dfake-gen-update-ratio",
+            type=int,
+            default=TrainingArgs.dfake_gen_update_ratio,
+            help=
+            "Self-forcing: How often to train generator vs critic (train generator every N steps)."
+        )
         parser.add_argument("--min-timestep-ratio",
                             type=float,
                             default=TrainingArgs.min_timestep_ratio,
@@ -1014,6 +1078,11 @@ class TrainingArgs(FastVideoArgs):
                             type=float,
                             default=TrainingArgs.fake_score_learning_rate,
                             help="Learning rate for fake score transformer")
+        parser.add_argument(
+            "--fake-score-betas",
+            type=str,
+            default=TrainingArgs.fake_score_betas,
+            help="Betas for fake score optimizer (format: 'beta1,beta2')")
         parser.add_argument(
             "--fake-score-lr-scheduler",
             type=str,
@@ -1029,7 +1098,46 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument(
             "--warp-denoising-step",
             action=StoreBoolean,
-            help="Whether to warp denoising step according to the scheduler time shift")
+            help=
+            "Whether to warp denoising step according to the scheduler time shift"
+        )
+
+        # Self-forcing specific arguments
+        parser.add_argument(
+            "--num-frame-per-block",
+            type=int,
+            default=TrainingArgs.num_frame_per_block,
+            help="Number of frames per block for causal generation")
+        parser.add_argument(
+            "--independent-first-frame",
+            action=StoreBoolean,
+            help="Whether the first frame is independent in causal generation")
+        parser.add_argument(
+            "--enable-gradient-masking",
+            action=StoreBoolean,
+            help="Whether to enable frame-level gradient masking")
+        parser.add_argument(
+            "--gradient-mask-last-n-frames",
+            type=int,
+            default=TrainingArgs.gradient_mask_last_n_frames,
+            help="Number of last frames to enable gradients for")
+        parser.add_argument(
+            "--validate-cache-structure",
+            action=StoreBoolean,
+            help="Whether to validate KV cache structure (debug flag)")
+        parser.add_argument(
+            "--same-step-across-blocks",
+            action=StoreBoolean,
+            help="Whether to use the same exit timestep for all blocks")
+        parser.add_argument(
+            "--last-step-only",
+            action=StoreBoolean,
+            help="Whether to only use the last timestep for training")
+        parser.add_argument(
+            "--context-noise",
+            type=int,
+            default=TrainingArgs.context_noise,
+            help="Context noise level for cache updates")
 
         return parser
 
@@ -1037,4 +1145,4 @@ class TrainingArgs(FastVideoArgs):
 def parse_int_list(value: str) -> list[int]:
     if not value:
         return []
-    return [int(x.strip()) for x in value.split(",")]
+    return [int(x.strip()) for x in value.split(",")]

fastvideo/layers/layernorm.py

@@ -9,9 +9,6 @@ import torch.nn.functional as F
 from fastvideo.layers.custom_op import CustomOp
 from fastvideo.platforms import current_platform
 
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
 
 @CustomOp.register("rms_norm")
 class RMSNorm(CustomOp):
@@ -103,13 +100,16 @@ class ScaleResidual(nn.Module):
     def forward(self, residual: torch.Tensor, x: torch.Tensor,
                 gate: torch.Tensor) -> torch.Tensor:
         """Apply gated residual connection."""
-        # logger.info("x.shape: %s", x.shape)
-        # if isinstance(gate, torch.Tensor):
-            # logger.info("gate.shape: %s", gate.shape)
-
-        num_frames = gate.shape[1]
-        frame_seqlen = x.shape[1] // num_frames
-        return residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)
+        # x.shape: [batch_size, seq_len, inner_dim]
+        if gate.dim() == 4:
+            # gate.shape: [batch_size, num_frames, 1, inner_dim]
+            num_frames = gate.shape[1]
+            frame_seqlen = x.shape[1] // num_frames
+            return residual + (x.unflatten(
+                dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)
+        else:
+            # gate.shape: [batch_size, 1, inner_dim]
+            return residual + x * gate
 
 
 # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py
@@ -168,7 +168,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
             raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
     def forward(self, residual: torch.Tensor, x: torch.Tensor,
-                gate: torch.Tensor, shift: torch.Tensor,
+                gate: torch.Tensor | int, shift: torch.Tensor,
                 scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Apply gated residual connection, followed by layernorm and 
@@ -180,36 +180,41 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
             - residual value (value after residual connection 
               but before normalization)
         """
+        # x.shape: [batch_size, seq_len, inner_dim]
         # Apply residual connection with gating
-        # logger.info("x.shape: %s", x.shape)
-        if isinstance(gate, int): 
+        if isinstance(gate, int):
             # used by cross-attention, should be 1
             assert gate == 1
-            residual_output = residual + x * gate
+            residual_output = residual + x
         elif isinstance(gate, torch.Tensor):
-            # logger.info("gate.shape: %s", gate.shape)
-            if gate.dim() == 3:
-                # used by bidirectional self attention
-                residual_output = residual + x * gate
-            else:
-                assert gate.dim() == 4 
+            if gate.dim() == 4:
+                # gate.shape: [batch_size, num_frames, 1, inner_dim]
                 num_frames = gate.shape[1]
                 frame_seqlen = x.shape[1] // num_frames
-                residual_output = residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)
-                # residual_output = residual + x * gate
+                residual_output = residual + (
+                    x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
+                    gate).flatten(1, 2)
+            else:
+                # used by bidirectional self attention
+                # gate.shape: [batch_size, 1, inner_dim]
+                residual_output = residual + x * gate
         else:
             raise ValueError(f"Gate type {type(gate)} not supported")
-        # logger.info("residual_output.shape: %s", residual_output.shape)
+        # residual_output.shape: [batch_size, seq_len, inner_dim]
 
         # Apply normalization
         normalized = self.norm(residual_output)
         # Apply scale and shift
         if isinstance(scale, torch.Tensor) and scale.dim() == 4:
+            # scale.shape: [batch_size, num_frames, 1, inner_dim]
+            # shift.shape: [batch_size, num_frames, 1, inner_dim]
             num_frames = scale.shape[1]
             frame_seqlen = normalized.shape[1] // num_frames
-            modulated = (normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2)
+            modulated = (
+                normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
+                (1 + scale) + shift).flatten(1, 2)
         else:
-            modulated = normalized * (1.0 + scale) + shift
+            modulated = normalized * (1 + scale) + shift
         return modulated, residual_output
 
 
@@ -251,16 +256,24 @@ class LayerNormScaleShift(nn.Module):
     def forward(self, x: torch.Tensor, shift: torch.Tensor,
                 scale: torch.Tensor) -> torch.Tensor:
         """Apply ln followed by scale and shift in a single fused operation."""
+        # x.shape: [batch_size, seq_len, inner_dim]
         normalized = self.norm(x)
+        if self.compute_dtype == torch.float32:
+            normalized = normalized.float()
+
         if scale.dim() == 4:
+            # scale.shape: [batch_size, num_frames, 1, inner_dim]
             num_frames = scale.shape[1]
             frame_seqlen = normalized.shape[1] // num_frames
-            if self.compute_dtype == torch.float32:
-                return (normalized.float().unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2).to(x.dtype)
-            else:
-                return (normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2)
+            output = (
+                normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
+                (1 + scale) + shift).flatten(1, 2)
         else:
-            if self.compute_dtype == torch.float32:
-                return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
-            else:
-                return normalized * (1.0 + scale) + shift
+            # scale.shape: [batch_size, 1, inner_dim]
+            # shift.shape: [batch_size, 1, inner_dim]
+            output = normalized * (1 + scale) + shift
+
+        if self.compute_dtype == torch.float32:
+            output = output.to(x.dtype)
+
+        return output

fastvideo/layers/lora/linear.py

@@ -63,7 +63,7 @@ class BaseLayerWithLoRA(nn.Module):
                             device=self.base_layer.weight.device,
                             dtype=self.base_layer.weight.dtype))
             torch.nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
-            torch.nn.init.kaiming_uniform_(self.lora_B, a=math.sqrt(5))
+            torch.nn.init.zeros_(self.lora_B)
         else:
             self.lora_A = None
             self.lora_B = None

fastvideo/models/dits/causal_wanvideo.py

@@ -147,6 +147,9 @@ class CausalWanSelfAttention(nn.Module):
                 # Assign new keys/values directly up to current_end
                 local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
                 local_start_index = local_end_index - num_new_tokens
+                # kv_cache["k"] = kv_cache["k"].detach()
+                # kv_cache["v"] = kv_cache["v"].detach()
+                # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None)
                 kv_cache["k"][:, local_start_index:local_end_index] = roped_key
                 kv_cache["v"][:, local_start_index:local_end_index] = v
             x = self.attn(
@@ -176,7 +179,7 @@ class CausalWanTransformerBlock(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
         self.to_q = ReplicatedLinear(dim, dim, bias=True)
         self.to_k = ReplicatedLinear(dim, dim, bias=True)
         self.to_v = ReplicatedLinear(dim, dim, bias=True)
@@ -209,8 +212,7 @@ class CausalWanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=True,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            dtype=torch.float32)
 
         # 2. Cross-attention
         # Only T2V for now
@@ -223,8 +225,7 @@ class CausalWanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=False,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -244,40 +245,39 @@ class CausalWanTransformerBlock(nn.Module):
         current_start: int = 0,
         cache_start: int | None = None,
     ) -> torch.Tensor:
-        # logger.info("temb.shape: %s", temb.shape)
-        num_frames = temb.shape[1]
-        # logger.info("first hidden_states.shape: %s", hidden_states.shape)
-        # logger.info("num_frames: %s", num_frames)
+        # hidden_states.shape: [batch_size, seq_length, inner_dim]
+        # temb.shape: [batch_size, num_frames, 6, inner_dim]
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
-        frame_seqlen = hidden_states.shape[1] // temb.shape[1]
-        # logger.info("frame_seqlen: %s", frame_seqlen)
+        num_frames = temb.shape[1]
+        frame_seqlen = hidden_states.shape[1] // num_frames    
         bs, seq_length, _ = hidden_states.shape
         orig_dtype = hidden_states.dtype
         # assert orig_dtype != torch.float32
-        e = self.scale_shift_table + temb.float()
-        # logger.info("e.shape: %s", e.shape)
+        e = self.scale_shift_table + temb
+        # e.shape: [batch_size, num_frames, 6, inner_dim]
+        assert e.shape == (bs, num_frames, 6, self.hidden_dim)
         shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
             6, dim=2)
-        assert shift_msa.dtype == torch.float32
+        # *_msa.shape: [batch_size, num_frames, 1, inner_dim]
+        # assert shift_msa.dtype == torch.float32
+
+        # logger.info("temb sum: %s, dtype: %s", temb.float().sum().item(), temb.dtype)
+        # logger.info("scale_msa sum: %s, dtype: %s", scale_msa.float().sum().item(), scale_msa.dtype)
+        # logger.info("shift_msa sum: %s, dtype: %s", shift_msa.float().sum().item(), shift_msa.dtype)
 
         # 1. Self-attention
-        # print(f"hidden_states: {hidden_states.shape}")
-        # print(f"hidden_states: {scale_msa.shape}")
-        # print(f"hidden_states: {shift_msa.shape}")
-        # norm_hidden_states = (self.norm1(hidden_states.float()) *
-        #                       (1 + scale_msa) + shift_msa).to(orig_dtype)
-        norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
-                              (1 + scale_msa) + shift_msa).flatten(1, 2).to(orig_dtype)
-        # logger.info("1 norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
+                        (1 + scale_msa) + shift_msa).flatten(1, 2)
+        # logger.info("norm_hidden_states sum: %s, shape: %s", norm_hidden_states.float().sum().item(), norm_hidden_states.shape)
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
 
         if self.norm_q is not None:
-            query = self.norm_q(query)
+            query = self.norm_q.forward_native(query)
         if self.norm_k is not None:
-            key = self.norm_k(key)
+            key = self.norm_k.forward_native(key)
 
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
@@ -291,10 +291,6 @@ class CausalWanTransformerBlock(nn.Module):
         null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
         norm_hidden_states, hidden_states = self.self_attn_residual_norm(
             hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        # logger.info("after self_attn_residual_norm norm_hidden_states.shape: %s", norm_hidden_states.shape)
-        # logger.info("after self_attn_residual_norm hidden_states.shape: %s", hidden_states.shape)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
 
         # 2. Cross-attention
         attn_output = self.attn2(norm_hidden_states,
@@ -303,17 +299,10 @@ class CausalWanTransformerBlock(nn.Module):
                                  crossattn_cache=crossattn_cache)
         norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
             hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        # logger.info("after cross_attn_residual_norm norm_hidden_states.shape: %s", norm_hidden_states.shape)
-        # logger.info("after cross_attn_residual_norm hidden_states.shape: %s", hidden_states.shape)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
         hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
-        # logger.info("after mlp_residual norm_hidden_states.shape: %s", norm_hidden_states.shape)
-        # logger.info("after mlp_residual hidden_states.shape: %s", hidden_states.shape)
-        hidden_states = hidden_states.to(orig_dtype)
 
         return hidden_states
 
@@ -376,12 +365,9 @@ class CausalWanTransformer3DModel(BaseDiT):
                                             norm_type="layer",
                                             eps=config.eps,
                                             elementwise_affine=False,
-                                            dtype=torch.float32,
-                                            compute_dtype=torch.float32)
-        # Debug: Log configuration values
-        proj_out_dim = config.out_channels * math.prod(config.patch_size)
-        
-        self.proj_out = nn.Linear(inner_dim, proj_out_dim)
+                                            dtype=torch.float32)
+        self.proj_out = nn.Linear(
+            inner_dim, config.out_channels * math.prod(config.patch_size))
         self.scale_shift_table = nn.Parameter(
             torch.randn(1, 2, inner_dim) / inner_dim**0.5)
 
@@ -389,7 +375,7 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         # Causal-specific
         self.block_mask = None
-        self.num_frame_per_block = 1
+        self.num_frame_per_block = 3
         self.independent_first_frame = False
 
         self.__post_init__()
@@ -470,7 +456,6 @@ class CausalWanTransformer3DModel(BaseDiT):
         This function will be run for num_frame times.
         Process the latent frames one by one (1560 tokens each)
         """
-        # logger.info("forward inference hidden_states.shape: %s", hidden_states.shape)
 
         orig_dtype = hidden_states.dtype
         if not isinstance(encoder_hidden_states, torch.Tensor):
@@ -502,15 +487,18 @@ class CausalWanTransformer3DModel(BaseDiT):
         )
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos.float(),
-                     freqs_sin.float()) if freqs_cos is not None else None
+        freqs_cis = (freqs_cos,
+                     freqs_sin) if freqs_cos is not None else None
 
         hidden_states = self.patch_embedding(hidden_states)
+        grid_sizes = torch.stack(
+            [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)])
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
-        # logger.info("forward inference flattened and transposed hidden_states.shape: %s", hidden_states.shape)
+
+        encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
 
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
+                        timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
         timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
 
         if encoder_hidden_states_image is not None:
@@ -549,22 +537,15 @@ class CausalWanTransformer3DModel(BaseDiT):
                                         **causal_kwargs)
 
         # 5. Output norm, projection & unpatchify
-        # logger.info("===== INFERENCE 5. Output norm, projection & unpatchify")
-        # logger.info("hidden_states.shape: %s", hidden_states.shape)
-        # logger.info("temb.shape: %s", temb.shape)
-        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
-                                                                          dim=1)
+        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
+        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
+                                                                    dim=2)
         hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
 
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
-        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
-        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+        output = self.unpatchify(hidden_states, grid_sizes)
 
-        return output
+        return torch.stack(output)
 
     def _forward_train(self,
                 hidden_states: torch.Tensor,
@@ -575,8 +556,6 @@ class CausalWanTransformer3DModel(BaseDiT):
                 start_frame: int = 0,
                 **kwargs) -> torch.Tensor:
 
-        # logger.info("===== forward train hidden_states.shape: %s", hidden_states.shape)
-        # logger.info("===== forward train timestep.shape: %s", timestep.shape)
         orig_dtype = hidden_states.dtype
         if not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
@@ -607,8 +586,8 @@ class CausalWanTransformer3DModel(BaseDiT):
         )
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos.float(),
-                     freqs_sin.float()) if freqs_cos is not None else None
+        freqs_cis = (freqs_cos,
+                     freqs_sin) if freqs_cos is not None else None
 
         # Construct blockwise causal attn mask
         if self.block_mask is None:
@@ -621,14 +600,14 @@ class CausalWanTransformer3DModel(BaseDiT):
             )
 
         hidden_states = self.patch_embedding(hidden_states)
+        grid_sizes = torch.stack(
+            [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)])
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
-        # logger.info("forward train flattened and transposed hidden_states.shape: %s", hidden_states.shape)
+
+        encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
 
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
-        # logger.info("forward train timestep_proj.shape: %s", timestep_proj.shape)
-        # logger.info("forward train timestep.shape: %s", timestep.shape)
-        # logger.info("forward train temb.shape: %s", temb.shape)
+                        timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
         timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
 
         if encoder_hidden_states_image is not None:
@@ -649,44 +628,21 @@ class CausalWanTransformer3DModel(BaseDiT):
                     timestep_proj, freqs_cis,
                     block_mask=self.block_mask)
         else:
-            for block_index, block in enumerate(self.blocks):
-                # logger.info("===== TRAIN block %d", block_index)
-                # logger.info("hidden_states.shape: %s", hidden_states.shape)
-                # logger.info("encoder_hidden_states.shape: %s", encoder_hidden_states.shape)
-                # logger.info("timestep_proj.shape: %s", timestep_proj.shape)
-                # logger.info("freqs_cis.shape: %s", freqs_cis.shape)
-                # logger.info("block_mask.shape: %s", self.block_mask.shape)
+            for block in self.blocks:
                 hidden_states = block(hidden_states, encoder_hidden_states,
                                         timestep_proj, freqs_cis,
                                         block_mask=self.block_mask)
 
         # 5. Output norm, projection & unpatchify
-        # logger.info("===== TRAIN 5. Output norm, projection & unpatchify")
-        # logger.info("hidden_states.shape: %s", hidden_states.shape)
-        # logger.info("temb.shape: %s", temb.shape)
-        # shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
         temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
-        # logger.info("WTFWTF train temb.shape: %s", temb.shape)
-        # logger.info("WTFWTF train self.scale_shift_table.shape: %s", self.scale_shift_table.shape)
         shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
-                                                                          dim=2)
-        # logger.info("DEBUG scale.shape: %s", scale.shape)
-        # logger.info("DEBUG shift.shape: %s", shift.shape)
+                                                                    dim=2)
         hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
-        # logger.info("DEBUG after proj_out hidden_states.shape: %s", hidden_states.shape)
-        # logger.info(f"DEBUG reshape dimensions: batch_size={batch_size}, post_patch_num_frames={post_patch_num_frames}")
-        # logger.info(f"DEBUG reshape dimensions: post_patch_height={post_patch_height}, post_patch_width={post_patch_width}")
-        # logger.info(f"DEBUG patch dimensions: p_t={p_t}, p_h={p_h}, p_w={p_w}")
 
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
-        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
-        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+        output = self.unpatchify(hidden_states, grid_sizes)
 
-        return output
+        return torch.stack(output)
 
     def forward(
         self,
@@ -697,3 +653,30 @@ class CausalWanTransformer3DModel(BaseDiT):
             return self._forward_inference(*args, **kwargs)
         else:
             return self._forward_train(*args, **kwargs)
+
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+
+
+        Returns:
+            Tensor:
+                Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_channels
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = u.permute(6, 0, 3, 1, 4, 2, 5)
+            # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous())
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out

fastvideo/models/dits/wanvideo.py

@@ -1,3 +1,5 @@
+import torch
+import torch.nn as nn
 # SPDX-License-Identifier: Apache-2.0
 
 import math
@@ -37,16 +39,14 @@ class WanImageEmbedding(torch.nn.Module):
     def __init__(self, in_features: int, out_features: int):
         super().__init__()
 
-        self.norm1 = FP32LayerNorm(in_features)
+        self.norm1 = nn.LayerNorm(in_features)
         self.ff = MLP(in_features, in_features, out_features, act_type="gelu")
-        self.norm2 = FP32LayerNorm(out_features)
+        self.norm2 = nn.LayerNorm(out_features)
 
-    def forward(self,
-                encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
-        dtype = encoder_hidden_states_image.dtype
+    def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
         hidden_states = self.norm1(encoder_hidden_states_image)
         hidden_states = self.ff(hidden_states)
-        hidden_states = self.norm2(hidden_states).to(dtype)
+        hidden_states = self.norm2(hidden_states)
         return hidden_states
 
 
@@ -62,7 +62,7 @@ class WanTimeTextImageEmbedding(nn.Module):
         super().__init__()
 
         self.time_embedder = TimestepEmbedder(
-            dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
+            dim, frequency_embedding_size=time_freq_dim, act_layer="silu", freq_dtype=torch.float64)
         self.time_modulation = ModulateProjection(dim,
                                                   factor=6,
                                                   act_layer="silu")
@@ -156,12 +156,12 @@ class WanT2VCrossAttention(WanSelfAttention):
         b, n, d = x.size(0), self.num_heads, self.head_dim
 
         # compute query, key, value
-        q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
+        q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d)
 
         if crossattn_cache is not None:
             if not crossattn_cache["is_init"]:
                 crossattn_cache["is_init"] = True
-                k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
+                k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
                 v = self.to_v(context)[0].view(b, -1, n, d)
                 crossattn_cache["k"] = k
                 crossattn_cache["v"] = v
@@ -169,7 +169,7 @@ class WanT2VCrossAttention(WanSelfAttention):
                 k = crossattn_cache["k"]
                 v = crossattn_cache["v"]
         else:
-            k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
+            k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
             v = self.to_v(context)[0].view(b, -1, n, d)
 
         # compute attention
@@ -213,10 +213,10 @@ class WanI2VCrossAttention(WanSelfAttention):
         b, n, d = x.size(0), self.num_heads, self.head_dim
 
         # compute query, key, value
-        q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
-        k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
+        q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d)
+        k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
         v = self.to_v(context)[0].view(b, -1, n, d)
-        k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(
+        k_img = self.norm_added_k.forward_native(self.add_k_proj(context_img)[0]).view(
             b, -1, n, d)
         v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
         img_x = self.attn(q, k_img, v_img)
@@ -247,7 +247,7 @@ class WanTransformerBlock(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
         self.to_q = ReplicatedLinear(dim, dim, bias=True)
         self.to_k = ReplicatedLinear(dim, dim, bias=True)
         self.to_v = ReplicatedLinear(dim, dim, bias=True)
@@ -278,29 +278,29 @@ class WanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=True,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            dtype=torch.float32)
 
         # 2. Cross-attention
         if added_kv_proj_dim is not None:
             # I2V
-            self.attn2 = WanI2VCrossAttention(dim,
-                                              num_heads,
-                                              qk_norm=qk_norm,
+            self.attn2 = WanI2VCrossAttention(dim, 
+                                              num_heads, 
+                                              qk_norm=qk_norm, 
                                               eps=eps)
+
         else:
             # T2V
-            self.attn2 = WanT2VCrossAttention(dim,
-                                              num_heads,
-                                              qk_norm=qk_norm,
+            self.attn2 = WanT2VCrossAttention(dim, 
+                                              num_heads, 
+                                              qk_norm=qk_norm, 
                                               eps=eps)
+
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
-            dim,
-            norm_type="layer",
-            eps=eps,
-            elementwise_affine=False,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            dim, 
+            norm_type="layer", 
+            eps=eps, 
+            elementwise_affine=False, 
+            dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -319,12 +319,11 @@ class WanTransformerBlock(nn.Module):
             hidden_states = hidden_states.squeeze(1)
         bs, seq_length, _ = hidden_states.shape
         orig_dtype = hidden_states.dtype
-        # assert orig_dtype != torch.float32
 
         if temb.dim() == 4:
             # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
             shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table.unsqueeze(0) + temb.float()
+                self.scale_shift_table.unsqueeze(0) + temb
             ).chunk(6, dim=2)
             # batch_size, seq_len, 1, inner_dim
             shift_msa = shift_msa.squeeze(2)
@@ -335,22 +334,20 @@ class WanTransformerBlock(nn.Module):
             c_gate_msa = c_gate_msa.squeeze(2)
         else:
             # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
-            e = self.scale_shift_table + temb.float()
+            e = self.scale_shift_table + temb
             shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
                 6, dim=1)
-        assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()) *
-                              (1 + scale_msa) + shift_msa).to(orig_dtype)
+        norm_hidden_states = self.norm1(hidden_states) * (1 + scale_msa) + shift_msa
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
 
         if self.norm_q is not None:
-            query = self.norm_q(query)
+            query = self.norm_q.forward_native(query)
         if self.norm_k is not None:
-            key = self.norm_k(key)
+            key = self.norm_k.forward_native(key)
 
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
@@ -370,26 +367,20 @@ class WanTransformerBlock(nn.Module):
         null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
         norm_hidden_states, hidden_states = self.self_attn_residual_norm(
             hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
 
         # 2. Cross-attention
-        attn_output = self.attn2(norm_hidden_states,
-                                 context=encoder_hidden_states,
+        attn_output = self.attn2(norm_hidden_states, 
+                                 context=encoder_hidden_states, 
                                  context_lens=None)
         norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
             hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
         hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
-        hidden_states = hidden_states.to(orig_dtype)
 
         return hidden_states
 
-
 class WanTransformerBlock_VSA(nn.Module):
 
     def __init__(self,
@@ -406,7 +397,7 @@ class WanTransformerBlock_VSA(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
         self.to_q = ReplicatedLinear(dim, dim, bias=True)
         self.to_k = ReplicatedLinear(dim, dim, bias=True)
         self.to_v = ReplicatedLinear(dim, dim, bias=True)
@@ -438,8 +429,7 @@ class WanTransformerBlock_VSA(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=True,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            dtype=torch.float32)
 
         # 2. Cross-attention
         if added_kv_proj_dim is not None:
@@ -459,8 +449,7 @@ class WanTransformerBlock_VSA(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=False,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -480,23 +469,22 @@ class WanTransformerBlock_VSA(nn.Module):
         bs, seq_length, _ = hidden_states.shape
         orig_dtype = hidden_states.dtype
         # assert orig_dtype != torch.float32
-        e = self.scale_shift_table + temb.float()
+        e = self.scale_shift_table + temb
         shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
             6, dim=1)
-        assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()) *
-                              (1 + scale_msa) + shift_msa).to(orig_dtype)
+        norm_hidden_states = (self.norm1(hidden_states) *
+                              (1 + scale_msa) + shift_msa)
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
         gate_compress, _ = self.to_gate_compress(norm_hidden_states)
 
         if self.norm_q is not None:
-            query = self.norm_q(query)
+            query = self.norm_q.forward_native(query)
         if self.norm_k is not None:
-            key = self.norm_k(key)
+            key = self.norm_k.forward_native(key)
 
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
@@ -521,8 +509,6 @@ class WanTransformerBlock_VSA(nn.Module):
         null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
         norm_hidden_states, hidden_states = self.self_attn_residual_norm(
             hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
 
         # 2. Cross-attention
         attn_output = self.attn2(norm_hidden_states,
@@ -530,17 +516,15 @@ class WanTransformerBlock_VSA(nn.Module):
                                  context_lens=None)
         norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
             hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
         hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
-        hidden_states = hidden_states.to(orig_dtype)
 
         return hidden_states
 
 
+
 class WanTransformer3DModel(CachableDiT):
     _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
     _compile_conditions = WanVideoConfig()._compile_conditions
@@ -598,8 +582,7 @@ class WanTransformer3DModel(CachableDiT):
                                             norm_type="layer",
                                             eps=config.eps,
                                             elementwise_affine=False,
-                                            dtype=torch.float32,
-                                            compute_dtype=torch.float32)
+                                            dtype=torch.float32)
         self.proj_out = nn.Linear(
             inner_dim, config.out_channels * math.prod(config.patch_size))
         self.scale_shift_table = nn.Parameter(
@@ -659,10 +642,12 @@ class WanTransformer3DModel(CachableDiT):
             rope_theta=10000)
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos.float(),
-                     freqs_sin.float()) if freqs_cos is not None else None
+        freqs_cis = (freqs_cos,
+                     freqs_sin) if freqs_cos is not None else None
 
         hidden_states = self.patch_embedding(hidden_states)
+        grid_sizes = torch.stack(
+            [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)])
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
         # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
@@ -672,6 +657,8 @@ class WanTransformer3DModel(CachableDiT):
         else:
             ts_seq_len = None
 
+        encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
+
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
             timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
         if ts_seq_len is not None:
@@ -728,14 +715,35 @@ class WanTransformer3DModel(CachableDiT):
         hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
 
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
-        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
-        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+        output = self.unpatchify(hidden_states, grid_sizes)
 
-        return output
+        return torch.stack(output)
+    
+    def unpatchify(self, x, grid_sizes):
+        r"""
+
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+
+
+        Returns:
+            Tensor:
+                Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_channels
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = u.permute(6, 0, 3, 1, 4, 2, 5)
+            # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous())
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out
 
     def maybe_cache_states(self, hidden_states: torch.Tensor,
                            original_hidden_states: torch.Tensor) -> None:
@@ -827,5 +835,4 @@ class WanTransformer3DModel(CachableDiT):
         if self.is_even:
             return hidden_states + self.previous_residual_even
         else:
-            return hidden_states + self.previous_residual_odd
-        
+            return hidden_states + self.previous_residual_odd

fastvideo/models/loader/component_loader.py

@@ -415,10 +415,6 @@ class TransformerLoader(ComponentLoader):
             raise ValueError(
                 "Model config does not contain a _class_name attribute. "
                 "Only diffusers format is supported.")
-        logger.info("transformer cls_name: %s", cls_name)
-        if fastvideo_args.override_transformer_cls_name is not None:
-            cls_name = fastvideo_args.override_transformer_cls_name
-            logger.info("Overriding transformer cls_name to %s", cls_name)
 
         fastvideo_args.model_paths["transformer"] = model_path
 
@@ -434,6 +430,16 @@ class TransformerLoader(ComponentLoader):
         if not safetensors_list:
             raise ValueError(f"No safetensors files found in {model_path}")
 
+        # Check if we should use custom initialization weights
+        custom_weights_path = getattr(fastvideo_args, 'init_weights_from_safetensors', None)
+        use_custom_weights = (custom_weights_path and os.path.exists(custom_weights_path) and 
+                            fastvideo_args.training_mode and 
+                            not hasattr(fastvideo_args, '_loading_teacher_critic_model'))
+
+        if use_custom_weights:
+            logger.info("Using custom initialization weights from: %s", custom_weights_path)
+            safetensors_list = [custom_weights_path]
+
         logger.info("Loading model from %s safetensors files in %s",
                     len(safetensors_list), model_path)
 

fastvideo/models/registry.py

@@ -61,6 +61,9 @@ _SCHEDULERS = {
      "FlowMatchEulerDiscreteScheduler"),
     "UniPCMultistepScheduler":
     ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"),
+    "SelfForcingFlowMatchScheduler":
+    ("schedulers", "scheduling_self_forcing_flow_match",
+     "SelfForcingFlowMatchScheduler"),
 }
 
 _FAST_VIDEO_MODELS = {

fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -635,8 +635,31 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
         noise: torch.Tensor,
         timestep: torch.IntTensor,
     ) -> torch.Tensor:
+
+        """
+        Args:
+            clean_latent: the clean latent with shape [B, C, H, W],
+                where B is batch_size or batch_size * num_frames
+            noise: the noise with shape [B, C, H, W]
+            timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
+
+        Returns:
+            the corrupted latent with shape [B, C, H, W]
+        """
+        # If timestep is [bs, num_frames]
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+            assert timestep.numel() == clean_latent.shape[0]
+        elif timestep.ndim == 1:
+            # If timestep is [1]
+            if timestep.shape[0] == 1:
+                timestep = timestep.expand(clean_latent.shape[0])
+            else:
+                assert timestep.numel() == clean_latent.shape[0]
+        else:
+            raise ValueError(f"[add_noise] Invalid timestep shape: {timestep.shape}")
+        # timestep shape should be [B]
         self.sigmas = self.sigmas.to(noise.device)
-        timestep = timestep.expand(clean_latent.shape[0])
         self.timesteps = self.timesteps.to(noise.device)
         timestep_id = torch.argmin(
             (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)

fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+import torch
+
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.base import BaseScheduler
+
+logger = init_logger(__name__)
+
+
+class SelfForcingFlowMatchSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+    prev_sample: torch.FloatTensor
+
+class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
+    
+    config_name = "scheduler_config.json"
+    order = 1
+    @register_to_config
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False, training=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps, training=training)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, return_dict=False, **kwargs):
+        sigma_start = self.sigma_min + \
+            (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / \
+            (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) /
+                          num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * \
+                (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self, model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final=False, return_dict=False, **kwargs):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            sigma_ = 1 if (
+                self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        if isinstance(prev_sample, torch.Tensor | float) and not return_dict:
+            return (prev_sample, )
+        return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(self, original_samples, noise, timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B*T, C, H, W]
+            - noise: the noise with shape [B*T, C, H, W]
+            - timestep: the timestep with shape [B*T]
+        Output: the corrupted latent with shape [B*T, C, H, W]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        """
+        Input:
+            - timestep: the timestep with shape [B*T]
+        Output: the corresponding weighting [B*T]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights
+    
+    def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
+        return sample
+
+    def set_shift(self, shift: float) -> None:
+        self.shift = shift
+    

fastvideo/models/utils.py

@@ -6,9 +6,6 @@ from typing import Any
 import torch
 
 
-from fastvideo.logger import init_logger
-logger = init_logger(__name__)
-
 # TODO(PY): move it elsewhere
 def auto_attributes(init_func):
     """
@@ -148,8 +145,30 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
                              scheduler: Any) -> torch.Tensor:
     """
     Convert predicted noise to clean latent.
+
+    Args:
+    pred_noise: the predicted noise with shape [B, C, H, W]
+        where B is batch_size or batch_size * num_frames
+    noise_input_latent: the noisy latent with shape [B, C, H, W],
+    timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
+    scheduler: the scheduler
+
+    Returns:
+        the predicted video with shape [B, C, H, W]
     """
-    timestep = timestep.expand(noise_input_latent.shape[0])
+    # If timestep is [bs, num_frames]
+    if timestep.ndim == 2:
+        timestep = timestep.flatten(0, 1)
+        assert timestep.numel() == noise_input_latent.shape[0]
+    elif timestep.ndim == 1:
+        # If timestep is [1]
+        if timestep.shape[0] == 1:
+            timestep = timestep.expand(noise_input_latent.shape[0])
+        else:
+            assert timestep.numel() == noise_input_latent.shape[0]
+    else:
+        raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}")
+    # timestep shape should be [B]
     dtype = pred_noise.dtype
     device = pred_noise.device
     pred_noise = pred_noise.float().to(device)

fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py

@@ -16,8 +16,7 @@ from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
                                         CausalDMDDenosingStage,
                                         InputValidationStage,
                                         LatentPreparationStage,
-                                        TextEncodingStage,
-                                        TimestepPreparationStage)
+                                        TextEncodingStage)
 # isort: on
 
 logger = init_logger(__name__)
@@ -29,10 +28,6 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
         "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
     ]
 
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift)
-
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
         """Set up pipeline stages with proper dependency injection."""
 
@@ -48,10 +43,6 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
         self.add_stage(stage_name="conditioning_stage",
                        stage=ConditioningStage())
 
-        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"),

fastvideo/pipelines/composed_pipeline_base.py

@@ -79,7 +79,7 @@ class ComposedPipelineBase(ABC):
             for name, module in self.modules.items():
                 if not isinstance(module, torch.nn.Module):
                     continue
-                if name == "transformer":
+                if "transformer" in name:
                     module.requires_grad_(True)
                 else:
                     module.requires_grad_(False)
@@ -121,25 +121,14 @@ class ComposedPipelineBase(ABC):
                         model_path: str,
                         device: str | None = None,
                         torch_dtype: torch.dtype | None = None,
-                        pipeline_config: PipelineConfig | None = None,
+                        pipeline_config: str | PipelineConfig | None = None,
                         args: argparse.Namespace | None = None,
                         required_config_modules: list[str] | None = None,
                         loaded_modules: dict[str, torch.nn.Module]
                         | None = None,
                         **kwargs) -> "ComposedPipelineBase":
         """
-        Load a pipeline from a pretrained model. 
-        Few different patterns are supported:
-        - Only provide model_path:
-            - This will load the pipeline in inference mode.
-            - The pipeline will be initialized with the default config.
-            - The pipeline will be initialized with the default modules.
-            - The pipeline will be initialized with the default stages.
-            - The pipeline will be initialized with the default stages.
-        - override the default config using pipeline_config or args or kwargs
-        - override the default modules using loaded_modules
-        - override the pipelineconfig
-
+        Load a pipeline from a pretrained model.
         loaded_modules: Optional[Dict[str, torch.nn.Module]] = None,
         If provided, loaded_modules will be used instead of loading from config/pretrained weights.
         """
@@ -147,18 +136,9 @@ class ComposedPipelineBase(ABC):
 
             kwargs['model_path'] = model_path
             fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
-            if pipeline_config is not None:
-                fastvideo_args.pipeline_config = pipeline_config
-            if fastvideo_args.override_transformer_cls_name is not None:
-                pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers")
-                fastvideo_args.pipeline_config = pipeline_config
         else:
             assert args is not None, "args must be provided for training mode"
             fastvideo_args = TrainingArgs.from_cli_args(args)
-            if fastvideo_args.override_transformer_cls_name is not None:
-                pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers")
-                fastvideo_args.pipeline_config = pipeline_config
-                logger.info("in 2 Overriding transformer cls name to %s", fastvideo_args.override_transformer_cls_name)
             # TODO(will): fix this so that its not so ugly
             fastvideo_args.model_path = model_path
             for key, value in kwargs.items():
@@ -169,8 +149,7 @@ class ComposedPipelineBase(ABC):
             # model is loaded with the correct precision. Subsequently we will
             # use FSDP2's MixedPrecisionPolicy to set the precision for the
             # fwd, bwd, and other operations' precision.
-            fastvideo_args.pipeline_config.dit_precision = 'fp32'
-            # assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
+            assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
 
         logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
 

fastvideo/pipelines/lora_pipeline.py

@@ -7,7 +7,7 @@ import torch
 import torch.distributed as dist
 import torch.nn as nn
 from safetensors.torch import load_file
-from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.tensor import DTensor
 
 from fastvideo.distributed import get_local_torch_device
@@ -32,6 +32,7 @@ class LoRAPipeline(ComposedPipelineBase):
     cur_adapter_name: str = ""
     cur_adapter_path: str = ""
     lora_layers: dict[str, BaseLayerWithLoRA] = {}
+    lora_layers_critic: dict[str, BaseLayerWithLoRA] = {}
     fastvideo_args: FastVideoArgs | TrainingArgs
     exclude_lora_layers: list[str] = []
     device: torch.device = get_local_torch_device()
@@ -81,6 +82,17 @@ class LoRAPipeline(ComposedPipelineBase):
 
     def set_trainable(self) -> None:
 
+        def set_lora_grads(lora_layers: dict[str, BaseLayerWithLoRA],
+                           device_mesh: DeviceMesh):
+            for name, layer in lora_layers.items():
+                layer.lora_A.requires_grad_(True)
+                layer.lora_B.requires_grad_(True)
+                layer.base_layer.requires_grad_(False)
+                layer.lora_A = nn.Parameter(
+                    DTensor.from_local(layer.lora_A, device_mesh=device_mesh))
+                layer.lora_B = nn.Parameter(
+                    DTensor.from_local(layer.lora_B, device_mesh=device_mesh))
+
         is_lora_training = self.training_mode and getattr(
             self.fastvideo_args, "lora_training", False)
         if not is_lora_training:
@@ -88,18 +100,12 @@ class LoRAPipeline(ComposedPipelineBase):
             return
 
         self.modules["transformer"].requires_grad_(False)
+        if "fake_score_transformer" in self.modules:
+            self.modules["fake_score_transformer"].requires_grad_(False)
         device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1),
                                        mesh_dim_names=["fake", "replicate"])
-        for name, layer in self.lora_layers.items():
-            # Enable grads for lora weights only
-            # Must convert to DTensor for compatibility with other FSDP modules in grad calculation
-            layer.lora_A.requires_grad_(True)
-            layer.lora_B.requires_grad_(True)
-            layer.base_layer.requires_grad_(False)
-            layer.lora_A = nn.Parameter(
-                DTensor.from_local(layer.lora_A, device_mesh=device_mesh))
-            layer.lora_B = nn.Parameter(
-                DTensor.from_local(layer.lora_B, device_mesh=device_mesh))
+        set_lora_grads(self.lora_layers, device_mesh)
+        set_lora_grads(self.lora_layers_critic, device_mesh)
 
     def convert_to_lora_layers(self) -> None:
         """
@@ -131,6 +137,24 @@ class LoRAPipeline(ComposedPipelineBase):
                 converted_count += 1
         logger.info("Converted %d layers to LoRA layers", converted_count)
 
+        if "fake_score_transformer" in self.modules:
+            for name, layer in self.modules[
+                    "fake_score_transformer"].named_modules():
+                if not self.is_target_layer(name):
+                    continue
+                layer = get_lora_layer(layer,
+                                       lora_rank=self.lora_rank,
+                                       lora_alpha=self.lora_alpha,
+                                       training_mode=self.training_mode)
+                if layer is not None:
+                    self.lora_layers_critic[name] = layer
+                    replace_submodule(self.modules["fake_score_transformer"],
+                                      name, layer)
+                    converted_count += 1
+            logger.info(
+                "Converted %d layers to LoRA layers in the critic model",
+                converted_count)
+
     def set_lora_adapter(self,
                          lora_nickname: str,
                          lora_path: str | None = None):  # type: ignore
@@ -224,4 +248,4 @@ class LoRAPipeline(ComposedPipelineBase):
 
     def unmerge_lora_weights(self) -> None:
         for name, layer in self.lora_layers.items():
-            layer.unmerge_lora_weights()
+            layer.unmerge_lora_weights()

fastvideo/pipelines/pipeline_batch_info.py

@@ -147,12 +147,7 @@ class ForwardBatch:
     modules: dict[str, Any] = field(default_factory=dict)
 
     # Final output (after pipeline completion)
-    output: torch.Tensor | None = None
-    return_trajectory_latents: bool = False
-    return_trajectory_decoded: bool = False
-    trajectory_timesteps: list[int] | None = None
-    trajectory_latents: torch.Tensor | None = None
-    trajectory_decoded: list[torch.Tensor] | None = None
+    output: Any = None
 
     # Extra parameters that might be needed by specific pipeline implementations
     extra: dict[str, Any] = field(default_factory=dict)
@@ -211,10 +206,6 @@ class TrainingBatch:
     infos: list[dict[str, Any]] | None = None
     mask_lat_size: torch.Tensor | None = None
 
-    # ODE trajectory supervision
-    trajectory_latents: torch.Tensor | None = None
-    trajectory_timesteps: torch.Tensor | None = None
-
     # Transformer inputs
     noisy_model_input: torch.Tensor | None = None
     timesteps: torch.Tensor | None = None
@@ -250,5 +241,5 @@ class TrainingBatch:
 
 @dataclass
 class PreprocessBatch(ForwardBatch):
-    video_loader: list["VideoDecoder"] = field(default_factory=list)
+    video_loader: list["VideoDecoder"] | list[str] = field(default_factory=list)
     video_file_name: list[str] = field(default_factory=list)

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -1,654 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-ODE Trajectory Data Preprocessing pipeline implementation.
-
-This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline
-using the modular pipeline architecture.
-
-Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772
-"""
-
-import os
-from collections.abc import Iterator
-from typing import Any
-
-import numpy as np
-import pyarrow as pa
-import torch
-from PIL import Image
-from torch.utils.data import DataLoader
-from torchdata.stateful_dataloader import StatefulDataLoader
-from tqdm import tqdm
-
-from fastvideo.configs.sample import SamplingParam
-from fastvideo.dataset import getdataset
-from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
-    BasePreprocessPipeline)
-from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage,
-                                        ImageVAEEncodingStage,
-                                        InputValidationStage,
-                                        LatentPreparationStage,
-                                        TextEncodingStage,
-                                        TimestepPreparationStage)
-from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict
-
-logger = init_logger(__name__)
-
-
-class FlowMatchScheduler:
-
-    order = 1
-
-    def __init__(self,
-                 num_inference_steps=100,
-                 num_train_timesteps=1000,
-                 shift=3.0,
-                 sigma_max=1.0,
-                 sigma_min=0.003 / 1.002,
-                 inverse_timesteps=False,
-                 extra_one_step=False,
-                 reverse_sigmas=False):
-        self.num_train_timesteps = num_train_timesteps
-        self.shift = shift
-        self.sigma_max = sigma_max
-        self.sigma_min = sigma_min
-        self.inverse_timesteps = inverse_timesteps
-        self.extra_one_step = extra_one_step
-        self.reverse_sigmas = reverse_sigmas
-        self.set_timesteps(num_inference_steps)
-
-    def set_timesteps(self,
-                      num_inference_steps=100,
-                      denoising_strength=1.0,
-                      training=False,
-                      device=None):
-        sigma_start = self.sigma_min + \
-            (self.sigma_max - self.sigma_min) * denoising_strength
-        if self.extra_one_step:
-            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
-                                         num_inference_steps + 1)[:-1]
-        else:
-            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
-                                         num_inference_steps)
-        if self.inverse_timesteps:
-            self.sigmas = torch.flip(self.sigmas, dims=[0])
-        self.sigmas = self.shift * self.sigmas / \
-            (1 + (self.shift - 1) * self.sigmas)
-        if self.reverse_sigmas:
-            self.sigmas = 1 - self.sigmas
-        self.timesteps = self.sigmas * self.num_train_timesteps
-        if training:
-            x = self.timesteps
-            y = torch.exp(
-                -2 * ((x - num_inference_steps / 2) / num_inference_steps)**2)
-            y_shifted = y - y.min()
-            bsmntw_weighing = y_shifted * \
-                (num_inference_steps / y_shifted.sum())
-            self.linear_timesteps_weights = bsmntw_weighing
-
-    def step(self,
-             model_output,
-             timestep,
-             sample,
-             to_final=False,
-             return_dict=False,
-             **kwargs):
-        assert return_dict is False
-        assert kwargs == {}
-        self.sigmas = self.sigmas.to(model_output.device)
-        self.timesteps = self.timesteps.to(model_output.device)
-        logger.info('step timestep: %s', timestep)
-        logger.info('step timestep: %s', timestep.shape)
-        # timestep  is [num_frames]
-        # timestep_id = torch.argmin(
-        #     (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        # assert timestep.ndim == 1
-        # assert timestep.shape[0] == 1
-        timestep_id = torch.argmin((self.timesteps - timestep).abs(), dim=0)
-        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
-        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
-            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
-        else:
-            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
-        prev_sample = sample + model_output * (sigma_ - sigma)
-        return (prev_sample, )
-
-    def scale_model_input(self, sample: torch.Tensor, *args,
-                          **kwargs) -> torch.Tensor:
-        """
-        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
-        current timestep.
-
-        Args:
-            sample (`torch.Tensor`):
-                The input sample.
-
-        Returns:
-            `torch.Tensor`:
-                A scaled input sample.
-        """
-        return sample
-
-    def add_noise(self, original_samples, noise, timestep):
-        """
-        Diffusion forward corruption process.
-        Input:
-            - clean_latent: the clean latent with shape [B, C, H, W]
-            - noise: the noise with shape [B, C, H, W]
-            - timestep: the timestep with shape [B]
-        Output: the corrupted latent with shape [B, C, H, W]
-        """
-        self.sigmas = self.sigmas.to(noise.device)
-        self.timesteps = self.timesteps.to(noise.device)
-        timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
-        sample = (1 - sigma) * original_samples + sigma * noise
-        return sample.type_as(noise)
-
-    def training_target(self, sample, noise, timestep):
-        target = noise - sample
-        return target
-
-    def training_weight(self, timestep):
-        timestep_id = torch.argmin(
-            (self.timesteps - timestep.to(self.timesteps.device)).abs())
-        weights = self.linear_timesteps_weights[timestep_id]
-        return weights
-
-
-class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
-    """ODE Trajectory preprocessing pipeline implementation."""
-
-    _required_config_modules = [
-        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
-    ]
-    preprocess_dataloader: StatefulDataLoader
-    preprocess_loader_iter: Iterator[dict[str, Any]]
-
-    def get_schema_fields(self):
-        """Get the schema fields for ODE Trajectory pipeline."""
-        return [f.name for f in pyarrow_schema_ode_trajectory]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        """Set up pipeline stages with proper dependency injection."""
-        fastvideo_args.pipeline_config.flow_shift = 5
-        logger.info('WTF flow_shift: %s',
-                    fastvideo_args.pipeline_config.flow_shift)
-
-        assert fastvideo_args.pipeline_config.flow_shift == 5
-        # self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
-        # shift=fastvideo_args.pipeline_config.flow_shift)
-        self.modules["scheduler"] = FlowMatchScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift,
-            sigma_min=0.0,
-            extra_one_step=True)
-        self.modules["scheduler"].set_timesteps(num_inference_steps=48,
-                                                denoising_strength=1.0)
-        logger.info('WTF scheduler timesteps: %s',
-                    self.modules["scheduler"].timesteps)
-
-        self.add_stage(stage_name="input_validation_stage",
-                       stage=InputValidationStage())
-        self.add_stage(stage_name="prompt_encoding_stage",
-                       stage=TextEncodingStage(
-                           text_encoders=[self.get_module("text_encoder")],
-                           tokenizers=[self.get_module("tokenizer")],
-                       ))
-        self.add_stage(stage_name="vae_encoding_stage",
-                       stage=ImageVAEEncodingStage(
-                           vae=self.get_module("vae"), ))
-        self.add_stage(stage_name="timestep_preparation_stage",
-                       stage=TimestepPreparationStage(
-                           scheduler=self.get_module("scheduler")))
-        self.add_stage(stage_name="latent_preparation_stage",
-                       stage=LatentPreparationStage(
-                           scheduler=self.get_module("scheduler"),
-                           transformer=self.get_module("transformer", None)))
-        self.add_stage(stage_name="denoising_stage",
-                       stage=DenoisingStage(
-                           transformer=self.get_module("transformer"),
-                           transformer_2=self.get_module("transformer_2", None),
-                           scheduler=self.get_module("scheduler"),
-                           pipeline=self,
-                       ))
-        self.add_stage(stage_name="decoding_stage",
-                       stage=DecodingStage(vae=self.get_module("vae")))
-
-    def preprocess_video_and_text_and_trajectory(self,
-                                                 fastvideo_args: FastVideoArgs,
-                                                 args):
-
-        for batch_idx, data in enumerate(self.pbar):
-            if data is None:
-                continue
-
-            with torch.inference_mode():
-                # Filter out invalid samples (those with all zeros)
-                valid_indices = []
-                for i, pixel_values in enumerate(data["pixel_values"]):
-                    if not torch.all(
-                            pixel_values == 0):  # Check if all values are zero
-                        valid_indices.append(i)
-                self.num_processed_samples += len(valid_indices)
-
-                if not valid_indices:
-                    continue
-
-                # Create new batch with only valid samples
-                valid_data = {
-                    "pixel_values":
-                    torch.stack(
-                        [data["pixel_values"][i] for i in valid_indices]),
-                    "text": [data["text"][i] for i in valid_indices],
-                    "path": [data["path"][i] for i in valid_indices],
-                    "fps": [data["fps"][i] for i in valid_indices],
-                    "duration": [data["duration"][i] for i in valid_indices],
-                }
-
-                # VAE
-                with torch.autocast("cuda", dtype=torch.float32):
-                    latents = self.get_module("vae").encode(
-                        valid_data["pixel_values"].to(
-                            get_local_torch_device())).mean
-
-                # Get extra features if needed
-                extra_features = self.get_extra_features(
-                    valid_data, fastvideo_args)
-
-                batch_captions = valid_data["text"]
-                logger.info(f"===== batch_captions: {batch_captions}")
-                # Encode text using the standalone TextEncodingStage API
-                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
-                    batch_captions,
-                    fastvideo_args,
-                    encoder_index=[0],
-                    return_attention_mask=True,
-                )
-                prompt_embeds = prompt_embeds_list[0]
-                prompt_attention_masks = prompt_masks_list[0]
-                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
-
-                # # Get sequence lengths from attention masks (number of 1s)
-                # seq_lens = prompt_attention_mask.sum(dim=1)
-
-                # non_padded_embeds = []
-                # non_padded_masks = []
-
-                # # Process each item in the batch
-                # for i in range(prompt_embeds.size(0)):
-                #     seq_len = seq_lens[i].item()
-                #     # Slice the embeddings and masks to keep only non-padding parts
-                #     non_padded_embeds.append(prompt_embeds[i, :seq_len])
-                #     non_padded_masks.append(prompt_attention_mask[i, :seq_len])
-
-                # Update the tensors with non-padded versions
-                # prompt_embeds = non_padded_embeds
-                # prompt_attention_masks = non_padded_masks
-                # prompt_embeds = prompt_embeds
-
-                # logger.info(f"===== prompt_embeds: {prompt_embeds[0].shape}")
-                # logger.info(f"===== prompt_attention_masks: {prompt_attention_masks[0].shape}")
-
-                sampling_params = SamplingParam.from_pretrained(args.model_path)
-
-                # encode negative prompt for trajectory collection
-                if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None:
-                    negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text(
-                        sampling_params.negative_prompt,
-                        fastvideo_args,
-                        encoder_index=[0],
-                        return_attention_mask=True,
-                    )
-                    negative_prompt_embed = negative_prompt_embeds_list[0][0]
-                    negative_prompt_attention_mask = negative_prompt_masks_list[
-                        0][0]
-                else:
-                    negative_prompt_embed = None
-                    negative_prompt_attention_mask = None
-
-                trajectory_latents = []
-                trajectory_timesteps = []
-                trajectory_decoded = []
-                for i, (prompt_embed, prompt_attention_mask) in enumerate(
-                        zip(prompt_embeds, prompt_attention_masks, strict=False)):
-                    prompt_embed = prompt_embed.unsqueeze(0)
-                    prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
-                    logger.info("what")
-                    logger.info(f"===== prompt_embed: {prompt_embed.shape}")
-                    logger.info(
-                        f"===== prompt_attention_mask: {prompt_attention_mask.shape}"
-                    )
-                    # Collect the trajectory data
-                    batch = ForwardBatch(
-                        **shallow_asdict(sampling_params),
-                        # data_type="video",
-                        # seed=args.seed,
-                        # prompt=batch_captions[i],
-                        # prompt_embeds=[prompt_embed],
-                        # prompt_attention_mask=[prompt_attention_mask],
-                        # height=args.max_height,
-                        # width=args.max_width,
-                        # num_frames=81,
-                        # fps=args.train_fps,
-                        # return_trajectory_latents=True,
-                        # guidance_scale=3.0,
-                        # do_classifier_free_guidance=True,
-                    )
-                    batch.prompt_embeds = [prompt_embed]
-                    batch.prompt_attention_mask = [prompt_attention_mask]
-                    batch.negative_prompt_embeds = [negative_prompt_embed]
-                    batch.negative_attention_mask = [
-                        negative_prompt_attention_mask
-                    ]
-                    batch.return_trajectory_latents = True
-                    batch.return_trajectory_decoded = False
-                    batch.height = args.max_height
-                    batch.width = args.max_width
-                    batch.num_inference_steps = 48
-                    # batch.num_frames = 81
-                    batch.fps = args.train_fps
-                    batch.guidance_scale = 6.0
-                    batch.do_classifier_free_guidance = True
-                    # fastvideo_args.pipeline_config.ti2v_task = True
-
-                    result_batch = self.input_validation_stage(
-                        batch, fastvideo_args)
-                    # result_batch = self.prompt_encoding_stage(result_batch, fastvideo_args)
-                    # result_batch = self.vae_encoding_stage(result_batch, fastvideo_args)
-                    result_batch = self.timestep_preparation_stage(
-                        batch, fastvideo_args)
-                    result_batch = self.latent_preparation_stage(
-                        result_batch, fastvideo_args)
-                    result_batch = self.denoising_stage(result_batch,
-                                                        fastvideo_args)
-                    result_batch = self.decoding_stage(result_batch,
-                                                       fastvideo_args)
-                    # trajectory_latents = result_batch.trajectory_latents
-                    trajectory_latents.append(
-                        result_batch.trajectory_latents.cpu())
-                    trajectory_timesteps.append(
-                        result_batch.trajectory_timesteps.cpu())
-                    trajectory_decoded.append(result_batch.trajectory_decoded)
-
-            extra_features["trajectory_latents"] = trajectory_latents
-            extra_features["trajectory_timesteps"] = trajectory_timesteps
-            logger.info(
-                f"===== trajectory_latents: {trajectory_latents[0].shape}")
-            logger.info(
-                f"===== trajectory_latents len: {len(trajectory_latents)}")
-            logger.info(f"===== trajectory_timesteps: {trajectory_timesteps}")
-            logger.info(
-                f"===== trajectory_timesteps len: {len(trajectory_timesteps)}")
-
-            if batch.return_trajectory_decoded:
-                logger.info("===== SAVING TRAJECTORY DECODED")
-                for i, decoded_frames in enumerate(trajectory_decoded):
-                    for j, decoded_frame in enumerate(decoded_frames):
-                        logger.info(
-                            f"===== SAVING TRAJECTORY DECODED {i} for prompt {batch_captions[i]}"
-                        )
-                        save_decoded_latents_as_video(
-                            decoded_frame,
-                            f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
-                            args.train_fps)
-            # assert False
-            # Prepare batch data for Parquet dataset
-            batch_data = []
-
-            # Add progress bar for saving outputs
-            save_pbar = tqdm(enumerate(valid_data["path"]),
-                             desc="Saving outputs",
-                             unit="item",
-                             leave=False)
-            for idx, video_path in save_pbar:
-                # Get the corresponding latent and info using video name
-                latent = latents[idx].cpu()
-                video_name = os.path.basename(video_path).split(".")[0]
-
-                # Convert tensors to numpy arrays
-                vae_latent = latent.cpu().numpy()
-                text_embedding = prompt_embeds[idx].cpu().numpy()
-
-                # Get extra features for this sample if needed
-                sample_extra_features = {}
-                if extra_features:
-                    for key, value in extra_features.items():
-                        logger.info(f"===== key: {key}")
-                        if isinstance(value, torch.Tensor):
-                            logger.info(f"===== value: {value[idx].shape}")
-                            sample_extra_features[key] = value[idx].cpu().numpy(
-                            )
-                        else:
-                            assert isinstance(value, list)
-                            if isinstance(value[idx], torch.Tensor):
-                                logger.info(
-                                    f"===== value in list: {value[idx].shape}")
-                                sample_extra_features[key] = value[idx].cpu(
-                                ).float().numpy()
-                            else:
-                                logger.info("===== value in list: not tensor")
-                                sample_extra_features[key] = value[idx]
-                            # logger.info(f"===== value: not tensor")
-                            # sample_extra_features[key] = value[idx]
-
-                # Create record for Parquet dataset
-                record = self.create_record(
-                    video_name=video_name,
-                    vae_latent=vae_latent,
-                    text_embedding=text_embedding,
-                    valid_data=valid_data,
-                    idx=idx,
-                    extra_features=sample_extra_features)
-                batch_data.append(record)
-
-            if batch_data:
-                # Add progress bar for writing to Parquet dataset
-                write_pbar = tqdm(total=1,
-                                  desc="Writing to Parquet dataset",
-                                  unit="batch")
-                # Convert batch data to PyArrow arrays
-                arrays = []
-                for field in self.get_schema_fields():
-                    if field.endswith('_bytes'):
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data],
-                                     type=pa.binary()))
-                    elif field.endswith('_shape'):
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data],
-                                     type=pa.list_(pa.int32())))
-                    elif field in ['width', 'height', 'num_frames']:
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data],
-                                     type=pa.int32()))
-                    elif field in ['duration_sec', 'fps']:
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data],
-                                     type=pa.float32()))
-                    else:
-                        arrays.append(
-                            pa.array([record[field] for record in batch_data]))
-
-                table = pa.Table.from_arrays(arrays,
-                                             names=self.get_schema_fields())
-                write_pbar.update(1)
-                write_pbar.close()
-
-                # Store the table in a list for later processing
-                if not hasattr(self, 'all_tables'):
-                    self.all_tables = []
-                self.all_tables.append(table)
-
-                logger.info("Collected batch with %s samples", len(table))
-
-            if self.num_processed_samples >= args.flush_frequency:
-                self._flush_tables(self.num_processed_samples, args,
-                                   self.combined_parquet_dir)
-                self.num_processed_samples = 0
-                self.all_tables = []
-
-    def get_extra_features(self, valid_data: dict[str, Any],
-                           fastvideo_args: FastVideoArgs) -> dict[str, Any]:
-
-        # TODO(will): move these to cpu at some point
-        self.get_module("vae").to(get_local_torch_device())
-
-        # generator = torch.Generator(device=get_local_torch_device(), seed=42)
-        generator = torch.Generator("cpu").manual_seed(42)
-
-        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)
-        _, _, 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
-
-        unprocessed_images = []
-        pil_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))
-            pil_images.append(frame_pil)
-            # processed_img = self.get_module("image_processor")(
-            #     images=frame_pil, return_tensors="pt")
-            unprocessed_images.append(frame_pil)
-        """Get VAE features from the first frame of each video"""
-        video_conditions = []
-        for frame in unprocessed_images:
-
-            latent = self.vae_encoding_stage.encode_image(
-                frame, height, width, fastvideo_args, generator)
-            video_conditions.append(latent)
-
-        features["image_condition_latents"] = video_conditions
-        features["pil_images"] = pil_images
-        return features
-
-    def create_record(
-            self,
-            video_name: str,
-            vae_latent: np.ndarray,
-            text_embedding: np.ndarray,
-            valid_data: dict[str, Any],
-            idx: int,
-            extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
-        """Create a record for the Parquet dataset with CLIP features."""
-        record = super().create_record(video_name=video_name,
-                                       vae_latent=vae_latent,
-                                       text_embedding=text_embedding,
-                                       valid_data=valid_data,
-                                       idx=idx,
-                                       extra_features=extra_features)
-
-        if extra_features and "image_condition_latents" in extra_features:
-            image_condition_latents = extra_features["image_condition_latents"]
-            record.update({
-                "image_condition_latents_bytes":
-                image_condition_latents.tobytes(),
-                "image_condition_latents_shape":
-                list(image_condition_latents.shape),
-                "image_condition_latents_dtype":
-                str(image_condition_latents.dtype),
-            })
-        else:
-            record.update({
-                "image_condition_latents_bytes": b"",
-                "image_condition_latents_shape": [],
-                "image_condition_latents_dtype": "",
-            })
-
-        if extra_features and "trajectory_latents" in extra_features:
-            trajectory_latents = extra_features["trajectory_latents"]
-            record.update({
-                "trajectory_latents_bytes":
-                trajectory_latents.tobytes(),
-                "trajectory_latents_shape":
-                list(trajectory_latents.shape),
-                "trajectory_latents_dtype":
-                str(trajectory_latents.dtype),
-            })
-        else:
-            record.update({
-                "trajectory_latents_bytes": b"",
-                "trajectory_latents_shape": [],
-                "trajectory_latents_dtype": "",
-            })
-
-        if extra_features and "trajectory_timesteps" in extra_features:
-            trajectory_timesteps = extra_features["trajectory_timesteps"]
-            record.update({
-                "trajectory_timesteps_bytes":
-                trajectory_timesteps.tobytes(),
-                "trajectory_timesteps_shape":
-                list(trajectory_timesteps.shape),
-                "trajectory_timesteps_dtype":
-                str(trajectory_timesteps.dtype),
-            })
-        else:
-            record.update({
-                "trajectory_timesteps_bytes": b"",
-                "trajectory_timesteps_shape": [],
-                "trajectory_timesteps_dtype": "",
-            })
-
-        if extra_features and "pil_image" in extra_features:
-            pil_image = extra_features["pil_image"]
-            record.update({
-                "pil_image_bytes": pil_image.tobytes(),
-                "pil_image_shape": list(pil_image.shape),
-                "pil_image_dtype": str(pil_image.dtype),
-            })
-        else:
-            record.update({
-                "pil_image_bytes": b"",
-                "pil_image_shape": [],
-                "pil_image_dtype": "",
-            })
-
-        return record
-
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
-        if not self.post_init_called:
-            self.post_init()
-
-        self.local_rank = int(os.getenv("RANK", 0))
-        os.makedirs(args.output_dir, exist_ok=True)
-        # Create directory for combined data
-        self.combined_parquet_dir = os.path.join(args.output_dir,
-                                                 "combined_parquet_dataset")
-        os.makedirs(self.combined_parquet_dir, exist_ok=True)
-
-        # Loading dataset
-        train_dataset = getdataset(args)
-
-        self.preprocess_dataloader = DataLoader(
-            train_dataset,
-            batch_size=args.preprocess_video_batch_size,
-            num_workers=args.dataloader_num_workers,
-        )
-
-        self.preprocess_loader_iter = iter(self.preprocess_dataloader)
-
-        self.num_processed_samples = 0
-        # Add progress bar for video preprocessing
-        self.pbar = tqdm(self.preprocess_loader_iter,
-                         desc="Processing videos",
-                         unit="batch",
-                         disable=self.local_rank != 0)
-
-        # Initialize class variables for data sharing
-        self.video_data: dict[str, Any] = {}  # Store video metadata and paths
-        self.latent_data: dict[str, Any] = {}  # Store latent tensors
-        self.preprocess_video_and_text_and_trajectory(fastvideo_args, args)
-
-
-EntryClass = PreprocessPipeline_ODE_Trajectory

fastvideo/pipelines/preprocess/preprocess_stages.py

@@ -4,9 +4,11 @@ from typing import cast
 
 import numpy as np
 import torch
+import torchvision
 from einops import rearrange
 from torchvision import transforms
 
+from fastvideo.configs.configs import VideoLoaderType
 from fastvideo.dataset.transform import (CenterCropResizeVideo,
                                          TemporalRandomCrop)
 from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
@@ -61,7 +63,16 @@ class VideoTransformStage(PipelineStage):
                 else:
                     frame_indices = frame_indices[:self.num_frames]
 
-            video = batch.video_loader[i].get_frames_at(frame_indices).data
+            if fastvideo_args.preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
+                video = batch.video_loader[i].get_frames_at(frame_indices).data
+            elif fastvideo_args.preprocess_config.video_loader_type == VideoLoaderType.TORCHVISION:
+                video, _, _ = torchvision.io.read_video(batch.video_loader[i],
+                                                        output_format="TCHW")
+                video = video[frame_indices]
+            else:
+                raise ValueError(
+                    f"Invalid video loader type: {fastvideo_args.preprocess_config.video_loader_type}"
+                )
             video = self.video_transform(video)
             video_pixel_batch.append(video)
 

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -9,8 +9,6 @@ from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import (
     PreprocessPipeline_I2V)
-from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import (
-    PreprocessPipeline_ODE_Trajectory)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import (
     PreprocessPipeline_T2V)
 from fastvideo.utils import maybe_download_model
@@ -26,8 +24,7 @@ def main(args) -> None:
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
     kwargs = {
         "vae_precision": "fp32",
-        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
-        "flow_shift": 5,
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
     }
     pipeline_config.update_config_from_dict(kwargs)
     fastvideo_args = FastVideoArgs(
@@ -38,19 +35,7 @@ def main(args) -> None:
         text_encoder_cpu_offload=False,
         pipeline_config=pipeline_config,
     )
-    if args.preprocess_task == "t2v":
-        PreprocessPipeline = PreprocessPipeline_T2V
-    elif args.preprocess_task == "i2v":
-        PreprocessPipeline = PreprocessPipeline_I2V
-    elif args.preprocess_task == "ode_trajectory":
-        PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
-    else:
-        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}")
-
-    logger.info(
-        f"Preprocess task: {args.preprocess_task} using {PreprocessPipeline.__name__}"
-    )
-
+    PreprocessPipeline = PreprocessPipeline_I2V if args.preprocess_task == "i2v" else PreprocessPipeline_T2V
     pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 

fastvideo/pipelines/stages/causal_denoising.py

@@ -4,11 +4,11 @@ from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler)
 from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.denoising import DenoisingStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
 
 try:
     from fastvideo.attention.backends.sliding_tile_attn import (
@@ -36,7 +36,6 @@ class CausalDMDDenosingStage(DenoisingStage):
 
     def __init__(self, transformer, scheduler) -> None:
         super().__init__(transformer, scheduler)
-        self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
         # KV and cross-attention cache state (initialized on first forward)
         self.kv_cache1: list | None = None
         self.crossattn_cache: list | None = None
@@ -71,8 +70,16 @@ class CausalDMDDenosingStage(DenoisingStage):
         # Timesteps for DMD
         timesteps = torch.tensor(
             fastvideo_args.pipeline_config.dmd_denoising_steps,
-            dtype=torch.long,
-            device=get_local_torch_device())
+            dtype=torch.long).cpu()
+
+        if fastvideo_args.pipeline_config.warp_denoising_step:
+            logger.info("Warping timesteps...")
+            scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(),
+                                             torch.tensor([0],
+                                                          dtype=torch.float32)))
+            timesteps = scheduler_timesteps[1000 - timesteps]
+        timesteps = timesteps.to(get_local_torch_device())
+        logger.info("Using timesteps: %s", timesteps)
 
         # Image kwargs (kept empty unless caller provides compatible args)
         image_kwargs: dict = {}
@@ -228,10 +235,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                                                        dim=2)
 
                     # Prepare inputs
-                    t_expand = t_cur.expand(latent_model_input.shape[0])
-                    # t_expand = t_cur * torch.ones((latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long)
-                    # t_expand = t_expand.repeat(1, self.sliding_window_num_frames)
-
+                    t_expand = t_cur.repeat(latent_model_input.shape[0])
 
                     # Attention metadata if needed
                     if (vsa_available and self.attn_backend
@@ -265,7 +269,10 @@ class CausalDMDDenosingStage(DenoisingStage):
                                             attn_metadata=attn_metadata,
                                             forward_batch=batch):
                         # Run transformer; follow DMD stage pattern
-                        t_expanded_noise= t_cur * torch.ones((latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long)
+                        t_expanded_noise = t_cur * torch.ones(
+                            (latent_model_input.shape[0], 1),
+                            device=latent_model_input.device,
+                            dtype=torch.long)
                         pred_noise_btchw = self.transformer(
                             latent_model_input,
                             prompt_embeds,
@@ -330,7 +337,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                     set_forward_context(current_timestep=0,
                                         attn_metadata=attn_metadata,
                                         forward_batch=batch):
-                    t_expanded_context = t_context * torch.ones((context_bcthw.shape[0], 1), device=context_bcthw.device, dtype=torch.long)
+                    t_expanded_context = t_context.unsqueeze(1)
                     _ = self.transformer(
                         context_bcthw,
                         prompt_embeds,
@@ -412,3 +419,27 @@ class CausalDMDDenosingStage(DenoisingStage):
                 False,
             })
         self.crossattn_cache = crossattn_cache
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify denoising stage inputs."""
+        result = VerificationResult()
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        result.add_check("image_embeds", batch.image_embeds, V.is_list)
+        result.add_check("image_latent", batch.image_latent,
+                         V.none_or_tensor_with_dims(5))
+        result.add_check("num_inference_steps", batch.num_inference_steps,
+                         V.positive_int)
+        result.add_check("guidance_scale", batch.guidance_scale,
+                         V.positive_float)
+        result.add_check("eta", batch.eta, V.non_negative_float)
+        result.add_check("generator", batch.generator,
+                         V.generator_or_list_generators)
+        result.add_check("do_classifier_free_guidance",
+                         batch.do_classifier_free_guidance, V.bool_value)
+        result.add_check(
+            "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
+            not batch.do_classifier_free_guidance or V.list_not_empty(x))
+        return result

fastvideo/pipelines/stages/decoding.py

@@ -50,50 +50,6 @@ class DecodingStage(PipelineStage):
         result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
         return result
 
-    @torch.no_grad()
-    def decode(self, latents: torch.Tensor,
-               fastvideo_args: FastVideoArgs) -> torch.Tensor:
-        """Decode latents into pixel space."""
-        self.vae = self.vae.to(get_local_torch_device())
-        latents = latents.to(get_local_torch_device())
-
-        # Setup VAE precision
-        vae_dtype = PRECISION_TO_TYPE[
-            fastvideo_args.pipeline_config.vae_precision]
-        vae_autocast_enabled = (
-            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
-
-        if isinstance(self.vae.scaling_factor, torch.Tensor):
-            latents = latents / self.vae.scaling_factor.to(
-                latents.device, latents.dtype)
-        else:
-            latents = latents / self.vae.scaling_factor
-
-        # Apply shifting if needed
-        if (hasattr(self.vae, "shift_factor")
-                and self.vae.shift_factor is not None):
-            if isinstance(self.vae.shift_factor, torch.Tensor):
-                latents += self.vae.shift_factor.to(latents.device,
-                                                    latents.dtype)
-            else:
-                latents += self.vae.shift_factor
-
-        # Decode latents
-        with torch.autocast(device_type="cuda",
-                            dtype=vae_dtype,
-                            enabled=vae_autocast_enabled):
-            if fastvideo_args.pipeline_config.vae_tiling:
-                self.vae.enable_tiling()
-            # if fastvideo_args.vae_sp:
-            #     self.vae.enable_parallel()
-            if not vae_autocast_enabled:
-                latents = latents.to(vae_dtype)
-            image = self.vae.decode(latents)
-
-        # Normalize image to [0, 1] range
-        image = (image / 2 + 0.5).clamp(0, 1)
-        return image
-
     @torch.no_grad()
     def forward(
         self,
@@ -110,7 +66,6 @@ class DecodingStage(PipelineStage):
         Returns:
             The batch with decoded outputs.
         """
-        # load vae if not already loaded (used for memory constrained devices)
         pipeline = self.pipeline() if self.pipeline else None
         if not fastvideo_args.model_loaded["vae"]:
             loader = VAELoader()
@@ -120,31 +75,58 @@ class DecodingStage(PipelineStage):
                 pipeline.add_module("vae", self.vae)
             fastvideo_args.model_loaded["vae"] = True
 
-        if fastvideo_args.output_type == "latent":
-            frames = batch.latents
-        else:
-            frames = self.decode(batch.latents, fastvideo_args)
+        self.vae = self.vae.to(get_local_torch_device())
 
-        # decode trajectory latents if needed
-        if batch.return_trajectory_decoded:
-            batch.trajectory_decoded = []
-            logger.info(f"batch.trajectory_latents.shape: {batch.trajectory_latents.shape}")
-            assert batch.trajectory_latents is not None, "batch should have trajectory latents"
-            for idx in range(batch.trajectory_latents.shape[1]):
-                # bathc.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
-                cur_latent = batch.trajectory_latents[:, idx, :, :, :, :]
-                logger.info(f"cur_latent.shape: {cur_latent.shape}")
-                cur_timestep = batch.trajectory_timesteps[idx]
-                logger.info(
-                    f"decoding trajectory latent for timestep: {cur_timestep}")
-                decoded_frames = self.decode(cur_latent, fastvideo_args)
-                batch.trajectory_decoded.append(decoded_frames.cpu().float())
+        latents = batch.latents
+        # TODO(will): remove this once we add input/output validation for stages
+        if latents is None:
+            raise ValueError("Latents must be provided")
+
+        # Skip decoding if output type is latent
+        if fastvideo_args.output_type == "latent":
+            image = latents
+        else:
+            # Setup VAE precision
+            vae_dtype = PRECISION_TO_TYPE[
+                fastvideo_args.pipeline_config.vae_precision]
+            vae_autocast_enabled = (vae_dtype != torch.float32
+                                    ) and not fastvideo_args.disable_autocast
+
+            if isinstance(self.vae.scaling_factor, torch.Tensor):
+                latents = latents / self.vae.scaling_factor.to(
+                    latents.device, latents.dtype)
+            else:
+                latents = latents / self.vae.scaling_factor
+
+            # Apply shifting if needed
+            if (hasattr(self.vae, "shift_factor")
+                    and self.vae.shift_factor is not None):
+                if isinstance(self.vae.shift_factor, torch.Tensor):
+                    latents += self.vae.shift_factor.to(latents.device,
+                                                        latents.dtype)
+                else:
+                    latents += self.vae.shift_factor
+
+            # Decode latents
+            with torch.autocast(device_type="cuda",
+                                dtype=vae_dtype,
+                                enabled=vae_autocast_enabled):
+                if fastvideo_args.pipeline_config.vae_tiling:
+                    self.vae.enable_tiling()
+                # if fastvideo_args.vae_sp:
+                #     self.vae.enable_parallel()
+                if not vae_autocast_enabled:
+                    latents = latents.to(vae_dtype)
+                image = self.vae.decode(latents)
+
+        # Normalize image to [0, 1] range
+        image = (image / 2 + 0.5).clamp(0, 1)
 
         # Convert to CPU float32 for compatibility
-        frames = frames.cpu().float()
+        image = image.cpu().float()
 
         # Update batch with decoded image
-        batch.output = frames
+        batch.output = image
 
         # Offload models if needed
         if hasattr(self, 'maybe_free_model_hooks'):

fastvideo/pipelines/stages/denoising.py

@@ -40,6 +40,13 @@ try:
 except ImportError:
     st_attn_available = False
 
+try:
+    from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend
+    from fastvideo.utils import is_vmoba_available
+    vmoba_attn_available = is_vmoba_available()
+except ImportError:
+    vmoba_attn_available = False
+
 try:
     from fastvideo.attention.backends.video_sparse_attn import (
         VideoSparseAttentionBackend)
@@ -77,6 +84,7 @@ class DenoisingStage(PipelineStage):
             supported_attention_backends=(
                 AttentionBackendEnum.SLIDING_TILE_ATTN,
                 AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+                AttentionBackendEnum.VMOBA_ATTN,
                 AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA
             )  # hack
         )
@@ -132,12 +140,11 @@ class DenoisingStage(PipelineStage):
             latents = latents[:, :, rank_in_sp_group, :, :, :]
             batch.latents = latents
             if batch.image_latent is not None:
-                if not fastvideo_args.pipeline_config.ti2v_task and not fastvideo_args.pipeline_config.t2v_as_i2v_task:
-                    image_latent = rearrange(batch.image_latent,
-                                             "b c (n t) h w -> b c n t h w",
-                                             n=sp_world_size).contiguous()
-                    image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
-                    batch.image_latent = image_latent
+                image_latent = rearrange(batch.image_latent,
+                                         "b c (n t) h w -> b c n t h w",
+                                         n=sp_world_size).contiguous()
+                image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
+                batch.image_latent = image_latent
         # Get timesteps and calculate warmup steps
         timesteps = batch.timesteps
         # TODO(will): remove this once we add input/output validation for stages
@@ -150,7 +157,8 @@ class DenoisingStage(PipelineStage):
         # Prepare image latents and embeddings for I2V generation
         image_embeds = batch.image_embeds
         if len(image_embeds) > 0:
-            assert torch.isnan(image_embeds[0]).sum() == 0
+            assert not torch.isnan(
+                image_embeds[0]).any(), "image_embeds contains nan"
             image_embeds = [
                 image_embed.to(target_dtype) for image_embed in image_embeds
             ]
@@ -187,11 +195,13 @@ class DenoisingStage(PipelineStage):
         # Get latents and embeddings
         latents = batch.latents
         prompt_embeds = batch.prompt_embeds
-        assert torch.isnan(prompt_embeds[0]).sum() == 0
+        assert not torch.isnan(
+            prompt_embeds[0]).any(), "prompt_embeds contains nan"
         if batch.do_classifier_free_guidance:
             neg_prompt_embeds = batch.negative_prompt_embeds
             assert neg_prompt_embeds is not None
-            assert torch.isnan(neg_prompt_embeds[0]).sum() == 0
+            assert not torch.isnan(
+                neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
 
         # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
         if fastvideo_args.boundary_ratio is not None:
@@ -237,9 +247,6 @@ class DenoisingStage(PipelineStage):
                                                              patch_size[2])
             seq_len = int(math.ceil(seq_len / sp_world_size)) * sp_world_size
 
-        trajectory_timesteps: list[int] = []
-        trajectory_latents: list[torch.Tensor] = []
-
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -267,24 +274,14 @@ class DenoisingStage(PipelineStage):
 
                 # Expand latents for I2V
                 latent_model_input = latents.to(target_dtype)
-                if batch.image_latent is not None and not fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                if batch.image_latent is not None:
                     assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
                     latent_model_input = torch.cat(
                         [latent_model_input, batch.image_latent],
                         dim=1).to(target_dtype)
-                elif batch.image_latent is not None and fastvideo_args.pipeline_config.t2v_as_i2v_task:
-                    assert batch.image_latent is not None, "image latents should be provided for T2V to I2V task"
-                    if rank_in_sp_group == 0:
-                        logger.info("latent_model_input.shape: %s",
-                                    latent_model_input.shape)
-                        latent_model_input = torch.cat([
-                            batch.image_latent,
-                            latent_model_input[:, :, 1:, :, :],
-                        ],
-                                                       dim=2).to(target_dtype)
-                        logger.info("latent_model_input.shape: %s",
-                                    latent_model_input.shape)
 
+                assert not torch.isnan(
+                    latent_model_input).any(), "latent_model_input contains nan"
                 if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
                     timestep = torch.stack([t]).to(get_local_torch_device())
                     temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
@@ -297,16 +294,8 @@ class DenoisingStage(PipelineStage):
                 else:
                     t_expand = t.repeat(latent_model_input.shape[0])
 
-                assert torch.isnan(latent_model_input).sum() == 0
                 latent_model_input = self.scheduler.scale_model_input(
                     latent_model_input, t)
-                if fastvideo_args.pipeline_config.t2v_as_i2v_task:
-                    if rank_in_sp_group == 0:
-                        latent_model_input = torch.cat([
-                            batch.image_latent,
-                            latent_model_input[:, :, 1:, :, :],
-                        ],
-                                                       dim=2).to(target_dtype)
 
                 # Prepare inputs for transformer
                 guidance_expand = (
@@ -349,6 +338,31 @@ class DenoisingStage(PipelineStage):
                             assert attn_metadata is not None, "attn_metadata cannot be None"
                         else:
                             attn_metadata = None
+                    elif (vmoba_attn_available
+                          and self.attn_backend == VMOBAAttentionBackend):
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
+                        )
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
+                            )
+                            # Prepare V-MoBA parameters from config
+                            moba_params = fastvideo_args.moba_config.copy()
+                            moba_params.update({
+                                "current_timestep":
+                                i,
+                                "raw_latent_shape":
+                                batch.raw_latent_shape[2:5],
+                                "patch_size":
+                                fastvideo_args.pipeline_config.dit_config.
+                                patch_size,
+                                "device":
+                                get_local_torch_device(),
+                            })
+                            attn_metadata = self.attn_metadata_builder.build(
+                                **moba_params)
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
                     else:
                         attn_metadata = None
                     # TODO(will): finalize the interface. vLLM uses this to
@@ -413,12 +427,6 @@ class DenoisingStage(PipelineStage):
                         latents = (1. - mask2[0]) * z + mask2[0] * latents
                         # latents = latents.unsqueeze(0)
 
-                # save trajectory latents if needed
-                if batch.return_trajectory_latents:
-                    trajectory_timesteps.append(t)
-                    # trajectory_latents.append(latents.cpu())
-                    trajectory_latents.append(latents)
-
                 # Update progress bar
                 if i == len(timesteps) - 1 or (
                     (i + 1) > num_warmup_steps and
@@ -427,32 +435,8 @@ class DenoisingStage(PipelineStage):
                     progress_bar.update()
 
         # Gather results if using sequence parallelism
-        trajectory_tensor: torch.Tensor | None = None
-        if trajectory_latents:
-            trajectory_tensor = torch.stack(trajectory_latents, dim=1)
-        else:
-            trajectory_tensor = None
-
         if sp_group:
             latents = sequence_model_parallel_all_gather(latents, dim=2)
-            if batch.return_trajectory_latents:
-                # logger.info("before stack trajectory_latents.shape: %s", trajectory_latents[0].shape)
-                logger.info("after stack trajectory_latents.shape: %s", trajectory_tensor.shape)
-                trajectory_tensor = trajectory_tensor.to(
-                    get_local_torch_device())
-                trajectory_tensor = sequence_model_parallel_all_gather(
-                    trajectory_tensor, dim=3)
-
-        if trajectory_tensor is not None:
-            batch.trajectory_timesteps = torch.tensor(trajectory_timesteps).cpu()
-            batch.trajectory_latents = trajectory_tensor.cpu()
-
-        if fastvideo_args.pipeline_config.t2v_as_i2v_task:
-            latents = torch.cat([
-                batch.image_latent,
-                latents[:, :, 1:, :, :],
-            ],
-                                dim=2)
 
         # Update batch with final latents
         batch.latents = latents
@@ -791,7 +775,8 @@ class DmdDenoisingStage(DenoisingStage):
 
         video_raw_latent_shape = latents.shape
         prompt_embeds = batch.prompt_embeds
-        assert torch.isnan(prompt_embeds[0]).sum() == 0
+        assert not torch.isnan(
+            prompt_embeds[0]).any(), "prompt_embeds contains nan"
         timesteps = torch.tensor(
             fastvideo_args.pipeline_config.dmd_denoising_steps,
             dtype=torch.long,
@@ -830,7 +815,8 @@ class DmdDenoisingStage(DenoisingStage):
                         batch.image_latent.permute(0, 2, 1, 3, 4)
                     ],
                                                    dim=2).to(target_dtype)
-                assert torch.isnan(latent_model_input).sum() == 0
+                assert not torch.isnan(
+                    latent_model_input).any(), "latent_model_input contains nan"
 
                 # Prepare inputs for transformer
                 t_expand = t.repeat(latent_model_input.shape[0])

fastvideo/pipelines/stages/encoding.py

@@ -92,24 +92,6 @@ class EncodingStage(PipelineStage):
                 latents = latents.to(vae_dtype)
             latents = self.vae.encode(latents).mean
 
-        # Apply shifting if needed (reverse of decoding)
-        if (hasattr(self.vae, "shift_factor")
-                and self.vae.shift_factor is not None):
-            if isinstance(self.vae.shift_factor, torch.Tensor):
-                latents -= self.vae.shift_factor.to(latents.device,
-                                                    latents.dtype)
-            else:
-                latents -= self.vae.shift_factor
-
-        # Apply scaling factor
-        if (hasattr(self.vae, "scaling_factor")
-                and self.vae.scaling_factor is not None):
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                latents = latents * self.vae.scaling_factor.to(
-                    latents.device, latents.dtype)
-            else:
-                latents = latents * self.vae.scaling_factor
-
         # Update batch with encoded latents
         batch.latents = latents
 

fastvideo/pipelines/stages/image_encoding.py

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

fastvideo/pipelines/stages/input_validation.py

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

fastvideo/pipelines/stages/latent_preparation.py

@@ -59,10 +59,6 @@ class LatentPreparationStage(PipelineStage):
 
         # Adjust batch size for number of videos per prompt
         batch_size *= batch.num_videos_per_prompt
-        logger.info(f"===== batch_size: {batch_size}")
-        logger.info(f"===== batch.prompt: {batch.prompt}")
-        logger.info(f"===== batch.prompt_embeds: {batch.prompt_embeds}")
-        logger.info(f"===== batch.prompt_attention_mask: {batch.prompt_attention_mask}")
 
         # Get required parameters
         dtype = batch.prompt_embeds[0].dtype

fastvideo/platforms/__init__.py

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

fastvideo/platforms/cuda.py

@@ -159,6 +159,20 @@ class CudaPlatformBase(Platform):
                     str(e))
                 raise ImportError(
                     "Video Sparse Attention backend is not installed. ") from e
+        elif selected_backend == AttentionBackendEnum.VMOBA_ATTN:
+            try:
+                from csrc.attn.vmoba_attn.vmoba import (  # noqa: F401
+                    moba_attn_varlen)
+                from fastvideo.attention.backends.vmoba import (  # noqa: F401
+                    VMOBAAttentionBackend)
+                logger.info("Using Video MOBA Attention backend.")
+
+                return "fastvideo.attention.backends.vmoba.VMOBAAttentionBackend"
+            except ImportError as e:
+                logger.error(
+                    "Failed to import Video MoBA Attention backend: %s", str(e))
+                raise ImportError(
+                    "Video MoBA Attention backend is not installed. ") from e
         elif selected_backend == AttentionBackendEnum.TORCH_SDPA:
             logger.info("Using Torch SDPA backend.")
             return "fastvideo.attention.backends.sdpa.SDPABackend"

fastvideo/platforms/interface.py

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

fastvideo/tests/inference/vmoba/test_vmoba_inference.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import subprocess
+from pathlib import Path
+
+def test_inference_vmoba():
+    """Test FastVideo VMOBA_ATTN inference pipeline"""
+
+    num_gpus = "1"
+    model_base = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    output_dir = Path("outputs_video/vmoba_1.3B/")
+    moba_config = "fastvideo/configs/backend/vmoba/wan_1.3B_77_480_832.json"
+
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VMOBA_ATTN"
+
+    cmd = [
+        "fastvideo", "generate",
+        "--model-path", model_base,
+        "--sp-size", num_gpus,
+        "--tp-size", "1",
+        "--num-gpus", num_gpus,
+        "--dit-cpu-offload", "False",
+        "--vae-cpu-offload", "False",
+        "--text-encoder-cpu-offload", "True",
+        "--pin-cpu-memory", "False",
+        "--height", "480",
+        "--width", "832",
+        "--num-frames", "77",
+        "--num-inference-steps", "50",
+        "--moba-config-path", moba_config,
+        "--fps", "16",
+        "--guidance-scale", "6.0",
+        "--flow-shift", "8.0",
+        "--prompt", "A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.",
+        "--negative-prompt", (
+            "Bright tones, overexposed, static, blurred details, subtitles, style, "
+            "works, paintings, images, static, overall gray, worst quality, low quality, "
+            "JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, "
+            "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, "
+            "still picture, messy background, three legs, many people in the background, walking backwards"
+        ),
+        "--seed", "1024",
+        "--output-path", str(output_dir),
+    ]
+
+    subprocess.run(cmd, check=True)
+
+    assert output_dir.exists(), f"Output directory {output_dir} does not exist"
+
+    video_files = list(output_dir.glob("*.mp4"))
+    assert len(video_files) > 0, "No video files were generated"
+
+    for video_file in video_files:
+        assert video_file.stat().st_size > 0, f"Video file {video_file} is empty"
+
+if __name__ == "__main__":
+    test_inference_vmoba()

fastvideo/tests/modal/pr_test.py

@@ -102,10 +102,18 @@ def run_precision_tests_STA():
 def run_precision_tests_VSA():
     run_test("python csrc/attn/tests/test_vsa.py")
 
+@app.function(gpu="L40S:1", image=image, timeout=900)
+def run_precision_tests_vmoba():
+    run_test("pytest csrc/attn/vmoba_attn/tests/test_vmoba_attn.py")
+
+@app.function(gpu="L40S:1", image=image, timeout=900)
+def run_inference_tests_vmoba():
+    run_test('python fastvideo/tests/inference/vmoba/test_vmoba_inference.py')
+
 @app.function(gpu="L40S:1", image=image, timeout=3600)
 def run_inference_lora_tests():
     run_test("pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs")
 
 @app.function(gpu="L40S:2", image=image, timeout=900)
 def run_distill_dmd_tests():
-    run_test("pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
+    run_test("pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")

fastvideo/tests/ssim/README.md

@@ -4,7 +4,8 @@ The reference videos in the `*_reference_videos` directory are used as part of a
 
 run `bash update_reference_videos.sh` from inside the `fastvideo/tests/ssim/` directory after running `test_inference_similarity.py` to update reference videos. Note: make sure to update the path to the corresponding device.
 
-all reference videos are were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
+reference videos were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
+causal videos were generated on commit b318063c0a4618f1d5d99ea82ca67a06aad0d19d
 
 ## Generation Details
 
@@ -76,4 +77,4 @@ Wan2.1-I2V-14B-480P-Diffusers: {
 ### Image-to-Video Prompts
 
 1. "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
-    Image path: "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    Image path: "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"

fastvideo/tests/ssim/test_causal_similarity.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+
+import torch
+import pytest
+
+from fastvideo import VideoGenerator
+from fastvideo.logger import init_logger
+from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
+from fastvideo.worker.multiproc_executor import MultiprocExecutor
+
+logger = init_logger(__name__)
+
+device_name = torch.cuda.get_device_name()
+device_reference_folder_suffix = '_reference_videos'
+
+if "A40" in device_name:
+    device_reference_folder = "A40" + device_reference_folder_suffix
+elif "L40S" in device_name:
+    device_reference_folder = "L40S" + device_reference_folder_suffix
+
+# Base parameters from the shell script
+
+SF_WAN_T2V_PARAMS = {
+    "num_gpus": 1,
+    "model_path": "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
+    "height": 480,
+    "width": 832,
+    "num_frames": 81,
+    "num_inference_steps": 4,
+    "seed": 1024,
+    "sp_size": 1,
+    "tp_size": 1,
+}
+
+
+MODEL_TO_PARAMS = {
+    "SFWan2.1-T2V-1.3B-Diffusers": SF_WAN_T2V_PARAMS,
+}
+
+I2V_MODEL_TO_PARAMS = {
+}
+
+TEST_PROMPTS = [
+    "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
+    # "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature."
+]
+
+I2V_TEST_PROMPTS = [
+    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
+]
+
+I2V_IMAGE_PATHS = [
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg",
+]
+
+@pytest.mark.parametrize("prompt", TEST_PROMPTS)
+@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
+@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys()))
+def test_causal_similarity(prompt, ATTENTION_BACKEND, model_id):
+    """
+    Test that runs inference with different parameters and compares the output
+    to reference videos using SSIM.
+    """
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
+
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+
+    base_output_dir = os.path.join(script_dir, 'generated_videos', model_id)
+    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
+    output_video_name = f"{prompt[:100]}.mp4"
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    BASE_PARAMS = MODEL_TO_PARAMS[model_id]
+    num_inference_steps = BASE_PARAMS["num_inference_steps"]
+
+    init_kwargs = {
+        "num_gpus": BASE_PARAMS["num_gpus"],
+        "sp_size": BASE_PARAMS["sp_size"],
+        "tp_size": BASE_PARAMS["tp_size"],
+        "dit_cpu_offload": True,
+    }
+    if BASE_PARAMS.get("vae_sp"):
+        init_kwargs["vae_sp"] = True
+        init_kwargs["vae_tiling"] = True
+    #if "text-encoder-precision" in BASE_PARAMS:
+    #    init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"]
+
+    generation_kwargs = {
+        "num_inference_steps": num_inference_steps,
+        "output_path": output_dir,
+        "height": BASE_PARAMS["height"],
+        "width": BASE_PARAMS["width"],
+        "num_frames": BASE_PARAMS["num_frames"],
+        "seed": BASE_PARAMS["seed"],
+    }
+    if "neg_prompt" in BASE_PARAMS:
+        generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]
+
+    generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs)
+    generator.generate_video(prompt, **generation_kwargs)
+
+    if isinstance(generator.executor, MultiprocExecutor):
+        generator.executor.shutdown()
+
+    assert os.path.exists(
+        output_dir), f"Output video was not generated at {output_dir}"
+
+    reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND)
+
+    if not os.path.exists(reference_folder):
+        logger.error("Reference folder missing")
+        raise FileNotFoundError(
+            f"Reference video folder does not exist: {reference_folder}")
+
+    # Find the matching reference video based on the prompt
+    reference_video_name = None
+
+    for filename in os.listdir(reference_folder):
+        if filename.endswith('.mp4') and prompt[:100] in filename:
+            reference_video_name = filename
+            break
+
+    if not reference_video_name:
+        logger.error(f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}")
+        raise FileNotFoundError(f"Reference video missing")
+
+    reference_video_path = os.path.join(reference_folder, reference_video_name)
+    generated_video_path = os.path.join(output_dir, output_video_name)
+
+    logger.info(
+        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
+    )
+    ssim_values = compute_video_ssim_torchvision(reference_video_path,
+                                                 generated_video_path,
+                                                 use_ms_ssim=True)
+
+    mean_ssim = ssim_values[0]
+    logger.info(f"SSIM mean value: {mean_ssim}")
+    logger.info(f"Writing SSIM results to directory: {output_dir}")
+
+    success = write_ssim_results(output_dir, ssim_values, reference_video_path,
+                                 generated_video_path, num_inference_steps,
+                                 prompt)
+
+    if not success:
+        logger.error("Failed to write SSIM results to file")
+
+    min_acceptable_ssim = 0.98
+    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"

fastvideo/tests/ssim/test_inference_similarity.py

@@ -101,7 +101,7 @@ I2V_IMAGE_PATHS = [
 
 
 @pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS)
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
+@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
 @pytest.mark.parametrize("model_id", list(I2V_MODEL_TO_PARAMS.keys()))
 def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     """

fastvideo/tests/transformers/test_ori_causal_wanvideo.py

@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import math
+
+import numpy as np
+import pytest
+import torch
+from fastvideo.wan.modules.causal_model import CausalWanModel
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.forward_context import set_forward_context
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.models.dits.causal_wanvideo import CausalWanTransformer3DModel
+from fastvideo.utils import maybe_download_model
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_ori_causal_wan_transformer():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
+                         dit_cpu_offload=True,
+                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
+    args.device = device
+
+    loader = TransformerLoader()
+    model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision)
+
+    model1 = CausalWanModel.from_pretrained(
+        "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+    causal_state_dict = torch.load("/mnt/weka/home/hao.zhang/wei/Self-Forcing/checkpoints/self_forcing_dmd.pt")["generator_ema"]
+    new_state_dict = {}
+    for k, v in causal_state_dict.items():
+        if k.startswith("model."):
+            new_state_dict[k.replace("model.", "")] = v
+    causal_state_dict = new_state_dict
+    model1.load_state_dict(causal_state_dict)
+
+    total_params = sum(p.numel() for p in model1.parameters())
+    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
+    weight_sum_model1 = sum(
+        p.to(torch.float64).sum().item() for p in model1.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model1 = weight_sum_model1 / total_params
+    logger.info("Model 1 weight sum: %s", weight_sum_model1)
+    logger.info("Model 1 weight mean: %s", weight_mean_model1)
+
+    # Calculate weight sum for model2 (converting to float64 to avoid overflow)
+    total_params_model2 = sum(p.numel() for p in model2.parameters())
+    weight_sum_model2 = sum(
+        p.to(torch.float64).sum().item() for p in model2.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model2 = weight_sum_model2 / total_params_model2
+    logger.info("Model 2 weight sum: %s", weight_sum_model2)
+    logger.info("Model 2 weight mean: %s", weight_mean_model2)
+
+    weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2)
+    logger.info("Weight sum difference: %s", weight_sum_diff)
+    weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2)
+    logger.info("Weight mean difference: %s", weight_mean_diff)
+
+    # Set both models to eval mode
+    model1 = model1.eval()
+    model2 = model2.eval()
+
+    # Create identical inputs for both models
+    batch_size = 1
+    text_seq_len = 30
+
+    # Video latents [B, C, T, H, W]
+    hidden_states = torch.randn(batch_size,
+                                16,
+                                12,
+                                160,
+                                90,
+                                device=device,
+                                dtype=precision)
+    block_sizes = [3 for _ in range(4)]
+    timesteps = [1000, 750, 500, 250]
+
+    # Text embeddings [B, L, D] (including global token)
+    encoder_hidden_states = torch.randn(batch_size,
+                                        text_seq_len + 1,
+                                        4096,
+                                        device=device,
+                                        dtype=precision)
+
+    output1 = _causal_inference(model1, hidden_states.clone(), encoder_hidden_states.clone(), block_sizes, timesteps, precision)
+    logger.info("Finish inference for model1")
+    output2 = _causal_inference(model2, hidden_states.clone(), encoder_hidden_states.clone(), block_sizes, timesteps, precision)
+
+    # Check if outputs have the same shape
+    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
+
+    logger.info("Output 1 Sum: %s", output1.float().sum().item())
+    logger.info("Output 2 Sum: %s", output2.float().sum().item())
+
+    # Check if outputs are similar (allowing for small numerical differences)
+    max_diff = torch.max(torch.abs(output1 - output2))
+    mean_diff = torch.mean(torch.abs(output1 - output2))
+    logger.info("Max Diff: %s", max_diff.item())
+    logger.info("Mean Diff: %s", mean_diff.item())
+    assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}"
+    # mean diff
+    assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}"
+
+def _causal_inference(transformer, latents, prompt_embeds, block_sizes, timesteps, target_dtype):
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+    start_index = 0
+    pos_start_base = 0
+    frame_seq_length = latents.shape[-1] * latents.shape[-2] // (WanVideoConfig().arch_config.patch_size[-1] * WanVideoConfig().arch_config.patch_size[-2])
+    seq_len = frame_seq_length * latents.shape[2]
+    kv_cache1 = _initialize_kv_cache(transformer, batch_size=latents.shape[0],
+                                      kv_cache_size=frame_seq_length * latents.shape[2],
+                                      dtype=target_dtype,
+                                      device=latents.device)
+    crossattn_cache = _initialize_crossattn_cache(
+                transformer,
+                batch_size=latents.shape[0],
+                max_text_len=WanVideoConfig().arch_config.text_len,
+                dtype=target_dtype,
+                device=latents.device)
+    for current_num_frames, t_cur in zip(block_sizes, timesteps):
+        # logger.info(f"Current frame idx: {start_index}, Current timestep: {t_cur}")
+        # logger.info(f"k cache sum: {sum(kv_cache['k'].float().sum().item() for kv_cache in kv_cache1)}, v cache sum: {sum(kv_cache['v'].float().sum().item() for kv_cache in kv_cache1)}")
+        # logger.info(f"latents sum: {latents.float().sum().item()}, encoder_hidden_states sum: {prompt_embeds.float().sum().item()}")
+        current_latents = latents[:, :, start_index:start_index +
+                                    current_num_frames, :, :]
+
+        attn_metadata = None
+
+        with set_forward_context(current_timestep=0,
+                                attn_metadata=attn_metadata,
+                                forward_batch=forward_batch):
+            # Run transformer; follow DMD stage pattern
+            t_expanded_noise = t_cur * torch.ones(
+                (current_latents.shape[0], 1),
+                device=current_latents.device,
+                dtype=torch.long)
+            if isinstance(transformer, CausalWanModel):
+                pred_noise_btchw = transformer(
+                    x=current_latents,
+                    context=prompt_embeds,
+                    t=t_expanded_noise,
+                    seq_len=seq_len,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
+                    current_start=(pos_start_base + start_index) *
+                    frame_seq_length
+                )
+            elif isinstance(transformer, CausalWanTransformer3DModel):
+                pred_noise_btchw = transformer(
+                    current_latents,
+                    prompt_embeds,
+                    t_expanded_noise,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
+                    current_start=(pos_start_base + start_index) *
+                    frame_seq_length,
+                    start_frame=start_index
+                )
+
+        # Write back and advance
+        latents[:, :, start_index:start_index +
+                current_num_frames, :, :] = pred_noise_btchw.clone()
+
+        # Re-run with context timestep to update KV cache using clean context
+        context_noise = 0
+        t_context = torch.ones([latents.shape[0]],
+                                device=latents.device,
+                                dtype=torch.long) * int(context_noise)
+        context_bcthw = pred_noise_btchw.to(target_dtype)
+        with set_forward_context(current_timestep=0,
+                                attn_metadata=attn_metadata,
+                                forward_batch=forward_batch):
+            t_expanded_context = t_context.unsqueeze(1)
+            if isinstance(transformer, CausalWanModel):
+                _ = transformer(
+                    x=context_bcthw,
+                    context=prompt_embeds,
+                    t=t_expanded_context,
+                    seq_len=seq_len,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
+                    current_start=(pos_start_base + start_index) *
+                    frame_seq_length
+                )
+            elif isinstance(transformer, CausalWanTransformer3DModel):
+                _ = transformer(
+                    context_bcthw,
+                    prompt_embeds,
+                    t_expanded_context,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
+                    current_start=(pos_start_base + start_index) *
+                    frame_seq_length,
+                    start_frame=start_index
+                )
+        start_index += current_num_frames
+
+    return latents
+
+def _initialize_kv_cache(transformer, batch_size, kv_cache_size, dtype, device) -> None:
+    """
+    Initialize a Per-GPU KV cache aligned with the Wan model assumptions.
+    """
+    kv_cache1 = []
+    if isinstance(transformer, CausalWanModel):
+        num_attention_heads = transformer.num_heads
+        attention_head_dim = transformer.dim // transformer.num_heads
+    elif isinstance(transformer, CausalWanTransformer3DModel):
+        num_attention_heads = transformer.num_attention_heads
+        attention_head_dim = transformer.attention_head_dim
+
+    for _ in range(len(transformer.blocks)):
+        kv_cache1.append({
+            "k":
+            torch.zeros([
+                batch_size, kv_cache_size, num_attention_heads,
+                attention_head_dim
+            ],
+                        dtype=dtype,
+                        device=device),
+            "v":
+            torch.zeros([
+                batch_size, kv_cache_size, num_attention_heads,
+                attention_head_dim
+            ],
+                        dtype=dtype,
+                        device=device),
+            "global_end_index":
+            torch.tensor([0], dtype=torch.long, device=device),
+            "local_end_index":
+            torch.tensor([0], dtype=torch.long, device=device),
+        })
+
+    return kv_cache1
+
+def _initialize_crossattn_cache(transformer, batch_size, max_text_len, dtype,
+                                    device) -> None:
+    """
+    Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions.
+    """
+    crossattn_cache = []
+    if isinstance(transformer, CausalWanModel):
+        num_attention_heads = transformer.num_heads
+        attention_head_dim = transformer.dim // transformer.num_heads
+    elif isinstance(transformer, CausalWanTransformer3DModel):
+        num_attention_heads = transformer.num_attention_heads
+        attention_head_dim = transformer.attention_head_dim
+
+    for _ in range(len(transformer.blocks)):
+        crossattn_cache.append({
+            "k":
+            torch.zeros([
+                batch_size, max_text_len, num_attention_heads,
+                attention_head_dim
+            ],
+                        dtype=dtype,
+                        device=device),
+            "v":
+            torch.zeros([
+                batch_size, max_text_len, num_attention_heads,
+                attention_head_dim
+            ],
+                        dtype=dtype,
+                        device=device),
+            "is_init":
+            False,
+        })
+    return crossattn_cache

fastvideo/tests/transformers/test_ori_wanvideo.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import math
+
+import numpy as np
+import pytest
+import torch
+from fastvideo.wan.modules.model import WanModel
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.forward_context import set_forward_context
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.utils import maybe_download_model
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_ori_wan_transformer():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
+                         dit_cpu_offload=True,
+                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
+    args.device = device
+
+    loader = TransformerLoader()
+    model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision)
+
+    model1 = WanModel.from_pretrained(
+        "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+
+    total_params = sum(p.numel() for p in model1.parameters())
+    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
+    weight_sum_model1 = sum(
+        p.to(torch.float64).sum().item() for p in model1.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model1 = weight_sum_model1 / total_params
+    logger.info("Model 1 weight sum: %s", weight_sum_model1)
+    logger.info("Model 1 weight mean: %s", weight_mean_model1)
+
+    # Calculate weight sum for model2 (converting to float64 to avoid overflow)
+    total_params_model2 = sum(p.numel() for p in model2.parameters())
+    weight_sum_model2 = sum(
+        p.to(torch.float64).sum().item() for p in model2.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model2 = weight_sum_model2 / total_params_model2
+    logger.info("Model 2 weight sum: %s", weight_sum_model2)
+    logger.info("Model 2 weight mean: %s", weight_mean_model2)
+
+    weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2)
+    logger.info("Weight sum difference: %s", weight_sum_diff)
+    weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2)
+    logger.info("Weight mean difference: %s", weight_mean_diff)
+
+    # Set both models to eval mode
+    model1 = model1.eval()
+    model2 = model2.eval()
+
+    # Create identical inputs for both models
+    batch_size = 1
+    text_seq_len = 30
+    seq_len = math.ceil((160 * 90) /
+                            (2 * 2) *
+                            21)
+
+    # Video latents [B, C, T, H, W]
+    hidden_states = torch.randn(batch_size,
+                                16,
+                                21,
+                                160,
+                                90,
+                                device=device,
+                                dtype=precision)
+
+    # Text embeddings [B, L, D] (including global token)
+    encoder_hidden_states = torch.randn(batch_size,
+                                        text_seq_len + 1,
+                                        4096,
+                                        device=device,
+                                        dtype=precision)
+
+    # Timestep
+    timestep = torch.tensor([500], device=device, dtype=precision)
+
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+
+    # with torch.amp.autocast('cuda', dtype=precision):
+    output1 = model1(
+        x=hidden_states,
+        context=encoder_hidden_states,
+        t=timestep,
+        seq_len=seq_len,
+    )
+    with set_forward_context(
+            current_timestep=0,
+            attn_metadata=None,
+            forward_batch=forward_batch,
+    ):
+        output2 = model2(hidden_states=hidden_states,
+                            encoder_hidden_states=encoder_hidden_states,
+                            timestep=timestep)
+
+    # Check if outputs have the same shape
+    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
+
+    # Check if outputs are similar (allowing for small numerical differences)
+    max_diff = torch.max(torch.abs(output1 - output2))
+    mean_diff = torch.mean(torch.abs(output1 - output2))
+    logger.info("Max Diff: %s", max_diff.item())
+    logger.info("Mean Diff: %s", mean_diff.item())
+    assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}"
+    # mean diff
+    assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}"

fastvideo/tests/transformers/test_train_ori_causal_wanvideo.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import math
+
+import numpy as np
+import pytest
+import torch
+from fastvideo.wan.modules.causal_model import CausalWanModel
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.forward_context import set_forward_context
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.utils import maybe_download_model
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_train_ori_causal_wan_transformer():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
+                         dit_cpu_offload=True,
+                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
+    args.device = device
+
+    loader = TransformerLoader()
+    model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision)
+
+    model1 = CausalWanModel.from_pretrained(
+        "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+    causal_state_dict = torch.load("/mnt/weka/home/hao.zhang/wei/Self-Forcing/checkpoints/self_forcing_dmd.pt")["generator_ema"]
+    new_state_dict = {}
+    for k, v in causal_state_dict.items():
+        if k.startswith("model."):
+            new_state_dict[k.replace("model.", "")] = v
+    causal_state_dict = new_state_dict
+    model1.load_state_dict(causal_state_dict)
+
+    model1.num_frame_per_block = 3
+    model2.num_frame_per_block = 3
+
+    total_params = sum(p.numel() for p in model1.parameters())
+    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
+    weight_sum_model1 = sum(
+        p.to(torch.float64).sum().item() for p in model1.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model1 = weight_sum_model1 / total_params
+    logger.info("Model 1 weight sum: %s", weight_sum_model1)
+    logger.info("Model 1 weight mean: %s", weight_mean_model1)
+
+    # Calculate weight sum for model2 (converting to float64 to avoid overflow)
+    total_params_model2 = sum(p.numel() for p in model2.parameters())
+    weight_sum_model2 = sum(
+        p.to(torch.float64).sum().item() for p in model2.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model2 = weight_sum_model2 / total_params_model2
+    logger.info("Model 2 weight sum: %s", weight_sum_model2)
+    logger.info("Model 2 weight mean: %s", weight_mean_model2)
+
+    weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2)
+    logger.info("Weight sum difference: %s", weight_sum_diff)
+    weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2)
+    logger.info("Weight mean difference: %s", weight_mean_diff)
+
+    # Set both models to eval mode
+    model1 = model1.eval()
+    model2 = model2.eval()
+
+    # Create identical inputs for both models
+    batch_size = 1
+    text_seq_len = 30
+    seq_len = math.ceil((160 * 90) /
+                            (2 * 2) *
+                            21)
+
+    # Video latents [B, C, T, H, W]
+    hidden_states = torch.randn(batch_size,
+                                16,
+                                21,
+                                160,
+                                90,
+                                device=device,
+                                dtype=precision)
+
+    # Text embeddings [B, L, D] (including global token)
+    encoder_hidden_states = torch.randn(batch_size,
+                                        text_seq_len + 1,
+                                        4096,
+                                        device=device,
+                                        dtype=precision)
+
+    # Timestep
+    timestep = torch.randint(0, 1000, (batch_size, 21), device=device, dtype=torch.long)
+    logger.info("timestep: %s", timestep)
+
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+
+    # with torch.amp.autocast('cuda', dtype=precision):
+    output1 = model1(
+        x=hidden_states,
+        context=encoder_hidden_states,
+        t=timestep,
+        seq_len=seq_len,
+    )
+    with set_forward_context(
+            current_timestep=0,
+            attn_metadata=None,
+            forward_batch=forward_batch,
+    ):
+        output2 = model2(hidden_states=hidden_states,
+                            encoder_hidden_states=encoder_hidden_states,
+                            timestep=timestep)
+
+    # Check if outputs have the same shape
+    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
+
+    # Check if outputs are similar (allowing for small numerical differences)
+    max_diff = torch.max(torch.abs(output1 - output2))
+    mean_diff = torch.mean(torch.abs(output1 - output2))
+    logger.info("Max Diff: %s", max_diff.item())
+    logger.info("Mean Diff: %s", mean_diff.item())
+    assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}"
+    # mean diff
+    assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}"

fastvideo/training/distillation_pipeline.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import copy
 import gc
+import json
 import os
 import time
 from abc import abstractmethod
@@ -11,6 +12,7 @@ from typing import Any
 import imageio
 import numpy as np
 import torch
+import torch.distributed as dist
 import torch.nn.functional as F
 import torchvision
 from einops import rearrange
@@ -36,9 +38,11 @@ from fastvideo.training.activation_checkpoint import (
     apply_activation_checkpointing)
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
-    clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler,
-    load_distillation_checkpoint, save_distillation_checkpoint, shift_timestep)
-from fastvideo.utils import is_vsa_available, set_random_seed
+    EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
+    get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint,
+    shift_timestep, compute_density_for_timestep_sampling, get_sigmas)
+from fastvideo.utils import (is_vsa_available, maybe_download_model,
+                             set_random_seed, verify_model_config_and_directory)
 
 import wandb  # isort: skip
 
@@ -87,9 +91,30 @@ class DistillationPipeline(TrainingPipeline):
         self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
             shift=self.timestep_shift)
 
-        # self.transformer is the generator model
-        self.real_score_transformer = self.get_module("real_score_transformer")
-        self.fake_score_transformer = self.get_module("fake_score_transformer")
+        self.transformer_2 = self.get_module("transformer_2", None)
+        self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
+        
+        if training_args.real_score_model_path:
+            logger.info(
+                f"Loading real score transformer from: {training_args.real_score_model_path}"
+            )
+            self.real_score_transformer = self.load_module_from_path(
+                training_args.real_score_model_path, "transformer",
+                training_args)
+        else:
+            self.real_score_transformer = self.get_module(
+                "real_score_transformer")
+
+        if training_args.fake_score_model_path:
+            logger.info(
+                f"Loading fake score transformer from: {training_args.fake_score_model_path}"
+            )
+            self.fake_score_transformer = self.load_module_from_path(
+                training_args.fake_score_model_path, "transformer",
+                training_args)
+        else:
+            self.fake_score_transformer = self.get_module(
+                "fake_score_transformer")
 
         self.real_score_transformer.requires_grad_(False)
         self.real_score_transformer.eval()
@@ -105,6 +130,39 @@ class DistillationPipeline(TrainingPipeline):
                 self.real_score_transformer,
                 checkpointing_type=training_args.
                 enable_gradient_checkpointing_type)
+            if self.transformer_2 is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2,
+                    checkpointing_type=training_args.
+                    enable_gradient_checkpointing_type)
+
+        if self.transformer_2 is not None:
+            self.transformer_2.train()
+            self.transformer_2.requires_grad_(True)
+            params_to_optimize_2 = self.transformer_2.parameters()
+            params_to_optimize_2 = list(
+                filter(lambda p: p.requires_grad, params_to_optimize_2))
+            
+            betas_str = training_args.betas
+            betas = tuple(float(x.strip()) for x in betas_str.split(","))
+            
+            self.optimizer_2 = torch.optim.AdamW(
+                params_to_optimize_2,
+                lr=training_args.learning_rate,
+                betas=betas,
+                weight_decay=training_args.weight_decay,
+                eps=1e-8,
+            )
+            self.lr_scheduler_2 = get_scheduler(
+                training_args.lr_scheduler,
+                optimizer=self.optimizer_2,
+                num_warmup_steps=training_args.lr_warmup_steps,
+                num_training_steps=training_args.max_train_steps,
+                num_cycles=training_args.lr_num_cycles,
+                power=training_args.lr_power,
+                min_lr_ratio=training_args.min_lr_ratio,
+                last_epoch=self.init_steps - 1,
+            )
 
         # Initialize optimizers
         fake_score_params = list(
@@ -116,10 +174,13 @@ class DistillationPipeline(TrainingPipeline):
         if fake_score_lr == 0.0:
             fake_score_lr = training_args.learning_rate
 
+        betas_str = training_args.fake_score_betas
+        betas = tuple(float(x.strip()) for x in betas_str.split(","))
+
         self.fake_score_optimizer = torch.optim.AdamW(
             fake_score_params,
             lr=fake_score_lr,
-            betas=(0.9, 0.999),
+            betas=betas,
             weight_decay=training_args.weight_decay,
             eps=1e-8,
         )
@@ -147,8 +208,19 @@ class DistillationPipeline(TrainingPipeline):
             self.training_args.pipeline_config.dmd_denoising_steps,
             dtype=torch.long,
             device=get_local_torch_device())
-        logger.info("Distillation generator model to %s denoising steps",
-                    len(self.denoising_step_list))
+
+        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
+            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
+                                   torch.tensor([0],
+                                                dtype=torch.float32))).cuda()
+            self.denoising_step_list = timesteps[1000 -
+                                                 self.denoising_step_list]
+            logger.info("Warping denoising_step_list")
+
+        self.denoising_step_list = self.denoising_step_list.to(
+            get_local_torch_device())
+        logger.info("Distillation generator model to %s denoising steps: %s",
+                    len(self.denoising_step_list), self.denoising_step_list)
         self.num_train_timestep = self.noise_scheduler.num_train_timesteps
 
         self.min_timestep = int(self.training_args.min_timestep_ratio *
@@ -158,6 +230,82 @@ class DistillationPipeline(TrainingPipeline):
 
         self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
 
+        self.generator_ema = None
+        if (self.training_args.ema_decay
+                is not None) and (self.training_args.ema_decay > 0.0):
+            self.generator_ema = EMA_FSDP(self.transformer,
+                                          decay=self.training_args.ema_decay)
+            logger.info(
+                f"Initialized generator EMA with decay={self.training_args.ema_decay}"
+            )
+        else:
+            logger.info("Generator EMA disabled (ema_decay <= 0.0)")
+
+    def load_module_from_path(self, model_path: str, module_type: str,
+                              training_args: "TrainingArgs"):
+        """
+        Load a module from a specific path using the same loading logic as the pipeline.
+        
+        Args:
+            model_path: Path to the model
+            module_type: Type of module to load (e.g., "transformer")
+            training_args: Training arguments
+            
+        Returns:
+            The loaded module
+        """
+        logger.info(f"Loading {module_type} from custom path: {model_path}")
+        # Set flag to prevent custom weight loading for teacher/critic models
+        training_args._loading_teacher_critic_model = True
+
+        try:
+            from fastvideo.models.loader.component_loader import (
+                PipelineComponentLoader)
+
+            # Download the model if it's a Hugging Face model ID
+            local_model_path = maybe_download_model(model_path)
+            logger.info(f"Model downloaded/found at: {local_model_path}")
+            config = verify_model_config_and_directory(local_model_path)
+
+            if module_type not in config:
+                if hasattr(self, '_extra_config_module_map'
+                           ) and module_type in self._extra_config_module_map:
+                    extra_module = self._extra_config_module_map[module_type]
+                    if extra_module in config:
+                        module_type = extra_module
+                        logger.info(f"Using {extra_module} for {module_type}")
+                    else:
+                        raise ValueError(
+                            f"Module {module_type} not found in config at {local_model_path}"
+                        )
+                else:
+                    raise ValueError(
+                        f"Module {module_type} not found in config at {local_model_path}"
+                    )
+
+            module_info = config[module_type]
+            if module_info is None:
+                raise ValueError(
+                    f"Module {module_type} has null value in config at {local_model_path}"
+                )
+
+            transformers_or_diffusers, architecture = module_info
+            component_path = os.path.join(local_model_path, module_type)
+            module = PipelineComponentLoader.load_module(
+                module_name=module_type,
+                component_model_path=component_path,
+                transformers_or_diffusers=transformers_or_diffusers,
+                fastvideo_args=training_args,
+            )
+
+            logger.info(
+                f"Successfully loaded {module_type} from {component_path}")
+            return module
+        finally:
+            # Always clean up the flag
+            if hasattr(training_args, '_loading_teacher_critic_model'):
+                delattr(training_args, '_loading_teacher_critic_model')
+
     @abstractmethod
     def initialize_validation_pipeline(self, training_args: TrainingArgs):
         """Initialize validation pipeline - must be implemented by subclasses."""
@@ -169,11 +317,118 @@ class DistillationPipeline(TrainingPipeline):
         """Prepare training environment for distillation."""
         self.transformer.requires_grad_(True)
         self.transformer.train()
+        if self.transformer_2 is not None:
+            self.transformer_2.requires_grad_(True)
+            self.transformer_2.train()
         self.fake_score_transformer.requires_grad_(True)
         self.fake_score_transformer.train()
 
         return training_batch
 
+    def apply_ema_to_model(self, model):
+        """Apply EMA weights to the model for validation or inference."""
+        if self.generator_ema is not None:
+            with self.generator_ema.apply_to_model(model):
+                return model
+        return model
+
+    def get_ema_model_copy(self):
+        """Get a copy of the model with EMA weights applied."""
+        if self.generator_ema is not None:
+            ema_model = copy.deepcopy(self.transformer)
+            self.generator_ema.copy_to_unwrapped(ema_model)
+            return ema_model
+        return None
+
+    def is_ema_ready(self, current_step: int = None):
+        """Check if EMA is ready for use (after ema_start_step)."""
+        if current_step is None:
+            current_step = getattr(self, 'current_trainstep', 0)
+        return (self.generator_ema is not None
+                and current_step >= self.training_args.ema_start_step)
+
+    def save_ema_weights(self, output_dir: str, step: int):
+        """Save EMA weights separately for inference purposes."""
+        if self.generator_ema is None:
+            logger.warning("Cannot save EMA weights: EMA not initialized")
+            return
+
+        if not self.is_ema_ready():
+            logger.warning(
+                "Cannot save EMA weights: EMA not ready yet (step < ema_start_step)"
+            )
+            return
+
+        try:
+            ema_model = self.get_ema_model_copy()
+            if ema_model is None:
+                logger.warning("Failed to create EMA model copy")
+                return
+
+            ema_save_dir = os.path.join(output_dir, f"ema_checkpoint-{step}")
+            os.makedirs(ema_save_dir, exist_ok=True)
+
+            # save as diffusers format
+            from safetensors.torch import save_file
+
+            from fastvideo.training.training_utils import (
+                custom_to_hf_state_dict, gather_state_dict_on_cpu_rank0)
+            cpu_state = gather_state_dict_on_cpu_rank0(ema_model, device=None)
+
+            if self.global_rank == 0:
+                weight_path = os.path.join(
+                    ema_save_dir, "diffusion_pytorch_model.safetensors")
+                diffusers_state_dict = custom_to_hf_state_dict(
+                    cpu_state, ema_model.reverse_param_names_mapping)
+                save_file(diffusers_state_dict, weight_path)
+
+                config_dict = ema_model.hf_config
+                if "dtype" in config_dict:
+                    del config_dict["dtype"]
+                config_path = os.path.join(ema_save_dir, "config.json")
+                with open(config_path, "w") as f:
+                    json.dump(config_dict, f, indent=4)
+
+                logger.info(f"EMA weights saved to {weight_path}")
+
+            del ema_model
+
+        except Exception as e:
+            logger.error(f"Failed to save EMA weights: {str(e)}")
+
+    def get_ema_stats(self):
+        """Get EMA statistics for monitoring."""
+        if self.generator_ema is None:
+            return {
+                "ema_enabled": False,
+                "ema_decay": None,
+                "ema_start_step": self.training_args.ema_start_step,
+                "ema_ready": False,
+                "ema_step": self.current_trainstep,
+            }
+
+        return {
+            "ema_enabled": True,
+            "ema_decay": self.training_args.ema_decay,
+            "ema_start_step": self.training_args.ema_start_step,
+            "ema_ready": self.is_ema_ready(),
+            "ema_step": self.current_trainstep,
+        }
+
+    def reset_ema(self):
+        """Reset EMA to current model weights."""
+        if self.generator_ema is not None:
+            logger.info("Resetting EMA to current model weights")
+            self.generator_ema.update(self.transformer)
+            # Force update to current weights by setting decay to 0 temporarily
+            original_decay = self.generator_ema.decay
+            self.generator_ema.decay = 0.0
+            self.generator_ema.update(self.transformer)
+            self.generator_ema.decay = original_decay
+            logger.info("EMA reset completed")
+        else:
+            logger.warning("Cannot reset EMA: EMA not initialized")
+
     def _build_distill_input_kwargs(
             self, noise_input: torch.Tensor, timestep: torch.Tensor,
             text_dict: dict[str, torch.Tensor] | None,
@@ -220,7 +475,9 @@ class DistillationPipeline(TrainingPipeline):
         training_batch = self._build_distill_input_kwargs(
             noisy_latent, timestep, training_batch.conditional_dict,
             training_batch)
-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+        
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
+        pred_noise = current_model(**training_batch.input_kwargs).permute(
             0, 2, 1, 3, 4)
         pred_video = pred_noise_to_pred_video(
             pred_noise=pred_noise.flatten(0, 1),
@@ -262,6 +519,7 @@ class DistillationPipeline(TrainingPipeline):
         max_target_idx = len(self.denoising_step_list) - 1
         noise_latents = []
         noise_latent_index = target_timestep_idx_int - 1
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
         if max_target_idx > 0:
             # Run student model for all steps before the target timestep
             with torch.no_grad():
@@ -273,7 +531,7 @@ class DistillationPipeline(TrainingPipeline):
                     training_batch_temp = self._build_distill_input_kwargs(
                         current_noise_latents, current_timestep_tensor,
                         training_batch.conditional_dict, training_batch)
-                    pred_flow = self.transformer(
+                    pred_flow = current_model(
                         **training_batch_temp.input_kwargs).permute(
                             0, 2, 1, 3, 4)
                     pred_clean = pred_noise_to_pred_video(
@@ -316,7 +574,7 @@ class DistillationPipeline(TrainingPipeline):
         training_batch = self._build_distill_input_kwargs(
             noisy_input, target_timestep, training_batch.conditional_dict,
             training_batch)
-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+        pred_noise = current_model(**training_batch.input_kwargs).permute(
             0, 2, 1, 3, 4)
         pred_video = pred_noise_to_pred_video(
             pred_noise=pred_noise.flatten(0, 1),
@@ -330,6 +588,7 @@ class DistillationPipeline(TrainingPipeline):
     def _dmd_forward(self, generator_pred_video: torch.Tensor,
                      training_batch: TrainingBatch) -> torch.Tensor:
         """Compute DMD (Diffusion Model Distillation) loss."""
+        original_latent = generator_pred_video
         with torch.no_grad():
             timestep = torch.randint(0,
                                      self.num_train_timestep, [1],
@@ -354,7 +613,7 @@ class DistillationPipeline(TrainingPipeline):
 
             noisy_latent = self.noise_scheduler.add_noise(
                 generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
-                timestep).unflatten(0, (1, generator_pred_video.shape[1]))
+                timestep).detach().unflatten(0, (1, generator_pred_video.shape[1]))
 
             # fake_score_transformer forward
             training_batch = self._build_distill_input_kwargs(
@@ -403,24 +662,24 @@ class DistillationPipeline(TrainingPipeline):
                 pred_real_video_uncond) * self.real_score_guidance_scale
 
             grad = (faker_score_pred_video - real_score_pred_video) / torch.abs(
-                generator_pred_video - real_score_pred_video).mean()
+                original_latent - real_score_pred_video).mean()
             grad = torch.nan_to_num(grad)
 
         dmd_loss = 0.5 * F.mse_loss(
-            generator_pred_video.float(),
-            (generator_pred_video.float() - grad.float()).detach())
+            original_latent.float(),
+            (original_latent.float() - grad.float()).detach())
 
         training_batch.dmd_latent_vis_dict.update({
             "training_batch_dmd_fwd_clean_latent":
             training_batch.latents,
             "generator_pred_video":
-            generator_pred_video,
+            original_latent.detach(),
             "real_score_pred_video":
-            real_score_pred_video,
+            real_score_pred_video.detach(),
             "faker_score_pred_video":
-            faker_score_pred_video,
+            faker_score_pred_video.detach(),
             "dmd_timestep":
-            timestep,
+            timestep.detach(),
         })
 
         return dmd_loss
@@ -513,16 +772,17 @@ class DistillationPipeline(TrainingPipeline):
             "encoder_hidden_states": training_batch.encoder_hidden_states,
             "encoder_attention_mask": training_batch.encoder_attention_mask,
         }
-        unconditional_dict = {
-            "encoder_hidden_states": self.negative_prompt_embeds,
-            "encoder_attention_mask": self.negative_prompt_attention_mask,
-        }
+        if getattr(self, "negative_prompt_embeds", None) is not None:
+            unconditional_dict = {
+                "encoder_hidden_states": self.negative_prompt_embeds,
+                "encoder_attention_mask": self.negative_prompt_attention_mask,
+            }
+            training_batch.unconditional_dict = unconditional_dict
 
         training_batch.dmd_latent_vis_dict = {}
         training_batch.fake_score_latent_vis_dict = {}
 
         training_batch.conditional_dict = conditional_dict
-        training_batch.unconditional_dict = unconditional_dict
         training_batch.raw_latent_shape = training_batch.latents.shape
         training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4)
         self.video_latent_shape = training_batch.latents.shape
@@ -556,6 +816,9 @@ class DistillationPipeline(TrainingPipeline):
             batches.append(batch)
 
         self.optimizer.zero_grad()
+        # TODO: confirm this
+        if self.transformer_2 is not None:
+            self.optimizer_2.zero_grad()
         total_dmd_loss = 0.0
         dmd_latent_vis_dict = {}
         fake_score_latent_vis_dict = {}
@@ -584,9 +847,32 @@ class DistillationPipeline(TrainingPipeline):
                         attn_metadata=batch_gen.attn_metadata_vsa):
                     (dmd_loss / gradient_accumulation_steps).backward()
                 total_dmd_loss += dmd_loss.detach().item()
-            self._clip_model_grad_norm_(batch_gen, self.transformer)
-            self.optimizer.step()
-            self.optimizer.zero_grad(set_to_none=True)
+            
+            # Only clip gradients for the model that is currently training
+            if self.train_transformer_2 and self.transformer_2 is not None:
+                self._clip_model_grad_norm_(batch_gen, self.transformer_2)
+                for param in self.transformer_2.parameters():
+                    # check if the gradient is not None and not zero
+                    assert param.grad is not None and param.grad.abs().sum() > 0
+                self.optimizer_2.step()
+                self.optimizer_2.zero_grad(set_to_none=True)
+            else:
+                self._clip_model_grad_norm_(batch_gen, self.transformer)
+                for param in self.transformer.parameters():
+                    # check if the gradient is not None and not zero
+                    assert param.grad is not None and param.grad.abs().sum() > 0
+                self.optimizer.step()
+                self.optimizer.zero_grad(set_to_none=True)
+
+            if self.generator_ema is not None:
+                # TODO: support EMA for transformer_2?
+                if self.train_transformer_2 and self.transformer_2 is not None:
+                    # Note: EMA currently only supports the main transformer
+                    # Could be extended to support transformer_2 in the future
+                    pass
+                else:
+                    self.generator_ema.update(self.transformer)
+
             avg_dmd_loss = torch.tensor(total_dmd_loss /
                                         gradient_accumulation_steps,
                                         device=self.device)
@@ -610,9 +896,18 @@ class DistillationPipeline(TrainingPipeline):
             fake_score_latent_vis_dict.update(
                 batch_fake.fake_score_latent_vis_dict)
         self._clip_model_grad_norm_(batch_fake, self.fake_score_transformer)
+        for param in self.fake_score_transformer.parameters():
+            # check if the gradient is not None and not zero
+            assert param.grad is not None and param.grad.abs().sum() > 0
         self.fake_score_optimizer.step()
         self.fake_score_lr_scheduler.step()
-        self.lr_scheduler.step()
+        
+        # Step the appropriate scheduler
+        if self.train_transformer_2 and self.transformer_2 is not None:
+            self.lr_scheduler_2.step()
+        else:
+            self.lr_scheduler.step()
+            
         self.fake_score_optimizer.zero_grad(set_to_none=True)
         avg_fake_score_loss = torch.tensor(total_fake_score_loss /
                                            gradient_accumulation_steps,
@@ -637,7 +932,8 @@ class DistillationPipeline(TrainingPipeline):
             self.transformer, self.fake_score_transformer, self.global_rank,
             self.training_args.resume_from_checkpoint, self.optimizer,
             self.fake_score_optimizer, self.train_dataloader, self.lr_scheduler,
-            self.fake_score_lr_scheduler, self.noise_random_generator)
+            self.fake_score_lr_scheduler, self.noise_random_generator,
+            self.generator_ema)
 
         if resumed_step > 0:
             self.init_steps = resumed_step
@@ -668,6 +964,14 @@ class DistillationPipeline(TrainingPipeline):
             sum(p.numel()
                 for p in self.fake_score_transformer.parameters()) / 1e9)
 
+        if self.generator_ema is not None:
+            logger.info("  Generator EMA enabled with decay: %s",
+                        self.training_args.ema_decay)
+            logger.info("  Generator EMA start step: %s",
+                        self.training_args.ema_start_step)
+        else:
+            logger.info("  Generator EMA disabled")
+
     @torch.no_grad()
     def _log_validation(self, transformer, training_args, global_step) -> None:
         training_args.inference_mode = True
@@ -699,6 +1003,18 @@ class DistillationPipeline(TrainingPipeline):
 
         transformer.eval()
 
+        # Optionally use EMA model for validation if available and ready
+        use_ema_for_validation = (self.training_args.use_ema
+                                  and self.is_ema_ready(global_step))
+        if use_ema_for_validation:
+            logger.info("Using EMA model for validation")
+            validation_transformer = self.transformer
+            ema_context = self.generator_ema.apply_to_model(
+                validation_transformer)
+        else:
+            validation_transformer = transformer
+            ema_context = None
+
         validation_steps = training_args.validation_sampling_steps.split(",")
         validation_steps = [int(step) for step in validation_steps]
         validation_steps = [step for step in validation_steps if step > 0]
@@ -714,50 +1030,98 @@ class DistillationPipeline(TrainingPipeline):
             step_videos: list[np.ndarray] = []
             step_captions: list[str] = []
 
-            for validation_batch in validation_dataloader:
-                batch = self._prepare_validation_batch(sampling_param,
-                                                       training_args,
-                                                       validation_batch,
-                                                       num_inference_steps)
+            if ema_context is not None:
+                with ema_context:
+                    for validation_batch in validation_dataloader:
+                        batch = self._prepare_validation_batch(
+                            sampling_param, training_args, validation_batch,
+                            num_inference_steps)
 
-                negative_prompt = batch.negative_prompt
-                batch_negative = ForwardBatch(
-                    data_type="video",
-                    prompt=negative_prompt,
-                    prompt_embeds=[],
-                    prompt_attention_mask=[],
-                )
-                result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
-                    batch_negative, training_args)
-                self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
-                    0], result_batch.prompt_attention_mask[0]
+                        negative_prompt = batch.negative_prompt
+                        batch_negative = ForwardBatch(
+                            data_type="video",
+                            prompt=negative_prompt,
+                            prompt_embeds=[],
+                            prompt_attention_mask=[],
+                        )
+                        result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                            batch_negative, training_args)
+                        self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                            0], result_batch.prompt_attention_mask[0]
 
-                logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                        logger.info(
+                            "rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
                             self.global_rank,
                             self.rank_in_sp_group,
                             batch.prompt,
                             local_main_process_only=False)
 
-                assert batch.prompt is not None and isinstance(
-                    batch.prompt, str)
-                step_captions.append(batch.prompt)
+                        assert batch.prompt is not None and isinstance(
+                            batch.prompt, str)
+                        step_captions.append(batch.prompt)
 
-                # Run validation inference
-                with torch.no_grad():
-                    output_batch = self.validation_pipeline.forward(
-                        batch, training_args)
-                samples = output_batch.output
-                if self.rank_in_sp_group != 0:
-                    continue
+                        # Run validation inference
+                        with torch.no_grad():
+                            output_batch = self.validation_pipeline.forward(
+                                batch, training_args)
+                        samples = output_batch.output
+                        if self.rank_in_sp_group != 0:
+                            continue
 
-                # 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))
-                step_videos.append(frames)
+                        # 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))
+                        step_videos.append(frames)
+            else:
+                # Use original transformer without EMA
+                for validation_batch in validation_dataloader:
+                    batch = self._prepare_validation_batch(
+                        sampling_param, training_args, validation_batch,
+                        num_inference_steps)
+
+                    negative_prompt = batch.negative_prompt
+                    batch_negative = ForwardBatch(
+                        data_type="video",
+                        prompt=negative_prompt,
+                        prompt_embeds=[],
+                        prompt_attention_mask=[],
+                    )
+                    result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                        batch_negative, training_args)
+                    self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                        0], result_batch.prompt_attention_mask[0]
+
+                    logger.info(
+                        "rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                        self.global_rank,
+                        self.rank_in_sp_group,
+                        batch.prompt,
+                        local_main_process_only=False)
+
+                    assert batch.prompt is not None and isinstance(
+                        batch.prompt, str)
+                    step_captions.append(batch.prompt)
+
+                    # Run validation inference
+                    with torch.no_grad():
+                        output_batch = self.validation_pipeline.forward(
+                            batch, training_args)
+                    samples = output_batch.output
+                    if self.rank_in_sp_group != 0:
+                        continue
+
+                    # 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))
+                    step_videos.append(frames)
 
             # Log validation results for this step
             world_group = get_world_group()
@@ -834,16 +1198,16 @@ class DistillationPipeline(TrainingPipeline):
                                                         latents.dtype)
                 else:
                     latents += self.vae.shift_factor
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                video = self.vae.decode(latents)
-            video = (video / 2 + 0.5).clamp(0, 1)
-            video = video.cpu().float()
-            video = video.permute(0, 2, 1, 3, 4)
-            video = (video * 255).numpy().astype(np.uint8)
-            wandb_loss_dict[latent_key] = wandb.Video(
-                video, fps=24, format="mp4")  # change to 16 for Wan2.1
-            # Clean up references
-            del video, latents
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    video = self.vae.decode(latents)
+                video = (video / 2 + 0.5).clamp(0, 1)
+                video = video.cpu().float()
+                video = video.permute(0, 2, 1, 3, 4)
+                video = (video * 255).numpy().astype(np.uint8)
+                wandb_loss_dict[latent_key] = wandb.Video(
+                    video, fps=24, format="mp4")  # change to 16 for Wan2.1
+                # Clean up references
+                del video, latents
 
         # Process DMD training data if available - use decode_stage instead of self.vae.decode
         if 'generator_pred_video' in dmd_latents_vis_dict:
@@ -904,6 +1268,10 @@ class DistillationPipeline(TrainingPipeline):
             device="cpu").manual_seed(self.seed)
         logger.info("Initialized random seeds with seed: %s", seed)
 
+        # Initialize current_trainstep for EMA ready checks
+        #TODO: check if needed
+        self.current_trainstep = self.init_steps
+
         # Resume from checkpoint if specified (this will restore random states)
         if self.training_args.resume_from_checkpoint:
             self._resume_from_checkpoint()
@@ -947,6 +1315,14 @@ class DistillationPipeline(TrainingPipeline):
             self.current_trainstep = step
             training_batch.current_vsa_sparsity = current_vsa_sparsity
 
+            if (step >= self.training_args.ema_start_step) and \
+                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+                self.generator_ema = EMA_FSDP(
+                    self.transformer, decay=self.training_args.ema_decay)
+                logger.info(
+                    f"Created generator EMA at step {step} with decay={self.training_args.ema_decay}"
+                )
+
             with torch.autocast("cuda", dtype=torch.bfloat16):
                 training_batch = self.train_one_step(training_batch)
 
@@ -960,11 +1336,19 @@ class DistillationPipeline(TrainingPipeline):
             avg_step_time = sum(step_times) / len(step_times)
 
             progress_bar.set_postfix({
-                "total_loss": f"{total_loss:.4f}",
-                "generator_loss": f"{generator_loss:.4f}",
-                "fake_score_loss": f"{fake_score_loss:.4f}",
-                "step_time": f"{step_time:.2f}s",
-                "grad_norm": grad_norm,
+                "total_loss":
+                f"{total_loss:.4f}",
+                "generator_loss":
+                f"{generator_loss:.4f}",
+                "fake_score_loss":
+                f"{fake_score_loss:.4f}",
+                "step_time":
+                f"{step_time:.2f}s",
+                "grad_norm":
+                grad_norm,
+                "ema":
+                "✓" if (self.generator_ema is not None and self.is_ema_ready())
+                else "✗",
             })
             progress_bar.update(1)
 
@@ -992,6 +1376,15 @@ class DistillationPipeline(TrainingPipeline):
                 if use_vsa:
                     log_data["VSA_train_sparsity"] = current_vsa_sparsity
 
+                if self.generator_ema is not None:
+                    log_data["ema_enabled"] = True
+                    log_data["ema_decay"] = self.training_args.ema_decay
+                else:
+                    log_data["ema_enabled"] = False
+
+                ema_stats = self.get_ema_stats()
+                log_data.update(ema_stats)
+
                 if training_batch.dmd_latent_vis_dict:
                     dmd_additional_logs = {
                         "generator_timestep":
@@ -1023,7 +1416,8 @@ class DistillationPipeline(TrainingPipeline):
                     self.global_rank, self.training_args.output_dir, step,
                     self.optimizer, self.fake_score_optimizer,
                     self.train_dataloader, self.lr_scheduler,
-                    self.fake_score_lr_scheduler, self.noise_random_generator)
+                    self.fake_score_lr_scheduler, self.noise_random_generator,
+                    self.generator_ema)
 
                 if self.transformer:
                     self.transformer.train()
@@ -1040,7 +1434,11 @@ class DistillationPipeline(TrainingPipeline):
                                              self.global_rank,
                                              self.training_args.output_dir,
                                              f"{step}_weight_only",
-                                             only_save_generator_weight=True)
+                                             only_save_generator_weight=True,
+                                             generator_ema=self.generator_ema)
+
+                if self.training_args.use_ema and self.is_ema_ready():
+                    self.save_ema_weights(self.training_args.output_dir, step)
 
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
                 if self.training_args.log_visualization:
@@ -1060,7 +1458,11 @@ class DistillationPipeline(TrainingPipeline):
             self.training_args.output_dir, self.training_args.max_train_steps,
             self.optimizer, self.fake_score_optimizer, self.train_dataloader,
             self.lr_scheduler, self.fake_score_lr_scheduler,
-            self.noise_random_generator)
+            self.noise_random_generator, self.generator_ema)
+
+        if self.training_args.use_ema and self.is_ema_ready():
+            self.save_ema_weights(self.training_args.output_dir,
+                                  self.training_args.max_train_steps)
 
         if get_sp_group():
-            cleanup_dist_env_and_memory()
+            cleanup_dist_env_and_memory()

fastvideo/training/ode_causal_pipeline.py

@@ -1,392 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import sys
-from copy import deepcopy
-from typing import cast
-
-import torch
-import torch.nn.functional as F
-
-from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler)
-from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import (
-    WanCausalDMDPipeline)
-from fastvideo.training.training_pipeline import TrainingPipeline
-from fastvideo.training.training_utils import (
-    clip_grad_norm_while_handling_failing_dtensor_cases)
-
-logger = init_logger(__name__)
-
-
-class ODEInitTrainingPipeline(TrainingPipeline):
-    """
-    Training pipeline for ODE-init using precomputed denoising trajectories.
-
-    Supervision: predict the next latent in the stored trajectory by
-    - feeding current latent at timestep t into the transformer to predict noise
-    - stepping the scheduler with the predicted noise
-    - minimizing MSE to the stored next latent at timestep t_next
-    """
-
-    _required_config_modules = ["scheduler", "transformer", "vae"]
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        # Match the preprocess/generation scheduler for consistent stepping
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift)
-
-    def set_schemas(self):
-        self.train_dataset_schema = pyarrow_schema_ode_trajectory
-
-    def initialize_training_pipeline(self, training_args: TrainingArgs):
-        super().initialize_training_pipeline(training_args)
-
-        self.noise_scheduler = self.get_module("scheduler")
-        self.vae = self.get_module("vae")
-        self.vae.requires_grad_(False)
-
-        self.timestep_shift = self.training_args.pipeline_config.flow_shift
-        assert self.timestep_shift == 5.0, "flow_shift must be 5.0"
-        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
-            shift=self.timestep_shift)
-
-        # logger.info(f"ARG dmd_denoising_steps: {training_args.pipeline_config.dmd_denoising_steps}")
-        logger.info(
-            f"ARG dmd_denoising_steps: {self.training_args.pipeline_config.dmd_denoising_steps}"
-        )
-        self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 0],
-                                                dtype=torch.long,
-                                                device=get_local_torch_device())
-        # self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], dtype=torch.long, device=get_local_torch_device())
-        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
-            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
-                                   torch.tensor([0],
-                                                dtype=torch.float32))).cuda()
-            logger.info(f"timesteps: {timesteps}")
-            self.dmd_denoising_steps = timesteps[1000 -
-                                                 self.dmd_denoising_steps]
-            logger.info(
-                f"warped self.dmd_denoising_steps: {self.dmd_denoising_steps}")
-            assert False, "warp_denoising_step must be false"
-        else:
-            # assert False, "warp_denoising_step must be true"
-            logger.info("not warped")
-        self.dmd_denoising_steps = self.dmd_denoising_steps.to(
-            get_local_torch_device())
-
-        logger.info(f"denoising_step_list: {self.dmd_denoising_steps}")
-
-        logger.info(
-            "Initialized ODE-init training pipeline with %s denoising steps",
-            len(self.dmd_denoising_steps))
-        # Cache for nearest trajectory index per DMD step (computed lazily on first batch)
-        self._cached_closest_idx_per_dmd = None
-        self.num_train_timestep = self.noise_scheduler.num_train_timesteps
-        # self.min_timestep = int(self.training_args.min_timestep_ratio *
-        #                         self.num_train_timestep)
-        # self.max_timestep = int(self.training_args.max_timestep_ratio *
-        #                         self.num_train_timestep)
-        # self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
-
-    def initialize_validation_pipeline(self, training_args: TrainingArgs):
-        logger.info("Initializing validation pipeline...")
-        args_copy = deepcopy(training_args)
-        args_copy.inference_mode = True
-        # Warm start validation with current transformer
-        self.validation_pipeline = WanCausalDMDPipeline.from_pretrained(
-            training_args.model_path,
-            args=args_copy,  # type: ignore
-            inference_mode=True,
-            loaded_modules={
-                "transformer": self.get_module("transformer"),
-            },
-            tp_size=training_args.tp_size,
-            sp_size=training_args.sp_size,
-            num_gpus=training_args.num_gpus,
-            pin_cpu_memory=training_args.pin_cpu_memory,
-            dit_cpu_offload=True)
-
-    def _get_next_batch(self, training_batch):  # type: ignore[override]
-        batch = next(self.train_loader_iter, None)  # type: ignore
-        if batch is None:
-            self.current_epoch += 1
-            logger.info("Starting epoch %s", self.current_epoch)
-            self.train_loader_iter = iter(self.train_dataloader)
-            batch = next(self.train_loader_iter)
-
-        # Required fields from parquet (ODE trajectory schema)
-        encoder_hidden_states = batch['text_embedding']
-        encoder_attention_mask = batch['text_attention_mask']
-        infos = batch['info_list']
-
-        # Trajectory tensors may include a leading singleton batch dim per row
-        trajectory_latents = batch['trajectory_latents']
-        if trajectory_latents.dim() == 7:
-            # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W]
-            trajectory_latents = trajectory_latents[:, 0]
-        elif trajectory_latents.dim() == 6:
-            # already [B, S, C, T, H, W]
-            pass
-        else:
-            raise ValueError(
-                f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}"
-            )
-
-        trajectory_timesteps = batch['trajectory_timesteps']
-        if trajectory_timesteps.dim() == 3:
-            # [B, 1, S] -> [B, S]
-            trajectory_timesteps = trajectory_timesteps[:, 0]
-        elif trajectory_timesteps.dim() == 2:
-            # [B, S]
-            pass
-        else:
-            raise ValueError(
-                f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}"
-            )
-        # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
-        trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)
-
-        # Move to device
-        device = get_local_torch_device()
-        training_batch.encoder_hidden_states = encoder_hidden_states.to(
-            device, dtype=torch.bfloat16)
-        training_batch.encoder_attention_mask = encoder_attention_mask.to(
-            device, dtype=torch.bfloat16)
-        training_batch.infos = infos
-
-        return training_batch, trajectory_latents.to(
-            device, dtype=torch.bfloat16), trajectory_timesteps.to(device)
-
-    def _get_timestep(self,
-                      min_timestep: int,
-                      max_timestep: int,
-                      batch_size: int,
-                      num_frame: int,
-                      num_frame_per_block: int,
-                      uniform_timestep: bool = False) -> torch.Tensor:
-        if uniform_timestep:
-            timestep = torch.randint(min_timestep,
-                                     max_timestep, [batch_size, 1],
-                                     device=self.device,
-                                     dtype=torch.long).repeat(1, num_frame)
-            return timestep
-        else:
-            timestep = torch.randint(min_timestep,
-                                     max_timestep, [batch_size, num_frame],
-                                     device=self.device,
-                                     dtype=torch.long)
-            # logger.info(f"individual timestep: {timestep}")
-            # make the noise level the same within every block
-            timestep = timestep.reshape(timestep.shape[0], -1,
-                                        num_frame_per_block)
-            timestep[:, :, 1:] = timestep[:, :, 0:1]
-            timestep = timestep.reshape(timestep.shape[0], -1)
-            return timestep
-
-    def _step_predict_next_latent(
-        self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_attention_mask: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        device = get_local_torch_device()
-        target_latent = traj_latents[:, -1]
-
-        # logger.info(f"traj_latents: {traj_latents.shape}")
-        # logger.info(f"traj_timesteps: {traj_timesteps.shape}")
-
-        # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
-        B, S, num_frames, num_channels, height, width = traj_latents.shape
-
-        # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps
-        if self._cached_closest_idx_per_dmd is None:
-            # Use the first sample's trajectory timesteps; assumed identical across batches
-            # s_steps = traj_timesteps[0].to(torch.long)  # [S]
-            # dmd = cast(torch.Tensor, self.dmd_denoising_steps).to(s_steps.device)  # [K]
-            # distances_ks: [K, S] = |s_steps - dmd|
-            # distances_ks = (s_steps.unsqueeze(0) - dmd.unsqueeze(1)).abs()
-            # self._cached_closest_idx_per_dmd = distances_ks.argmin(dim=1).to(torch.long).cpu()  # [K]
-            self._cached_closest_idx_per_dmd = torch.tensor(
-                [0, 12, 24, 36], dtype=torch.long).cpu()
-            # logger.info(f"self._cached_closest_idx_per_dmd: {self._cached_closest_idx_per_dmd}")
-            # logger.info(f"corresponding timesteps: {self.noise_scheduler.timesteps[self._cached_closest_idx_per_dmd]}")
-
-        # logger.info(f"traj_latents: {traj_latents.shape}")
-        # Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd
-        # traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K]
-        # Output: [B, K, C, T, H, W]
-        relevant_traj_latents = torch.index_select(
-            traj_latents,
-            dim=1,
-            index=self._cached_closest_idx_per_dmd.to(traj_latents.device))
-        # logger.info(f"relevant_traj_latents: {relevant_traj_latents.shape}")
-        # assert relevant_traj_latents.shape[0] == 1
-
-        indexes = self._get_timestep(  # [B, num_frames]
-            0,
-            len(self.dmd_denoising_steps),
-            B,
-            num_frames,
-            3,
-            uniform_timestep=False)
-        # noisy_input = relevant_traj_latents[indexes]
-        # logger.info(f"indexes: {indexes.shape}")
-        noisy_input = torch.gather(
-            relevant_traj_latents,
-            dim=1,
-            index=indexes.reshape(B, 1, num_frames, 1, 1,
-                                  1).expand(-1, -1, -1, num_channels, height,
-                                            width).to(self.device)).squeeze(1)
-        # noisy_input = noisy_input.unsqueeze(0)
-
-        # # Sample a single DMD step for the whole batch and fetch its cached nearest S-index
-        # K = len(self.dmd_denoising_steps)
-        # dmd_idx = torch.randint(0, K, (1,), device=device)
-        # logger.info(f"dmd_idx: {dmd_idx}")
-        # assert self._cached_closest_idx_per_dmd is not None
-        # nearest_s_idx = int(self._cached_closest_idx_per_dmd[int(dmd_idx.item())])
-        # nearest_idx = torch.full((B,), nearest_s_idx, device=device, dtype=torch.long)
-
-        # batch_indices = torch.arange(B, device=device)
-        # noisy_input = traj_latents[batch_indices, nearest_idx]  # [B, C, T, H, W]
-        # target_latent = traj_latents[batch_indices, -1]  # [B, C, T, H, W]
-        # t = traj_timesteps[batch_indices, nearest_idx]  # [B]
-
-        # Scale model input as in inference for consistency with stored trajectories
-        # noisy_input = self.modules["scheduler"].scale_model_input(noisy_input, t)
-        # logger.info(f"indexes: {indexes.shape}")
-        # logger.info(f"indexes: {indexes}")
-        timestep = self.dmd_denoising_steps[indexes]
-        # logger.info(f"timestep: {timestep.shape}")
-        # logger.info(f"timestep: {timestep}")
-
-        # Prepare inputs for transformer
-        input_kwargs = {
-            "hidden_states": noisy_input.permute(0, 2, 1, 3, 4),
-            "encoder_hidden_states": encoder_hidden_states,
-            "timestep": timestep.to(device, dtype=torch.bfloat16),
-            "encoder_attention_mask": encoder_attention_mask,
-            "return_dict": False,
-        }
-        # Predict noise and step the scheduler to obtain next latent
-        with set_forward_context(current_timestep=timestep,
-                                 attn_metadata=None,
-                                 forward_batch=None):
-            noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
-            # logger.info(f"noise_pred: {noise_pred.shape}")
-        if isinstance(noise_pred, (tuple, list)):
-            noise_pred = noise_pred[0]
-
-        from fastvideo.models.utils import pred_noise_to_pred_video
-        noise_pred = pred_noise_to_pred_video(
-            pred_noise=noise_pred.flatten(0, 1),
-            noise_input_latent=noisy_input.flatten(0, 1),
-            timestep=timestep.flatten(0, 1),
-            scheduler=self.modules["scheduler"]).unflatten(
-                0, noise_pred.shape[:2])
-
-        # noisy_input = pred_noise_to_pred_video(noise_pred, noisy_input, t, self.modules["scheduler"])
-        # next_latent_pred = self.modules["scheduler"].step(
-        #     noise_pred, t, current_latents, return_dict=False)[0]
-        return noise_pred, target_latent, timestep
-
-    def train_one_step(self, training_batch):  # type: ignore[override]
-        self.transformer.train()
-        self.optimizer.zero_grad()
-        training_batch.total_loss = 0.0
-        args = cast(TrainingArgs, self.training_args)
-
-        # Using cached nearest index per DMD step; computation happens in _step_predict_next_latent
-
-        for _ in range(args.gradient_accumulation_steps):
-            training_batch, traj_latents, traj_timesteps = self._get_next_batch(
-                training_batch)
-            text_embeds = training_batch.encoder_hidden_states
-            text_attention_mask = training_batch.encoder_attention_mask
-            assert traj_latents.shape[0] == 1
-
-            # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
-            B, S = traj_latents.shape[0], traj_latents.shape[1]
-            if S < 2:
-                raise ValueError("Trajectory must contain at least 2 steps")
-
-            # Sample per-sample current step i in [0, S-2]
-
-            # idx = torch.randint(low=0, high=S - 1, size=(B, ),
-            #                     device=traj_latents.device)
-
-            # Gather current latents and next latents
-            # batch_indices = torch.arange(B, device=traj_latents.device)
-            # current_latents = traj_latents[batch_indices, idx]  # [B, C, T,H,W]
-            # current_latent = traj_timesteps[:, -1, :, :, :, :]
-            # target_latents = traj_latents[:, -1, :, :, :, :]
-
-            # Corresponding timesteps t (long) -> cast per sample
-            # t = traj_timesteps[:, -1, :, :, :, :]
-            # if t.dtype != torch.long:
-            #     t = t.long()
-
-            # Forward to predict next latent by stepping scheduler with predicted noise
-            noise_pred, target_latent, t = self._step_predict_next_latent(
-                traj_latents, traj_timesteps, text_embeds, text_attention_mask)
-
-            mask = t != 0
-
-            # Compute loss
-            loss = F.mse_loss(noise_pred[mask],
-                              target_latent[mask],
-                              reduction="mean")
-            loss = loss / args.gradient_accumulation_steps
-
-            with set_forward_context(current_timestep=t,
-                                     attn_metadata=None,
-                                     forward_batch=None):
-                loss.backward()
-            avg_loss = loss.detach().clone()
-            training_batch.total_loss += avg_loss.item()
-
-        # Clip grad and step optimizers
-        grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
-            [p for p in self.transformer.parameters() if p.requires_grad],
-            args.max_grad_norm if args.max_grad_norm is not None else 0.0)
-
-        self.optimizer.step()
-        self.lr_scheduler.step()
-
-        if grad_norm is None:
-            grad_value = 0.0
-        else:
-            try:
-                if isinstance(grad_norm, torch.Tensor):
-                    grad_value = float(grad_norm.detach().float().item())
-                else:
-                    grad_value = float(grad_norm)
-            except Exception:
-                grad_value = 0.0
-        training_batch.grad_norm = grad_value
-        return training_batch
-
-
-def main(args) -> None:
-    logger.info("Starting ODE-init training pipeline...")
-    logger.info(f"ARG dmd_denoising_steps: {args.dmd_denoising_steps}")
-    pipeline = ODEInitTrainingPipeline.from_pretrained(
-        args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("ODE-init training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/training/training_pipeline.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+import gc
 import dataclasses
 import math
 import os
@@ -22,6 +23,7 @@ from tqdm.auto import tqdm
 import fastvideo.envs as envs
 from fastvideo.attention.backends.video_sparse_attn import (
     VideoSparseAttentionMetadataBuilder)
+# from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder
 from fastvideo.configs.sample import SamplingParam
 from fastvideo.dataset import build_parquet_map_style_dataloader
 from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
@@ -41,11 +43,15 @@ from fastvideo.training.training_utils import (
     compute_density_for_timestep_sampling, get_scheduler, get_sigmas,
     load_checkpoint, normalize_dit_input, save_checkpoint,
     shard_latents_across_sp)
-from fastvideo.utils import is_vsa_available, set_random_seed, shallow_asdict
+# from fastvideo.utils import (is_vmoba_available, is_vsa_available,
+#                              set_random_seed, shallow_asdict)
+from fastvideo.utils import (is_vsa_available,
+                             set_random_seed, shallow_asdict)
 
 import wandb  # isort: skip
 
 vsa_available = is_vsa_available()
+# vmoba_available = is_vmoba_available()
 
 logger = init_logger(__name__)
 
@@ -64,6 +70,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
     train_dataloader: StatefulDataLoader
     train_loader_iter: Iterator[dict[str, Any]]
     current_epoch: int = 0
+    train_transformer_2: bool = False
 
     def __init__(
             self,
@@ -99,6 +106,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         self.sp_world_size = self.sp_group.world_size
         self.local_rank = world_group.local_rank
         self.transformer = self.get_module("transformer")
+        self.transformer_2 = self.get_module("transformer_2", None)
         self.seed = training_args.seed
         self.set_schemas()
 
@@ -111,16 +119,25 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 self.transformer,
                 checkpointing_type=training_args.
                 enable_gradient_checkpointing_type)
+            if self.transformer_2 is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2,
+                    checkpointing_type=training_args.
+                    enable_gradient_checkpointing_type)
 
         noise_scheduler = self.modules["scheduler"]
         self.set_trainable()
         params_to_optimize = self.transformer.parameters()
         params_to_optimize = list(
             filter(lambda p: p.requires_grad, params_to_optimize))
+        # Parse betas from string format "beta1,beta2"
+        betas_str = training_args.betas
+        betas = tuple(float(x.strip()) for x in betas_str.split(","))
+        
         self.optimizer = torch.optim.AdamW(
             params_to_optimize,
             lr=training_args.learning_rate,
-            betas=(0.9, 0.999),
+            betas=betas,
             weight_decay=training_args.weight_decay,
             eps=1e-8,
         )
@@ -138,6 +155,30 @@ class TrainingPipeline(LoRAPipeline, ABC):
             min_lr_ratio=training_args.min_lr_ratio,
             last_epoch=self.init_steps - 1,
         )
+        if self.transformer_2 is not None:
+            # Ensure transformer_2 has trainable parameters before creating optimizer
+            self.transformer_2.train()
+            self.transformer_2.requires_grad_(True)
+            params_to_optimize_2 = self.transformer_2.parameters()
+            params_to_optimize_2 = list(
+                filter(lambda p: p.requires_grad, params_to_optimize_2))
+            self.optimizer_2 = torch.optim.AdamW(
+                params_to_optimize_2,
+                lr=training_args.learning_rate,
+                betas=(0.9, 0.999),
+                weight_decay=training_args.weight_decay,
+                eps=1e-8,
+            )
+            self.lr_scheduler_2 = get_scheduler(
+                training_args.lr_scheduler,
+                optimizer=self.optimizer_2,
+                num_warmup_steps=training_args.lr_warmup_steps,
+                num_training_steps=training_args.max_train_steps,
+                num_cycles=training_args.lr_num_cycles,
+                power=training_args.lr_power,
+                min_lr_ratio=training_args.min_lr_ratio,
+                last_epoch=self.init_steps - 1,
+            )
 
         self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
             training_args.data_path,
@@ -152,7 +193,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
             seed=self.seed)
 
         self.noise_scheduler = noise_scheduler
-
+        self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
         self.num_update_steps_per_epoch = math.ceil(
             len(self.train_dataloader) /
             training_args.gradient_accumulation_steps * training_args.sp_size /
@@ -178,9 +219,25 @@ class TrainingPipeline(LoRAPipeline, ABC):
     def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
         self.transformer.train()
         self.optimizer.zero_grad()
+        if self.transformer_2 is not None:
+            self.transformer_2.train()
+            self.optimizer_2.zero_grad()
         training_batch.total_loss = 0.0
         return training_batch
 
+    def _enable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        """Enable training mode and gradients for the specified model."""
+        for param in model.parameters():
+            param.requires_grad = True
+        model.train()
+        optimizer.zero_grad()
+
+    def _disable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        """Disable training mode and gradients for the specified model."""
+        for param in model.parameters():
+            param.requires_grad = False
+        optimizer.zero_grad(set_to_none=True)
+
     def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
         batch = next(self.train_loader_iter, None)  # type: ignore
         if batch is None:
@@ -224,17 +281,17 @@ class TrainingPipeline(LoRAPipeline, ABC):
                             generator=self.noise_gen_cuda,
                             device=latents.device,
                             dtype=latents.dtype)
-        u = compute_density_for_timestep_sampling(
-            weighting_scheme=self.training_args.weighting_scheme,
-            batch_size=batch_size,
-            generator=self.noise_random_generator,
-            logit_mean=self.training_args.logit_mean,
-            logit_std=self.training_args.logit_std,
-            mode_scale=self.training_args.mode_scale,
-        )
-        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
-        timesteps = self.noise_scheduler.timesteps[indices].to(
-            device=latents.device)
+        timesteps = self._sample_timesteps(batch_size, latents.device)
+
+        # Enable training for the model that will be trained next and disable the other
+        if self.train_transformer_2:
+            self._enable_training(self.transformer_2, self.optimizer_2)
+            self._disable_training(self.transformer, self.optimizer)
+        else:
+            self._enable_training(self.transformer, self.optimizer)
+            if self.transformer_2 is not None:
+                self._disable_training(self.transformer_2, self.optimizer_2)
+
         if self.training_args.sp_size > 1:
             # Make sure that the timesteps are the same across all sp processes.
             sp_group = get_sp_group()
@@ -257,6 +314,38 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         return training_batch
 
+    def _sample_timesteps(self, batch_size, device):
+        # Determine which model to train based on the boundary timestep
+        if (self.transformer_2 is not None and self.boundary_timestep is not None and
+                torch.rand(1, generator=self.noise_random_generator).item() <= self.training_args.boundary_ratio):
+            self.train_transformer_2 = True
+        else:
+            self.train_transformer_2 = False
+
+        # Broadcast the decision to all processes
+        decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0, device=self.device)
+        dist.broadcast(decision, src=0)
+        self.train_transformer_2 = decision.item() == 1.0
+
+        # Sample u from the appropriate range
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=self.training_args.weighting_scheme,
+            batch_size=batch_size,
+            generator=self.noise_random_generator,
+            logit_mean=self.training_args.logit_mean,
+            logit_std=self.training_args.logit_std,
+            mode_scale=self.training_args.mode_scale,
+        )
+
+        boundary_ratio = self.training_args.boundary_ratio
+        if self.train_transformer_2:
+            u = (1 - boundary_ratio) + u * boundary_ratio # min: 1 - boundary_ratio, max: 1
+        else:
+            u = u * (1 - boundary_ratio) # min: 0, max: 1 - boundary_ratio
+
+        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        return self.noise_scheduler.timesteps[indices].to(device=device)
+        
     def _build_attention_metadata(
             self, training_batch: TrainingBatch) -> TrainingBatch:
         latents_shape = training_batch.raw_latent_shape
@@ -307,11 +396,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
         #         [1000.0],
         #         device=training_batch.noisy_model_input.device,
         #         dtype=torch.bfloat16)
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
 
         with set_forward_context(
                 current_timestep=training_batch.current_timestep,
                 attn_metadata=training_batch.attn_metadata):
-            model_pred = self.transformer(**input_kwargs)
+            model_pred = current_model(**input_kwargs)
             if self.training_args.precondition_outputs:
                 assert training_batch.sigmas is not None
                 model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
@@ -342,7 +432,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
         # the following:
         # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
         if max_grad_norm is not None:
-            model_parts = [self.transformer]
+            # Only clip gradients for the model that is currently training
+            if self.train_transformer_2 and self.transformer_2 is not None:
+                model_parts = [self.transformer_2]
+            else:
+                model_parts = [self.transformer]
+            
             grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
                 [p for m in model_parts for p in m.parameters()],
                 max_grad_norm,
@@ -387,9 +482,14 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         training_batch = self._clip_grad_norm(training_batch)
 
-        self.optimizer.step()
-        self.lr_scheduler.step()
-
+        # Only step the optimizer and scheduler for the model that is currently training
+        if self.train_transformer_2 and self.transformer_2 is not None:
+            self.optimizer_2.step()
+            self.lr_scheduler_2.step()
+        else:
+            self.optimizer.step()
+            self.lr_scheduler.step()
+            
         training_batch.total_loss = training_batch.total_loss
         training_batch.grad_norm = training_batch.grad_norm
         return training_batch
@@ -421,6 +521,11 @@ class TrainingPipeline(LoRAPipeline, ABC):
         logger.info("Starting training with %s B trainable parameters",
                     round(num_trainable_params / 1e9, 3))
 
+        if getattr(self, "transformer_2", None) is not None:
+            num_trainable_params = _get_trainable_params(self.transformer_2)
+            logger.info("Transformer 2: Starting training with %s B trainable parameters",
+                        round(num_trainable_params / 1e9, 3))
+
         # Set random seeds for deterministic training
         self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
             self.seed)
@@ -441,7 +546,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         self._log_training_info()
 
-        self._log_validation(self.transformer, self.training_args,
+        self._log_validation(self.training_args,
                              self.init_steps)
 
         # Train!
@@ -503,7 +608,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 self.transformer.train()
                 self.sp_group.barrier()
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                self._log_validation(self.transformer, self.training_args, step)
+                self._log_validation(self.training_args, step)
                 gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
                 trainable_params = round(
                     _get_trainable_params(self.transformer) / 1e9, 3)
@@ -584,12 +689,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
         return batch
 
     @torch.no_grad()
-    def _log_validation(self, transformer, training_args, global_step) -> None:
+    def _log_validation(self, training_args, global_step) -> None:
         """
         Generate a validation video and log it to wandb to check the quality during training.
         """
         training_args.inference_mode = True
-        training_args.dit_cpu_offload = True
+        training_args.dit_cpu_offload = False
         if not training_args.log_validation:
             return
         if self.validation_pipeline is None:
@@ -611,7 +716,9 @@ class TrainingPipeline(LoRAPipeline, ABC):
                                            batch_size=None,
                                            num_workers=0)
 
-        transformer.eval()
+        self.transformer.eval()
+        if getattr(self, "transformer_2", None) is not None:
+            self.transformer_2.eval()
 
         validation_steps = training_args.validation_sampling_steps.split(",")
         validation_steps = [int(step) for step in validation_steps]
@@ -703,4 +810,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         # Re-enable gradients for training
         training_args.inference_mode = False
-        transformer.train()
+        self.transformer.train()
+        if getattr(self, "transformer_2", None) is not None:
+            self.transformer_2.train()

fastvideo/training/training_utils.py

@@ -202,6 +202,7 @@ def save_distillation_checkpoint(generator_transformer,
                                  generator_scheduler=None,
                                  fake_score_scheduler=None,
                                  noise_generator=None,
+                                 generator_ema=None,
                                  only_save_generator_weight=False) -> None:
     """
     Save distillation checkpoint with both generator and fake_score models.
@@ -233,6 +234,8 @@ def save_distillation_checkpoint(generator_transformer,
         if generator_scheduler is not None:
             generator_states["scheduler"] = SchedulerWrapper(
                 generator_scheduler)
+        if generator_ema is not None:
+            generator_states["ema"] = generator_ema.state_dict()
 
         generator_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
                                          "generator")
@@ -346,14 +349,10 @@ def load_checkpoint(transformer,
     """
     if not os.path.exists(checkpoint_path):
         logger.warning("Checkpoint path %s does not exist", checkpoint_path)
-        assert False
         return 0
 
     # Extract step number from checkpoint path
-    try:
-        step = int(os.path.basename(checkpoint_path).split('-')[-1])
-    except:
-        step = 1
+    step = int(os.path.basename(checkpoint_path).split('-')[-1])
 
     if rank == 0:
         logger.info("Loading checkpoint from step %s", step)
@@ -406,7 +405,8 @@ def load_distillation_checkpoint(generator_transformer,
                                  dataloader=None,
                                  generator_scheduler=None,
                                  fake_score_scheduler=None,
-                                 noise_generator=None) -> int:
+                                 noise_generator=None,
+                                 generator_ema=None) -> int:
     """
     Load distillation checkpoint with both generator and fake_score models.
     Returns the step number from which training should resume.
@@ -460,6 +460,18 @@ def load_distillation_checkpoint(generator_transformer,
         end_time - begin_time,
         local_main_process_only=False)
 
+    # Load EMA state if available and generator_ema is provided
+    if generator_ema is not None:
+        try:
+            ema_state = generator_states.get("ema")
+            if ema_state is not None:
+                generator_ema.load_state_dict(ema_state)
+                logger.info("rank: %s, generator EMA state loaded successfully", rank)
+            else:
+                logger.info("rank: %s, no EMA state found in checkpoint", rank)
+        except Exception as e:
+            logger.warning("rank: %s, failed to load EMA state: %s", rank, str(e))
+
     # Load critic distributed checkpoint
     critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
                                   "critic")
@@ -1282,3 +1294,156 @@ def get_scheduler(
                          num_warmup_steps=num_warmup_steps,
                          num_training_steps=num_training_steps,
                          last_epoch=last_epoch)
+
+
+class EMA_FSDP:
+    """
+    FSDP2-friendly EMA with two modes:
+      - mode="local_shard" (default): maintain float32 CPU EMA of local parameter shards on every rank.
+        Provides a context manager to temporarily swap EMA weights into the live model for teacher forward.
+      - mode="rank0_full": maintain a consolidated float32 CPU EMA of full parameters on rank 0 only
+        using gather_state_dict_on_cpu_rank0(). Useful for checkpoint export; not for teacher forward.
+
+    Usage (local_shard for CM teacher):
+      ema = EMA_FSDP(model, decay=0.999, mode="local_shard")
+      for step in ...:
+          ema.update(model)
+      with ema.apply_to_model(model):
+          with torch.no_grad():
+              y_teacher = model(...)
+
+    Usage (rank0_full for export):
+      ema = EMA_FSDP(model, decay=0.999, mode="rank0_full")
+      ema.update(model)
+      ema.state_dict()  # on rank 0
+    """
+    def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"):
+        self.decay = float(decay)
+        self.mode = mode
+        self.shadow: dict[str, torch.Tensor] = {}
+        self.rank = dist.get_rank() if dist.is_initialized() else 0
+        if self.mode not in {"local_shard", "rank0_full"}:
+            raise ValueError(f"Unsupported EMA_FSDP mode: {self.mode}")
+        self._init_shadow(module)
+
+    @staticmethod
+    def _to_local_tensor(t: torch.Tensor) -> torch.Tensor:
+        # DTensor-aware to_local fetch; fall back to raw tensor
+        try:
+            from torch.distributed.tensor import DTensor  # type: ignore
+            if isinstance(t, DTensor):
+                return t.to_local()
+        except Exception:
+            pass
+        return t
+
+    @torch.no_grad()
+    def _init_shadow(self, module):
+        if self.mode == "rank0_full":
+            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
+            if self.rank == 0:
+                self.shadow = {k: v.detach().clone().float().cpu() for k, v in cpu_state.items()}
+            else:
+                self.shadow = {}
+            return
+
+        # local_shard: maintain EMA of local shards for requires_grad params
+        self.shadow = {}
+        for name, p in module.named_parameters():
+            if not p.requires_grad:
+                continue
+            local = self._to_local_tensor(p.detach())
+            self.shadow[name] = local.clone().float().cpu()
+
+    @torch.no_grad()
+    def update(self, module):
+        d = self.decay
+        if self.mode == "rank0_full":
+            if self.rank != 0:
+                return
+            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
+            for n, v in cpu_state.items():
+                v_cpu = v.detach().float().cpu()
+                if n not in self.shadow:
+                    self.shadow[n] = v_cpu.clone()
+                else:
+                    self.shadow[n].mul_(d).add_(v_cpu, alpha=1.0 - d)
+            return
+
+        # local_shard: update local shard EMA on every rank
+        for name, p in module.named_parameters():
+            if not p.requires_grad:
+                continue
+            local = self._to_local_tensor(p.detach())
+            v_cpu = local.float().cpu()
+            if name not in self.shadow:
+                self.shadow[name] = v_cpu.clone()
+            else:
+                self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
+
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        if self.mode == "rank0_full":
+            return {k: v.clone() for k, v in self.shadow.items()} if self.rank == 0 else {}
+        return {k: v.clone() for k, v in self.shadow.items()}
+
+    def load_state_dict(self, sd: dict[str, torch.Tensor]):
+        self.shadow = {k: v.clone() for k, v in sd.items()}
+
+    @torch.no_grad()
+    def copy_to_unwrapped(self, module) -> None:
+        """
+        Copy EMA weights into a non-sharded (unwrapped) module. Intended for export/eval.
+        For mode="rank0_full", only rank 0 has the full EMA state.
+        """
+        if self.mode == "rank0_full" and self.rank != 0:
+            return
+        name_to_param = dict(module.named_parameters())
+        for n, w in self.shadow.items():
+            if n in name_to_param:
+                p = name_to_param[n]
+                p.data.copy_(w.to(dtype=p.dtype, device=p.device))
+
+    class _ApplyEMACtx:
+        def __init__(self, ema: "EMA_FSDP", module):
+            self.ema = ema
+            self.module = module
+            self.saved: dict[str, torch.Tensor] = {}
+
+        def __enter__(self):
+            if self.ema.mode != "local_shard":
+                raise RuntimeError("EMA apply_to_model is only supported for mode='local_shard'")
+            with torch.no_grad():
+                for name, p in self.module.named_parameters():
+                    if not p.requires_grad:
+                        continue
+                    # Save local shard
+                    p_local = EMA_FSDP._to_local_tensor(p.detach())
+                    if p_local.numel() == 0:
+                        # Nothing to swap on this rank for this param
+                        continue
+                    self.saved[name] = p_local.clone().to(device=p_local.device, dtype=p_local.dtype)
+                    if name in self.ema.shadow:
+                        ema_cpu = self.ema.shadow[name]
+                        if ema_cpu.numel() != p_local.numel():
+                            # Shard shape mismatch (e.g., empty shard here), skip
+                            continue
+                        # Copy EMA shard into local param shard
+                        p_local.copy_(ema_cpu.to(dtype=p_local.dtype, device=p_local.device))
+            return self.module
+
+        def __exit__(self, exc_type, exc, tb):
+            with torch.no_grad():
+                for name, p in self.module.named_parameters():
+                    if name in self.saved:
+                        p_local = EMA_FSDP._to_local_tensor(p.detach())
+                        if p_local.numel() == 0:
+                            continue
+                        saved_local = self.saved[name]
+                        if saved_local.numel() != p_local.numel():
+                            continue
+                        p_local.copy_(saved_local)
+            self.saved.clear()
+            return False
+
+    def apply_to_model(self, module):
+        return EMA_FSDP._ApplyEMACtx(self, module)

fastvideo/training/wan_self_forcing_distillation_pipeline.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import WanCausalDMDPipeline
+from fastvideo.training.self_forcing_distillation_pipeline import SelfForcingDistillationPipeline
+from fastvideo.utils import is_vsa_available
+
+vsa_available = is_vsa_available()
+
+logger = init_logger(__name__)
+
+
+class WanSelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
+    """
+    A self-forcing distillation pipeline for Wan that uses the self-forcing methodology
+    with DMD for video generation.
+    """
+    _required_config_modules = [
+        "scheduler", "transformer", "vae", "real_score_transformer",
+        "fake_score_transformer"
+    ]
+
+    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
+        validation_pipeline = WanCausalDMDPipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,  # type: ignore
+            inference_mode=True,
+            loaded_modules={"transformer": self.get_module("transformer")},
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True)
+
+        self.validation_pipeline = validation_pipeline
+
+
+def main(args) -> None:
+    logger.info("Starting Wan self-forcing distillation pipeline...")
+
+    pipeline = WanSelfForcingDistillationPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, args=args)
+
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Wan self-forcing distillation pipeline completed")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    main(args) 

fastvideo/training/wan_t2v_i2v_training_pipeline.py

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

fastvideo/utils.py

@@ -2,10 +2,6 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py
 
 import argparse
-from einops import rearrange
-import torchvision
-import numpy as np
-import imageio
 import ctypes
 import hashlib
 import importlib
@@ -818,6 +814,17 @@ def is_vsa_available() -> bool:
     return importlib.util.find_spec("vsa") is not None
 
 
+@lru_cache(maxsize=1)
+def is_vmoba_available() -> bool:
+    if importlib.util.find_spec("csrc.attn.vmoba_attn.vmoba") is None:
+        return False
+    try:
+        import flash_attn
+        return flash_attn.__version__ >= "2.7.4"
+    except Exception:
+        return False
+
+
 # adapted from: https://github.com/Wan-Video/Wan2.2/blob/main/wan/utils/utils.py
 def masks_like(tensor,
                zero=False,
@@ -879,16 +886,3 @@ def best_output_size(w, h, dw, dh, expected_area):
         return ow1, oh1
     else:
         return ow2, oh2
-
-
-def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor], output_path: str, fps: int):
-    # Process outputs
-    videos = rearrange(decoded_latents, "b c t h w -> t b c h w")
-    frames = []
-    for x in videos:
-        x = torchvision.utils.make_grid(x, nrow=6)
-        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-        frames.append((x * 255).numpy().astype(np.uint8))
-
-    os.makedirs(os.path.dirname(output_path), exist_ok=True)
-    imageio.mimsave(output_path, frames, fps=fps, format="mp4")

fastvideo/workflow/preprocess/components.py

@@ -13,7 +13,8 @@ import pyarrow.parquet as pq
 import torch
 from datasets import Dataset, Video, load_dataset
 
-from fastvideo.configs.configs import DatasetType, PreprocessConfig
+from fastvideo.configs.configs import (DatasetType, PreprocessConfig,
+                                       VideoLoaderType)
 from fastvideo.distributed.parallel_state import get_world_rank, get_world_size
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
@@ -431,7 +432,8 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str,
             return item
 
         dataset = dataset.map(add_video_column)
-        dataset = dataset.cast_column("video", Video())
+        if preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
+            dataset = dataset.cast_column("video", Video())
     else:
         raise ValueError(
             f"Invalid dataset type: {preprocess_config.dataset_type}")

scripts/checkpoint_conversion/wan_to_diffusers.py

@@ -135,8 +135,8 @@ _self_forcing_to_diffusers_param_names_mapping: dict = {
     r"blocks.\1.norm2.\2",
 }
 
-state_dict = load_state_dict_from_file("checkpoints/ode_init.pt")
-state_dict = state_dict["generator"]
+state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt")
+state_dict = state_dict["generator_ema"]
 new_state_dict = OrderedDict()
 for k, v in state_dict.items():
     new_key = k
@@ -156,6 +156,6 @@ del state_dict
 
 save_torch_state_dict(
     new_state_dict,
-    "ode_init_diffusers/",
+    "new2/",
     max_shard_size="10GB"
 )

scripts/inference/v1_inference_wan_1.3B_VMoba.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+num_gpus=1
+export FASTVIDEO_ATTENTION_BACKEND=VMOBA_ATTN
+export MODEL_BASE=FastVideo/Wan2.1-T2V-1.3B-Diffusers
+# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# You can either use --prompt or --prompt-txt, but not both.
+fastvideo generate \
+    --model-path $MODEL_BASE \
+    --sp-size $num_gpus \
+    --tp-size 1 \
+    --num-gpus $num_gpus \
+    --dit-cpu-offload False \
+    --vae-cpu-offload False \
+    --text-encoder-cpu-offload True \
+    --pin-cpu-memory False \
+    --height 480 \
+    --width 832 \
+    --num-frames 77 \
+    --num-inference-steps 50 \
+    --moba-config-path fastvideo/configs/backend/vmoba/wan_1.3B_77_480_832.json \
+    --fps 16 \
+    --guidance-scale 6.0 \
+    --flow-shift 8.0 \
+    --prompt-txt assets/prompt.txt \
+    --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
+    --seed 1024 \
+    --output-path outputs_video/

wan2.2_train.sh

@@ -0,0 +1,57 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export WANDB_API_KEY='2f25ad37933894dbf0966c838c0b8494987f9f2f'
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+export TRITON_CACHE_DIR=/tmp/triton_cache
+# DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/
+DATA_DIR=/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol_processed_t2v/combined_parquet_dataset
+# VALIDATION_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/mixkit/validation_8.json
+VALIDATION_DIR=/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol-single_processed_t2v/validation.json
+NUM_GPUS=8
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+CHECKPOINT_PATH="outputs_train_test/wan_finetune/checkpoint-10"
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_training_pipeline.py \
+    --model_path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+    --inference_mode False\
+    --pretrained_model_name_or_path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+    --cache_dir "/home/ray/.cache" \
+    --data_path "$DATA_DIR" \
+    --validation_dataset_file  "$VALIDATION_DIR" \
+    --train_batch_size 1 \
+    --num_latent_t 16\
+    --sp_size 4 \
+    --tp_size 1 \
+    --num_gpus $NUM_GPUS \
+    --hsdp_replicate_dim 1 \
+    --hsdp-shard-dim 8 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps 1 \
+    --max_train_steps 30000 \
+    --learning_rate 1e-5 \
+    --mixed_precision "bf16" \
+    --checkpointing_steps 1000 \
+    --validation_steps 30 \
+    --validation_sampling_steps "40" \
+    --log_validation True \
+    --checkpoints_total_limit 3 \
+    --ema_start_step 0 \
+    --training_cfg_rate 0.1 \
+    --seed 1024 \
+    --output_dir "outputs_train_test/wan_finetune_v1" \
+    --tracker_project_name VSA_finetune \
+    --num_height 448 \
+    --num_width 832 \
+    --num_frames 61 \
+    --flow_shift 5 \
+    --validation_guidance_scale "5.0" \
+    --num_euler_timesteps 50 \
+    --master_weight_type "fp32" \
+    --dit_precision "fp32" \
+    --weight_decay 0.01 \
+    --max_grad_norm 1.0