Matthew/causal (#787)

Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com>
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: RandNMR73 <quantumnium@gmail.com>
Co-authored-by: RandNMR73 <notomatthew31@gmail.com>

.buildkite/pipeline.yml

@@ -176,26 +176,3 @@ 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,15 +109,6 @@ 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,4 +64,5 @@ 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-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/S7HLCSTh" 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-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/rG0QpZdw" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">

csrc/attn/vmoba_attn/README.md

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

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

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

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

@@ -1,860 +0,0 @@
-# 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/distill/SFWan2.1-I2V/distill_dmd_i2v_1.3B.slurm

@@ -10,59 +10,72 @@
 #SBATCH --output=dmd_t2v_output/t2v_%j.out
 #SBATCH --error=dmd_t2v_output/t2v_%j.err
 #SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
 
 # Basic Info
+export WANDB_MODE="online"
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 # different cache dir for different processes
 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29503
+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_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
 
 # Configs
 NUM_GPUS=1
 
-# Model paths for Self-Forcing DMD distillation:
+# Model paths for 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/"
+DATA_DIR="data/crush-smol-single_processed_t2v/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
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd
+  --output_dir "checkpoints/SFwan_t2v_finetune"
+  --max_train_steps 500
   --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 ✓)
+  --num_frames 81
   --enable_gradient_checkpointing_type "full"
   --log_visualization
   --simulate_generator_forward
-  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --num_frame_per_block 3
   --enable_gradient_masking
   --gradient_mask_last_n_frames 21
 )
 
 # Parallel arguments
 parallel_args=(
-  --num_gpus $NUM_GPUS # 64
+  --num_gpus 1 # 64
   --sp_size 1
   --tp_size 1
   --hsdp_replicate_dim 1 # 64
-  --hsdp_shard_dim $NUM_GPUS
+  --hsdp_shard_dim 1
 )
 
 # Model arguments
@@ -84,7 +97,7 @@ dataset_args=(
 validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
+  --validation_steps 10
   --validation_sampling_steps "4"
   --validation_guidance_scale "6.0" # not used for dmd inference
 )
@@ -93,8 +106,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 1e-5
   --mixed_precision "bf16"
-  --training_state_checkpointing_steps 500
-  --weight_only_checkpointing_steps 500
+  --training_state_checkpointing_steps 50
+  --weight_only_checkpointing_steps 50
   --weight_decay 0.01
   --betas '0.0,0.999'
   --max_grad_norm 1.0
@@ -123,22 +136,14 @@ dmd_args=(
   --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 \
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
 --nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
     fastvideo/training/wan_self_forcing_distillation_pipeline.py \
     "${parallel_args[@]}" \
     "${model_args[@]}" \
@@ -147,5 +152,4 @@ torchrun \
     "${optimizer_args[@]}" \
     "${validation_args[@]}" \
     "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}" \
-    "${self_forcing_args[@]}"
+    "${dmd_args[@]}"

examples/distill/SFWan2.1-I2V/download_single_example.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+
+# Download the full dataset first
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+
+# Create a single-example dataset for debugging
+SINGLE_EXAMPLE_DIR="data/crush-smol-single"
+mkdir -p "$SINGLE_EXAMPLE_DIR/videos"
+
+# Copy the specific video that matches the validation.json style (macaron crushing)
+cp "data/crush-smol/videos/7P02AihYkCU-Scene-005.mp4" "$SINGLE_EXAMPLE_DIR/videos/"
+
+# Create a single-line videos.txt
+echo "videos/7P02AihYkCU-Scene-005.mp4" > "$SINGLE_EXAMPLE_DIR/videos.txt"
+
+# Create a single-line prompt.txt with the macaron crushing prompt
+echo "PIKA_CRUSH A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out." > "$SINGLE_EXAMPLE_DIR/prompt.txt"
+
+# Generate the JSON file and merge.txt for the single example
+python scripts/dataset_preparation/prepare_json_file.py --data_folder "$SINGLE_EXAMPLE_DIR" --output "videos2caption.json"
+
+# Create a validation.json that uses the same example for consistency
+cat > "$SINGLE_EXAMPLE_DIR/validation.json" << 'EOF'
+{
+    "data": [
+      {
+        "caption": "A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out.",
+        "image_path": null,
+        "video_path": null,
+        "num_inference_steps": 50,
+        "height": 480,
+        "width": 832,
+        "num_frames": 81
+      }
+    ]
+  }
+EOF
+
+echo "Single example dataset created at $SINGLE_EXAMPLE_DIR"
+echo "Contains:"
+echo "- 1 video: $(cat $SINGLE_EXAMPLE_DIR/videos.txt)"
+echo "- 1 prompt: $(cat $SINGLE_EXAMPLE_DIR/prompt.txt)"
+echo "- Validation file created with the same example for consistency" 

examples/distill/SFWan2.1-I2V/preprocess_single_example.sh

@@ -0,0 +1,29 @@
+#!/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-single_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 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 "t2v"
+
+# Copy the validation.json to the output directory for consistency
+cp "data/crush-smol-single/validation.json" "$OUTPUT_DIR/"
+
+echo "Preprocessing completed. Validation file copied to $OUTPUT_DIR/" 

examples/distill/SFWan2.1-I2V/validation.json

@@ -0,0 +1,31 @@
+{
+    "data": [
+      {
+        "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
+        "image_path": null,
+        "video_path": null,
+        "num_inference_steps": 50,
+        "height": 480,
+        "width": 832,
+        "num_frames": 77
+      },
+      {
+        "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
+        "image_path": null,
+        "video_path": null,
+        "num_inference_steps": 50,
+        "height": 480,
+        "width": 832,
+        "num_frames": 77
+      },
+      {
+        "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
+        "image_path": null,
+        "video_path": null,
+        "num_inference_steps": 50,
+        "height": 480,
+        "width": 832,
+        "num_frames": 77
+      }
+    ]
+  }

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

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

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

@@ -9,9 +9,9 @@ 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-14B-Diffusers",
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         # FastVideo will automatically handle distributed setup
-        num_gpus=4,
+        num_gpus=1,
         use_fsdp_inference=True,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
@@ -25,7 +25,9 @@ def main():
     # 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 watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open."
+        "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)
     # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
@@ -33,7 +35,11 @@ def main():
     # Generate another video with a different prompt, without reloading the
     # model!
     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")
+        "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)
 
 

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

@@ -10,7 +10,6 @@ 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 \

fastvideo/attention/backends/vmoba.py

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

@@ -1,16 +0,0 @@
-{
-    "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

@@ -1,16 +0,0 @@
-{
-    "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,29 +32,6 @@ 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."""
@@ -74,7 +51,6 @@ 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
@@ -144,12 +120,6 @@ 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,
@@ -204,10 +174,6 @@ 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,13 +15,9 @@ 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.VMOBA_ATTN,
-    )
+        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
+        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
+        AttentionBackendEnum.VIDEO_SPARSE_ATTN)
 
     hidden_size: int = 0
     num_attention_heads: int = 0

fastvideo/configs/models/dits/wanvideo.py

@@ -109,4 +109,4 @@ class WanVideoArchConfig(DiTArchConfig):
 class WanVideoConfig(DiTConfig):
     arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig)
 
-    prefix: str = "Wan"
+    prefix: str = "Wan"

fastvideo/configs/pipelines/registry.py

@@ -7,14 +7,12 @@ from collections.abc import Callable
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
+from fastvideo.configs.pipelines.wan import (FastWan2_1_T2V_480P_Config,
+                                             FastWan2_2_TI2V_5B_Config,
+                                             SelfForcingWanT2V480PConfig,
+                                             WanI2V480PConfig, WanI2V720PConfig,
+                                             WanT2V480PConfig, WanT2V720PConfig)
 
-# isort: off
-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)
-# isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
                              verify_model_config_and_directory)
@@ -35,10 +33,10 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
     "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
+    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig,
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig,
+    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig,
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
-    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
-    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
-    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
     # Add other specific weight variants
 }
 

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 = 3.0
+    flow_shift: int = 3
 
     # 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: float | None = 5.0
+    flow_shift: int = 5
 
 
 @dataclass
@@ -94,7 +94,7 @@ class WanI2V720PConfig(WanI2V480PConfig):
     # WanConfig-specific parameters with defaults
 
     # Denoising stage
-    flow_shift: float | None = 5.0
+    flow_shift: int = 5
 
 
 @dataclass
@@ -104,7 +104,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
     # WanConfig-specific parameters with defaults
 
     # Denoising stage
-    flow_shift: float | None = 8.0
+    flow_shift: int = 8
     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: float | None = 5.0
+    flow_shift: int = 5
     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: float | None = 5.0
+    flow_shift: int = 5
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 757, 522])
 
@@ -146,7 +146,6 @@ class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
 @dataclass
 class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
     is_causal: bool = True
-    flow_shift: float | None = 5.0
+    flow_shift: int = 5
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 750, 500, 250])
-    warp_denoising_step: bool = True

fastvideo/configs/sample/base.py

@@ -191,13 +191,6 @@ class SamplingParam:
             default=SamplingParam.image_path,
             help="Path to input image for image-to-video generation",
         )
-        parser.add_argument(
-            "--moba-config-path",
-            type=str,
-            default=None,
-            help=
-            "Path to a JSON file containing V-MoBA specific configurations.",
-        )
         return parser
 
 

fastvideo/configs/sample/registry.py

@@ -37,8 +37,6 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
-    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
-    Wan2_1_Fun_1_3B_InP_SamplingParam,
 
     # Wan2.2
     "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,

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
@@ -119,7 +119,7 @@ class FastVideoArgs:
     output_type: str = "pil"
 
     # CPU offload parameters
-    dit_cpu_offload: bool = True
+    dit_cpu_offload: bool = False
     use_fsdp_inference: bool = True
     text_encoder_cpu_offload: bool = True
     image_encoder_cpu_offload: bool = True
@@ -139,10 +139,6 @@ 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
 
@@ -170,16 +166,6 @@ 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
@@ -705,9 +691,6 @@ class TrainingArgs(FastVideoArgs):
     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":
@@ -1041,15 +1024,6 @@ 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,
@@ -1125,19 +1099,6 @@ class TrainingArgs(FastVideoArgs):
             "--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
 
@@ -1145,4 +1106,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

@@ -100,16 +100,7 @@ class ScaleResidual(nn.Module):
     def forward(self, residual: torch.Tensor, x: torch.Tensor,
                 gate: torch.Tensor) -> torch.Tensor:
         """Apply gated residual connection."""
-        # 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
+        return residual + x * gate
 
 
 # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py
@@ -168,7 +159,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 | int, shift: torch.Tensor,
+                gate: torch.Tensor, shift: torch.Tensor,
                 scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Apply gated residual connection, followed by layernorm and 
@@ -180,41 +171,12 @@ 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
-        if isinstance(gate, int):
-            # used by cross-attention, should be 1
-            assert gate == 1
-            residual_output = residual + x
-        elif isinstance(gate, torch.Tensor):
-            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)
-            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")
-        # residual_output.shape: [batch_size, seq_len, inner_dim]
-
+        residual_output = residual + x * gate
         # 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 + scale) + shift).flatten(1, 2)
-        else:
-            modulated = normalized * (1 + scale) + shift
+        modulated = normalized * (1.0 + scale) + shift
         return modulated, residual_output
 
 
@@ -256,24 +218,8 @@ 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
-            output = (
-                normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
-                (1 + scale) + shift).flatten(1, 2)
+            return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
         else:
-            # 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
+            return normalized * (1.0 + scale) + shift

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.zeros_(self.lora_B)
+            torch.nn.init.kaiming_uniform_(self.lora_B, a=math.sqrt(5))
         else:
             self.lora_A = None
             self.lora_B = None

fastvideo/layers/rotary_embedding.py

@@ -33,7 +33,6 @@ from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
 
-
 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
     x1 = x[..., :x.shape[-1] // 2]
     x2 = x[..., x.shape[-1] // 2:]

fastvideo/models/dits/causal_wanvideo.py

@@ -147,9 +147,6 @@ 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(
@@ -179,7 +176,7 @@ class CausalWanTransformerBlock(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = FP32LayerNorm(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)
@@ -212,7 +209,8 @@ class CausalWanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=True,
-            dtype=torch.float32)
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 2. Cross-attention
         # Only T2V for now
@@ -225,7 +223,8 @@ class CausalWanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=False,
-            dtype=torch.float32)
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -245,39 +244,27 @@ class CausalWanTransformerBlock(nn.Module):
         current_start: int = 0,
         cache_start: int | None = None,
     ) -> torch.Tensor:
-        # 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)
-        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
-        # e.shape: [batch_size, num_frames, 6, inner_dim]
-        assert e.shape == (bs, num_frames, 6, self.hidden_dim)
+        e = self.scale_shift_table + temb.float()
         shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
-            6, dim=2)
-        # *_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)
+            6, dim=1)
+        assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
-        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)
+        norm_hidden_states = (self.norm1(hidden_states.float()) *
+                              (1 + scale_msa) + shift_msa).to(orig_dtype)
         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.forward_native(query)
+            query = self.norm_q(query)
         if self.norm_k is not None:
-            key = self.norm_k.forward_native(key)
+            key = self.norm_k(key)
 
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
@@ -291,6 +278,8 @@ 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)
+        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,
@@ -299,10 +288,13 @@ 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)
+        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
 
@@ -365,7 +357,8 @@ class CausalWanTransformer3DModel(BaseDiT):
                                             norm_type="layer",
                                             eps=config.eps,
                                             elementwise_affine=False,
-                                            dtype=torch.float32)
+                                            dtype=torch.float32,
+                                            compute_dtype=torch.float32)
         self.proj_out = nn.Linear(
             inner_dim, config.out_channels * math.prod(config.patch_size))
         self.scale_shift_table = nn.Parameter(
@@ -375,7 +368,7 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         # Causal-specific
         self.block_mask = None
-        self.num_frame_per_block = 3
+        self.num_frame_per_block = 1
         self.independent_first_frame = False
 
         self.__post_init__()
@@ -487,19 +480,15 @@ class CausalWanTransformer3DModel(BaseDiT):
         )
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos,
-                     freqs_sin) if freqs_cos is not None else None
+        freqs_cis = (freqs_cos.float(),
+                     freqs_sin.float()) 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)
 
-        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_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
+            timestep, encoder_hidden_states, encoder_hidden_states_image)
+        timestep_proj = timestep_proj.unflatten(1, (6, -1))
 
         if encoder_hidden_states_image is not None:
             encoder_hidden_states = torch.concat(
@@ -537,15 +526,19 @@ class CausalWanTransformer3DModel(BaseDiT):
                                         **causal_kwargs)
 
         # 5. Output norm, projection & unpatchify
-        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
-        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
-                                                                    dim=2)
+        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
+                                                                          dim=1)
         hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
 
-        output = self.unpatchify(hidden_states, grid_sizes)
+        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)
 
-        return torch.stack(output)
+        return output
 
     def _forward_train(self,
                 hidden_states: torch.Tensor,
@@ -586,8 +579,8 @@ class CausalWanTransformer3DModel(BaseDiT):
         )
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos,
-                     freqs_sin) if freqs_cos is not None else None
+        freqs_cis = (freqs_cos.float(),
+                     freqs_sin.float()) if freqs_cos is not None else None
 
         # Construct blockwise causal attn mask
         if self.block_mask is None:
@@ -600,15 +593,11 @@ 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)
 
-        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_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
+            timestep, encoder_hidden_states, encoder_hidden_states_image)
+        timestep_proj = timestep_proj.unflatten(1, (6, -1))
 
         if encoder_hidden_states_image is not None:
             encoder_hidden_states = torch.concat(
@@ -634,15 +623,19 @@ class CausalWanTransformer3DModel(BaseDiT):
                                         block_mask=self.block_mask)
 
         # 5. Output norm, projection & unpatchify
-        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
-        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
-                                                                    dim=2)
+        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
+                                                                          dim=1)
         hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
 
-        output = self.unpatchify(hidden_states, grid_sizes)
+        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)
 
-        return torch.stack(output)
+        return output
 
     def forward(
         self,
@@ -653,30 +646,3 @@ 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,5 +1,3 @@
-import torch
-import torch.nn as nn
 # SPDX-License-Identifier: Apache-2.0
 
 import math
@@ -39,14 +37,16 @@ class WanImageEmbedding(torch.nn.Module):
     def __init__(self, in_features: int, out_features: int):
         super().__init__()
 
-        self.norm1 = nn.LayerNorm(in_features)
+        self.norm1 = FP32LayerNorm(in_features)
         self.ff = MLP(in_features, in_features, out_features, act_type="gelu")
-        self.norm2 = nn.LayerNorm(out_features)
+        self.norm2 = FP32LayerNorm(out_features)
 
-    def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
+    def forward(self,
+                encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
+        dtype = encoder_hidden_states_image.dtype
         hidden_states = self.norm1(encoder_hidden_states_image)
         hidden_states = self.ff(hidden_states)
-        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.norm2(hidden_states).to(dtype)
         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", freq_dtype=torch.float64)
+            dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
         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.forward_native(self.to_q(x)[0]).view(b, -1, n, d)
+        q = self.norm_q(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.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
+                k = self.norm_k(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.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
+            k = self.norm_k(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.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)
+        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)
         v = self.to_v(context)[0].view(b, -1, n, d)
-        k_img = self.norm_added_k.forward_native(self.add_k_proj(context_img)[0]).view(
+        k_img = self.norm_added_k(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 = nn.LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = FP32LayerNorm(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)
+            dtype=torch.float32,
+            compute_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)
+            dim,
+            norm_type="layer",
+            eps=eps,
+            elementwise_affine=False,
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -319,11 +319,12 @@ 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
+                self.scale_shift_table.unsqueeze(0) + temb.float()
             ).chunk(6, dim=2)
             # batch_size, seq_len, 1, inner_dim
             shift_msa = shift_msa.squeeze(2)
@@ -334,20 +335,22 @@ 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
+            e = self.scale_shift_table + temb.float()
             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) * (1 + scale_msa) + shift_msa
+        norm_hidden_states = (self.norm1(hidden_states.float()) *
+                              (1 + scale_msa) + shift_msa).to(orig_dtype)
         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.forward_native(query)
+            query = self.norm_q(query)
         if self.norm_k is not None:
-            key = self.norm_k.forward_native(key)
+            key = self.norm_k(key)
 
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
@@ -367,20 +370,26 @@ 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,
@@ -397,7 +406,7 @@ class WanTransformerBlock_VSA(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = FP32LayerNorm(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)
@@ -429,7 +438,8 @@ class WanTransformerBlock_VSA(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=True,
-            dtype=torch.float32)
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 2. Cross-attention
         if added_kv_proj_dim is not None:
@@ -449,7 +459,8 @@ class WanTransformerBlock_VSA(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=False,
-            dtype=torch.float32)
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -469,22 +480,23 @@ 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
+        e = self.scale_shift_table + temb.float()
         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) *
-                              (1 + scale_msa) + shift_msa)
+        norm_hidden_states = (self.norm1(hidden_states.float()) *
+                              (1 + scale_msa) + shift_msa).to(orig_dtype)
         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.forward_native(query)
+            query = self.norm_q(query)
         if self.norm_k is not None:
-            key = self.norm_k.forward_native(key)
+            key = self.norm_k(key)
 
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
@@ -509,6 +521,8 @@ 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,
@@ -516,15 +530,17 @@ 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
@@ -582,7 +598,8 @@ class WanTransformer3DModel(CachableDiT):
                                             norm_type="layer",
                                             eps=config.eps,
                                             elementwise_affine=False,
-                                            dtype=torch.float32)
+                                            dtype=torch.float32,
+                                            compute_dtype=torch.float32)
         self.proj_out = nn.Linear(
             inner_dim, config.out_channels * math.prod(config.patch_size))
         self.scale_shift_table = nn.Parameter(
@@ -642,12 +659,10 @@ 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,
-                     freqs_sin) if freqs_cos is not None else None
+        freqs_cis = (freqs_cos.float(),
+                     freqs_sin.float()) 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)
@@ -657,8 +672,6 @@ 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:
@@ -715,35 +728,14 @@ class WanTransformer3DModel(CachableDiT):
         hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
 
-        output = self.unpatchify(hidden_states, grid_sizes)
+        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)
 
-        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
+        return output
 
     def maybe_cache_states(self, hidden_states: torch.Tensor,
                            original_hidden_states: torch.Tensor) -> None:
@@ -835,4 +827,5 @@ 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

@@ -435,7 +435,7 @@ class TransformerLoader(ComponentLoader):
         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]

fastvideo/models/loader/fsdp_load.py

@@ -248,6 +248,13 @@ def load_model_from_full_model_state_dict(
     sharded_sd = {}
     custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
         full_sd_iterator, param_names_mapping)  # type: ignore
+    print(custom_param_sd.keys())
+    print("--------------------------------")
+    print("--------------------------------")
+    print("--------------------------------")
+    print("--------------------------------")
+    print("--------------------------------")
+    print(meta_sd.keys())
     for target_param_name, full_tensor in custom_param_sd.items():
         meta_sharded_param = meta_sd.get(target_param_name)
         if meta_sharded_param is None:

fastvideo/models/registry.py

@@ -61,9 +61,6 @@ _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,31 +635,8 @@ 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

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

@@ -145,30 +145,8 @@ 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]
     """
-    # 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]
+    timestep = timestep.expand(noise_input_latent.shape[0])
     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,7 +16,8 @@ from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
                                         CausalDMDDenosingStage,
                                         InputValidationStage,
                                         LatentPreparationStage,
-                                        TextEncodingStage)
+                                        TextEncodingStage,
+                                        TimestepPreparationStage)
 # isort: on
 
 logger = init_logger(__name__)
@@ -28,6 +29,10 @@ 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."""
 
@@ -43,6 +48,10 @@ 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 "transformer" in name:
+                if name == "transformer":
                     module.requires_grad_(True)
                 else:
                     module.requires_grad_(False)
@@ -136,9 +136,11 @@ class ComposedPipelineBase(ABC):
 
             kwargs['model_path'] = model_path
             fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
+            logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
         else:
             assert args is not None, "args must be provided for training mode"
             fastvideo_args = TrainingArgs.from_cli_args(args)
+            logger.info("training args in from_pretrained: %s", fastvideo_args)
             # TODO(will): fix this so that its not so ugly
             fastvideo_args.model_path = model_path
             for key, value in kwargs.items():

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 DeviceMesh, init_device_mesh
+from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.tensor import DTensor
 
 from fastvideo.distributed import get_local_torch_device
@@ -32,7 +32,6 @@ 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()
@@ -82,17 +81,6 @@ 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:
@@ -100,12 +88,18 @@ 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"])
-        set_lora_grads(self.lora_layers, device_mesh)
-        set_lora_grads(self.lora_layers_critic, device_mesh)
+        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))
 
     def convert_to_lora_layers(self) -> None:
         """
@@ -137,24 +131,6 @@ 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
@@ -248,4 +224,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

@@ -241,5 +241,5 @@ class TrainingBatch:
 
 @dataclass
 class PreprocessBatch(ForwardBatch):
-    video_loader: list["VideoDecoder"] | list[str] = field(default_factory=list)
-    video_file_name: list[str] = field(default_factory=list)
+    video_loader: list["VideoDecoder"] = field(default_factory=list)
+    video_file_name: list[str] = field(default_factory=list)

fastvideo/pipelines/preprocess/preprocess_stages.py

@@ -4,11 +4,9 @@ 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
@@ -63,16 +61,7 @@ class VideoTransformStage(PipelineStage):
                 else:
                     frame_indices = frame_indices[:self.num_frames]
 
-            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 = batch.video_loader[i].get_frames_at(frame_indices).data
             video = self.video_transform(video)
             video_pixel_batch.append(video)
 

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,6 +36,7 @@ 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
@@ -70,16 +71,8 @@ class CausalDMDDenosingStage(DenoisingStage):
         # Timesteps for DMD
         timesteps = torch.tensor(
             fastvideo_args.pipeline_config.dmd_denoising_steps,
-            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)
+            dtype=torch.long,
+            device=get_local_torch_device())
 
         # Image kwargs (kept empty unless caller provides compatible args)
         image_kwargs: dict = {}
@@ -269,14 +262,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)
                         pred_noise_btchw = self.transformer(
                             latent_model_input,
                             prompt_embeds,
-                            t_expanded_noise,
+                            t_expand,
                             kv_cache=self.kv_cache1,
                             crossattn_cache=self.crossattn_cache,
                             current_start=(pos_start_base + start_index) *
@@ -337,11 +326,10 @@ class CausalDMDDenosingStage(DenoisingStage):
                     set_forward_context(current_timestep=0,
                                         attn_metadata=attn_metadata,
                                         forward_batch=batch):
-                    t_expanded_context = t_context.unsqueeze(1)
                     _ = self.transformer(
                         context_bcthw,
                         prompt_embeds,
-                        t_expanded_context,
+                        t_context,
                         kv_cache=self.kv_cache1,
                         crossattn_cache=self.crossattn_cache,
                         current_start=(pos_start_base + start_index) *
@@ -419,27 +407,3 @@ 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/denoising.py

@@ -40,13 +40,6 @@ 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)
@@ -84,7 +77,6 @@ class DenoisingStage(PipelineStage):
             supported_attention_backends=(
                 AttentionBackendEnum.SLIDING_TILE_ATTN,
                 AttentionBackendEnum.VIDEO_SPARSE_ATTN,
-                AttentionBackendEnum.VMOBA_ATTN,
                 AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA
             )  # hack
         )
@@ -157,8 +149,7 @@ class DenoisingStage(PipelineStage):
         # Prepare image latents and embeddings for I2V generation
         image_embeds = batch.image_embeds
         if len(image_embeds) > 0:
-            assert not torch.isnan(
-                image_embeds[0]).any(), "image_embeds contains nan"
+            assert torch.isnan(image_embeds[0]).sum() == 0
             image_embeds = [
                 image_embed.to(target_dtype) for image_embed in image_embeds
             ]
@@ -195,13 +186,11 @@ class DenoisingStage(PipelineStage):
         # Get latents and embeddings
         latents = batch.latents
         prompt_embeds = batch.prompt_embeds
-        assert not torch.isnan(
-            prompt_embeds[0]).any(), "prompt_embeds contains nan"
+        assert torch.isnan(prompt_embeds[0]).sum() == 0
         if batch.do_classifier_free_guidance:
             neg_prompt_embeds = batch.negative_prompt_embeds
             assert neg_prompt_embeds is not None
-            assert not torch.isnan(
-                neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
+            assert torch.isnan(neg_prompt_embeds[0]).sum() == 0
 
         # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
         if fastvideo_args.boundary_ratio is not None:
@@ -279,9 +268,6 @@ class DenoisingStage(PipelineStage):
                     latent_model_input = torch.cat(
                         [latent_model_input, batch.image_latent],
                         dim=1).to(target_dtype)
-
-                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()
@@ -294,6 +280,7 @@ 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)
 
@@ -338,31 +325,6 @@ 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
@@ -775,8 +737,7 @@ class DmdDenoisingStage(DenoisingStage):
 
         video_raw_latent_shape = latents.shape
         prompt_embeds = batch.prompt_embeds
-        assert not torch.isnan(
-            prompt_embeds[0]).any(), "prompt_embeds contains nan"
+        assert torch.isnan(prompt_embeds[0]).sum() == 0
         timesteps = torch.tensor(
             fastvideo_args.pipeline_config.dmd_denoising_steps,
             dtype=torch.long,
@@ -815,8 +776,7 @@ class DmdDenoisingStage(DenoisingStage):
                         batch.image_latent.permute(0, 2, 1, 3, 4)
                     ],
                                                    dim=2).to(target_dtype)
-                assert not torch.isnan(
-                    latent_model_input).any(), "latent_model_input contains nan"
+                assert torch.isnan(latent_model_input).sum() == 0
 
                 # Prepare inputs for transformer
                 t_expand = t.repeat(latent_model_input.shape[0])
@@ -878,6 +838,7 @@ class DmdDenoisingStage(DenoisingStage):
                             **pos_cond_kwargs,
                         ).permute(0, 2, 1, 3, 4)
 
+                    from fastvideo.models.utils import pred_noise_to_pred_video
                     pred_video = pred_noise_to_pred_video(
                         pred_noise=pred_noise.flatten(0, 1),
                         noise_input_latent=noise_latents.flatten(0, 1),

fastvideo/pipelines/stages/encoding.py

@@ -92,6 +92,24 @@ 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/text_encoding.py

@@ -59,37 +59,58 @@ class TextEncodingStage(PipelineStage):
         assert len(self.text_encoders) == len(
             fastvideo_args.pipeline_config.text_encoder_configs)
 
-        # Encode positive prompt with all available encoders
-        assert batch.prompt is not None
-        prompt_text: str | list[str] = batch.prompt
-        all_indices: list[int] = list(range(len(self.text_encoders)))
-        prompt_embeds_list, prompt_masks_list = self.encode_text(
-            prompt_text,
-            fastvideo_args,
-            encoder_index=all_indices,
-            return_attention_mask=True,
-        )
-        for pe in prompt_embeds_list:
-            batch.prompt_embeds.append(pe)
-        if batch.prompt_attention_mask is not None:
-            for am in prompt_masks_list:
-                batch.prompt_attention_mask.append(am)
+        for tokenizer, text_encoder, encoder_config, preprocess_func, postprocess_func in zip(
+                self.tokenizers,
+                self.text_encoders,
+                fastvideo_args.pipeline_config.text_encoder_configs,
+                fastvideo_args.pipeline_config.preprocess_text_funcs,
+                fastvideo_args.pipeline_config.postprocess_text_funcs,
+                strict=True):
 
-        # Encode negative prompt if CFG is enabled
-        if batch.do_classifier_free_guidance:
-            assert isinstance(batch.negative_prompt, str)
-            neg_embeds_list, neg_masks_list = self.encode_text(
-                batch.negative_prompt,
-                fastvideo_args,
-                encoder_index=all_indices,
-                return_attention_mask=True,
-            )
-            assert batch.negative_prompt_embeds is not None
-            for ne in neg_embeds_list:
-                batch.negative_prompt_embeds.append(ne)
-            if batch.negative_attention_mask is not None:
-                for nm in neg_masks_list:
-                    batch.negative_attention_mask.append(nm)
+            assert isinstance(batch.prompt, str | list)
+            if isinstance(batch.prompt, str):
+                batch.prompt = [batch.prompt]
+            texts = []
+            for prompt_str in batch.prompt:
+                texts.append(preprocess_func(prompt_str))
+            text_inputs = tokenizer(texts,
+                                    **encoder_config.tokenizer_kwargs).to(
+                                        get_local_torch_device())
+            input_ids = text_inputs["input_ids"]
+            attention_mask = text_inputs["attention_mask"]
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                outputs = text_encoder(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                )
+            prompt_embeds = postprocess_func(outputs)
+            batch.prompt_embeds.append(prompt_embeds)
+            if batch.prompt_attention_mask is not None:
+                batch.prompt_attention_mask.append(attention_mask)
+
+            if batch.do_classifier_free_guidance:
+                assert isinstance(batch.negative_prompt, str)
+                negative_text = preprocess_func(batch.negative_prompt)
+                negative_text_inputs = tokenizer(
+                    negative_text, **encoder_config.tokenizer_kwargs).to(
+                        get_local_torch_device())
+                negative_input_ids = negative_text_inputs["input_ids"]
+                negative_attention_mask = negative_text_inputs["attention_mask"]
+                with set_forward_context(current_timestep=0,
+                                         attn_metadata=None):
+                    negative_outputs = text_encoder(
+                        input_ids=negative_input_ids,
+                        attention_mask=negative_attention_mask,
+                        output_hidden_states=True,
+                    )
+                negative_prompt_embeds = postprocess_func(negative_outputs)
+
+                assert batch.negative_prompt_embeds is not None
+                batch.negative_prompt_embeds.append(negative_prompt_embeds)
+                if batch.negative_attention_mask is not None:
+                    batch.negative_attention_mask.append(
+                        negative_attention_mask)
 
         return batch
 
@@ -108,171 +129,6 @@ class TextEncodingStage(PipelineStage):
                          V.none_or_list)
         return result
 
-    @torch.no_grad()
-    def encode_text(
-        self,
-        text: str | list[str],
-        fastvideo_args: FastVideoArgs,
-        encoder_index: int | list[int] | None = None,
-        return_attention_mask: bool = False,
-        return_type: str = "list",  # one of: "list", "dict", "stack"
-        device: torch.device | str | None = None,
-        dtype: torch.dtype | None = None,
-        max_length: int | None = None,
-        truncation: bool | None = None,
-        padding: bool | str | None = None,
-    ):
-        """
-        Encode plain text using selected text encoder(s) and return embeddings.
-
-        Args:
-            text: A single string or a list of strings to encode.
-            fastvideo_args: The inference arguments providing pipeline config,
-                including tokenizer and encoder settings, preprocess and postprocess
-                functions.
-            encoder_index: Encoder selector by index. Accepts an int or list of ints.
-            return_attention_mask: If True, also return attention masks for each
-                selected encoder.
-            return_type: "list" (default) returns a list aligned with selection;
-                "dict" returns a dict keyed by encoder index as a string; "stack" stacks along a
-                new first dimension (requires matching shapes).
-            device: Optional device override for inputs; defaults to local torch device.
-            dtype: Optional dtype to cast returned embeddings to.
-            max_length: Optional per-call tokenizer override.
-            truncation: Optional per-call tokenizer override.
-            padding: Optional per-call tokenizer override.
-
-        Returns:
-            Depending on return_type and return_attention_mask:
-            - list: List[Tensor] or (List[Tensor], List[Tensor])
-            - dict: Dict[str, Tensor] or (Dict[str, Tensor], Dict[str, Tensor])
-            - stack: Tensor of shape [num_encoders, ...] or a tuple with stacked
-              attention masks
-        """
-
-        assert len(self.tokenizers) == len(self.text_encoders)
-        assert len(self.text_encoders) == len(
-            fastvideo_args.pipeline_config.text_encoder_configs)
-
-        # Resolve selection into indices
-        encoder_cfgs = fastvideo_args.pipeline_config.text_encoder_configs
-        if encoder_index is None:
-            indices: list[int] = [0]
-        elif isinstance(encoder_index, int):
-            indices = [encoder_index]
-        else:
-            indices = list(encoder_index)
-        # validate range
-        num_encoders = len(self.text_encoders)
-        for idx in indices:
-            if idx < 0 or idx >= num_encoders:
-                raise IndexError(
-                    f"encoder index {idx} out of range [0, {num_encoders-1}]")
-
-        # Validate indices are within range
-        num_encoders = len(self.text_encoders)
-
-        # Normalize input to list[str]
-        assert isinstance(text, str | list)
-        if isinstance(text, str):
-            texts: list[str] = [text]
-        else:
-            texts = text
-
-        embeds_list: list[torch.Tensor] = []
-        attn_masks_list: list[torch.Tensor] = []
-
-        preprocess_funcs = fastvideo_args.pipeline_config.preprocess_text_funcs
-        postprocess_funcs = fastvideo_args.pipeline_config.postprocess_text_funcs
-        encoder_cfgs = fastvideo_args.pipeline_config.text_encoder_configs
-
-        if return_type not in ("list", "dict", "stack"):
-            raise ValueError(
-                f"Invalid return_type '{return_type}'. Expected one of: 'list', 'dict', 'stack'"
-            )
-
-        target_device = device if device is not None else get_local_torch_device(
-        )
-
-        for i in indices:
-            tokenizer = self.tokenizers[i]
-            text_encoder = self.text_encoders[i]
-            encoder_config = encoder_cfgs[i]
-            preprocess_func = preprocess_funcs[i]
-            postprocess_func = postprocess_funcs[i]
-
-            processed_texts: list[str] = []
-            for prompt_str in texts:
-                processed_texts.append(preprocess_func(prompt_str))
-
-            tok_kwargs = dict(encoder_config.tokenizer_kwargs)
-            if max_length is not None:
-                tok_kwargs["max_length"] = max_length
-            if truncation is not None:
-                tok_kwargs["truncation"] = truncation
-            if padding is not None:
-                tok_kwargs["padding"] = padding
-
-            text_inputs = tokenizer(processed_texts,
-                                    **tok_kwargs).to(target_device)
-
-            input_ids = text_inputs["input_ids"]
-            attention_mask = text_inputs["attention_mask"]
-
-            with set_forward_context(current_timestep=0, attn_metadata=None):
-                outputs = text_encoder(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
-                )
-
-            prompt_embeds = postprocess_func(outputs)
-            if dtype is not None:
-                prompt_embeds = prompt_embeds.to(dtype=dtype)
-            embeds_list.append(prompt_embeds)
-            if return_attention_mask:
-                attn_masks_list.append(attention_mask)
-
-        # Shape results according to return_type
-        if return_type == "list":
-            if return_attention_mask:
-                return embeds_list, attn_masks_list
-            return embeds_list
-
-        if return_type == "dict":
-            key_strs = [str(i) for i in indices]
-            embeds_dict = {
-                k: v
-                for k, v in zip(key_strs, embeds_list, strict=False)
-            }
-            if return_attention_mask:
-                attn_dict = {
-                    k: v
-                    for k, v in zip(key_strs, attn_masks_list, strict=False)
-                }
-                return embeds_dict, attn_dict
-            return embeds_dict
-
-        # return_type == "stack"
-        # Validate shapes are compatible
-        base_shape = list(embeds_list[0].shape)
-        for t in embeds_list[1:]:
-            if list(t.shape) != base_shape:
-                raise ValueError(
-                    f"Cannot stack embeddings with differing shapes: {[list(t.shape) for t in embeds_list]}"
-                )
-        stacked_embeds = torch.stack(embeds_list, dim=0)
-        if return_attention_mask:
-            base_mask_shape = list(attn_masks_list[0].shape)
-            for m in attn_masks_list[1:]:
-                if list(m.shape) != base_mask_shape:
-                    raise ValueError(
-                        f"Cannot stack attention masks with differing shapes: {[list(m.shape) for m in attn_masks_list]}"
-                    )
-            stacked_masks = torch.stack(attn_masks_list, dim=0)
-            return stacked_embeds, stacked_masks
-        return stacked_embeds
-
     def verify_output(self, batch: ForwardBatch,
                       fastvideo_args: FastVideoArgs) -> VerificationResult:
         """Verify text encoding stage outputs."""

fastvideo/platforms/cuda.py

@@ -159,20 +159,6 @@ 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,7 +19,6 @@ 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

@@ -1,58 +0,0 @@
-# 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,18 +102,10 @@ 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,8 +4,7 @@ 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.
 
-reference videos were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
-causal videos were generated on commit b318063c0a4618f1d5d99ea82ca67a06aad0d19d
+all reference videos are were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
 
 ## Generation Details
 
@@ -77,4 +76,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

@@ -1,152 +0,0 @@
-# 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"])
+@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
 @pytest.mark.parametrize("model_id", list(I2V_MODEL_TO_PARAMS.keys()))
 def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     """

fastvideo/tests/stages/test_text_encoding.py

@@ -1,128 +0,0 @@
-import torch
-import types
-import pytest
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig, BaseEncoderOutput
-from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
-
-class TensorDict(dict):
-    def to(self, device):
-        return TensorDict({k: v.to(device) for k, v in self.items()})
-
-class FakeTokenizer:
-    def __call__(self, texts, **kwargs):
-        B = len(texts)
-        seq_len = int(kwargs.get("max_length", 4))
-        return TensorDict({
-            "input_ids": torch.arange(B * seq_len).view(B, seq_len),
-            "attention_mask": torch.ones(B, seq_len, dtype=torch.long),
-        })
-
-class FakeTextEncoder(torch.nn.Module):
-    def __init__(self, hidden_size=8):
-        super().__init__()
-        self.hidden_size = hidden_size
-
-    def forward(self, input_ids, attention_mask, output_hidden_states=True):
-        B, T = input_ids.shape
-        last_hidden_state = torch.arange(B * T * self.hidden_size, dtype=torch.float32).view(B, T, self.hidden_size)
-        return types.SimpleNamespace(last_hidden_state=last_hidden_state)
-
-def id_preprocess(x: str) -> str:
-    return x
-
-def take_mean_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
-    # [B, T, H] -> [B, H]
-    return outputs.last_hidden_state.mean(dim=1)
-
-def make_args(num_encoders=2, text_len=4, hidden_size=8):
-    enc_cfgs = []
-    preprocess_fns = []
-    postprocess_fns = []
-    for _ in range(num_encoders):
-        arch = TextEncoderArchConfig(text_len=text_len)
-        enc_cfgs.append(TextEncoderConfig(arch_config=arch))
-        preprocess_fns.append(id_preprocess)
-        postprocess_fns.append(take_mean_postprocess)
-    pipe_cfg = PipelineConfig(
-        text_encoder_configs=tuple(enc_cfgs),
-        text_encoder_precisions=tuple(["fp32"] * num_encoders),
-        preprocess_text_funcs=tuple(preprocess_fns),
-        postprocess_text_funcs=tuple(postprocess_fns),
-    )
-    return FastVideoArgs(model_path="", pipeline_config=pipe_cfg), hidden_size
-
-def make_stage(num_encoders=2, hidden_size=8):
-    tokenizers = [FakeTokenizer() for _ in range(num_encoders)]
-    encoders = [FakeTextEncoder(hidden_size=hidden_size) for _ in range(num_encoders)]
-    return TextEncodingStage(text_encoders=encoders, tokenizers=tokenizers)
-
-def test_encode_text_selection_and_shapes():
-    fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8)
-    stage = make_stage(num_encoders=2, hidden_size=hidden)
-
-    # list return, two encoders
-    embeds = stage.encode_text(["a", "b"], fastvideo_args, encoder_index=[0, 1])
-    assert isinstance(embeds, list) and len(embeds) == 2
-    for e in embeds:
-        assert e.shape == (2, hidden)
-
-    # with masks
-    embeds2, masks2 = stage.encode_text("a", fastvideo_args, encoder_index=[1], return_attention_mask=True)
-    assert len(embeds2) == 1 and len(masks2) == 1
-    assert embeds2[0].shape == (1, hidden)
-    assert masks2[0].shape == (1, 4)
-
-    # dict return
-    d = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="dict")
-    assert set(d.keys()) == {"0", "1"}
-    assert d["0"].shape == (2, hidden)
-
-    # stack return
-    s = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="stack")
-    assert s.shape == (2, 2, hidden)  # [encoders, batch, hidden]
-
-    # overrides: dtype + max_length
-    e3, m3 = stage.encode_text(["a"], fastvideo_args, encoder_index=[0], dtype=torch.float16, return_attention_mask=True, max_length=3)
-    assert e3[0].dtype == torch.float16
-    assert m3[0].shape[1] == 3
-
-def test_forward_integration_cfg_off_and_on():
-    fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8)
-    stage = make_stage(num_encoders=2, hidden_size=hidden)
-
-    # CFG off
-    batch = ForwardBatch(
-        data_type="video",
-        prompt="a cat",
-        negative_prompt="",
-        do_classifier_free_guidance=False,
-        prompt_embeds=[],
-        negative_prompt_embeds=None,
-        prompt_attention_mask=[],
-        negative_attention_mask=None,
-    )
-    out = stage.forward(batch, fastvideo_args)
-    assert len(out.prompt_embeds) == 2
-    for e in out.prompt_embeds:
-        assert e.shape[1] == hidden
-
-    # CFG on
-    batch2 = ForwardBatch(
-        data_type="video",
-        prompt=["a cat", "a dog"],
-        negative_prompt="bad picture",
-        do_classifier_free_guidance=True,
-        prompt_embeds=[],
-        negative_prompt_embeds=[],
-        prompt_attention_mask=[],
-        negative_attention_mask=[],
-    )
-    out2 = stage.forward(batch2, fastvideo_args)
-    assert len(out2.prompt_embeds) == 2
-    assert len(out2.negative_prompt_embeds) == 2
-    assert len(out2.prompt_attention_mask) == 2
-    assert len(out2.negative_attention_mask) == 2

fastvideo/tests/transformers/test_ori_causal_wanvideo.py

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

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

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

@@ -12,7 +12,6 @@ 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
@@ -40,7 +39,7 @@ from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
     EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
     get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint,
-    shift_timestep, compute_density_for_timestep_sampling, get_sigmas)
+    shift_timestep)
 from fastvideo.utils import (is_vsa_available, maybe_download_model,
                              set_random_seed, verify_model_config_and_directory)
 
@@ -91,9 +90,6 @@ class DistillationPipeline(TrainingPipeline):
         self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
             shift=self.timestep_shift)
 
-        self.transformer_2 = self.get_module("transformer_2", None)
-        self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
-        
         if training_args.real_score_model_path:
             logger.info(
                 f"Loading real score transformer from: {training_args.real_score_model_path}"
@@ -130,39 +126,6 @@ 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(
@@ -317,9 +280,6 @@ 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()
 
@@ -475,9 +435,7 @@ class DistillationPipeline(TrainingPipeline):
         training_batch = self._build_distill_input_kwargs(
             noisy_latent, timestep, training_batch.conditional_dict,
             training_batch)
-        
-        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
-        pred_noise = current_model(**training_batch.input_kwargs).permute(
+        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
             0, 2, 1, 3, 4)
         pred_video = pred_noise_to_pred_video(
             pred_noise=pred_noise.flatten(0, 1),
@@ -519,7 +477,6 @@ 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():
@@ -531,7 +488,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 = current_model(
+                    pred_flow = self.transformer(
                         **training_batch_temp.input_kwargs).permute(
                             0, 2, 1, 3, 4)
                     pred_clean = pred_noise_to_pred_video(
@@ -574,7 +531,7 @@ class DistillationPipeline(TrainingPipeline):
         training_batch = self._build_distill_input_kwargs(
             noisy_input, target_timestep, training_batch.conditional_dict,
             training_batch)
-        pred_noise = current_model(**training_batch.input_kwargs).permute(
+        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
             0, 2, 1, 3, 4)
         pred_video = pred_noise_to_pred_video(
             pred_noise=pred_noise.flatten(0, 1),
@@ -588,7 +545,6 @@ 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],
@@ -613,7 +569,7 @@ class DistillationPipeline(TrainingPipeline):
 
             noisy_latent = self.noise_scheduler.add_noise(
                 generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
-                timestep).detach().unflatten(0, (1, generator_pred_video.shape[1]))
+                timestep).unflatten(0, (1, generator_pred_video.shape[1]))
 
             # fake_score_transformer forward
             training_batch = self._build_distill_input_kwargs(
@@ -662,24 +618,24 @@ class DistillationPipeline(TrainingPipeline):
                 pred_real_video_uncond) * self.real_score_guidance_scale
 
             grad = (faker_score_pred_video - real_score_pred_video) / torch.abs(
-                original_latent - real_score_pred_video).mean()
+                generator_pred_video - real_score_pred_video).mean()
             grad = torch.nan_to_num(grad)
 
         dmd_loss = 0.5 * F.mse_loss(
-            original_latent.float(),
-            (original_latent.float() - grad.float()).detach())
+            generator_pred_video.float(),
+            (generator_pred_video.float() - grad.float()).detach())
 
         training_batch.dmd_latent_vis_dict.update({
             "training_batch_dmd_fwd_clean_latent":
             training_batch.latents,
             "generator_pred_video":
-            original_latent.detach(),
+            generator_pred_video,
             "real_score_pred_video":
-            real_score_pred_video.detach(),
+            real_score_pred_video,
             "faker_score_pred_video":
-            faker_score_pred_video.detach(),
+            faker_score_pred_video,
             "dmd_timestep":
-            timestep.detach(),
+            timestep,
         })
 
         return dmd_loss
@@ -772,17 +728,16 @@ class DistillationPipeline(TrainingPipeline):
             "encoder_hidden_states": training_batch.encoder_hidden_states,
             "encoder_attention_mask": training_batch.encoder_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
+        unconditional_dict = {
+            "encoder_hidden_states": self.negative_prompt_embeds,
+            "encoder_attention_mask": self.negative_prompt_attention_mask,
+        }
 
         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
@@ -816,9 +771,6 @@ 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 = {}
@@ -847,31 +799,12 @@ class DistillationPipeline(TrainingPipeline):
                         attn_metadata=batch_gen.attn_metadata_vsa):
                     (dmd_loss / gradient_accumulation_steps).backward()
                 total_dmd_loss += dmd_loss.detach().item()
-            
-            # 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)
+            self._clip_model_grad_norm_(batch_gen, self.transformer)
+            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)
+                self.generator_ema.update(self.transformer)
 
             avg_dmd_loss = torch.tensor(total_dmd_loss /
                                         gradient_accumulation_steps,
@@ -896,18 +829,9 @@ 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()
-        
-        # 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.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,
@@ -1465,4 +1389,4 @@ class DistillationPipeline(TrainingPipeline):
                                   self.training_args.max_train_steps)
 
         if get_sp_group():
-            cleanup_dist_env_and_memory()
+            cleanup_dist_env_and_memory()

fastvideo/training/training_pipeline.py

@@ -23,7 +23,6 @@ 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
@@ -43,15 +42,11 @@ 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_vmoba_available, is_vsa_available,
-#                              set_random_seed, shallow_asdict)
-from fastvideo.utils import (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__)
 
@@ -128,6 +123,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         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"
@@ -225,19 +221,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
         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:
@@ -283,15 +266,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
                             dtype=latents.dtype)
         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()
@@ -812,4 +786,4 @@ class TrainingPipeline(LoRAPipeline, ABC):
         training_args.inference_mode = False
         self.transformer.train()
         if getattr(self, "transformer_2", None) is not None:
-            self.transformer_2.train()
+            self.transformer_2.train()

fastvideo/training/training_pipeline_switch.py

@@ -0,0 +1,797 @@
+# SPDX-License-Identifier: Apache-2.0
+import dataclasses
+import math
+import os
+import time
+from abc import ABC, abstractmethod
+from collections import deque
+from collections.abc import Iterator
+from typing import Any
+
+import imageio
+import numpy as np
+import torch
+import torch.distributed as dist
+import torchvision
+from diffusers import FlowMatchEulerDiscreteScheduler
+from einops import rearrange
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm.auto import tqdm
+
+import fastvideo.envs as envs
+from fastvideo.attention.backends.video_sparse_attn import (
+    VideoSparseAttentionMetadataBuilder)
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset import build_parquet_map_style_dataloader
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.dataset.validation_dataset import ValidationDataset
+from fastvideo.distributed import (cleanup_dist_env_and_memory,
+                                   get_local_torch_device, get_sp_group,
+                                   get_world_group)
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import (ComposedPipelineBase, ForwardBatch,
+                                 LoRAPipeline, TrainingBatch)
+from fastvideo.training.activation_checkpoint import (
+    apply_activation_checkpointing)
+from fastvideo.training.training_utils import (
+    clip_grad_norm_while_handling_failing_dtensor_cases,
+    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
+
+import wandb  # isort: skip
+
+vsa_available = is_vsa_available()
+
+logger = init_logger(__name__)
+
+
+def _get_trainable_params(model: torch.nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+class TrainingPipeline(LoRAPipeline, ABC):
+    """
+    A pipeline for training a model. All training pipelines should inherit from this class.
+    All reusable components and code should be implemented in this class.
+    """
+    _required_config_modules = ["scheduler", "transformer"]
+    validation_pipeline: ComposedPipelineBase
+    train_dataloader: StatefulDataLoader
+    train_loader_iter: Iterator[dict[str, Any]]
+    current_epoch: int = 0
+    train_transformer_2: bool = False
+
+    def __init__(
+            self,
+            model_path: str,
+            fastvideo_args: TrainingArgs,
+            required_config_modules: list[str] | None = None,
+            loaded_modules: dict[str, torch.nn.Module] | None = None) -> None:
+        fastvideo_args.inference_mode = False
+        self.lora_training = fastvideo_args.lora_training
+        if self.lora_training and fastvideo_args.lora_rank is None:
+            raise ValueError("lora rank must be set when using lora training")
+
+        set_random_seed(fastvideo_args.seed)  # for lora param init
+        super().__init__(model_path, fastvideo_args, required_config_modules,
+                         loaded_modules)  # type: ignore
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        raise RuntimeError(
+            "create_pipeline_stages should not be called for training pipeline")
+
+    def set_schemas(self) -> None:
+        self.train_dataset_schema = pyarrow_schema_t2v
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing training pipeline...")
+        self.device = get_local_torch_device()
+        self.training_args = training_args
+        world_group = get_world_group()
+        self.world_size = world_group.world_size
+        self.global_rank = world_group.rank
+        self.sp_group = get_sp_group()
+        self.rank_in_sp_group = self.sp_group.rank_in_group
+        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()
+
+        # Set random seeds for deterministic training
+        assert self.seed is not None, "seed must be set"
+        set_random_seed(self.seed)
+        self.transformer.train()
+        if training_args.enable_gradient_checkpointing_type is not None:
+            self.transformer = apply_activation_checkpointing(
+                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))
+        self.optimizer = torch.optim.AdamW(
+            params_to_optimize,
+            lr=training_args.learning_rate,
+            betas=(0.9, 0.999),
+            weight_decay=training_args.weight_decay,
+            eps=1e-8,
+        )
+
+        self.init_steps = 0
+        logger.info("optimizer: %s", self.optimizer)
+
+        self.lr_scheduler = get_scheduler(
+            training_args.lr_scheduler,
+            optimizer=self.optimizer,
+            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,
+        )
+        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,
+            training_args.train_batch_size,
+            parquet_schema=self.train_dataset_schema,
+            num_data_workers=training_args.dataloader_num_workers,
+            cfg_rate=training_args.training_cfg_rate,
+            drop_last=True,
+            text_padding_length=training_args.pipeline_config.
+            text_encoder_configs[0].arch_config.
+            text_len,  # type: ignore[attr-defined]
+            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 /
+            training_args.train_sp_batch_size)
+        self.num_train_epochs = math.ceil(training_args.max_train_steps /
+                                          self.num_update_steps_per_epoch)
+
+        # TODO(will): is there a cleaner way to track epochs?
+        self.current_epoch = 0
+
+        if self.global_rank == 0:
+            project = training_args.tracker_project_name or "fastvideo"
+            wandb_config = dataclasses.asdict(training_args)
+            wandb.init(project=project,
+                       config=wandb_config,
+                       name=training_args.wandb_run_name)
+
+    @abstractmethod
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        raise NotImplementedError(
+            "Training pipelines must implement this method")
+
+    def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
+        # At the beginning, disable training for all models
+        self._disable_training(self.transformer, self.optimizer)
+        if self.transformer_2 is not None:
+            self._disable_training(self.transformer_2, self.optimizer_2)
+        
+        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:
+            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, encoder_hidden_states, encoder_attention_mask, infos = batch
+        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']
+        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.infos = infos
+
+        return training_batch
+
+    def _normalize_dit_input(self,
+                             training_batch: TrainingBatch) -> TrainingBatch:
+        # TODO(will): support other models
+        training_batch.latents = normalize_dit_input('wan',
+                                                     training_batch.latents,
+                                                     self.get_module("vae"))
+        return training_batch
+
+    def _prepare_dit_inputs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+        latents = training_batch.latents
+        batch_size = latents.shape[0]
+        noise = torch.randn(latents.shape,
+                            generator=self.noise_gen_cuda,
+                            device=latents.device,
+                            dtype=latents.dtype)
+        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()
+            sp_group.broadcast(timesteps, src=0)
+        sigmas = get_sigmas(
+            self.noise_scheduler,
+            latents.device,
+            timesteps,
+            n_dim=latents.ndim,
+            dtype=latents.dtype,
+        )
+        noisy_model_input = (1.0 -
+                             sigmas) * training_batch.latents + sigmas * noise
+
+        training_batch.noisy_model_input = noisy_model_input
+        training_batch.timesteps = timesteps
+        training_batch.sigmas = sigmas
+        training_batch.noise = noise
+        training_batch.raw_latent_shape = training_batch.latents.shape
+
+        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 = boundary_ratio + u * (1.0 - boundary_ratio)
+        else:
+            u = u * 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
+        patch_size = self.training_args.pipeline_config.dit_config.patch_size
+        current_vsa_sparsity = training_batch.current_vsa_sparsity
+        assert latents_shape is not None
+        assert training_batch.timesteps is not None
+        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
+            training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder(  # type: ignore
+            ).build(  # type: ignore
+                raw_latent_shape=latents_shape[2:5],
+                current_timestep=training_batch.timesteps,
+                patch_size=patch_size,
+                VSA_sparsity=current_vsa_sparsity,
+                device=get_local_torch_device())
+        else:
+            training_batch.attn_metadata = None
+
+        return training_batch
+
+    def _build_input_kwargs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+        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,
+            "return_dict":
+            False,
+        }
+        return training_batch
+
+    def _transformer_forward_and_compute_loss(
+            self, training_batch: TrainingBatch) -> TrainingBatch:
+        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
+            assert training_batch.attn_metadata is not None
+        else:
+            assert training_batch.attn_metadata is None
+        input_kwargs = training_batch.input_kwargs
+
+        # if 'hunyuan' in self.training_args.model_type:
+        #     input_kwargs["guidance"] = torch.tensor(
+        #         [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 = 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
+            assert training_batch.latents is not None
+            assert training_batch.noise is not None
+            target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
+
+            # make sure no implicit broadcasting happens
+            assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
+            loss = (torch.mean((model_pred.float() - target.float())**2) /
+                    self.training_args.gradient_accumulation_steps)
+
+            loss.backward()
+            avg_loss = loss.detach().clone()
+
+        # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}",
+        #             local_main_process_only=False)
+        world_group = get_world_group()
+        world_group.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
+        training_batch.total_loss += avg_loss.item()
+
+        return training_batch
+
+    def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch:
+        max_grad_norm = self.training_args.max_grad_norm
+
+        # TODO(will): perhaps move this into transformer api so that we can do
+        # the following:
+        # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+        if max_grad_norm is not None:
+            # 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,
+                foreach=None,
+            )
+            assert grad_norm is not float('nan') or grad_norm is not float(
+                'inf')
+            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
+        else:
+            grad_norm = 0.0
+        training_batch.grad_norm = grad_norm
+        return training_batch
+
+    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
+        training_batch = self._prepare_training(training_batch)
+
+        for _ in range(self.training_args.gradient_accumulation_steps):
+            training_batch = self._get_next_batch(training_batch)
+
+            # Normalize DIT input
+            training_batch = self._normalize_dit_input(training_batch)
+            # Create noisy model input
+            training_batch = self._prepare_dit_inputs(training_batch)
+
+            # Shard latents across sp groups
+            training_batch.latents = shard_latents_across_sp(
+                training_batch.latents,
+                num_latent_t=self.training_args.num_latent_t)
+            # shard noisy_model_input to match
+            training_batch.noisy_model_input = shard_latents_across_sp(
+                training_batch.noisy_model_input,
+                num_latent_t=self.training_args.num_latent_t)
+            # shard noise to match latents
+            training_batch.noise = shard_latents_across_sp(
+                training_batch.noise,
+                num_latent_t=self.training_args.num_latent_t)
+
+            training_batch = self._build_attention_metadata(training_batch)
+            training_batch = self._build_input_kwargs(training_batch)
+            training_batch = self._transformer_forward_and_compute_loss(
+                training_batch)
+
+        training_batch = self._clip_grad_norm(training_batch)
+
+        # 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
+
+    def _resume_from_checkpoint(self) -> None:
+        logger.info("Loading checkpoint from %s",
+                    self.training_args.resume_from_checkpoint)
+        resumed_step = load_checkpoint(
+            self.transformer, self.global_rank,
+            self.training_args.resume_from_checkpoint, self.optimizer,
+            self.train_dataloader, self.lr_scheduler,
+            self.noise_random_generator)
+        if resumed_step > 0:
+            self.init_steps = resumed_step
+            logger.info("Successfully resumed from step %s", resumed_step)
+        else:
+            logger.warning("Failed to load checkpoint, starting from step 0")
+            self.init_steps = 0
+
+    def train(self) -> None:
+        assert self.seed is not None, "seed must be set"
+        set_random_seed(self.seed + self.global_rank)
+        logger.info('rank: %s: start training',
+                    self.global_rank,
+                    local_main_process_only=False)
+        if not self.post_init_called:
+            self.post_init()
+        num_trainable_params = _get_trainable_params(self.transformer)
+        logger.info("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)
+        self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed(
+            self.seed)
+        self.validation_random_generator = torch.Generator(
+            device="cpu").manual_seed(self.seed)
+        logger.info("Initialized random seeds with seed: %s", self.seed)
+
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
+
+        if self.training_args.resume_from_checkpoint:
+            self._resume_from_checkpoint()
+
+        self.train_loader_iter = iter(self.train_dataloader)
+
+        step_times: deque[float] = deque(maxlen=100)
+
+        self._log_training_info()
+
+        self._log_validation(self.transformer, self.training_args,
+                             self.init_steps)
+
+        # Train!
+        progress_bar = tqdm(
+            range(0, self.training_args.max_train_steps),
+            initial=self.init_steps,
+            desc="Steps",
+            # Only show the progress bar once on each machine.
+            disable=self.local_rank > 0,
+        )
+        for step in range(self.init_steps + 1,
+                          self.training_args.max_train_steps + 1):
+            start_time = time.perf_counter()
+            if vsa_available:
+                vsa_sparsity = self.training_args.VSA_sparsity
+                vsa_decay_rate = self.training_args.VSA_decay_rate
+                vsa_decay_interval_steps = self.training_args.VSA_decay_interval_steps
+                current_decay_times = min(step // vsa_decay_interval_steps,
+                                          vsa_sparsity // vsa_decay_rate)
+                current_vsa_sparsity = current_decay_times * vsa_decay_rate
+            else:
+                current_vsa_sparsity = 0.0
+
+            training_batch = TrainingBatch()
+            training_batch.current_timestep = step
+            training_batch.current_vsa_sparsity = current_vsa_sparsity
+            training_batch = self.train_one_step(training_batch)
+
+            loss = training_batch.total_loss
+            grad_norm = training_batch.grad_norm
+
+            step_time = time.perf_counter() - start_time
+            step_times.append(step_time)
+            avg_step_time = sum(step_times) / len(step_times)
+
+            progress_bar.set_postfix({
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+            })
+            progress_bar.update(1)
+            if self.global_rank == 0:
+                wandb.log(
+                    {
+                        "train_loss": loss,
+                        "learning_rate": self.lr_scheduler.get_last_lr()[0],
+                        "step_time": step_time,
+                        "avg_step_time": avg_step_time,
+                        "grad_norm": grad_norm,
+                        "vsa_sparsity": current_vsa_sparsity,
+                    },
+                    step=step,
+                )
+            if step % self.training_args.checkpointing_steps == 0:
+                save_checkpoint(self.transformer, self.global_rank,
+                                self.training_args.output_dir, step,
+                                self.optimizer, self.train_dataloader,
+                                self.lr_scheduler, self.noise_random_generator)
+                self.transformer.train()
+                self.sp_group.barrier()
+            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
+                self._log_validation(self.transformer, self.training_args, step)
+                gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+                trainable_params = round(
+                    _get_trainable_params(self.transformer) / 1e9, 3)
+                logger.info(
+                    "GPU memory usage after validation: %s MB, trainable params: %sB",
+                    gpu_memory_usage, trainable_params)
+
+        wandb.finish()
+        save_checkpoint(self.transformer, self.global_rank,
+                        self.training_args.output_dir,
+                        self.training_args.max_train_steps, self.optimizer,
+                        self.train_dataloader, self.lr_scheduler,
+                        self.noise_random_generator)
+
+        if get_sp_group():
+            cleanup_dist_env_and_memory()
+
+    def _log_training_info(self) -> None:
+        total_batch_size = (self.world_size *
+                            self.training_args.gradient_accumulation_steps /
+                            self.training_args.sp_size *
+                            self.training_args.train_sp_batch_size)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %s", len(self.train_dataset))
+        logger.info("  Dataloader size = %s", len(self.train_dataloader))
+        logger.info("  Num Epochs = %s", self.num_train_epochs)
+        logger.info("  Resume training from step %s",
+                    self.init_steps)  # type: ignore
+        logger.info("  Instantaneous batch size per device = %s",
+                    self.training_args.train_batch_size)
+        logger.info(
+            "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
+            total_batch_size)
+        logger.info("  Gradient Accumulation steps = %s",
+                    self.training_args.gradient_accumulation_steps)
+        logger.info("  Total optimization steps = %s",
+                    self.training_args.max_train_steps)
+        logger.info("  Total training parameters per FSDP shard = %s B",
+                    round(_get_trainable_params(self.transformer) / 1e9, 3))
+        # print dtype
+        logger.info("  Master weight dtype: %s",
+                    self.transformer.parameters().__next__().dtype)
+
+        gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+        logger.info("GPU memory usage before train_one_step: %s MB",
+                    gpu_memory_usage)
+        logger.info("VSA validation sparsity: %s",
+                    self.training_args.VSA_sparsity)
+
+    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.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=self.validation_random_generator,
+            n_tokens=n_tokens,
+            eta=0.0,
+            VSA_sparsity=training_args.VSA_sparsity,
+        )
+
+        return batch
+
+    @torch.no_grad()
+    def _log_validation(self, transformer, 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
+        if not training_args.log_validation:
+            return
+        if self.validation_pipeline is None:
+            raise ValueError("Validation pipeline is not set")
+
+        logger.info("Starting validation")
+
+        # Create sampling parameters if not provided
+        sampling_param = SamplingParam.from_pretrained(training_args.model_path)
+
+        # Prepare validation prompts
+        logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s',
+                    self.global_rank,
+                    training_args.validation_dataset_file,
+                    local_main_process_only=False)
+        validation_dataset = ValidationDataset(
+            training_args.validation_dataset_file)
+        validation_dataloader = DataLoader(validation_dataset,
+                                           batch_size=None,
+                                           num_workers=0)
+
+        transformer.eval()
+
+        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]
+        # Log validation results for this step
+        world_group = get_world_group()
+        num_sp_groups = world_group.world_size // self.sp_group.world_size
+
+        # Process each validation prompt for each validation step
+        for num_inference_steps in validation_steps:
+            logger.info("rank: %s: num_inference_steps: %s",
+                        self.global_rank,
+                        num_inference_steps,
+                        local_main_process_only=False)
+            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)
+                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
+                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)
+
+            # Only sp_group leaders (rank_in_sp_group == 0) need to send their
+            # results to global rank 0
+            if self.rank_in_sp_group == 0:
+                if self.global_rank == 0:
+                    # Global rank 0 collects results from all sp_group leaders
+                    all_videos = step_videos  # Start with own results
+                    all_captions = step_captions
+
+                    # Receive from other sp_group leaders
+                    for sp_group_idx in range(1, num_sp_groups):
+                        src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
+                        recv_videos = world_group.recv_object(src=src_rank)
+                        recv_captions = world_group.recv_object(src=src_rank)
+                        all_videos.extend(recv_videos)
+                        all_captions.extend(recv_captions)
+
+                    video_filenames = []
+                    for i, (video, caption) in enumerate(
+                            zip(all_videos, all_captions, strict=True)):
+                        os.makedirs(training_args.output_dir, exist_ok=True)
+                        filename = os.path.join(
+                            training_args.output_dir,
+                            f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4"
+                        )
+                        imageio.mimsave(filename, video, fps=sampling_param.fps)
+                        video_filenames.append(filename)
+
+                    logs = {
+                        f"validation_videos_{num_inference_steps}_steps": [
+                            wandb.Video(filename, caption=caption)
+                            for filename, caption in zip(
+                                video_filenames, all_captions, strict=True)
+                        ]
+                    }
+                    wandb.log(logs, step=global_step)
+                else:
+                    # Other sp_group leaders send their results to global rank 0
+                    world_group.send_object(step_videos, dst=0)
+                    world_group.send_object(step_captions, dst=0)
+
+        # Re-enable gradients for training
+        training_args.inference_mode = False
+        transformer.train()

fastvideo/training/wan_self_forcing_distillation_pipeline.py

@@ -25,6 +25,11 @@ class WanSelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
         "fake_score_transformer"
     ]
 
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        """Initialize Wan-specific scheduler."""
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
     def create_training_stages(self, training_args: TrainingArgs):
         """
         May be used in future refactors.

fastvideo/training/wan_training_pipeline.py

@@ -34,7 +34,7 @@ class WanTrainingPipeline(TrainingPipeline):
     def initialize_validation_pipeline(self, training_args: TrainingArgs):
         logger.info("Initializing validation pipeline...")
         args_copy = deepcopy(training_args)
-
+        assert training_args.dit_cpu_offload == False
         args_copy.inference_mode = True
         validation_pipeline = WanPipeline.from_pretrained(
             training_args.model_path,
@@ -42,12 +42,13 @@ class WanTrainingPipeline(TrainingPipeline):
             inference_mode=True,
             loaded_modules={
                 "transformer": self.get_module("transformer"),
+                "transformer_2": self.get_module("transformer_2")
             },
             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)
+            dit_cpu_offload=training_args.dit_cpu_offload)
 
         self.validation_pipeline = validation_pipeline
 

fastvideo/utils.py

@@ -34,6 +34,7 @@ from torch.distributed.fsdp import MixedPrecisionPolicy
 import fastvideo.envs as envs
 from fastvideo.logger import init_logger
 
+
 logger = init_logger(__name__)
 
 T = TypeVar("T")
@@ -614,7 +615,6 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]:
             f"Failed to download or parse model_index.json for {model_name_or_path}: {e}"
         ) from e
 
-
 def update_environment_variables(envs: dict[str, str]):
     for k, v in envs.items():
         if k in os.environ and os.environ[k] != v:
@@ -814,17 +814,6 @@ 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,

fastvideo/v1/pipelines/stages/encoding.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Encoding stage for diffusion pipelines.
+"""
+from typing import Optional
+
+import PIL.Image
+import torch
+
+from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.vaes.common import ParallelTiledVAE
+from fastvideo.v1.models.vision_utils import (get_default_height_width,
+                                              normalize, numpy_to_pt,
+                                              pil_to_numpy, resize)
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines.stages.base import PipelineStage
+from fastvideo.v1.pipelines.stages.validators import V  # Import validators
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
+from fastvideo.v1.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class EncodingStage(PipelineStage):
+    """
+    Stage for encoding pixel representations into latent space.
+    
+    This stage handles the encoding of pixel representations into the final
+    input format (e.g., latents).
+    """
+
+    def __init__(self, vae: ParallelTiledVAE) -> None:
+        self.vae: ParallelTiledVAE = vae
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Encode pixel representations into latent space.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with encoded outputs.
+        """
+        self.vae = self.vae.to(get_torch_device())
+
+        assert batch.height is not None
+        assert batch.width is not None
+        latent_height = batch.height // self.vae.spatial_compression_ratio
+        latent_width = batch.width // self.vae.spatial_compression_ratio
+
+        image = batch.preprocessed_image
+        # TODO(will)
+        if image is None:
+            assert batch.pil_image is not None
+            image = batch.pil_image
+            image = self.preprocess(
+                image,
+                vae_scale_factor=self.vae.spatial_compression_ratio,
+                height=batch.height,
+                width=batch.width).to(get_torch_device(), dtype=torch.float32)
+
+            image = image.unsqueeze(2)
+        else:
+            # assumes image is loaded from parquet file and used for validation
+            torch.distributed.breakpoint()
+            image = image.transpose(1, 2)
+        logger.info("image: %s", image.shape)
+        video_condition = torch.cat([
+            image,
+            image.new_zeros(image.shape[0], image.shape[1],
+                            batch.num_frames - 1, batch.height, batch.width)
+        ],
+                                    dim=2)
+        video_condition = video_condition.to(device=get_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)
+
+        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
+
+        mask_lat_size = torch.ones(1, 1, batch.num_frames, latent_height,
+                                   latent_width)
+        mask_lat_size[:, :, list(range(1, batch.num_frames))] = 0
+        first_frame_mask = mask_lat_size[:, :, 0:1]
+        first_frame_mask = torch.repeat_interleave(
+            first_frame_mask,
+            dim=2,
+            repeats=self.vae.temporal_compression_ratio)
+        mask_lat_size = torch.concat(
+            [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
+        mask_lat_size = mask_lat_size.view(1, -1,
+                                           self.vae.temporal_compression_ratio,
+                                           latent_height, latent_width)
+        mask_lat_size = mask_lat_size.transpose(1, 2)
+        mask_lat_size = mask_lat_size.to(latent_condition.device)
+
+        batch.image_latent = torch.concat([mask_lat_size, latent_condition],
+                                          dim=1)
+
+        # Offload models if needed
+        if hasattr(self, 'maybe_free_model_hooks'):
+            self.maybe_free_model_hooks()
+
+        return batch
+
+    def retrieve_latents(self,
+                         encoder_output: torch.Tensor,
+                         generator: Optional[torch.Generator] = None,
+                         sample_mode: str = "sample"):
+        if sample_mode == "sample":
+            return encoder_output.sample(generator)
+        elif sample_mode == "argmax":
+            return encoder_output.mode()
+        else:
+            raise AttributeError(
+                "Could not access latents of provided encoder_output")
+
+    def preprocess(
+            self,
+            image: PIL.Image.Image,
+            vae_scale_factor: int,
+            height: Optional[int] = None,
+            width: Optional[int] = None,
+            resize_mode: str = "default",  # "default", "fill", "crop"
+    ) -> torch.Tensor:
+        image = [image]
+
+        height, width = get_default_height_width(image[0], vae_scale_factor,
+                                                 height, width)
+        image = [
+            resize(i, height, width, resize_mode=resize_mode) for i in image
+        ]
+        image = pil_to_numpy(image)  # to np
+        image = numpy_to_pt(image)  # to pt
+
+        do_normalize = True
+        if image.min() < 0:
+            do_normalize = False
+        if do_normalize:
+            image = normalize(image)
+
+        return image
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify encoding stage inputs."""
+        result = VerificationResult()
+        # result.add_check("pil_image", batch.pil_image)
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        result.add_check("generator", batch.generator,
+                         V.generator_or_list_generators)
+        result.add_check("num_frames", batch.num_frames, V.positive_int)
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify encoding stage outputs."""
+        result = VerificationResult()
+        result.add_check("image_latent", batch.image_latent,
+                         [V.is_tensor, V.with_dims(5)])
+        return result

fastvideo/workflow/preprocess/components.py

@@ -13,8 +13,7 @@ import pyarrow.parquet as pq
 import torch
 from datasets import Dataset, Video, load_dataset
 
-from fastvideo.configs.configs import (DatasetType, PreprocessConfig,
-                                       VideoLoaderType)
+from fastvideo.configs.configs import DatasetType, PreprocessConfig
 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
@@ -432,8 +431,7 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str,
             return item
 
         dataset = dataset.map(add_video_column)
-        if preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
-            dataset = dataset.cast_column("video", Video())
+        dataset = dataset.cast_column("video", Video())
     else:
         raise ValueError(
             f"Invalid dataset type: {preprocess_config.dataset_type}")

scripts/huggingface/upload_hf.py

@@ -3,7 +3,7 @@ from huggingface_hub import HfApi
 api = HfApi()
 
 api.upload_folder(
-    folder_path="Wan2.2-TI2V-5B-Diffusers",
-    repo_id="FastVideo/FastWan2.2-TI2V-5B-Diffusers",
+    folder_path="wow",
+    repo_id="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
     repo_type="model",
 )

scripts/inference/v1_inference_wan_1.3B_VMoba.sh

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

@@ -1,6 +1,6 @@
 export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export WANDB_API_KEY='2f25ad37933894dbf0966c838c0b8494987f9f2f'
+export WANDB_MODE=offline
+export WANDB_API_KEY='50632ebd88ffd970521cec9ab4a1a2d7e85bfc45'
 # 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/
@@ -38,7 +38,6 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --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 \