training refactor done

kv cache debugging
revert
2025-09-08 03:48:01 +00:00 · 2025-09-08 02:27:19 +00:00 · 2025-09-07 23:38:02 +00:00 · 2025-09-07 23:36:43 +00:00 · 2025-09-07 08:49:36 +00:00 · 2025-09-07 02:40:19 +00:00
98 changed files with 6062 additions and 701 deletions
@@ -18,6 +18,12 @@ on:
        required: false
        default: false
        type: boolean
+      python_3_12_cuda_12_9:
+        description: 'Build Python 3.12 image Cuda 12.9'
+        required: false
+        default: false
+        type: boolean
+

 permissions:
  contents: read
@@ -49,4 +55,13 @@ jobs:
      python_version: '3.12'
      dockerfile_path: docker/Dockerfile.python3.12
      tag_suffix: py3.12
+    secrets: inherit
+
+  build-python-3-12-cuda-12-9:
+    if: ${{ github.event.inputs.python_3_12_cuda_12_9 == 'true' }}
+    uses: ./.github/workflows/build-image-template.yml
+    with:
+      python_version: '3.12'
+      dockerfile_path: docker/Dockerfile.python3.12.cuda12.9.1
+      tag_suffix: py3.12-cuda12.9.1
    secrets: inherit
@@ -64,3 +64,5 @@ docs/source/distillation/examples/
 !docs/source/_static/images/**/*.png
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif
+
+dmd_t2v_output/
@@ -22,6 +22,7 @@ exclude: |
    examples/.*|
    .github/workflows/fastvideo-publish.yml|
    .github/workflows/sta-publish.yml|
+    .github/workflows/vsa-publish.yml|
    .github/workflows/build-image-template.yml|
    docs/source/inference/support_matrix.md
  )
@@ -1,5 +1,5 @@
 <div align="center">
-<img src=assets/logo.png width="30%"/>
+<img src=assets/logos/logo.svg width="30%"/>
 </div>

 **FastVideo is a unified post-training and inference framework for accelerated video generation.**
@@ -7,7 +7,7 @@
 FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.

 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/qqPzbrw" 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">
@@ -0,0 +1,6 @@
+<svg width="160" height="93" viewBox="0 0 160 93" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M28.8511 91.66L57.6319 1.86368H64.5394L35.7585 91.66H28.8511Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+<path d="M15.0376 91.66L43.8185 1.86368H46.1209L17.3401 91.66H15.0376Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+<path d="M1.22217 91.66L30.003 1.86366H31.1543L2.3734 91.66H1.22217Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.15122"/>
+<path d="M71.4465 1.86483L42.666 91.6599H69.144L78.3538 58.2746H123.251L129.007 39.855H84.1099L89.866 22.5868H152.032L157.788 1.86483H71.4465Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+</svg>
@@ -0,0 +1,18 @@
+<svg width="252" height="105" viewBox="0 0 252 105" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM100.768 13.1217L87.7028 29.4852H103.143L100.768 13.1217Z" fill="#356CFF"/>
+<path d="M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM109.081 90.697L116.802 65.8487C116.802 65.8487 120.959 65.8487 132.242 65.8487C143.525 65.8487 137.586 78.5759 135.211 84.0304C133.307 88.4021 127.491 90.697 122.74 90.697C117.989 90.697 109.081 90.697 109.081 90.697Z" fill="#356CFF"/>
+<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944H159.747C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273L125.188 48.273L124 37.97L147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852H131.836C120.142 29.4852 125.897 1.00043 141.337 1.00043L173.188 1.00056Z" fill="#356CFF"/>
+<path d="M179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056Z" fill="#356CFF"/>
+<path d="M161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM237.948 77.9692C239.984 70.6965 240.917 65.242 228.446 65.242C215.975 65.242 211.818 71.9087 210.037 77.9692C208.255 84.0298 208.255 91.3025 219.538 91.3025C230.821 91.3025 235.911 85.2419 237.948 77.9692Z" fill="#356CFF"/>
+<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944M173.188 1.00056C173.188 1.00056 156.777 1.00043 141.337 1.00043M173.188 1.00056L141.337 1.00043M141.337 20.3944C146.088 20.3944 150.839 20.3944 159.747 20.3944M141.337 20.3944H159.747M159.747 20.3944C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273M148.463 48.273C139.556 48.273 125.188 48.273 125.188 48.273M148.463 48.273L125.188 48.273M125.188 48.273L124 37.97M124 37.97C124 37.97 141.931 37.97 147.87 37.97M124 37.97L147.87 37.97M147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852M151.433 29.4852C146.682 29.4852 138.962 29.4852 131.836 29.4852M151.433 29.4852H131.836M131.836 29.4852C120.142 29.4852 125.897 1.00043 141.337 1.00043M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057ZM96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM87.7028 29.4852L100.768 13.1217L103.143 29.4852H87.7028ZM89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457ZM108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM116.802 65.8487L109.081 90.697C109.081 90.697 117.989 90.697 122.74 90.697C127.491 90.697 133.307 88.4021 135.211 84.0304C137.586 78.5759 143.525 65.8487 132.242 65.8487C120.959 65.8487 116.802 65.8487 116.802 65.8487ZM179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056ZM161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457ZM230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM228.446 65.242C240.917 65.242 239.984 70.6965 237.948 77.9692C235.911 85.2419 230.821 91.3025 219.538 91.3025C208.255 91.3025 208.255 84.0298 210.037 77.9692C211.818 71.9087 215.975 65.242 228.446 65.242Z" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M15.2524 55.5451L21.191 100.999L24.7541 100.999L18.8156 55.5451L15.2524 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M8.12646 55.5451L14.065 100.999L15.2527 100.999L9.31417 55.5451L8.12646 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M1 55.5451L6.93853 100.999L7.53239 100.999L1.59385 55.5451L1 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
+<path d="M15.2524 48.2724L30.0988 1H33.6619L18.8156 48.2724H15.2524Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M8.12646 48.2724L22.9728 1H24.1605L9.31417 48.2724H8.12646Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M1 48.2724L15.8463 1H16.4402L1.59385 48.2724H1Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
+<path d="M85.3271 55.5457H67.5116L87 12.7363L44.3513 68.2729H58.6038L43.1636 101L85.3271 55.5457Z" fill="#FDC717" stroke="#FDC717" stroke-width="1.18771" stroke-miterlimit="16"/>
+</svg>
@@ -0,0 +1,6 @@
+<svg width="160" height="93" viewBox="0 0 160 93" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M28.8511 91.66L57.6319 1.86368H64.5394L35.7585 91.66H28.8511Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+<path d="M15.0376 91.66L43.8185 1.86368H46.1209L17.3401 91.66H15.0376Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+<path d="M1.22217 91.66L30.003 1.86366H31.1543L2.3734 91.66H1.22217Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.15122"/>
+<path d="M71.4465 1.86483L42.666 91.6599H69.144L78.3538 58.2746H123.251L129.007 39.855H84.1099L89.866 22.5868H152.032L157.788 1.86483H71.4465Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+</svg>
@@ -13,9 +13,9 @@ BLOCK_M = 64
 BLOCK_N = 64

 def pytorch_test(Q, K, V, block_sparse_mask, dO):
-    q_ = Q.clone().requires_grad_()
-    k_ = K.clone().requires_grad_()
-    v_ = V.clone().requires_grad_()
+    q_ = Q.clone().float().requires_grad_()
+    k_ = K.clone().float().requires_grad_()
+    v_ = V.clone().float().requires_grad_()

    QK = torch.matmul(q_, k_.transpose(-2, -1))
    QK /= (q_.size(-1) ** 0.5)
@@ -35,9 +35,9 @@ def pytorch_test(Q, K, V, block_sparse_mask, dO):


 def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, non_pad_index, dO):
-    Q = Q.clone().requires_grad_()
-    K = K.clone().requires_grad_()
-    V = V.clone().requires_grad_()
+    Q = Q.detach().requires_grad_()
+    K = K.detach().requires_grad_()
+    V = V.detach().requires_grad_()
    
    q_padded = vsa_pad(Q, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
    k_padded = vsa_pad(K, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
@@ -60,11 +60,9 @@ def get_non_pad_index(

    return index_pad[index_mask]

-def generate_tensor(shape, mean, std, dtype, device):
+def generate_tensor(shape, dtype, device):
    tensor = torch.randn(shape, dtype=dtype, device=device)
-    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
-    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
-    return scaled_tensor.contiguous()
+    return tensor

 def generate_variable_block_sizes(num_blocks, min_size=32, max_size=64, device="cuda"):
    return torch.randint(min_size, max_size + 1, (num_blocks,), device=device, dtype=torch.int32)
@@ -75,7 +73,7 @@ def vsa_pad(x, non_pad_index, num_blocks, block_size):
    padded_x[:, :, non_pad_index, :] = x
    return padded_x

-def check_correctness(h, d, num_blocks, k, mean, std, num_iterations=20, error_mode='all'):
+def check_correctness(h, d, num_blocks, k,  num_iterations=20, error_mode='all'):
    results = {
        'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
        'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
@@ -91,10 +89,10 @@ def check_correctness(h, d, num_blocks, k, mean, std, num_iterations=20, error_m
    block_mask = generate_block_sparse_mask_for_function(h, num_blocks, k, device)
    full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, device)
    for _ in range(num_iterations):
-        Q = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
-        K = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
-        V = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
-        dO = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
+        Q = generate_tensor((1, h, S, d),  torch.bfloat16, device)
+        K = generate_tensor((1, h, S, d), torch.bfloat16, device)
+        V = generate_tensor((1, h, S, d),  torch.bfloat16, device)
+        dO = generate_tensor((1, h, S, d), torch.bfloat16, device)

        # dO_padded = torch.zeros_like(dO_padded)
        # dO_padded[:, :, non_pad_index, :] = dO
@@ -107,7 +105,8 @@ def check_correctness(h, d, num_blocks, k, mean, std, num_iterations=20, error_m
                abs_diff = torch.abs(diff)
                results[name]['sum_diff'] += torch.sum(abs_diff).item()
                results[name]['sum_abs'] += torch.sum(torch.abs(pt)).item()
-                results[name]['max_diff'] = max(results[name]['max_diff'], torch.max(abs_diff).item())
+                rel_max_diff = torch.max(abs_diff) / torch.mean(torch.abs(pt))
+                results[name]['max_diff'] = max(results[name]['max_diff'], rel_max_diff.item())
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

@@ -119,27 +118,27 @@ def check_correctness(h, d, num_blocks, k, mean, std, num_iterations=20, error_m

    return results

-def generate_error_graphs(h, d, mean, std, error_mode='all'):
+def generate_error_graphs(h, d, error_mode='all'):
    test_configs = [
        {"num_blocks": 16, "k": 2, "description": "Small sequence"},
        {"num_blocks": 32, "k": 4, "description": "Medium sequence"},
        {"num_blocks": 53, "k": 6, "description": "Large sequence"},
    ]
    
-    print(f"\nError Analysis for h={h}, d={d}, mean={mean}, std={std}, mode={error_mode}")
+    print(f"\nError Analysis for h={h}, d={d}, mode={error_mode}")
    print("=" * 150)
    print(f"{'Config':<20} {'Blocks':<8} {'K':<4} "
-          f"{'gQ Avg':<12} {'gQ Max':<12} "
-          f"{'gK Avg':<12} {'gK Max':<12} "
-          f"{'gV Avg':<12} {'gV Max':<12} "
-          f"{'gO Avg':<12} {'gO Max':<12}")
+          f"{'gQ Avg':<12} {'Rel gQ Max':<12} "
+          f"{'gK Avg':<12} {'Rel gK Max':<12} "
+          f"{'gV Avg':<12} {'Rel gV Max':<12} "
+          f"{'gO Avg':<12} {'Rel gO Max':<12}")
    print("-" * 150)

    for config in test_configs:
        num_blocks = config["num_blocks"]
        k = config["k"]
        description = config["description"]
-        results = check_correctness(h, d, num_blocks, k, mean, std, error_mode=error_mode)
+        results = check_correctness(h, d, num_blocks, k, error_mode=error_mode)
        print(f"{description:<20} {num_blocks:<8} {k:<4} "
              f"{results['gQ']['avg_diff']:<12.6e} {results['gQ']['max_diff']:<12.6e} "
              f"{results['gK']['avg_diff']:<12.6e} {results['gK']['max_diff']:<12.6e} "
@@ -150,10 +149,8 @@ def generate_error_graphs(h, d, mean, std, error_mode='all'):

 if __name__ == "__main__":
    h, d = 16, 128
-    mean = 0.0
-    std = 1
    print("Block Sparse Attention with Variable Block Sizes Analysis")
    print("=" * 60)
    for mode in ['backward']:
-        generate_error_graphs(h, d, mean, std, error_mode=mode)
+        generate_error_graphs(h, d, error_mode=mode)
    print("\nAnalysis completed for all modes.")
@@ -568,7 +568,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
        __syncthreads(); // wait for sd_smem shared memory write
        warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK
        warpgroup::mma_commit_group();
-        tma::store_async_wait();
        warpgroup::mma_async_wait();
        // store qg to shared memory
        warpgroup::store(qg_smem, qg_reg);
@@ -578,6 +577,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
        if (threadIdx.x / 32 == 0) {
            coord<qg_tile> tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0};
            tma::store_add_async(g.qg, qg_smem, tile_idx);
+            tma::store_async_wait();
        }
    }

@@ -624,7 +624,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
        __syncthreads(); // wait for sd_smem shared memory write
        warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK
        warpgroup::mma_commit_group();
-        tma::store_async_wait();
        warpgroup::mma_async_wait();
        // store qg to shared memory
        warpgroup::store(qg_smem, qg_reg);
@@ -634,13 +633,14 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
        if (threadIdx.x / 32 == 0) {
            coord<qg_tile> tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0};
            tma::store_add_async(g.qg, qg_smem, tile_idx);
+            tma::store_async_wait();
        }
    }

    // store kq and vq

    // ! the following two line seems unnecessary.
-    tma::store_async_wait(); // ensure qg is finished
+    // tma::store_async_wait(); // ensure qg is finished
    __syncthreads();

    warpgroup::store(kg_smem[0], kg_reg);
@@ -1174,4 +1174,4 @@ block_sparse_attention_backward(torch::Tensor q,

    return {qg, kg, vg};
    //cudadevicesynchronize();
-}
+}
@@ -1,7 +1,9 @@
-FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+FROM nvidia/cuda:12.8.0-devel-ubuntu22.04

 ENV DEBIAN_FRONTEND=noninteractive

+SHELL ["/bin/bash", "-c"]
+
 WORKDIR /FastVideo

 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -9,17 +11,25 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    ca-certificates \
    openssh-server \
+    zsh \
+    vim \
+    curl \
+    gcc-11 \
+    g++-11 \
+    clang-11 \
    && rm -rf /var/lib/apt/lists/*

-RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
-    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
-    rm Miniconda3-latest-Linux-x86_64.sh
+# Set up C++20 compilers for ThunderKittens
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11

-ENV PATH=/opt/conda/bin:$PATH
+# Set CUDA environment variables
+ENV CUDA_HOME=/usr/local/cuda-12.8
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH

-RUN conda create --name fastvideo-dev python=3.10.0 -y
-
-SHELL ["/bin/bash", "-c"]
+# Install uv and source its environment
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    echo 'source $HOME/.local/bin/env' >> /root/.bashrc

 # Copy just the pyproject.toml first to leverage Docker cache
 COPY pyproject.toml ./
@@ -27,22 +37,36 @@ COPY pyproject.toml ./
 # Create a dummy README to satisfy the installation
 RUN echo "# Placeholder" > README.md

-RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
-    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
-    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
-    conda clean -afy
+# Create and activate virtual environment with specific Python version and seed
+RUN source $HOME/.local/bin/env && \
+    uv venv --python 3.10 --seed /opt/venv && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir --upgrade pip && \
+    uv pip install --no-cache-dir .[dev] && \
+    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation

 COPY . .

-RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+# Install dependencies using uv and set up shell configuration
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir -e .[dev] && \
+    git config --unset-all http.https://github.com/.extraheader || true && \
+    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
+    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile

-# Remove authentication headers
-RUN git config --unset-all http.https://github.com/.extraheader || true
+# Install STA (Sliding Tile Attention)
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_sta.py install

-# Set up automatic conda environment activation for all shells
-RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
-    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
-    # Ensure .bashrc is sourced for SSH login shells
-    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+# Install VSA 
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_vsa.py install

 EXPOSE 22
@@ -1,7 +1,9 @@
-FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+FROM nvidia/cuda:12.8.0-devel-ubuntu22.04

 ENV DEBIAN_FRONTEND=noninteractive

+SHELL ["/bin/bash", "-c"]
+
 WORKDIR /FastVideo

 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -9,17 +11,25 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    ca-certificates \
    openssh-server \
+    zsh \
+    vim \
+    curl \
+    gcc-11 \
+    g++-11 \
+    clang-11 \
    && rm -rf /var/lib/apt/lists/*

-RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
-    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
-    rm Miniconda3-latest-Linux-x86_64.sh
+# Set up C++20 compilers for ThunderKittens
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11

-ENV PATH=/opt/conda/bin:$PATH
+# Set CUDA environment variables
+ENV CUDA_HOME=/usr/local/cuda-12.8
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH

-RUN conda create --name fastvideo-dev python=3.11.11 -y
-
-SHELL ["/bin/bash", "-c"]
+# Install uv and source its environment
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    echo 'source $HOME/.local/bin/env' >> /root/.bashrc

 # Copy just the pyproject.toml first to leverage Docker cache
 COPY pyproject.toml ./
@@ -27,22 +37,36 @@ COPY pyproject.toml ./
 # Create a dummy README to satisfy the installation
 RUN echo "# Placeholder" > README.md

-RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
-    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
-    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
-    conda clean -afy
+# Create and activate virtual environment with specific Python version and seed
+RUN source $HOME/.local/bin/env && \
+    uv venv --python 3.11 --seed /opt/venv && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir --upgrade pip && \
+    uv pip install --no-cache-dir .[dev] && \
+    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation

 COPY . .

-RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+# Install dependencies using uv and set up shell configuration
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir -e .[dev] && \
+    git config --unset-all http.https://github.com/.extraheader || true && \
+    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
+    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile

-# Remove authentication headers
-RUN git config --unset-all http.https://github.com/.extraheader || true
+# Install STA (Sliding Tile Attention)
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_sta.py install

-# Set up automatic conda environment activation for all shells
-RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
-    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
-    # Ensure .bashrc is sourced for SSH login shells
-    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+# Install VSA 
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_vsa.py install

 EXPOSE 22
@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir flash-attn==2.8.0.post2 --no-build-isolation
+    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation

 COPY . .

@@ -0,0 +1,72 @@
+FROM nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+SHELL ["/bin/bash", "-c"]
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    zsh \
+    vim \
+    curl \
+    gcc-11 \
+    g++-11 \
+    clang-11 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set up C++20 compilers for ThunderKittens
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+# Set CUDA environment variables
+ENV CUDA_HOME=/usr/local/cuda-12.9
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+
+# Install uv and source its environment
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject.toml ./
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+# Create and activate virtual environment with specific Python version and seed
+RUN source $HOME/.local/bin/env && \
+    uv venv --python 3.12 --seed /opt/venv && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir --upgrade pip && \
+    uv pip install --no-cache-dir .[dev] && \
+    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
+
+COPY . .
+
+# Install dependencies using uv and set up shell configuration
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir -e .[dev] && \
+    git config --unset-all http.https://github.com/.extraheader || true && \
+    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
+    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+# Install STA (Sliding Tile Attention)
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_sta.py install
+
+# Install VSA 
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_vsa.py install
+
+EXPOSE 22
@@ -96,8 +96,7 @@ copybutton_prompt_is_regexp = True
 #
 html_title = project
 html_theme = 'sphinx_book_theme'
-html_logo = '../../assets/logo.jpg'
-#html_favicon = 'assets/logos/vllm-logo-only-light.ico'
+html_logo = '../../assets/logos/icon_simple.svg'
 html_theme_options = {
    'path_to_docs': 'docs/source',
    'repository_url': 'https://github.com/hao-ai-lab/FastVideo/',
@@ -3,12 +3,12 @@

 If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:

-**Image:** [`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest`](https://ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev)
+**Images:** [`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest`](https://ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev)

 ## Starting the container

 ```bash
-docker run --gpus all -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
+docker run --gpus all -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
 ```

 This will:
@@ -6,7 +6,7 @@ You can easily use the FastVideo Docker image as a custom container on [RunPod](

 ## Creating a new pod

-Choose a GPU that supports CUDA 12.4
+Choose a GPU that supports CUDA 12.8

 Pick 1 or 2 L40S GPU(s)

@@ -22,10 +22,20 @@ source ~/.bashrc
 Create and activate a Conda environment for FastVideo:

 ```
-conda create -n fastvideo python=3.10 -y
+conda create -n fastvideo python=3.12 -y
 conda activate fastvideo
 ```

+Install `uv` (optional, but recommended):
+
+From instructions on [uv](https://astral.sh/uv/):
+
+```
+curl -LsSf https://astral.sh/uv/install.sh | sh
+# or 
+wget -qO- https://astral.sh/uv/install.sh | sh
+```
+
 Clone the FastVideo repository and go to the FastVideo directory:

 ```
@@ -36,10 +46,10 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 Now you can install FastVideo and setup git hooks for running linting. By using `pre-commit`, the linters will run and have to pass before you'll be able to make a commit.

 ```bash
-pip install -e .[dev]
+uv pip install -e .[dev]

 # Can also install flash-attn (optional)
-pip install flash-attn==2.7.4.post1 --no-build-isolation 
+uv pip install flash-attn --no-build-isolation 

 # Linting, formatting and static type checking
 pre-commit install --hook-type pre-commit --hook-type commit-msg
@@ -50,3 +60,14 @@ pre-commit run --all-files
 # Unit tests
 pytest tests/
 ```
+
+If you are on a Hopper GPU, you should also install [FA3](https://github.com/Dao-AILab/flash-attention) for much better performance:
+
+```
+git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention/hopper
+
+# make sure you have ninja installed
+uv pip install ninja
+
+python setup.py install
+```
@@ -6,7 +6,7 @@ Instructions to install FastVideo for NVIDIA CUDA GPUs.

 - **OS: Linux or Windows WSL**
 - **Python: 3.10-3.12**
- **CUDA 12.4**
+- **CUDA 12.8**
 - **At least 1 NVIDIA GPU**

 ## Set up using Python
@@ -38,6 +38,7 @@ conda activate fastvideo

 :::{tip}
 We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
+Note that you can also use `uv` to install FastVideo in a Conda environment.
 :::

 Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:
@@ -60,7 +61,7 @@ uv pip install fastvideo
 Also optionally install flash-attn:

 ```bash
-pip install flash-attn==2.7.4.post1 --no-build-isolation
+pip install flash-attn --no-build-isolation
 ```

 ### Installation from Source
@@ -87,7 +88,7 @@ uv pip install -e .
 #### Flash Attention

 ```bash
-pip install flash-attn==2.7.4.post1 --no-build-isolation
+pip install flash-attn --no-build-isolation
 ```

 ## Set up using Docker
@@ -102,7 +103,7 @@ If you're planning to contribute to FastVideo please see the following page:
 ## Hardware Requirements

 ### For Basic Inference
- NVIDIA GPU with CUDA 12.4 support
+- NVIDIA GPU with CUDA 12.8 support

 ### For Lora Finetuning
 - 40GB GPU memory each for 2 GPUs with lora
@@ -39,6 +39,7 @@ conda activate fastvideo

 :::{tip}
 We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
+Note that you can also use `uv` to install FastVideo in a Conda environment.
 :::

 Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:
@@ -1,6 +1,6 @@
 # Welcome to FastVideo

-:::{figure} ../../assets/logo.png
+:::{figure} ../../assets/logos/logo.svg
 :align: center
 :alt: FastVideo
 :class: no-scaled-link
@@ -101,7 +101,6 @@ sliding_tile_attention/demo
 :maxdepth: 1

 video_sparse_attention/installation
-video_sparse_attention/demo
 :::

 :::{toctree}
@@ -5,7 +5,7 @@ This page contains step-by-step instructions to get you quickly started with vid
 ## Requirements
 - **OS**: Linux (Tested on Ubuntu 22.04+)
 - **Python**: 3.10-3.12
- **CUDA**: 12.4
+- **CUDA**: 12.8
 - **GPU**: At least one NVIDIA GPU

 ## Installation
@@ -6,6 +6,7 @@ The symbols used have the following meanings:

 - ✅ = Full compatibility
 - ❌ = No compatibility
+- ⭕ = Does not apply to this model

 ## Models x Optimization
 The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods and FastVideo will use the optimal default parameters when initializing and generating videos.
@@ -37,51 +38,94 @@ The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods a
  * TeaCache
  * Sliding Tile Attn
  * Sage Attn
+  * Video Sparse Attention (VSA)
+- * FastWan2.1 T2V 1.3B
+  * `FastVideo/FastWan2.1-T2V-1.3B-Diffusers`
+  * 480P
+  * ⭕
+  * ⭕
+  * ⭕
+  * ✅
+- * FastWan2.2 TI2V 5B Full Attn*
+  * `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers`
+  * 720P
+  * ⭕
+  * ⭕
+  * ⭕
+  * ✅
+- * Wan2.2 TI2V 5B 
+  * `Wan-AI/Wan2.2-TI2V-5B-Diffusers`
+  * 720P
+  * ⭕
+  * ⭕
+  * ✅
+  * ⭕
+- * Wan2.2 T2V A14B
+  * `Wan-AI/Wan2.2-T2V-A14B-Diffusers`
+  * 480P<br>720P
+  * ❌
+  * ❌
+  * ✅
+  * ⭕
+- * Wan2.2 I2V A14B
+  * `Wan-AI/Wan2.2-I2V-A14B-Diffusers`
+  * 480P<br>720P
+  * ❌
+  * ❌
+  * ✅
+  * ⭕
 - * HunyuanVideo
  * `hunyuanvideo-community/HunyuanVideo`
  * 720px1280p<br>544px960p
  * ❌
  * ✅
  * ✅
+  * ⭕
 - * FastHunyuan
  * `FastVideo/FastHunyuan-diffusers`
  * 720px1280p<br>544px960p
  * ❌
  * ✅
  * ✅
- * Wan T2V 1.3B
+  * ⭕
+- * Wan2.1 T2V 1.3B
  * `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`
  * 480P
  * ✅
  * ✅*
  * ✅
- * Wan T2V 14B
+  * ⭕
+- * Wan2.1 T2V 14B
  * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
  * 480P, 720P
  * ✅
  * ✅*
  * ✅
- * Wan I2V 480P
+  * ⭕
+- * Wan2.1 I2V 480P
  * `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
  * 480P
  * ✅
  * ✅*
  * ✅
- * Wan I2V 720P
+  * ⭕
+- * Wan2.1 I2V 720P
  * `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers`
  * 720P
  * ✅
-  * ✅*
  * ✅
+  * ✅
+  * ⭕
 - * StepVideo T2V
  * `FastVideo/stepvideo-t2v-diffusers`
  * 768px768px204f<br>544px992px204f<br>544px992px136f
  * ❌
  * ❌
  * ✅
+  * ⭕
 :::

-**Note**: there are some known quality issues with Wan2.1 + Sliding Tile Attn. We are working on fixing this issue.
+**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.

 ## Special requirements

@@ -1,3 +0,0 @@
-(vsa-demo)=
-
-# 🎬 Demo
@@ -23,10 +23,10 @@ sudo apt update
 sudo apt install clang-11
 ```

-Set up CUDA environment (if using CUDA 12.4):
+Set up CUDA environment (if using CUDA 12.8):

 ```bash
-export CUDA_HOME=/usr/local/cuda-12.4
+export CUDA_HOME=/usr/local/cuda-12.8
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 ```
@@ -0,0 +1,165 @@
+#!/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
+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=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=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:
+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/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 "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 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus 1 # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim 1
+)
+
+# 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 10
+  --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 50
+  --weight_only_checkpointing_steps 50
+  --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'
+)
+
+# Self-forcing specific arguments
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks False  # 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
+)
+
+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[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -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" 
@@ -0,0 +1,24 @@
+#!/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/merge.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 8 \
+    --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 8 \
+    --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "t2v" 
@@ -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/" 
@@ -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
+      }
+    ]
+  }
@@ -0,0 +1,31 @@
+import os
+import time
+from fastvideo import VideoGenerator, SamplingParam
+
+OUTPUT_PATH = "video_samples_causal"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        text_encoder_cpu_offload=False,
+        dit_cpu_offload=False,
+    )
+
+    sampling_param = SamplingParam.from_pretrained(model_name)
+
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+
+if __name__ == "__main__":
+    main()
@@ -1,59 +0,0 @@
-# FastVideo Gradio Demo
-
-This is a Gradio-based web interface for generating videos using the FastVideo framework. The demo allows users to create videos from text prompts with various customization options.
-
-## Overview
-
-The demo uses the FastVideo framework to generate videos based on text prompts. It provides a simple web interface built with Gradio that allows users to:
-
- Enter text prompts to generate videos
- Customize video parameters (dimensions, number of frames, etc.)
- Use negative prompts to guide the generation process
- Set or randomize seeds for reproducibility
-
---
-
-## Usage
-
-Run the demo with:
-
-```bash
-python examples/inference/gradio/gradio_demo.py
-```
-
-This will start a web server at `http://0.0.0.0:7860` where you can access the interface.
-
---
-
-## Model Initialization
-
-This demo initializes a `VideoGenerator` with the minimum required arguments for inference. Users can seamlessly adjust inference options between generations, including prompts, resolution, video length, or even the number of inference steps, *without ever needing to reload the model*.
-
-## Video Generation
-
-The core functionality is in the `generate_video` function, which:
-1. Processes user inputs
-2. Uses the FastVideo VideoGenerator from earlier to run inference (`generator.generate_video()`)
-3. Returns an output path that Gradio uses to display the generated video
-
-## Gradio Interface
-
-The interface is built with several components:
- A text input for the prompt
- A video display for the result
- Inference options in a collapsible accordion:
-  - Height and width sliders
-  - Number of frames slider
-  - Guidance scale slider
-  - Inference steps slider
-  - Negative prompt options
-  - Seed controls
-
-### Inference Options
-
- **Height/Width**: Control the resolution of the generated video
- **Number of Frames**: Set how many frames to generate
- **Guidance Scale**: Control how closely the generation follows the prompt
- **Inference Steps**: More steps can improve quality but take longer
- **Negative Prompt**: Specify what you don't want to see in the video
- **Seed**: Control randomness for reproducible results
@@ -1,169 +0,0 @@
-import argparse
-import os
-from copy import deepcopy
-
-import gradio as gr
-import torch
-
-from fastvideo import VideoGenerator
-from fastvideo.configs.sample.base import SamplingParam
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="FastVideo Gradio Demo")
-    parser.add_argument("--model_path",
-                        type=str,
-                        default="FastVideo/FastHunyuan-diffusers",
-                        help="Path to the model")
-    parser.add_argument("--num_gpus",
-                        type=int,
-                        default=1,
-                        help="Number of GPUs to use")
-    parser.add_argument("--output_path",
-                        type=str,
-                        default="outputs",
-                        help="Path to save generated videos")
-    parsed_args = parser.parse_args()
-
-    # args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
-
-    generator = VideoGenerator.from_pretrained(
-        model_path=parsed_args.model_path, num_gpus=parsed_args.num_gpus)
-
-    default_params = SamplingParam.from_pretrained(parsed_args.model_path)
-
-    def generate_video(
-        prompt,
-        negative_prompt,
-        use_negative_prompt,
-        seed,
-        guidance_scale,
-        num_frames,
-        height,
-        width,
-        num_inference_steps,
-        randomize_seed=False,
-    ):
-        params = deepcopy(default_params)
-        params.prompt = prompt
-        params.negative_prompt = negative_prompt
-        params.seed = seed
-        params.guidance_scale = guidance_scale
-        params.num_frames = num_frames
-        params.height = height
-        params.width = width
-        params.num_inference_steps = num_inference_steps
-
-        if randomize_seed:
-            params.seed = torch.randint(0, 1000000, (1, )).item()
-
-        if not use_negative_prompt:
-            params.negative_prompt = None
-
-        generator.generate_video(prompt=prompt, sampling_param=params)
-
-        output_path = os.path.join(parsed_args.output_path,
-                                   f"{params.prompt[:100]}.mp4")
-
-        return output_path, params.seed
-
-    examples = [
-        "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand’s movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough.",
-        "A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks. The cars glint in the late afternoon sun, their deep crimson and gold accents lending a touch of elegance. The tracks carve a precarious path along the cliffside, revealing glimpses of a roaring river far below. Inside, passengers peer out the large windows, their faces lit with awe as the landscape unfolds.",
-        "A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background. Strings of fairy lights hang above, casting a warm, golden glow over the scene. Groups of people gather around high tables, their laughter blending with the soft rhythm of live jazz. The aroma of freshly mixed cocktails and charred appetizers wafts through the air, mingling with the cool night breeze.",
-    ]
-
-    with gr.Blocks() as demo:
-        gr.Markdown("# FastVideo Inference Demo")
-
-        with gr.Group():
-            with gr.Row():
-                prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    max_lines=1,
-                    placeholder="Enter your prompt",
-                    container=False,
-                )
-                run_button = gr.Button("Run", scale=0)
-            result = gr.Video(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced options", open=False):
-            with gr.Group():
-                with gr.Row():
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=256,
-                        maximum=1024,
-                        step=32,
-                        value=default_params.height,
-                    )
-                    width = gr.Slider(label="Width",
-                                      minimum=256,
-                                      maximum=1024,
-                                      step=32,
-                                      value=default_params.width)
-
-                with gr.Row():
-                    num_frames = gr.Slider(
-                        label="Number of Frames",
-                        minimum=21,
-                        maximum=163,
-                        value=default_params.num_frames,
-                    )
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=1,
-                        maximum=12,
-                        value=default_params.guidance_scale,
-                    )
-                    num_inference_steps = gr.Slider(
-                        label="Inference Steps",
-                        minimum=4,
-                        maximum=100,
-                        value=default_params.num_inference_steps,
-                    )
-
-                with gr.Row():
-                    use_negative_prompt = gr.Checkbox(
-                        label="Use negative prompt", value=False)
-                    negative_prompt = gr.Text(
-                        label="Negative prompt",
-                        max_lines=1,
-                        placeholder="Enter a negative prompt",
-                        visible=False,
-                    )
-
-                seed = gr.Slider(label="Seed",
-                                 minimum=0,
-                                 maximum=1000000,
-                                 step=1,
-                                 value=default_params.seed)
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                seed_output = gr.Number(label="Used Seed")
-
-        gr.Examples(examples=examples, inputs=prompt)
-
-        use_negative_prompt.change(
-            fn=lambda x: gr.update(visible=x),
-            inputs=use_negative_prompt,
-            outputs=default_params.negative_prompt,
-        )
-
-        run_button.click(
-            fn=generate_video,
-            inputs=[
-                prompt,
-                negative_prompt,
-                use_negative_prompt,
-                seed,
-                guidance_scale,
-                num_frames,
-                height,
-                width,
-                num_inference_steps,
-                randomize_seed,
-            ],
-            outputs=[result, seed_output],
-        )
-
-    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)
@@ -0,0 +1,730 @@
+import argparse
+import os
+import requests
+import base64
+import time
+
+import gradio as gr
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+MODEL_PATH_MAPPING = {
+    "FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+    "FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+}
+
+
+class RayServeClient:
+    def __init__(self, backend_url: str):
+        self.backend_url = backend_url
+        self.session = requests.Session()
+    
+    def check_health(self) -> bool:
+        try:
+            response = self.session.get(f"{self.backend_url}/health", timeout=5)
+            return response.status_code == 200
+        except requests.exceptions.RequestException:
+            return False
+    
+    def generate_video(self, request_data: dict) -> dict:
+        start_time = time.time()
+        
+        try:
+            headers = {"Content-Type": "application/json"}
+            
+            response = self.session.post(
+                f"{self.backend_url}/generate_video",
+                json=request_data,
+                headers=headers,
+                timeout=300
+            )
+            
+            round_trip_time = time.time() - start_time
+            
+            if response.status_code == 200:
+                result = response.json()
+                backend_total = result.get("total_time", 0)
+                network_time = round_trip_time - backend_total
+                result["network_time"] = network_time
+                return result
+            else:
+                return {"success": False, "error_message": f"HTTP {response.status_code}: {response.text}"}
+                
+        except requests.exceptions.RequestException as e:
+            return {"success": False, "error_message": f"Request failed: {str(e)}"}
+
+
+def save_video_from_base64(video_data: str, output_dir: str, prompt: str) -> str:
+    if not video_data:
+        return None
+    
+    try:
+        if video_data.startswith('data:video/'):
+            video_data = video_data.split(',')[1]
+        
+        video_bytes = base64.b64decode(video_data)
+        
+        safe_prompt = prompt[:50].replace(' ', '_').replace('/', '_').replace('\\', '_')
+        video_filename = f"{safe_prompt}.mp4"
+        video_path = os.path.join(output_dir, video_filename)
+        
+        os.makedirs(output_dir, exist_ok=True)
+        
+        with open(video_path, 'wb') as f:
+            f.write(video_bytes)
+        
+        return video_path
+        
+    except Exception as e:
+        print(f"Failed to save video: {e}")
+        return None
+
+
+def create_timing_display(inference_time, encoding_time, network_time, total_time, stage_execution_times, num_frames):
+    dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A"
+    
+    timing_html = f"""
+    <div style="margin: 10px 0;">
+        <h3 style="text-align: center; margin-bottom: 10px;">⏱️ Timing Breakdown</h3>
+        <div style="display: grid; grid-template-columns: repeat(5, 1fr); gap: 10px; margin-bottom: 10px;">
+            <div class="timing-card timing-card-highlight">
+                <div style="font-size: 20px;">🚀</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">DiT Denoising</div>
+                <div style="font-size: 16px; color: #ffa200; font-weight: bold;">{dit_denoising_time}</div>
+            </div>
+            <div class="timing-card">
+                <div style="font-size: 20px;">🧠</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">E2E (w. vae/text encoder)</div>
+                <div style="font-size: 16px; color: #2563eb;">{inference_time:.2f}s</div>
+            </div>
+            <div class="timing-card">
+                <div style="font-size: 20px;">🎬</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Video Encoding</div>
+                <div style="font-size: 16px; color: #dc2626;">{encoding_time:.2f}s</div>
+            </div>
+            <div class="timing-card">
+                <div style="font-size: 20px;">🌐</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Network Transfer</div>
+                <div style="font-size: 16px; color: #059669;">{network_time:.2f}s</div>
+            </div>
+            <div class="timing-card">
+                <div style="font-size: 20px;">📊</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Total Processing</div>
+                <div style="font-size: 18px; color: #0277bd;">{total_time:.2f}s</div>
+            </div>
+        </div>"""
+    
+    if inference_time > 0:
+        fps = num_frames / inference_time
+        timing_html += f"""
+        <div class="performance-card" style="margin-top: 15px;">
+            <span style="font-weight: bold;">Generation Speed: </span>
+            <span style="font-size: 18px; color: #6366f1; font-weight: bold;">{fps:.1f} frames/second</span>
+        </div>"""
+    
+    return timing_html + "</div>"
+
+
+def load_example_prompts():
+    def contains_chinese(text):
+        return any('\u4e00' <= char <= '\u9fff' for char in text)
+    
+    def load_from_file(filepath):
+        prompts, labels = [], []
+        try:
+            with open(filepath, "r", encoding='utf-8') as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not contains_chinese(line):
+                        label = line[:100] + "..." if len(line) > 100 else line
+                        labels.append(label)
+                        prompts.append(line)
+        except Exception as e:
+            print(f"Warning: Could not read {filepath}: {e}")
+        return prompts, labels
+    
+    examples, example_labels = load_from_file("prompts/prompts_final.txt")
+    
+    if not examples:
+        examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."]
+        example_labels = ["Crowded rooftop bar at night"]
+    
+    return examples, example_labels
+
+
+def create_gradio_interface(backend_url: str, default_params: dict[str, SamplingParam]):
+    
+    client = RayServeClient(backend_url)
+    
+    def generate_video(
+        prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,
+        num_frames, height, width, randomize_seed, model_selection, progress
+    ):
+        if not client.check_health():
+            return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", ""
+        
+        # Validate dimensions
+        max_pixels = 720 * 1280
+        if height * width > max_pixels:
+            return None, f"Video dimensions too large. Maximum: 720x1280 pixels", ""
+        
+        if progress:
+            progress(0.1, desc="Checking backend health...")
+        
+        request_data = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "use_negative_prompt": use_negative_prompt,
+            "seed": seed,
+            "guidance_scale": guidance_scale,
+            "num_frames": num_frames,
+            "height": height,
+            "width": width,
+            "randomize_seed": randomize_seed,
+            "return_frames": False,
+            "image_path": None,
+            "model_path": MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
+        }
+        
+        if progress:
+            progress(0.4, desc="Generating video...")
+        
+        response = client.generate_video(request_data)
+        
+        if progress:
+            progress(0.8, desc="Processing response...")
+        
+        if response.get("success", False):
+            video_data = response.get("video_data", "")
+            used_seed = response.get("seed", seed)
+            inference_time = response.get("inference_time", 0.0)
+            encoding_time = response.get("encoding_time", 0.0)
+            total_time = response.get("total_time", 0.0)
+            network_time = response.get("network_time", 0.0)
+            stage_execution_times = response.get("stage_execution_times", [])
+            
+            timing_details = create_timing_display(
+                inference_time, encoding_time, network_time, total_time, 
+                stage_execution_times, num_frames
+            )
+            
+            if video_data:
+                if progress:
+                    progress(0.9, desc="Saving video...")
+                
+                video_path = save_video_from_base64(video_data, "outputs", prompt)
+                
+                if progress:
+                    progress(1.0, desc="Generation complete!")
+                
+                if video_path and os.path.exists(video_path):
+                    return video_path, used_seed, timing_details
+                else:
+                    return None, "Failed to save video", ""
+            else:
+                return None, "No video data received from backend", ""
+        else:
+            error_msg = response.get("error_message", "Unknown error occurred")
+            return None, f"Generation failed: {error_msg}", ""
+    
+    examples, example_labels = load_example_prompts()
+    
+    theme = gr.themes.Base().set(
+        button_primary_background_fill="#2563eb",
+        button_primary_background_fill_hover="#1d4ed8",
+        button_primary_text_color="white",
+        slider_color="#2563eb",
+        checkbox_background_color_selected="#2563eb",
+    )
+    
+    def get_default_values(model_name):
+        model_path = MODEL_PATH_MAPPING.get(model_name)
+        if model_path and model_path in default_params:
+            params = default_params[model_path]
+            return {
+                'height': params.height,
+                'width': params.width,
+                'num_frames': params.num_frames,
+                'guidance_scale': params.guidance_scale,
+                'seed': params.seed,
+            }
+        
+        return {
+            'height': 448,
+            'width': 832,
+            'num_frames': 61,
+            'guidance_scale': 3.0,
+            'seed': 1024,
+        }
+    
+    initial_values = get_default_values("FastWan2.1-T2V-1.3B")
+    
+    with gr.Blocks(title="FastWan", theme=theme) as demo:
+        gr.Image("assets/logos/logo.svg", show_label=False, container=False, height=80)
+        gr.HTML("""
+        <div style="text-align: center; margin-bottom: 10px;">
+            <p style="font-size: 18px;"> Make Video Generation Go Blurrrrrrr </p>
+            <p style="font-size: 18px;"> <a href="https://github.com/hao-ai-lab/FastVideo/tree/main" target="_blank">Code</a> | <a href="https://hao-ai-lab.github.io/blogs/fastvideo_post_training/" target="_blank">Blog</a> | <a href="https://hao-ai-lab.github.io/FastVideo/" target="_blank">Docs</a>  </p>
+        </div>
+        """)
+        
+        with gr.Accordion("🎥 What Is FastVideo?", open=False):
+            gr.HTML("""
+            <div style="padding: 20px; line-height: 1.6;">
+                <p style="font-size: 16px; margin-bottom: 15px;">
+                    FastVideo is an inference and post-training framework for diffusion models. It 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>
+            </div>
+            """)
+        
+        with gr.Row():
+            model_selection = gr.Dropdown(
+                choices=list(MODEL_PATH_MAPPING.keys()),
+                value="FastWan2.1-T2V-1.3B",
+                label="Select Model",
+                interactive=True
+            )
+
+        with gr.Row():
+            example_dropdown = gr.Dropdown(
+                choices=example_labels,
+                label="Example Prompts",
+                value=None,
+                interactive=True,
+                allow_custom_value=False
+            )
+        
+        with gr.Row():
+            with gr.Column(scale=6):
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=3,
+                    placeholder="Describe your scene...",
+                    container=False,
+                    lines=3,
+                    autofocus=True,
+                )
+            with gr.Column(scale=1, min_width=120, elem_classes="center-button"):
+                run_button = gr.Button("Run", variant="primary", size="lg")
+        
+        with gr.Row():
+            with gr.Column():
+                error_output = gr.Text(label="Error", visible=False)
+                timing_display = gr.Markdown(label="Timing Breakdown", visible=False)
+
+        with gr.Row(equal_height=True, elem_classes="main-content-row"):
+            with gr.Column(scale=1, elem_classes="advanced-options-column"):
+                with gr.Group():
+                    gr.HTML("<div style='margin: 0 0 15px 0; text-align: center; font-size: 16px;'>Advanced Options</div>")
+                    with gr.Row():
+                        height = gr.Number(
+                            label="Height",
+                            value=initial_values['height'],
+                            interactive=False,
+                            container=True
+                        )
+                        width = gr.Number(
+                            label="Width",
+                            value=initial_values['width'],
+                            interactive=False,
+                            container=True
+                        )
+                    
+                    with gr.Row():
+                        num_frames = gr.Number(
+                            label="Number of Frames",
+                            value=initial_values['num_frames'],
+                            interactive=False,
+                            container=True
+                        )
+                        guidance_scale = gr.Slider(
+                            label="Guidance Scale",
+                            minimum=1,
+                            maximum=12,
+                            value=initial_values['guidance_scale'],
+                        )
+                    
+                    with gr.Row():
+                        use_negative_prompt = gr.Checkbox(
+                            label="Use negative prompt", value=False)
+                        negative_prompt = gr.Text(
+                            label="Negative prompt",
+                            max_lines=3,
+                            lines=3,
+                            placeholder="Enter a negative prompt",
+                            visible=False,
+                        )
+
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=1000000,
+                        step=1,
+                        value=initial_values['seed'],
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                    seed_output = gr.Number(label="Used Seed")
+        
+            with gr.Column(scale=1, elem_classes="video-column"):
+                result = gr.Video(
+                    label="Generated Video", 
+                    show_label=True,
+                    height=466,
+                    width=600,
+                    container=True,
+                    elem_classes="video-component"
+                )
+        
+        gr.HTML("""
+        <style>
+        .center-button {
+            display: flex !important;
+            justify-content: center !important;
+            height: 100% !important;
+            padding-top: 1.4em !important;
+        }
+        
+        .gradio-container {
+            max-width: 1200px !important;
+            margin: 0 auto !important;
+        }
+        
+        .main {
+            max-width: 1200px !important;
+            margin: 0 auto !important;
+        }
+        
+        .gr-form, .gr-box, .gr-group {
+            max-width: 1200px !important;
+        }
+        
+        .gr-video {
+            max-width: 500px !important;
+            margin: 0 auto !important;
+        }
+        
+        .main-content-row {
+            display: flex !important;
+            align-items: flex-start !important;
+            min-height: 500px !important;
+            gap: 20px !important;
+        }
+        
+        .advanced-options-column,
+        .video-column {
+            display: flex !important;
+            flex-direction: column !important;
+            flex: 1 !important;
+            min-height: 400px !important;
+            align-items: stretch !important;
+        }
+        
+        .video-column > * {
+            margin-top: 0 !important;
+        }
+        
+        .video-column .gr-video,
+        .video-component {
+            margin-top: 0 !important;
+            padding-top: 0 !important;
+        }
+        
+        .video-column .gr-video .gr-form {
+            margin-top: 0 !important;
+        }
+        
+        .advanced-options-column .gr-group,
+        .video-column .gr-video {
+            margin-top: 0 !important;
+            vertical-align: top !important;
+        }
+        
+        .advanced-options-column > *:last-child,
+        .video-column > *:last-child {
+            flex-grow: 0 !important;
+        }
+        
+        @media (max-width: 1400px) {
+            .main-content-row {
+                min-height: 600px !important;
+            }
+            
+            .advanced-options-column,
+            .video-column {
+                min-height: 600px !important;
+            }
+        }
+        
+        @media (max-width: 1200px) {
+            .main-content-row {
+                flex-direction: column !important;
+                align-items: stretch !important;
+            }
+            
+            .advanced-options-column,
+            .video-column {
+                min-height: auto !important;
+                width: 100% !important;
+            }
+        }
+        
+        .timing-card {
+            background: var(--background-fill-secondary) !important;
+            border: 1px solid var(--border-color-primary) !important;
+            color: var(--body-text-color) !important;
+            padding: 10px;
+            border-radius: 8px;
+            text-align: center;
+            min-height: 80px;
+            display: flex;
+            flex-direction: column;
+            justify-content: center;
+        }
+        
+        .timing-card-highlight {
+            background: var(--background-fill-primary) !important;
+            border: 2px solid var(--color-accent) !important;
+        }
+        
+        .performance-card {
+            background: var(--background-fill-secondary) !important;
+            border: 1px solid var(--border-color-primary) !important;
+            color: var(--body-text-color) !important;
+            padding: 10px;
+            border-radius: 6px;
+            text-align: center;
+        }
+        
+        .gr-number input[readonly] {
+            background-color: var(--background-fill-secondary) !important;
+            border: 1px solid var(--border-color-primary) !important;
+            color: var(--body-text-color-subdued) !important;
+            cursor: default !important;
+            text-align: center !important;
+            font-weight: 500 !important;
+        }
+        </style>
+        """)
+        
+        def on_example_select(example_label):
+            if example_label and example_label in example_labels:
+                index = example_labels.index(example_label)
+                return examples[index]
+            return ""
+        
+        example_dropdown.change(
+            fn=on_example_select,
+            inputs=example_dropdown,
+            outputs=prompt,
+        )
+        
+        gr.HTML("""
+        <div style="text-align: center; margin-top: 10px; margin-bottom: 15px;">
+            <p style="font-size: 16px; margin: 0;">The compute for this demo is generously provided by <a href="https://www.gmicloud.ai/" target="_blank">GMI Cloud</a>. Note that this demo is meant to showcase FastWan's quality and that under a large number of requests, generation speed may be affected. We are also rate-limiting users to 3 requests per minute.</p>
+        </div>
+        """)
+        
+        use_negative_prompt.change(
+            fn=lambda x: gr.update(visible=x),
+            inputs=use_negative_prompt,
+            outputs=negative_prompt,
+        )
+        
+        def on_model_selection_change(selected_model):
+            if not selected_model:
+                selected_model = "FastWan2.1-T2V-1.3B"
+            
+            model_path = MODEL_PATH_MAPPING.get(selected_model)
+            
+            if model_path and model_path in default_params:
+                params = default_params[model_path]
+                return (
+                    gr.update(value=params.height),
+                    gr.update(value=params.width),
+                    gr.update(value=params.num_frames),
+                    gr.update(value=params.guidance_scale),
+                    gr.update(value=params.seed),
+                )
+            
+            return (
+                gr.update(value=448),
+                gr.update(value=832),
+                gr.update(value=61),
+                gr.update(value=3.0),
+                gr.update(value=1024),
+            )
+        
+        model_selection.change(
+            fn=on_model_selection_change,
+            inputs=model_selection,
+            outputs=[height, width, num_frames, guidance_scale, seed],
+        )
+        
+        def handle_generation(*args, progress=None, request: gr.Request = None):
+            model_selection, prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, num_frames, height, width, randomize_seed = args
+            
+            result_path, seed_or_error, timing_details = generate_video(
+                prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, 
+                num_frames, height, width, randomize_seed, model_selection, progress
+            )
+            
+            if result_path and os.path.exists(result_path):
+                return (
+                    result_path, 
+                    seed_or_error, 
+                    gr.update(visible=False),
+                    gr.update(visible=True, value=timing_details),
+                )
+            else:
+                return (
+                    None, 
+                    seed_or_error, 
+                    gr.update(visible=True, value=seed_or_error),
+                    gr.update(visible=False),
+                )
+        
+        run_button.click(
+            fn=handle_generation,
+            inputs=[
+                model_selection,
+                prompt,
+                negative_prompt,
+                use_negative_prompt,
+                seed,
+                guidance_scale,
+                num_frames,
+                height,
+                width,
+                randomize_seed,
+            ],
+            outputs=[result, seed_output, error_output, timing_display],
+            concurrency_limit=20,
+        )
+    
+    return demo
+
+
+def main():
+    parser = argparse.ArgumentParser(description="FastVideo Gradio Frontend")
+    parser.add_argument("--backend_url", type=str, default="http://localhost:8000",
+                        help="URL of the Ray Serve backend")
+    parser.add_argument("--t2v_model_paths", type=str,
+                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.1-T2V-14B-Diffusers",
+                        help="Comma separated list of paths to the T2V model(s)")
+    parser.add_argument("--host", type=str, default="0.0.0.0",
+                        help="Host to bind to")
+    parser.add_argument("--port", type=int, default=7860,
+                        help="Port to bind to")
+    
+    args = parser.parse_args()
+    
+    default_params = {}
+    model_paths = args.t2v_model_paths.split(",")
+    for model_path in model_paths:
+        default_params[model_path] = SamplingParam.from_pretrained(model_path)
+    
+    demo = create_gradio_interface(args.backend_url, default_params)
+    
+    print(f"Starting Gradio frontend at http://{args.host}:{args.port}")
+    print(f"Backend URL: {args.backend_url}")
+    print(f"T2V Models: {args.t2v_model_paths}")
+    
+    from fastapi import FastAPI, Request, HTTPException
+    from fastapi.responses import HTMLResponse, FileResponse
+    import uvicorn
+    
+    app = FastAPI()
+    
+    @app.get("/logo.svg")
+    def get_logo():
+        return FileResponse(
+            "assets/logos/logo.svg", 
+            media_type="image/svg+xml",
+            headers={
+                "Cache-Control": "public, max-age=3600",
+                "Access-Control-Allow-Origin": "*"
+            }
+        )
+    
+    @app.get("/favicon.ico")
+    def get_favicon():
+        favicon_path = "assets/logos/icon_simple.svg"
+        
+        if os.path.exists(favicon_path):
+            return FileResponse(
+                favicon_path, 
+                media_type="image/svg+xml",
+                headers={
+                    "Cache-Control": "public, max-age=3600",
+                    "Access-Control-Allow-Origin": "*"
+                }
+            )
+        else:
+            raise HTTPException(status_code=404, detail="Favicon not found")
+    
+    @app.get("/", response_class=HTMLResponse)
+    def index(request: Request):
+        base_url = str(request.base_url).rstrip('/')
+        return f"""
+        <!DOCTYPE html>
+        <html lang="en">
+        <head>
+            <meta charset="UTF-8" />
+            <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+            
+            <title>FastWan</title>
+            <meta name="title" content="FastWan">
+            <meta name="description" content="Make video generation go blurrrrrrr">
+            <meta name="keywords" content="FastVideo, video generation, AI, machine learning, FastWan">
+            
+            <meta property="og:type" content="website">
+            <meta property="og:url" content="{base_url}/">
+            <meta property="og:title" content="FastWan">
+            <meta property="og:description" content="Make video generation go blurrrrrrr">
+            <meta property="og:image" content="{base_url}/logo.svg">
+            <meta property="og:image:width" content="1200">
+            <meta property="og:image:height" content="630">
+            <meta property="og:site_name" content="FastWan">
+            
+            <meta property="twitter:card" content="summary_large_image">
+            <meta property="twitter:url" content="{base_url}/">
+            <meta property="twitter:title" content="FastWan">
+            <meta property="twitter:description" content="Make video generation go blurrrrrrr">
+            <meta property="twitter:image" content="{base_url}/logo.svg">
+            <link rel="icon" type="image/png" sizes="32x32" href="/favicon.ico">
+            <link rel="icon" type="image/png" sizes="16x16" href="/favicon.ico">
+            <link rel="apple-touch-icon" href="/favicon.ico">
+            <style>
+                body, html {{
+                    margin: 0;
+                    padding: 0;
+                    height: 100%;
+                    overflow: hidden;
+                }}
+                iframe {{
+                    width: 100%;
+                    height: 100vh;
+                    border: none;
+                }}
+            </style>
+        </head>
+        <body>
+            <iframe src="/gradio" width="100%" height="100%" style="border: none;"></iframe>
+        </body>
+        </html>
+        """
+    
+    app = gr.mount_gradio_app(
+        app, 
+        demo, 
+        path="/gradio",
+        allowed_paths=[os.path.abspath("outputs"), os.path.abspath("fastvideo-logos")]
+    )
+    
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main() 
@@ -0,0 +1,379 @@
+import time
+import os
+import torch
+import base64
+import io
+from copy import deepcopy
+from typing import Dict, Any, Optional, List
+import signal
+import sys
+
+import ray
+from ray import serve
+from fastapi import FastAPI, Request, Response
+from pydantic import BaseModel
+import numpy as np
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+import imageio
+from ray.serve.handle import DeploymentHandle
+from prometheus_client import Counter, Histogram, generate_latest
+
+NUM_GPUS = 16
+DEFAULT_FPS = 16
+SEED_RANGE_MAX = 1_000_000
+SUPPORTED_MODELS = [
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", 
+    "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+]
+
+MODEL_CONFIGS = {
+    "1.3B": {
+        "num_cpus": 2,
+        "text_encoder_cpu_offload": False,
+        "dit_cpu_offload": False,
+        "vae_cpu_offload": False,
+        "VSA_sparsity": 0.8,
+    },
+    "14B": {
+        "num_cpus": 16,
+        "text_encoder_cpu_offload": True,
+        "dit_cpu_offload": True,
+        "vae_cpu_offload": False,
+        "VSA_sparsity": 0.9,
+    }
+}
+
+
+class VideoGenerationRequest(BaseModel):
+    prompt: str
+    negative_prompt: Optional[str] = None
+    use_negative_prompt: bool = False
+    seed: int = 42
+    guidance_scale: float = 7.5
+    num_frames: int = 21
+    height: int = 448
+    width: int = 832
+    randomize_seed: bool = False
+    return_frames: bool = False
+    model_path: Optional[str] = None
+
+
+class VideoGenerationResponse(BaseModel):
+    video_data: Optional[str] = None
+    seed: int
+    success: bool
+    error_message: Optional[str] = None
+    generation_time: Optional[float] = None
+    model_load_time: Optional[float] = None
+    inference_time: Optional[float] = None
+    encoding_time: Optional[float] = None
+    total_time: Optional[float] = None
+    stage_names: Optional[List[str]] = None
+    stage_execution_times: Optional[List[float]] = None
+
+
+def encode_video_to_base64(frames: List[np.ndarray], fps: int = DEFAULT_FPS) -> str:
+    if not frames:
+        return ""
+    
+    try:
+        buffer = io.BytesIO()
+        imageio.mimsave(buffer, frames, fps=fps, format="mp4")
+        buffer.seek(0)
+        
+        video_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+        return f"data:video/mp4;base64,{video_base64}"
+        
+    except Exception as e:
+        print(f"Warning: Failed to encode video: {e}")
+        return ""
+
+
+def setup_model_environment(model_path: str) -> None:
+    if "fullattn" in model_path.lower():
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
+    else:
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+    os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"
+
+
+def process_generation_result(result: Any) -> tuple[List[np.ndarray], float, List[str], List[float]]:
+    frames = result if isinstance(result, list) else result.get("frames", [])
+    generation_time = result.get("generation_time", 0.0) if isinstance(result, dict) else 0.0
+    
+    logging_info = result.get("logging_info", None)
+    if logging_info:
+        stage_names = logging_info.get_execution_order()
+        stage_execution_times = [
+            logging_info.get_stage_info(stage_name).get("execution_time", 0.0) 
+            for stage_name in stage_names
+        ]
+    else:
+        stage_names = []
+        stage_execution_times = []
+    
+    return frames, generation_time, stage_names, stage_execution_times
+
+
+def prepare_sampling_params(video_request: VideoGenerationRequest, default_params: Any) -> Any:
+    params = deepcopy(default_params)
+    params.prompt = video_request.prompt
+    
+    if video_request.use_negative_prompt:
+        params.negative_prompt = video_request.negative_prompt
+
+    params.seed = (video_request.seed if not video_request.randomize_seed 
+                  else torch.randint(0, SEED_RANGE_MAX, (1,)).item())
+    params.randomize_seed = video_request.randomize_seed
+    params.guidance_scale = video_request.guidance_scale
+    params.num_frames = video_request.num_frames
+    params.height = video_request.height
+    params.width = video_request.width
+    params.save_video = False
+    params.return_frames = False
+    
+    return params
+
+
+class BaseModelDeployment:
+    def __init__(self, model_path: str, output_path: str = "outputs"):
+        self.model_path = model_path
+        self.output_path = output_path
+        self.generator = None
+        self.default_params = None
+
+        os.makedirs(self.output_path, exist_ok=True)
+        setup_model_environment(self.model_path)
+
+    def _initialize_generator(self, config: Dict[str, Any]) -> None:
+        from fastvideo.entrypoints.video_generator import VideoGenerator
+        from fastvideo.configs.sample.base import SamplingParam
+
+        print(f"Initializing model: {self.model_path}")
+        self.generator = VideoGenerator.from_pretrained(
+            model_path=self.model_path,
+            num_gpus=1,
+            use_fsdp_inference=True,
+            text_encoder_cpu_offload=config["text_encoder_cpu_offload"],
+            dit_cpu_offload=config["dit_cpu_offload"],
+            vae_cpu_offload=config["vae_cpu_offload"],
+            VSA_sparsity=config["VSA_sparsity"],
+            enable_stage_verification=False,
+        )
+        self.default_params = SamplingParam.from_pretrained(self.model_path)
+
+    def generate_video(self, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
+        total_start_time = time.time()
+        
+        params = prepare_sampling_params(video_request, self.default_params)
+
+        inference_start_time = time.time()
+        result = self.generator.generate_video(
+            prompt=video_request.prompt,
+            sampling_param=params,
+            save_video=False,
+            return_frames=False,
+        )
+        inference_time = time.time() - inference_start_time
+
+        frames, generation_time, stage_names, stage_execution_times = process_generation_result(result)
+
+        encoding_start_time = time.time()
+        video_data = encode_video_to_base64(frames, fps=DEFAULT_FPS)
+        encoding_time = time.time() - encoding_start_time
+        
+        total_time = time.time() - total_start_time
+
+        return VideoGenerationResponse(
+            video_data=video_data,
+            seed=params.seed,
+            success=True,
+            generation_time=generation_time,
+            inference_time=inference_time,
+            encoding_time=encoding_time,
+            total_time=total_time,
+            stage_names=stage_names,
+            stage_execution_times=stage_execution_times,
+        )
+
+
+@serve.deployment(
+    ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+)
+class T2VModelDeployment(BaseModelDeployment):
+    def __init__(self, t2v_model_path: str, output_path: str = "outputs"):
+        super().__init__(t2v_model_path, output_path)
+        self._initialize_generator(MODEL_CONFIGS["1.3B"])
+        print("✅ T2V model initialized successfully")
+
+
+@serve.deployment(
+    ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+)
+class T2V14BModelDeployment(BaseModelDeployment):
+    def __init__(self, t2v_14b_model_path: str, output_path: str = "outputs"):
+        super().__init__(t2v_14b_model_path, output_path)
+        # Override environment for 14B model
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+        self._initialize_generator(MODEL_CONFIGS["14B"])
+        print("✅ T2V 14B model initialized successfully")
+
+
+app = FastAPI()
+limiter = Limiter(key_func=get_remote_address)
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+
+
+@serve.deployment(num_replicas=50, ray_actor_options={"num_cpus": 2})
+@serve.ingress(app)
+class FastVideoAPI:
+
+    def __init__(self, t2v_deployments: Dict[str, DeploymentHandle]):
+        self.t2v_deployments = t2v_deployments
+        
+        # Initialize Prometheus metrics
+        self.request_count = Counter('fastvideo_requests_total', 'Total FastVideo requests', ['model_type', 'status'])
+        self.request_duration = Histogram('fastvideo_request_duration_seconds', 'FastVideo request duration', ['model_type'])
+        self.video_generation_time = Histogram('fastvideo_video_generation_seconds', 'Video generation time', ['model_type'])
+    
+    def _get_model_name(self, model_path: Optional[str]) -> str:
+        return model_path.split('/')[-1] if model_path else "unknown"
+    
+    def _record_metrics(self, model_name: str, status: str, duration: float, response: Optional[VideoGenerationResponse] = None) -> None:
+        self.request_count.labels(model_type=model_name, status=status).inc()
+        self.request_duration.labels(model_type=model_name).observe(duration)
+        
+        if response and hasattr(response, 'generation_time') and response.generation_time:
+            self.video_generation_time.labels(model_type=model_name).observe(response.generation_time)
+    
+    @app.post("/generate_video", response_model=VideoGenerationResponse)
+    @limiter.limit("10/minute")
+    async def generate_video(self, request: Request, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
+        """Route the request to the appropriate model deployment based on model_path."""
+        start_time = time.time()
+        model_name = self._get_model_name(video_request.model_path)
+        
+        try:
+            if video_request.model_path not in self.t2v_deployments:
+                raise ValueError(f"Model {video_request.model_path} not found")
+            
+            response_ref = self.t2v_deployments[video_request.model_path].generate_video.remote(video_request)
+            response = await response_ref
+            
+            self._record_metrics(model_name, "success", time.time() - start_time, response)
+            return response
+
+        except Exception as e:
+            self._record_metrics(model_name, "error", time.time() - start_time)
+            
+            return VideoGenerationResponse(
+                video_data=None,
+                seed=video_request.seed,
+                success=False,
+                error_message=str(e),
+                generation_time=0,
+                inference_time=0,
+                encoding_time=0,
+                total_time=0,
+            )
+    
+    @app.get("/health")
+    @limiter.limit("10/minute")
+    async def health_check(self, request: Request) -> Dict[str, str]:
+        return {"status": "healthy"}
+    
+    @app.get("/metrics")
+    async def metrics(self) -> Response:
+        return Response(generate_latest(), media_type="text/plain")
+
+
+def validate_configuration(model_paths: List[str], replicas: List[int]) -> None:
+    assert len(model_paths) == len(replicas), "Number of models and replicas must match"
+    assert sum(replicas) <= NUM_GPUS, f"Total replicas ({sum(replicas)}) must be <= {NUM_GPUS}"
+    
+    for model, replica_count in zip(model_paths, replicas):
+        assert model in SUPPORTED_MODELS, f"Model {model} not supported"
+        assert replica_count > 0, f"Replicas must be greater than 0"
+
+
+def start_ray_serve(
+    *,
+    t2v_model_paths: str,
+    t2v_model_replicas: str,
+    output_path: str = "outputs",
+    host: str = "0.0.0.0",
+    port: int = 8000,
+) -> None:
+    if not ray.is_initialized():
+        ray.init()
+
+    model_paths = t2v_model_paths.split(",")
+    replicas = [int(r) for r in t2v_model_replicas.split(",")]
+    validate_configuration(model_paths, replicas)
+
+    t2v_deps = {}
+    for model_path, replica_count in zip(model_paths, replicas):
+        t2v_dep = T2VModelDeployment.options(num_replicas=replica_count).bind(model_path, output_path)
+        t2v_deps[model_path] = t2v_dep
+
+    api = FastVideoAPI.bind(t2v_deps)
+    serve.run(api, route_prefix="/", name="fast_video")
+
+    print(f"Ray Serve backend started at http://{host}:{port}")
+    for model_path, replica_count in zip(model_paths, replicas):
+        print(f"T2V Model: {model_path} | Replicas: {replica_count}")
+    print(f"Health check: http://{host}:{port}/health")
+    print(f"Video generation endpoint: http://{host}:{port}/generate_video")
+
+
+def setup_signal_handlers() -> None:
+    signal.signal(signal.SIGINT, lambda *_: sys.exit(0))
+    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
+
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="FastVideo Ray Serve Backend")
+    parser.add_argument("--t2v_model_paths",
+                        type=str,
+                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+                        help="Comma separated list of paths to the T2V model(s)")
+    parser.add_argument("--t2v_model_replicas",
+                        type=str,
+                        default="4,4",
+                        help="Comma separated list of number of replicas for the T2V model(s)")
+    parser.add_argument("--output_path",
+                        type=str,
+                        default="outputs",
+                        help="Path to save generated videos")
+    parser.add_argument("--host",
+                        type=str,
+                        default="0.0.0.0",
+                        help="Host to bind to")
+    parser.add_argument("--port",
+                        type=int,
+                        default=8000,
+                        help="Port to bind to")
+    
+    args = parser.parse_args()
+
+    model_paths = args.t2v_model_paths.split(",")
+    replicas = [int(r) for r in args.t2v_model_replicas.split(",")]
+    validate_configuration(model_paths, replicas)
+    
+    start_ray_serve(
+        t2v_model_paths=args.t2v_model_paths,
+        t2v_model_replicas=args.t2v_model_replicas,
+        output_path=args.output_path,
+        host=args.host,
+        port=args.port,
+    )
+
+    setup_signal_handlers()
+    print("✅ FastVideo backend is running. Press Ctrl-C to stop.")
+    while True:
+        time.sleep(3600) 
@@ -0,0 +1,3 @@
+python examples/inference/gradio/start_ray_serve_app.py \
+--t2v_model_paths "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers" \
+--t2v_model_replicas "4,4"
@@ -0,0 +1,257 @@
+"""
+Startup script for FastVideo with Ray Serve backend and Gradio frontend.
+This script starts both the backend and frontend services.
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+import time
+import threading
+import signal
+import requests
+from pathlib import Path
+from typing import Dict, Any, Optional
+
+
+DEFAULT_BACKEND_HOST = "0.0.0.0"
+DEFAULT_BACKEND_PORT = 8000
+DEFAULT_FRONTEND_HOST = "0.0.0.0"
+DEFAULT_FRONTEND_PORT = 7860
+DEFAULT_OUTPUT_PATH = "outputs"
+DEFAULT_T2V_MODELS = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
+DEFAULT_T2V_REPLICAS = "4,4"
+
+HEALTH_CHECK_TIMEOUT = 5
+HEALTH_CHECK_MAX_RETRIES = 100
+HEALTH_CHECK_INTERVAL = 2
+PROCESS_SHUTDOWN_TIMEOUT = 5
+PROCESS_MONITOR_INTERVAL = 1
+
+PROJECT_ROOT = Path(__file__).parent.parent.parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+
+class ServiceManager:
+    
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.backend_process: Optional[subprocess.Popen] = None
+        self.frontend_process: Optional[subprocess.Popen] = None
+        self.backend_url = f"http://{args.backend_host}:{args.backend_port}"
+    
+    def check_backend_health(self, max_retries: int = HEALTH_CHECK_MAX_RETRIES) -> bool:
+        health_url = f"{self.backend_url}/health"
+        
+        for attempt in range(max_retries):
+            try:
+                response = requests.get(health_url, timeout=HEALTH_CHECK_TIMEOUT)
+                if response.status_code == 200:
+                    print(f"✅ Backend is healthy at {self.backend_url}")
+                    return True
+            except requests.exceptions.RequestException:
+                pass
+            
+            if attempt < max_retries - 1:
+                print(f"⏳ Waiting for backend to start... ({attempt + 1}/{max_retries})")
+                time.sleep(HEALTH_CHECK_INTERVAL)
+        
+        print(f"❌ Backend failed to start within {max_retries * HEALTH_CHECK_INTERVAL} seconds")
+        return False
+    
+    def _create_monitor_thread(self, process: subprocess.Popen, service_name: str) -> threading.Thread:
+        def monitor():
+            if process.stdout:
+                for line in process.stdout:
+                    print(f"[{service_name}] {line.rstrip()}")
+        
+        thread = threading.Thread(target=monitor, daemon=True)
+        thread.start()
+        return thread
+    
+    def _start_service(self, script_name: str, args_dict: Dict[str, Any], service_name: str) -> subprocess.Popen:
+        script_path = Path(__file__).parent / script_name
+        
+        cmd = [sys.executable, str(script_path)]
+        for key, value in args_dict.items():
+            cmd.extend([f"--{key}", str(value)])
+        
+        print(f"🚀 Starting {service_name}...")
+        print(f"Command: {' '.join(cmd)}")
+        
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=True,
+            bufsize=1
+        )
+        
+        self._create_monitor_thread(process, service_name.upper())
+        return process
+    
+    def start_backend(self) -> subprocess.Popen:
+        backend_args = {
+            "t2v_model_paths": self.args.t2v_model_paths,
+            "t2v_model_replicas": self.args.t2v_model_replicas,
+            "output_path": self.args.output_path,
+            "host": self.args.backend_host,
+            "port": self.args.backend_port
+        }
+        
+        self.backend_process = self._start_service("ray_serve_backend.py", backend_args, "backend")
+        return self.backend_process
+    
+    def start_frontend(self) -> subprocess.Popen:
+        frontend_args = {
+            "backend_url": self.backend_url,
+            "t2v_model_paths": self.args.t2v_model_paths,
+            "host": self.args.frontend_host,
+            "port": self.args.frontend_port
+        }
+        
+        self.frontend_process = self._start_service("gradio_frontend.py", frontend_args, "frontend")
+        return self.frontend_process
+    
+    def shutdown_services(self) -> None:
+        print("\n🛑 Shutting down services...")
+        
+        processes = []
+        if self.frontend_process:
+            self.frontend_process.terminate()
+            processes.append(("frontend", self.frontend_process))
+        
+        if self.backend_process:
+            self.backend_process.terminate()
+            processes.append(("backend", self.backend_process))
+        
+        for name, process in processes:
+            try:
+                process.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT)
+                print(f"✅ {name.capitalize()} stopped gracefully")
+            except subprocess.TimeoutExpired:
+                print(f"⚠️  Force killing {name} process...")
+                process.kill()
+        
+        print("✅ All services stopped")
+    
+    def monitor_processes(self) -> None:
+        if not self.backend_process or not self.frontend_process:
+            print("❌ Processes not properly initialized")
+            return
+        
+        try:
+            while True:
+                if self.frontend_process.poll() is not None:
+                    print("❌ Frontend process died unexpectedly")
+                    break
+                
+                if self.backend_process.poll() is not None:
+                    print("❌ Backend process died unexpectedly")
+                    break
+                
+                time.sleep(PROCESS_MONITOR_INTERVAL)
+        
+        except KeyboardInterrupt:
+            pass
+        
+        self.shutdown_services()
+
+
+def setup_signal_handlers(service_manager: ServiceManager) -> None:
+    def signal_handler(signum: int, frame: Any) -> None:
+        service_manager.shutdown_services()
+        sys.exit(0)
+    
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
+
+def print_startup_info(args: argparse.Namespace) -> None:
+    print("🎬 FastVideo Ray Serve App")
+    print("=" * 50)
+    print(f"T2V Models: {args.t2v_model_paths}")
+    print(f"T2V Model Replicas: {args.t2v_model_replicas}")
+    print(f"Output: {args.output_path}")
+    print(f"Backend: http://{args.backend_host}:{args.backend_port}")
+    print(f"Frontend: http://{args.frontend_host}:{args.frontend_port}")
+    print("=" * 50)
+
+
+def parse_arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="FastVideo Ray Serve App")
+    
+    parser.add_argument("--t2v_model_paths", 
+                        type=str,
+                        default=DEFAULT_T2V_MODELS,
+                        help="Comma separated list of paths to the T2V model(s)")
+    parser.add_argument("--t2v_model_replicas", 
+                        type=str, 
+                        default=DEFAULT_T2V_REPLICAS,
+                        help="Comma separated list of number of replicas for the T2V model(s)")
+    parser.add_argument("--output_path", 
+                        type=str, 
+                        default=DEFAULT_OUTPUT_PATH,
+                        help="Path to save generated videos")
+    
+    parser.add_argument("--backend_host", 
+                        type=str, 
+                        default=DEFAULT_BACKEND_HOST,
+                        help="Backend host to bind to")
+    parser.add_argument("--backend_port", 
+                        type=int, 
+                        default=DEFAULT_BACKEND_PORT,
+                        help="Backend port to bind to")
+    
+    parser.add_argument("--frontend_host", 
+                        type=str, 
+                        default=DEFAULT_FRONTEND_HOST,
+                        help="Frontend host to bind to")
+    parser.add_argument("--frontend_port", 
+                        type=int, 
+                        default=DEFAULT_FRONTEND_PORT,
+                        help="Frontend port to bind to")
+    
+    parser.add_argument("--skip_backend_check", 
+                        action="store_true",
+                        help="Skip backend health check")
+    
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_arguments()
+    
+    os.makedirs(args.output_path, exist_ok=True)
+    print_startup_info(args)
+    
+    service_manager = ServiceManager(args)
+    setup_signal_handlers(service_manager)
+    
+    try:
+        service_manager.start_backend()
+        
+        if not args.skip_backend_check:
+            if not service_manager.check_backend_health():
+                print("❌ Backend failed to start. Terminating...")
+                service_manager.shutdown_services()
+                sys.exit(1)
+        
+        service_manager.start_frontend()
+        
+        print("\n🎉 Both services are starting up!")
+        print(f"📺 Frontend will be available at: http://{args.frontend_host}:{args.frontend_port}")
+        print(f"🔧 Backend API will be available at: http://{args.backend_host}:{args.backend_port}")
+        print("\nPress Ctrl+C to stop both services...")
+        
+        service_manager.monitor_processes()
+    
+    except Exception as e:
+        print(f"❌ Unexpected error: {e}")
+        service_manager.shutdown_services()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
@@ -14,9 +14,10 @@ def main():
        vae_cpu_offload=True,
        text_encoder_cpu_offload=True,
        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" 
-        lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1000/transformer",
+        lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-160/transformer",
        lora_nickname="crush_smol"
    )
+    generator.unmerge_lora_weights()
    kwargs = {
        "height": 480,
        "width": 832,
@@ -0,0 +1,24 @@
+#!/bin/bash
+
+GPU_NUM=2 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+DATASET_PATH="data/crush-smol/"
+OUTPUT_DIR="data/crush-smol_processed_i2v/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
+    --model_path $MODEL_PATH \
+    --mode preprocess \
+    --workload_type i2v \
+    --preprocess.dataset_type merged \
+    --preprocess.dataset_path $DATASET_PATH \
+    --preprocess.dataset_output_dir $OUTPUT_DIR \
+    --preprocess.preprocess_video_batch_size 2 \
+    --preprocess.dataloader_num_workers 0 \
+    --preprocess.max_height 480 \
+    --preprocess.max_width 832 \
+    --preprocess.num_frames 77 \
+    --preprocess.train_fps 16 \
+    --preprocess.samples_per_file 8 \
+    --preprocess.flush_frequency 8 \
+    --preprocess.video_length_tolerance_range 5
@@ -7,7 +7,7 @@ export WANDB_MODE=online
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=2
+NUM_GPUS=1
 # export CUDA_VISIBLE_DEVICES=4,5


@@ -76,6 +76,7 @@ miscellaneous_args=(
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
+  --resume_from_checkpoint "checkpoints/wan_t2v_finetune_lora/checkpoint-160"
 )

 torchrun \
@@ -1,19 +1,19 @@
 #!/bin/bash

-GPU_NUM=1 # 2,4,8
+GPU_NUM=2 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_TYPE="wan"
-DATASET_PATH="data/crush-smol-test/"
+DATASET_PATH="data/crush-smol/"
 OUTPUT_DIR="data/crush-smol_processed_t2v/"

 torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/pipelines/preprocess/v1_preprocessing_new.py \
+    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
    --model_path $MODEL_PATH \
    --mode preprocess \
    --workload_type t2v \
+    --preprocess.dataset_type merged \
    --preprocess.dataset_path $DATASET_PATH \
    --preprocess.dataset_output_dir $OUTPUT_DIR \
-    --preprocess.preprocess_video_batch_size 4 \
+    --preprocess.preprocess_video_batch_size 2 \
    --preprocess.dataloader_num_workers 0 \
    --preprocess.max_height 480 \
    --preprocess.max_width 832 \
@@ -1,9 +1,36 @@
 import dataclasses
+from enum import Enum
 from typing import Any, Optional

 from fastvideo.configs.utils import update_config_from_args
+from fastvideo.logger import init_logger
 from fastvideo.utils import FlexibleArgumentParser, StoreBoolean

+logger = init_logger(__name__)
+
+
+class DatasetType(str, Enum):
+    """
+    Enumeration for different dataset types.
+    """
+    HF = "hf"
+    MERGED = "merged"
+
+    @classmethod
+    def from_string(cls, value: str) -> "DatasetType":
+        """Convert string to DatasetType enum."""
+        try:
+            return cls(value.lower())
+        except ValueError:
+            raise ValueError(
+                f"Invalid dataset type: {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 [dataset_type.value for dataset_type in cls]
+

@dataclasses.dataclass
 class PreprocessConfig:
@@ -12,6 +39,7 @@ class PreprocessConfig:
    # Model and dataset configuration
    model_path: str = ""
    dataset_path: str = ""
+    dataset_type: DatasetType = DatasetType.HF
    dataset_output_dir: str = "./output"

    # Dataloader configuration
@@ -35,6 +63,9 @@ class PreprocessConfig:
    # Model configuration
    training_cfg_rate: float = 0.0

+    # framework configuration
+    seed: int = 42
+
    @staticmethod
    def add_cli_args(parser: FlexibleArgumentParser,
                     prefix: str = "preprocess") -> FlexibleArgumentParser:
@@ -52,6 +83,12 @@ class PreprocessConfig:
            type=str,
            default=PreprocessConfig.dataset_path,
            help="Path to the dataset directory for preprocessing")
+        preprocess_args.add_argument(
+            f"--{prefix_with_dot}dataset-type",
+            type=str,
+            choices=DatasetType.choices(),
+            default=PreprocessConfig.dataset_type.value,
+            help="Type of the dataset")
        preprocess_args.add_argument(
            f"--{prefix_with_dot}dataset-output-dir",
            type=str,
@@ -123,6 +160,10 @@ class PreprocessConfig:
                                     type=float,
                                     default=PreprocessConfig.training_cfg_rate,
                                     help="Training CFG rate")
+        preprocess_args.add_argument(f"--{prefix_with_dot}seed",
+                                     type=int,
+                                     default=PreprocessConfig.seed,
+                                     help="Seed for random number generator")

        return parser

@@ -130,6 +171,10 @@ class PreprocessConfig:
    def from_kwargs(cls, kwargs: dict[str,
                                      Any]) -> Optional["PreprocessConfig"]:
        """Create PreprocessConfig from keyword arguments."""
+        if 'dataset_type' in kwargs and isinstance(kwargs['dataset_type'], str):
+            kwargs['dataset_type'] = DatasetType.from_string(
+                kwargs['dataset_type'])
+
        preprocess_config = cls()
        if not update_config_from_args(
                preprocess_config, kwargs, prefix="preprocess", pop_args=True):
@@ -92,6 +92,12 @@ class WanVideoArchConfig(DiTArchConfig):
    pos_embed_seq_len: int | None = None
    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])

+    # Causal Wan
+    local_attn_size: int = -1  # Window size for temporal local attention (-1 indicates global attention)
+    sink_size: int = 0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
+    num_frames_per_block: int = 3
+    sliding_window_num_frames: int = 21
+
    def __post_init__(self):
        super().__post_init__()
        self.out_channels = self.out_channels or self.in_channels
@@ -103,4 +109,4 @@ class WanVideoArchConfig(DiTArchConfig):
 class WanVideoConfig(DiTConfig):
    arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig)

-    prefix: str = "Wan"
+    prefix: str = "Wan"
@@ -9,6 +9,7 @@ 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)
 from fastvideo.logger import init_logger
@@ -34,6 +35,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
    "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,
    # Add other specific weight variants
 }

@@ -57,7 +57,7 @@ class WanT2V480PConfig(PipelineConfig):
    # WanConfig-specific added parameters

    def __post_init__(self):
-        self.vae_config.load_encoder = False
+        self.vae_config.load_encoder = True
        self.vae_config.load_decoder = True


@@ -138,3 +138,14 @@ class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
@dataclass
 class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
    pass
+
+
+# =============================================
+# ============= Causal Self-Forcing =============
+# =============================================
+@dataclass
+class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
+    is_causal: bool = True
+    flow_shift: int = 5
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000, 750, 500, 250])
@@ -17,6 +17,7 @@ from fastvideo.configs.sample.wan import (
    WanI2V_14B_720P_SamplingParam,
    WanT2V_1_3B_SamplingParam,
    WanT2V_14B_SamplingParam,
+    SelfForcingWanT2V480PConfig,
 )
 # isort: on
 from fastvideo.logger import init_logger
@@ -28,16 +29,29 @@ logger = init_logger(__name__)
 SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
    "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
    "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
+
+    # Wan2.1
    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam,
    "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,
-    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
-    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
+
+    # Wan2.2
    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
-    "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
+    "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers":
+    Wan2_2_TI2V_5B_SamplingParam,
    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam,
    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam,
+
+    # FastWan2.1
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
+
+    # FastWan2.2
+    "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
+
+    # Causal Self-Forcing Wan2.1
+    "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
    # Add other specific weight variants
 }

@@ -141,3 +141,11 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
    guidance_scale_2: float = 3.5
    num_inference_steps: int = 40
    fps: int = 16
+
+
+# =============================================
+# ============= Causal Self-Forcing =============
+# =============================================
+@dataclass
+class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
+    pass
@@ -9,6 +9,7 @@ diffusion models.
 import math
 import os
 import time
+from copy import deepcopy
 from typing import Any

 import imageio
@@ -202,6 +203,8 @@ class VideoGenerator:
        if sampling_param is None:
            sampling_param = SamplingParam.from_pretrained(
                fastvideo_args.model_path)
+        else:
+            sampling_param = deepcopy(sampling_param)

        kwargs["prompt"] = prompt
        sampling_param.update(kwargs)
@@ -275,6 +278,7 @@ class VideoGenerator:
                       width: {target_width}
                video_length: {sampling_param.num_frames}
                      prompt: {prompt}
+                      image_path: {sampling_param.image_path}
                  neg_prompt: {sampling_param.negative_prompt}
                        seed: {sampling_param.seed}
                 infer_steps: {sampling_param.num_inference_steps}
@@ -304,7 +308,8 @@ class VideoGenerator:
        # Run inference
        start_time = time.perf_counter()
        output_batch = self.executor.execute_forward(batch, fastvideo_args)
-        samples = output_batch
+        samples = output_batch.output
+        logging_info = output_batch.logging_info

        gen_time = time.perf_counter() - start_time
        logger.info("Generated successfully in %.2f seconds", gen_time)
@@ -334,9 +339,11 @@ class VideoGenerator:
        else:
            return {
                "samples": samples,
+                "frames": frames,
                "prompts": prompt,
                "size": (target_height, target_width, batch.num_frames),
-                "generation_time": gen_time
+                "generation_time": gen_time,
+                "logging_info": logging_info,
            }

    def set_lora_adapter(self,
@@ -344,6 +351,16 @@ class VideoGenerator:
                         lora_path: str | None = None) -> None:
        self.executor.set_lora_adapter(lora_nickname, lora_path)

+    def unmerge_lora_weights(self) -> None:
+        """
+        Use unmerged weights for inference to produce videos that align with 
+        validation videos generated during training.
+        """
+        self.executor.unmerge_lora_weights()
+
+    def merge_lora_weights(self) -> None:
+        self.executor.merge_lora_weights()
+
    def shutdown(self):
        """
        Shutdown the video generator.
@@ -591,6 +591,11 @@ class TrainingArgs(FastVideoArgs):
    pretrained_model_name_or_path: str = ""
    dit_model_name_or_path: str = ""

+    # DMD model paths - separate paths for each network
+    generator_model_path: str = ""  # path for generator (student) model
+    real_score_model_path: str = ""  # path for real score (teacher) model
+    fake_score_model_path: str = ""  # path for fake score (critic) model
+
    # diffusion setting
    ema_decay: float = 0.0
    ema_start_step: int = 0
@@ -613,6 +618,7 @@ class TrainingArgs(FastVideoArgs):
    checkpoints_total_limit: int = 0
    checkpointing_steps: int = 0
    resume_from_checkpoint: str = ""  # specify the checkpoint folder to resume from
+    init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading

    # optimizer & scheduler
    num_train_epochs: int = 0
@@ -644,6 +650,7 @@ class TrainingArgs(FastVideoArgs):
    linear_quadratic_threshold: float = 0.0
    linear_range: float = 0.0
    weight_decay: float = 0.0
+    betas: str = "0.9,0.999"  # betas for optimizer, format: "beta1,beta2"
    use_ema: bool = False
    multi_phased_distill_schedule: str = ""
    pred_decay_weight: float = 0.0
@@ -664,16 +671,29 @@ class TrainingArgs(FastVideoArgs):

    # distillation args
    generator_update_interval: int = 5
+    dfake_gen_update_ratio: int = 5  # self-forcing: how often to train generator vs critic
    min_timestep_ratio: float = 0.2
    max_timestep_ratio: float = 0.98
    real_score_guidance_scale: float = 3.5
    fake_score_learning_rate: float = 0.0  # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
    fake_score_lr_scheduler: str = "constant"  # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
+    fake_score_betas: str = "0.9,0.999"  # betas for fake score optimizer, format: "beta1,beta2"
    training_state_checkpointing_steps: int = 0  # for resuming training
    weight_only_checkpointing_steps: int = 0  # for inference
    log_visualization: bool = False
    # simulate generator forward to match inference
    simulate_generator_forward: bool = False
+    warp_denoising_step: bool = False
+
+    # Self-forcing specific arguments
+    num_frame_per_block: int = 3
+    independent_first_frame: bool = False
+    enable_gradient_masking: bool = True
+    gradient_mask_last_n_frames: int = 21
+    validate_cache_structure: bool = False  # Debug flag for cache validation
+    same_step_across_blocks: bool = False  # Use same exit timestep for all blocks
+    last_step_only: bool = False  # Only use the last timestep for training
+    context_noise: int = 0  # Context noise level for cache updates

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -775,6 +795,20 @@ class TrainingArgs(FastVideoArgs):
                            type=str,
                            help="Directory to cache models")

+        # DMD model paths - separate paths for each network
+        parser.add_argument(
+            "--generator-model-path",
+            type=str,
+            help="Path to generator (student) model for DMD distillation")
+        parser.add_argument(
+            "--real-score-model-path",
+            type=str,
+            help="Path to real score (teacher) model for DMD distillation")
+        parser.add_argument(
+            "--fake-score-model-path",
+            type=str,
+            help="Path to fake score (critic) model for DMD distillation")
+
        # Diffusion settings
        parser.add_argument("--ema-decay",
                            type=float,
@@ -845,6 +879,10 @@ class TrainingArgs(FastVideoArgs):
        parser.add_argument("--resume-from-checkpoint",
                            type=str,
                            help="Path to checkpoint to resume from")
+        parser.add_argument(
+            "--init-weights-from-safetensors",
+            type=str,
+            help="Path to safetensors file for initial weight loading")
        parser.add_argument("--logging-dir",
                            type=str,
                            help="Directory for logging")
@@ -949,6 +987,10 @@ class TrainingArgs(FastVideoArgs):
                            help="Linear quadratic threshold")
        parser.add_argument("--linear-range", type=float, help="Linear range")
        parser.add_argument("--weight-decay", type=float, help="Weight decay")
+        parser.add_argument("--betas",
+                            type=str,
+                            default=TrainingArgs.betas,
+                            help="Betas for optimizer (format: 'beta1,beta2')")
        parser.add_argument("--use-ema",
                            action=StoreBoolean,
                            help="Whether to use EMA")
@@ -990,6 +1032,13 @@ class TrainingArgs(FastVideoArgs):
                            type=int,
                            default=TrainingArgs.generator_update_interval,
                            help="Ratio of student updates to critic updates.")
+        parser.add_argument(
+            "--dfake-gen-update-ratio",
+            type=int,
+            default=TrainingArgs.dfake_gen_update_ratio,
+            help=
+            "Self-forcing: How often to train generator vs critic (train generator every N steps)."
+        )
        parser.add_argument("--min-timestep-ratio",
                            type=float,
                            default=TrainingArgs.min_timestep_ratio,
@@ -1006,6 +1055,11 @@ class TrainingArgs(FastVideoArgs):
                            type=float,
                            default=TrainingArgs.fake_score_learning_rate,
                            help="Learning rate for fake score transformer")
+        parser.add_argument(
+            "--fake-score-betas",
+            type=str,
+            default=TrainingArgs.fake_score_betas,
+            help="Betas for fake score optimizer (format: 'beta1,beta2')")
        parser.add_argument(
            "--fake-score-lr-scheduler",
            type=str,
@@ -1018,6 +1072,49 @@ class TrainingArgs(FastVideoArgs):
            "--simulate-generator-forward",
            action=StoreBoolean,
            help="Whether to simulate generator forward to match inference")
+        parser.add_argument(
+            "--warp-denoising-step",
+            action=StoreBoolean,
+            help=
+            "Whether to warp denoising step according to the scheduler time shift"
+        )
+
+        # Self-forcing specific arguments
+        parser.add_argument(
+            "--num-frame-per-block",
+            type=int,
+            default=TrainingArgs.num_frame_per_block,
+            help="Number of frames per block for causal generation")
+        parser.add_argument(
+            "--independent-first-frame",
+            action=StoreBoolean,
+            help="Whether the first frame is independent in causal generation")
+        parser.add_argument(
+            "--enable-gradient-masking",
+            action=StoreBoolean,
+            help="Whether to enable frame-level gradient masking")
+        parser.add_argument(
+            "--gradient-mask-last-n-frames",
+            type=int,
+            default=TrainingArgs.gradient_mask_last_n_frames,
+            help="Number of last frames to enable gradients for")
+        parser.add_argument(
+            "--validate-cache-structure",
+            action=StoreBoolean,
+            help="Whether to validate KV cache structure (debug flag)")
+        parser.add_argument(
+            "--same-step-across-blocks",
+            action=StoreBoolean,
+            help="Whether to use the same exit timestep for all blocks")
+        parser.add_argument(
+            "--last-step-only",
+            action=StoreBoolean,
+            help="Whether to only use the last timestep for training")
+        parser.add_argument(
+            "--context-noise",
+            type=int,
+            default=TrainingArgs.context_noise,
+            help="Context noise level for cache updates")

        return parser

@@ -9,6 +9,9 @@ import torch.nn.functional as F
 from fastvideo.layers.custom_op import CustomOp
 from fastvideo.platforms import current_platform

+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)

@CustomOp.register("rms_norm")
 class RMSNorm(CustomOp):
@@ -100,7 +103,13 @@ class ScaleResidual(nn.Module):
    def forward(self, residual: torch.Tensor, x: torch.Tensor,
                gate: torch.Tensor) -> torch.Tensor:
        """Apply gated residual connection."""
-        return residual + x * gate
+        # logger.info("x.shape: %s", x.shape)
+        # if isinstance(gate, torch.Tensor):
+            # logger.info("gate.shape: %s", gate.shape)
+
+        num_frames = gate.shape[1]
+        frame_seqlen = x.shape[1] // num_frames
+        return residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)


 # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py
@@ -172,11 +181,35 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
              but before normalization)
        """
        # Apply residual connection with gating
-        residual_output = residual + x * gate
+        # logger.info("x.shape: %s", x.shape)
+        if isinstance(gate, int): 
+            # used by cross-attention, should be 1
+            assert gate == 1
+            residual_output = residual + x * gate
+        elif isinstance(gate, torch.Tensor):
+            # logger.info("gate.shape: %s", gate.shape)
+            if gate.dim() == 3:
+                # used by bidirectional self attention
+                residual_output = residual + x * gate
+            else:
+                assert gate.dim() == 4 
+                num_frames = gate.shape[1]
+                frame_seqlen = x.shape[1] // num_frames
+                residual_output = residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)
+                # residual_output = residual + x * gate
+        else:
+            raise ValueError(f"Gate type {type(gate)} not supported")
+        # logger.info("residual_output.shape: %s", residual_output.shape)
+
        # Apply normalization
        normalized = self.norm(residual_output)
        # Apply scale and shift
-        modulated = normalized * (1.0 + scale) + shift
+        if isinstance(scale, torch.Tensor) and scale.dim() == 4:
+            num_frames = scale.shape[1]
+            frame_seqlen = normalized.shape[1] // num_frames
+            modulated = (normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2)
+        else:
+            modulated = normalized * (1.0 + scale) + shift
        return modulated, residual_output


@@ -219,7 +252,15 @@ class LayerNormScaleShift(nn.Module):
                scale: torch.Tensor) -> torch.Tensor:
        """Apply ln followed by scale and shift in a single fused operation."""
        normalized = self.norm(x)
-        if self.compute_dtype == torch.float32:
-            return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
+        if scale.dim() == 4:
+            num_frames = scale.shape[1]
+            frame_seqlen = normalized.shape[1] // num_frames
+            if self.compute_dtype == torch.float32:
+                return (normalized.float().unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2).to(x.dtype)
+            else:
+                return (normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2)
        else:
-            return normalized * (1.0 + scale) + shift
+            if self.compute_dtype == torch.float32:
+                return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
+            else:
+                return normalized * (1.0 + scale) + shift
@@ -76,7 +76,7 @@ class BaseLayerWithLoRA(nn.Module):
            lora_B = self.lora_B.to_local()
            lora_A = self.lora_A.to_local()

-        if (self.training_mode or not self.merged) and not self.disable_lora:
+        if not self.merged and not self.disable_lora:
            delta = x @ (
                self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
                @ self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)))
@@ -29,7 +29,11 @@ import torch

 from fastvideo.distributed.parallel_state import get_sp_group
 from fastvideo.layers.custom_op import CustomOp
+from fastvideo.logger import init_logger

+logger = init_logger(__name__)
+
+logger = init_logger(__name__)

 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., :x.shape[-1] // 2]
@@ -267,6 +271,7 @@ def get_nd_rotary_pos_embed(
    sp_rank: int = 0,
    sp_world_size: int = 1,
    dtype: torch.dtype = torch.float32,
+    start_frame: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    """
    This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
@@ -292,6 +297,9 @@ def get_nd_rotary_pos_embed(
    full_grid = get_meshgrid_nd(
        start, *args, dim=len(rope_dim_list))  # [3, W, H, D] / [2, W, H]

+    if start_frame > 0:
+        full_grid[0] += start_frame
+
    # Shard the grid if using sequence parallelism (sp_world_size > 1)
    assert shard_dim < len(
        rope_dim_list
@@ -370,6 +378,7 @@ def get_rotary_pos_embed(
    interpolation_factor=1.0,
    shard_dim: int = 0,
    dtype: torch.dtype = torch.float32,
+    start_frame: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Generate rotary positional embeddings for the given sizes.
@@ -413,6 +422,7 @@ def get_rotary_pos_embed(
        sp_rank=sp_rank,
        sp_world_size=sp_world_size,
        dtype=dtype,
+        start_frame=start_frame,
    )
    return freqs_cos, freqs_sin

@@ -86,12 +86,16 @@ class TimestepEmbedder(nn.Module):
                       dtype=dtype)
        self.freq_dtype = freq_dtype

-    def forward(self, t: torch.Tensor) -> torch.Tensor:
+    def forward(self,
+                t: torch.Tensor,
+                timestep_seq_len: int | None = None) -> torch.Tensor:
        t_freq = timestep_embedding(t,
                                    self.frequency_embedding_size,
                                    self.max_period,
                                    dtype=self.freq_dtype).to(
                                        self.mlp.fc_in.weight.dtype)
+        if timestep_seq_len is not None:
+            t_freq = t_freq.unflatten(0, (1, timestep_seq_len))
        # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype)
        t_emb = self.mlp(t_freq)
        return t_emb
@@ -172,4 +176,4 @@ def unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor:
    x = torch.einsum("nthwcopq->nctohpwq", x)
    imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))

-    return imgs
+    return imgs
@@ -120,14 +120,14 @@ def _info(logger: Logger,

    if not _warned_local_main_process and local_main_process_only:
        logger.warning(
-            '%s is_local_main_process is set to True, logging only from the local main process.%s',
+            '%s By default, logger.info(..) will only log from the local main process. Set logger.info(..., is_local_main_process=False) to log from all processes.%s',
            GREEN,
            RESET,
        )
        _warned_local_main_process = True
    if not _warned_main_process and main_process_only:
        logger.warning(
-            '%s is_main_process_only is set to True, logging only from the main process.%s',
+            '%s is_main_process_only is set to True, logging only from the main (RANK==0) process.%s',
            GREEN,
            RESET,
        )
@@ -0,0 +1,712 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+from torch.nn.attention.flex_attention import BlockMask
+# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention
+# see https://github.com/pytorch/pytorch/issues/133254
+# change to default for other models
+flex_attention = torch.compile(
+    flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
+import torch.distributed as dist
+
+import fastvideo.envs as envs
+from fastvideo.attention import (DistributedAttention,
+                                 LocalAttention)
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.forward_context import get_forward_context
+from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift,
+                                        RMSNorm, ScaleResidual,
+                                        ScaleResidualLayerNormScaleShift)
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.mlp import MLP
+from fastvideo.layers.rotary_embedding import (_apply_rotary_emb,
+                                               get_rotary_pos_embed)
+from fastvideo.layers.visual_embedding import (PatchEmbed)
+from fastvideo.logger import init_logger
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImageEmbedding
+from fastvideo.platforms import AttentionBackendEnum, current_platform
+
+logger = init_logger(__name__)
+class CausalWanSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 local_attn_size: int = -1,
+                 sink_size: int = 0,
+                 qk_norm=True,
+                 eps=1e-6,
+                 parallel_attention=False) -> None:
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+        self.parallel_attention = parallel_attention
+        self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560
+
+        # Scaled dot product attention
+        self.attn = LocalAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
+                                          AttentionBackendEnum.TORCH_SDPA))
+
+    def forward(self, 
+                q: torch.Tensor,
+                k: torch.Tensor,
+                v: torch.Tensor,
+                freqs_cis: tuple[torch.Tensor, torch.Tensor],
+                block_mask: BlockMask,
+                kv_cache: dict | None = None,
+                current_start: int = 0,
+                cache_start: int | None = None):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        if cache_start is None:
+            cache_start = current_start
+
+        cos, sin = freqs_cis
+        roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v)
+        roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v)
+
+        if kv_cache is None:
+            # Padding for flex attention
+            padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
+            padded_roped_query = torch.cat(
+                [roped_query,
+                    torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
+                                device=q.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            padded_roped_key = torch.cat(
+                [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],
+                                        device=k.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            padded_v = torch.cat(
+                [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],
+                                device=v.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            x = flex_attention(
+                query=padded_roped_query.transpose(2, 1),
+                key=padded_roped_key.transpose(2, 1),
+                value=padded_v.transpose(2, 1),
+                block_mask=block_mask
+            )[:, :, :-padded_length].transpose(2, 1)
+        else:
+            frame_seqlen = q.shape[1]
+            current_end = current_start + roped_query.shape[1]
+            sink_tokens = self.sink_size * frame_seqlen
+            # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache
+            kv_cache_size = kv_cache["k"].shape[1]
+            num_new_tokens = roped_query.shape[1]
+            if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
+                    num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
+                # Calculate the number of new tokens added in this step
+                # Shift existing cache content left to discard oldest tokens
+                # Clone the source slice to avoid overlapping memory error
+                num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
+                num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
+                kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+                    kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+                kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+                    kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+                # Insert the new keys/values at the end
+                local_end_index = kv_cache["local_end_index"].item() + current_end - \
+                    kv_cache["global_end_index"].item() - num_evicted_tokens
+                local_start_index = local_end_index - num_new_tokens
+                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                kv_cache["v"][:, local_start_index:local_end_index] = v
+            else:
+                # 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"].clone()
+                kv_cache["v"] = kv_cache["v"].clone()
+
+                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                kv_cache["v"][:, local_start_index:local_end_index] = v
+            x = self.attn(
+                roped_query,
+                kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index],
+                kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index]
+            )
+            kv_cache["global_end_index"].fill_(current_end)
+            kv_cache["local_end_index"].fill_(local_end_index)
+
+        return x
+
+class CausalWanTransformerBlock(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 ffn_dim: int,
+                 num_heads: int,
+                 local_attn_size: int = -1,
+                 sink_size: int = 0,
+                 qk_norm: str = "rms_norm_across_heads",
+                 cross_attn_norm: bool = False,
+                 eps: float = 1e-6,
+                 added_kv_proj_dim: int | None = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+
+        # 1. Self-attention
+        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)
+
+        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        self.attn1 = CausalWanSelfAttention(
+            dim,
+            num_heads,
+            local_attn_size=local_attn_size,
+            sink_size=sink_size,
+            qk_norm=qk_norm,
+            eps=eps)
+        self.hidden_dim = dim
+        self.num_attention_heads = num_heads
+        self.local_attn_size = local_attn_size
+        dim_head = dim // num_heads
+        if qk_norm == "rms_norm":
+            self.norm_q = RMSNorm(dim_head, eps=eps)
+            self.norm_k = RMSNorm(dim_head, eps=eps)
+        elif qk_norm == "rms_norm_across_heads":
+            # LTX applies qk norm across all heads
+            self.norm_q = RMSNorm(dim, eps=eps)
+            self.norm_k = RMSNorm(dim, eps=eps)
+        else:
+            print("QK Norm type not supported")
+            raise Exception
+        assert cross_attn_norm is True
+        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim,
+            norm_type="layer",
+            eps=eps,
+            elementwise_affine=True,
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
+
+        # 2. Cross-attention
+        # Only T2V for now
+        self.attn2 = WanT2VCrossAttention(dim,
+                                            num_heads,
+                                            qk_norm=qk_norm,
+                                            eps=eps)
+        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim,
+            norm_type="layer",
+            eps=eps,
+            elementwise_affine=False,
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
+
+        # 3. Feed-forward
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.mlp_residual = ScaleResidual()
+
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        block_mask: BlockMask,
+        kv_cache: dict | None = None,
+        crossattn_cache: dict | None = None,
+        current_start: int = 0,
+        cache_start: int | None = None,
+    ) -> torch.Tensor:
+        logger.info("temb.shape: %s", temb.shape)
+        num_frames = temb.shape[1]
+        logger.info("first hidden_states.shape: %s", hidden_states.shape)
+        logger.info("num_frames: %s", num_frames)
+        if hidden_states.dim() == 4:
+            hidden_states = hidden_states.squeeze(1)
+        frame_seqlen = hidden_states.shape[1] // temb.shape[1]
+        logger.info("frame_seqlen: %s", frame_seqlen)
+        bs, seq_length, _ = hidden_states.shape
+        orig_dtype = hidden_states.dtype
+        # assert orig_dtype != torch.float32
+        e = self.scale_shift_table + temb.float()
+        logger.info("e.shape: %s", e.shape)
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
+            6, dim=2)
+        assert shift_msa.dtype == torch.float32
+
+        # 1. Self-attention
+        logger.info("hidden_states.shape: %s", hidden_states.shape)
+        logger.info("scale_msa.shape: %s", scale_msa.shape)
+        logger.info("shift_msa.shape: %s", shift_msa.shape)
+        
+        norm_hidden_states_unflattened = self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
+        # logger.info("norm_hidden_states_unflattened.shape: %s", norm_hidden_states_unflattened.shape)
+        
+        # norm_hidden_states = (self.norm1(hidden_states.float()) *
+        #                       (1 + scale_msa) + shift_msa).to(orig_dtype)
+        norm_hidden_states = (norm_hidden_states_unflattened *
+                              (1 + scale_msa) + shift_msa).flatten(1, 2).to(orig_dtype)
+        # logger.info("1 norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        query, _ = self.to_q(norm_hidden_states)
+        key, _ = self.to_k(norm_hidden_states)
+        value, _ = self.to_v(norm_hidden_states)
+
+        if self.norm_q is not None:
+            query = self.norm_q(query)
+        if self.norm_k is not None:
+            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))
+        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+
+        attn_output = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start)
+        attn_output = attn_output.flatten(2)
+        attn_output, _ = self.to_out(attn_output)
+        attn_output = attn_output.squeeze(1)
+
+        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
+        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
+            hidden_states, attn_output, gate_msa, null_shift, null_scale)
+        # logger.info("after self_attn_residual_norm norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        # logger.info("after self_attn_residual_norm hidden_states.shape: %s", hidden_states.shape)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(
+            orig_dtype), hidden_states.to(orig_dtype)
+
+        # 2. Cross-attention
+        attn_output = self.attn2(norm_hidden_states,
+                                 context=encoder_hidden_states,
+                                 context_lens=None,
+                                 crossattn_cache=crossattn_cache)
+        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
+            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
+        # logger.info("after cross_attn_residual_norm norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        # logger.info("after cross_attn_residual_norm hidden_states.shape: %s", hidden_states.shape)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(
+            orig_dtype), hidden_states.to(orig_dtype)
+
+        # 3. Feed-forward
+        ff_output = self.ffn(norm_hidden_states)
+        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
+        # logger.info("after mlp_residual norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        logger.info("after mlp_residual hidden_states.shape: %s", hidden_states.shape)
+        hidden_states = hidden_states.to(orig_dtype)
+
+        return hidden_states
+
+class CausalWanTransformer3DModel(BaseDiT):
+    _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
+    _compile_conditions = WanVideoConfig()._compile_conditions
+    _supported_attention_backends = WanVideoConfig(
+    )._supported_attention_backends
+    param_names_mapping = WanVideoConfig().param_names_mapping
+    reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
+    lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping
+
+    def __init__(self, config: WanVideoConfig, hf_config: dict[str,
+                                                               Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_dim = config.attention_head_dim
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels
+        self.num_channels_latents = config.num_channels_latents
+        self.patch_size = config.patch_size
+        self.text_len = config.text_len
+        self.local_attn_size = config.local_attn_size
+
+        # 1. Patch & position embedding
+        self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
+                                          embed_dim=inner_dim,
+                                          patch_size=config.patch_size,
+                                          flatten=False)
+
+        # 2. Condition embeddings
+        self.condition_embedder = WanTimeTextImageEmbedding(
+            dim=inner_dim,
+            time_freq_dim=config.freq_dim,
+            text_embed_dim=config.text_dim,
+            image_embed_dim=config.image_dim,
+        )
+
+        # 3. Transformer blocks
+        self.blocks = nn.ModuleList([
+            CausalWanTransformerBlock(inner_dim,
+                              config.ffn_dim,
+                              config.num_attention_heads,
+                              config.local_attn_size,
+                              config.sink_size,
+                              config.qk_norm,
+                              config.cross_attn_norm,
+                              config.eps,
+                              config.added_kv_proj_dim,
+                              self._supported_attention_backends,
+                              prefix=f"{config.prefix}.blocks.{i}")
+            for i in range(config.num_layers)
+        ])
+
+        # 4. Output norm & projection
+        self.norm_out = LayerNormScaleShift(inner_dim,
+                                            norm_type="layer",
+                                            eps=config.eps,
+                                            elementwise_affine=False,
+                                            dtype=torch.float32,
+                                            compute_dtype=torch.float32)
+        # Debug: Log configuration values
+        proj_out_dim = config.out_channels * math.prod(config.patch_size)
+        
+        self.proj_out = nn.Linear(inner_dim, proj_out_dim)
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+
+        self.gradient_checkpointing = False
+
+        # Causal-specific
+        self.block_mask = None
+        self.num_frame_per_block = 1
+        self.independent_first_frame = False
+
+        self.__post_init__()
+
+    @staticmethod
+    def _prepare_blockwise_causal_attn_mask(
+        device: torch.device | str, num_frames: int = 21,
+        frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1
+    ) -> BlockMask:
+        """
+        we will divide the token sequence into the following format
+        [1 latent frame] [1 latent frame] ... [1 latent frame]
+        We use flexattention to construct the attention mask
+        """
+        total_length = num_frames * frame_seqlen
+
+        # we do right padding to get to a multiple of 128
+        padded_length = math.ceil(total_length / 128) * 128 - total_length
+
+        ends = torch.zeros(total_length + padded_length,
+                           device=device, dtype=torch.long)
+
+        # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+        frame_indices = torch.arange(
+            start=0,
+            end=total_length,
+            step=frame_seqlen * num_frame_per_block,
+            device=device
+        )
+
+        for tmp in frame_indices:
+            ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \
+                frame_seqlen * num_frame_per_block
+
+        def attention_mask(b, h, q_idx, kv_idx):
+            if local_attn_size == -1:
+                return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)
+            else:
+                return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx)
+            # return ((kv_idx < total_length) & (q_idx < total_length))  | (q_idx == kv_idx) # bidirectional mask
+
+        block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,
+                                       KV_LEN=total_length + padded_length, _compile=False, device=device)
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                f" cache a block wise causal mask with block size of {num_frame_per_block} frames")
+            print(block_mask)
+
+        # import imageio
+        # import numpy as np
+        # from torch.nn.attention.flex_attention import create_mask
+
+        # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +
+        #                    padded_length, KV_LEN=total_length + padded_length, device=device)
+        # import cv2
+        # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))
+        # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))
+
+        return block_mask
+
+    def _forward_inference(
+                self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+                timestep: torch.LongTensor,
+                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
+                | None = None,
+                kv_cache: dict = None,
+                crossattn_cache: dict = None,
+                current_start: int = 0,
+                cache_start: int = 0,
+                start_frame: int = 0,
+                **kwargs) -> torch.Tensor:
+        r"""
+        Run the diffusion model with kv caching.
+        See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details.
+        This function will be run for num_frame times.
+        Process the latent frames one by one (1560 tokens each)
+        """
+        # logger.info("forward inference hidden_states.shape: %s", hidden_states.shape)
+
+        orig_dtype = hidden_states.dtype
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+        if isinstance(encoder_hidden_states_image,
+                      list) and len(encoder_hidden_states_image) > 0:
+            encoder_hidden_states_image = encoder_hidden_states_image[0]
+        else:
+            encoder_hidden_states_image = None
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # Get rotary embeddings
+        d = self.hidden_size // self.num_attention_heads
+        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (post_patch_num_frames * get_sp_world_size(), post_patch_height,
+             post_patch_width),
+            self.hidden_size,
+            self.num_attention_heads,
+            rope_dim_list,
+            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            rope_theta=10000,
+            start_frame=start_frame # Assume that start_frame is 0 when kv_cache is None
+        )
+        freqs_cos = freqs_cos.to(hidden_states.device)
+        freqs_sin = freqs_sin.to(hidden_states.device)
+        freqs_cis = (freqs_cos.float(),
+                     freqs_sin.float()) if freqs_cos is not None else None
+
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        # logger.info("forward inference flattened and transposed hidden_states.shape: %s", hidden_states.shape)
+
+        logger.info("timestep shape: %s", timestep.shape)
+
+        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)
+
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states = torch.concat(
+                [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+
+        encoder_hidden_states = encoder_hidden_states.to(
+            orig_dtype) if current_platform.is_mps(
+            ) else encoder_hidden_states  # cast to orig_dtype for MPS
+
+        assert encoder_hidden_states.dtype == orig_dtype
+
+        # 4. Transformer blocks
+        for block_index, block in enumerate(self.blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                causal_kwargs = {
+                    "kv_cache": kv_cache[block_index],
+                    "current_start": current_start,
+                    "cache_start": cache_start,
+                    "block_mask": self.block_mask
+                }
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, encoder_hidden_states,
+                    timestep_proj, freqs_cis,
+                    **causal_kwargs)
+            else:
+                causal_kwargs = {
+                    "kv_cache": kv_cache[block_index],
+                    "crossattn_cache": crossattn_cache[block_index],
+                    "current_start": current_start,
+                    "cache_start": cache_start,
+                    "block_mask": self.block_mask
+                }
+                hidden_states = block(hidden_states, encoder_hidden_states,
+                                        timestep_proj, freqs_cis,
+                                        **causal_kwargs)
+
+        # 5. Output norm, projection & unpatchify
+        # logger.info("===== INFERENCE 5. Output norm, projection & unpatchify")
+        # logger.info("hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("temb.shape: %s", temb.shape)
+        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
+        # logger.info("WTFWTF train temb.shape: %s", temb.shape)
+        # logger.info("WTFWTF train self.scale_shift_table.shape: %s", self.scale_shift_table.shape)
+        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
+                                                                          dim=2)
+        hidden_states = self.norm_out(hidden_states, shift, scale)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
+                                              post_patch_height,
+                                              post_patch_width, p_t, p_h, p_w,
+                                              -1)
+        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return output
+
+    def _forward_train(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+                timestep: torch.LongTensor,
+                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
+                | None = None,
+                start_frame: int = 0,
+                **kwargs) -> torch.Tensor:
+
+        # logger.info("===== forward train hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("===== forward train timestep.shape: %s", timestep.shape)
+        orig_dtype = hidden_states.dtype
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+        if isinstance(encoder_hidden_states_image,
+                      list) and len(encoder_hidden_states_image) > 0:
+            encoder_hidden_states_image = encoder_hidden_states_image[0]
+        else:
+            encoder_hidden_states_image = None
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # Get rotary embeddings
+        d = self.hidden_size // self.num_attention_heads
+        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (post_patch_num_frames * get_sp_world_size(), post_patch_height,
+             post_patch_width),
+            self.hidden_size,
+            self.num_attention_heads,
+            rope_dim_list,
+            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            rope_theta=10000,
+            start_frame=start_frame
+        )
+        freqs_cos = freqs_cos.to(hidden_states.device)
+        freqs_sin = freqs_sin.to(hidden_states.device)
+        freqs_cis = (freqs_cos.float(),
+                     freqs_sin.float()) if freqs_cos is not None else None
+
+        # Construct blockwise causal attn mask
+        if self.block_mask is None:
+            self.block_mask = self._prepare_blockwise_causal_attn_mask(
+                device=hidden_states.device,
+                num_frames=num_frames,
+                frame_seqlen=post_patch_height * post_patch_width,
+                num_frame_per_block=self.num_frame_per_block,
+                local_attn_size=self.local_attn_size
+            )
+
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        # logger.info("forward train flattened and transposed hidden_states.shape: %s", hidden_states.shape)
+
+        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
+            timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
+        # logger.info("forward train timestep_proj.shape: %s", timestep_proj.shape)
+        # logger.info("forward train timestep.shape: %s", timestep.shape)
+        # logger.info("forward train temb.shape: %s", temb.shape)
+        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
+
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states = torch.concat(
+                [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+
+        encoder_hidden_states = encoder_hidden_states.to(
+            orig_dtype) if current_platform.is_mps(
+            ) else encoder_hidden_states  # cast to orig_dtype for MPS
+
+        assert encoder_hidden_states.dtype == orig_dtype
+
+        # 4. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for block in self.blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, encoder_hidden_states,
+                    timestep_proj, freqs_cis,
+                    block_mask=self.block_mask)
+        else:
+            for block_index, block in enumerate(self.blocks):
+                logger.info("===== TRAIN block %d", block_index)
+                logger.info("hidden_states.shape: %s", hidden_states.shape)
+                # logger.info("encoder_hidden_states.shape: %s", encoder_hidden_states.shape)
+                logger.info("timestep_proj.shape: %s", timestep_proj.shape)
+                # logger.info("freqs_cis.shape: %s", freqs_cis.shape)
+                # logger.info("block_mask.shape: %s", self.block_mask.shape)
+                hidden_states = block(hidden_states, encoder_hidden_states,
+                                        timestep_proj, freqs_cis,
+                                        block_mask=self.block_mask)
+
+        # 5. Output norm, projection & unpatchify
+        # logger.info("===== TRAIN 5. Output norm, projection & unpatchify")
+        # logger.info("hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("temb.shape: %s", temb.shape)
+        # shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
+        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
+        # logger.info("WTFWTF train temb.shape: %s", temb.shape)
+        # logger.info("WTFWTF train self.scale_shift_table.shape: %s", self.scale_shift_table.shape)
+        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
+                                                                          dim=2)
+        # logger.info("DEBUG scale.shape: %s", scale.shape)
+        # logger.info("DEBUG shift.shape: %s", shift.shape)
+        hidden_states = self.norm_out(hidden_states, shift, scale)
+        hidden_states = self.proj_out(hidden_states)
+        # logger.info("DEBUG after proj_out hidden_states.shape: %s", hidden_states.shape)
+        # logger.info(f"DEBUG reshape dimensions: batch_size={batch_size}, post_patch_num_frames={post_patch_num_frames}")
+        # logger.info(f"DEBUG reshape dimensions: post_patch_height={post_patch_height}, post_patch_width={post_patch_width}")
+        # logger.info(f"DEBUG patch dimensions: p_t={p_t}, p_h={p_h}, p_w={p_w}")
+
+        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
+                                              post_patch_height,
+                                              post_patch_width, p_t, p_h, p_w,
+                                              -1)
+        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return output
+
+    def forward(
+        self,
+        *args,
+        **kwargs
+    ):
+        if kwargs.get('kv_cache', None) is not None:
+            return self._forward_inference(*args, **kwargs)
+        else:
+            return self._forward_train(*args, **kwargs)
@@ -81,8 +81,9 @@ class WanTimeTextImageEmbedding(nn.Module):
        timestep: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        encoder_hidden_states_image: torch.Tensor | None = None,
+        timestep_seq_len: int | None = None,
    ):
-        temb = self.time_embedder(timestep)
+        temb = self.time_embedder(timestep, timestep_seq_len)
        timestep_proj = self.time_modulation(temb)

        encoder_hidden_states = self.text_embedder(encoder_hidden_states)
@@ -145,7 +146,7 @@ class WanSelfAttention(nn.Module):

 class WanT2VCrossAttention(WanSelfAttention):

-    def forward(self, x, context, context_lens):
+    def forward(self, x, context, context_lens, crossattn_cache=None):
        r"""
        Args:
            x(Tensor): Shape [B, L1, C]
@@ -156,8 +157,20 @@ class WanT2VCrossAttention(WanSelfAttention):

        # compute query, key, value
        q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
-        k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
-        v = self.to_v(context)[0].view(b, -1, n, d)
+
+        if crossattn_cache is not None:
+            if not crossattn_cache["is_init"]:
+                crossattn_cache["is_init"] = True
+                k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
+                v = self.to_v(context)[0].view(b, -1, n, d)
+                crossattn_cache["k"] = k
+                crossattn_cache["v"] = v
+            else:
+                k = crossattn_cache["k"]
+                v = crossattn_cache["v"]
+        else:
+            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
        x = self.attn(q, k, v)
@@ -307,9 +320,24 @@ class WanTransformerBlock(nn.Module):
        bs, seq_length, _ = hidden_states.shape
        orig_dtype = hidden_states.dtype
        # assert orig_dtype != torch.float32
-        e = self.scale_shift_table + temb.float()
-        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
-            6, dim=1)
+
+        if temb.dim() == 4:
+            # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
+            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
+                self.scale_shift_table.unsqueeze(0) + temb.float()
+            ).chunk(6, dim=2)
+            # batch_size, seq_len, 1, inner_dim
+            shift_msa = shift_msa.squeeze(2)
+            scale_msa = scale_msa.squeeze(2)
+            gate_msa = gate_msa.squeeze(2)
+            c_shift_msa = c_shift_msa.squeeze(2)
+            c_scale_msa = c_scale_msa.squeeze(2)
+            c_gate_msa = c_gate_msa.squeeze(2)
+        else:
+            # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
+            e = self.scale_shift_table + temb.float()
+            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
@@ -637,9 +665,21 @@ class WanTransformer3DModel(CachableDiT):
        hidden_states = self.patch_embedding(hidden_states)
        hidden_states = hidden_states.flatten(2).transpose(1, 2)

+        # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
+        if timestep.dim() == 2:
+            ts_seq_len = timestep.shape[1]
+            timestep = timestep.flatten()  # batch_size * seq_len
+        else:
+            ts_seq_len = None
+
        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep, encoder_hidden_states, encoder_hidden_states_image)
-        timestep_proj = timestep_proj.unflatten(1, (6, -1))
+            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
+        if ts_seq_len is not None:
+            # batch_size, seq_len, 6, inner_dim
+            timestep_proj = timestep_proj.unflatten(2, (6, -1))
+        else:
+            # batch_size, 6, inner_dim
+            timestep_proj = timestep_proj.unflatten(1, (6, -1))

        if encoder_hidden_states_image is not None:
            encoder_hidden_states = torch.concat(
@@ -676,8 +716,15 @@ class WanTransformer3DModel(CachableDiT):
            if enable_teacache:
                self.maybe_cache_states(hidden_states, original_hidden_states)
        # 5. Output norm, projection & unpatchify
-        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
-                                                                          dim=1)
+        if temb.dim() == 3:
+            # batch_size, seq_len, inner_dim (wan 2.2 ti2v)
+            shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
+            shift = shift.squeeze(2)
+            scale = scale.squeeze(2)
+        else:
+            # batch_size, inner_dim
+            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)

@@ -781,3 +828,4 @@ class WanTransformer3DModel(CachableDiT):
            return hidden_states + self.previous_residual_even
        else:
            return hidden_states + self.previous_residual_odd
+        
@@ -430,6 +430,16 @@ class TransformerLoader(ComponentLoader):
        if not safetensors_list:
            raise ValueError(f"No safetensors files found in {model_path}")

+        # Check if we should use custom initialization weights
+        custom_weights_path = getattr(fastvideo_args, 'init_weights_from_safetensors', None)
+        use_custom_weights = (custom_weights_path and os.path.exists(custom_weights_path) and 
+                            fastvideo_args.training_mode and 
+                            not hasattr(fastvideo_args, '_loading_teacher_critic_model'))
+        
+        if use_custom_weights:
+            logger.info("Using custom initialization weights from: %s", custom_weights_path)
+            safetensors_list = [custom_weights_path]
+
        logger.info("Loading model from %s safetensors files in %s",
                    len(safetensors_list), model_path)

@@ -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:
@@ -25,12 +25,14 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
    "HunyuanVideoTransformer3DModel":
    ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+    "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
    "StepVideoModel": ("dits", "stepvideo", "StepVideoModel")
 }

 _IMAGE_TO_VIDEO_DIT_MODELS = {
    # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+    "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
 }

 _TEXT_ENCODER_MODELS = {
@@ -636,8 +636,14 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
        timestep: torch.IntTensor,
    ) -> torch.Tensor:
        self.sigmas = self.sigmas.to(noise.device)
-        timestep = timestep.expand(clean_latent.shape[0])
+        # TODO: hack
+        if timestep.ndim == 2 and timestep.shape[1] == 1:
+            timestep = timestep.expand(clean_latent.shape[0])
        self.timesteps = self.timesteps.to(noise.device)
+        logger.info("self.timesteps shape: %s", self.timesteps.shape)
+        logger.info("timestep shape: %s", timestep.shape)
+        if timestep.ndim > 1:
+            timestep = timestep.squeeze(0)
        timestep_id = torch.argmin(
            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
@@ -6,6 +6,9 @@ from typing import Any
 import torch


+from fastvideo.logger import init_logger
+logger = init_logger(__name__)
+
 # TODO(PY): move it elsewhere
 def auto_attributes(init_func):
    """
@@ -137,3 +140,27 @@ def modulate(x: torch.Tensor,
    else:
        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(
            1)  # type: ignore[union-attr]
+
+
+def pred_noise_to_pred_video(pred_noise: torch.Tensor,
+                             noise_input_latent: torch.Tensor,
+                             timestep: torch.Tensor,
+                             scheduler: Any) -> torch.Tensor:
+    """
+    Convert predicted noise to clean latent.
+    """
+    logger.info(f"timestep: {timestep.shape}")
+    logger.info(f"noise_input_latent: {noise_input_latent.shape}")
+    logger.info(f"pred_noise: {pred_noise.shape}")
+    timestep = timestep.expand(noise_input_latent.shape[0])
+    dtype = pred_noise.dtype
+    device = pred_noise.device
+    pred_noise = pred_noise.float().to(device)
+    noise_input_latent = noise_input_latent.float().to(device)
+    sigmas = scheduler.sigmas.float().to(device)
+    timesteps = scheduler.timesteps.float().to(device)
+    timestep_id = torch.argmin(
+        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+    pred_video = noise_input_latent - sigma_t * pred_noise
+    return pred_video.to(dtype)
@@ -27,6 +27,7 @@ from fastvideo.layers.activation import get_act_fn
 from fastvideo.models.vaes.common import (DiagonalGaussianDistribution,
                                          ParallelTiledVAE)
 from fastvideo.platforms import current_platform
+from fastvideo.logger import init_logger

 CACHE_T = 2

@@ -35,6 +36,7 @@ feat_cache = contextvars.ContextVar("feat_cache", default=None)
 feat_idx = contextvars.ContextVar("feat_idx", default=0)
 first_chunk = contextvars.ContextVar("first_chunk", default=None)

+logger = init_logger(__name__)

@contextmanager
 def forward_context(first_frame_arg=False,
@@ -1129,6 +1131,7 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
            self._conv_idx = 0
            self._feat_map = [None] * self._conv_num
        # cache encode
+        logger.info("self.config.load_encoder: %s", self.config.load_encoder)
        if self.config.load_encoder:
            self._enc_conv_num = _count_conv3d(self.encoder)
            self._enc_conv_idx = 0
@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Wan causal DMD pipeline implementation.
+
+This module wires the causal DMD denoising stage into the modular pipeline.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+
+# isort: off
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
+                                        CausalDMDDenosingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage)
+# isort: on
+
+logger = init_logger(__name__)
+
+
+class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
+
+    _required_config_modules = [
+        "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."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        self.add_stage(stage_name="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"),
+                           transformer=self.get_module("transformer", None)))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=CausalDMDDenosingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = WanCausalDMDPipeline
@@ -12,12 +12,10 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.pipelines.lora_pipeline import LoRAPipeline

 # isort: off
-from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage,
-                                        DecodingStage, DmdDenoisingStage,
-                                        EncodingStage, InputValidationStage,
-                                        LatentPreparationStage,
-                                        TextEncodingStage,
-                                        TimestepPreparationStage)
+from fastvideo.pipelines.stages import (
+    ImageEncodingStage, ConditioningStage, DecodingStage, DmdDenoisingStage,
+    ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage,
+    TextEncodingStage, TimestepPreparationStage)
 # isort: on
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
    FlowMatchEulerDiscreteScheduler)
@@ -67,7 +65,7 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
                           transformer=self.get_module("transformer")))

        self.add_stage(stage_name="image_latent_preparation_stage",
-                       stage=EncodingStage(vae=self.get_module("vae")))
+                       stage=ImageVAEEncodingStage(vae=self.get_module("vae")))

        self.add_stage(stage_name="denoising_stage",
                       stage=DmdDenoisingStage(
@@ -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():
@@ -282,7 +284,11 @@ class ComposedPipelineBase(ABC):
        for module_name, (transformers_or_diffusers,
                          architecture) in model_index.items():
            if transformers_or_diffusers is None:
-                self.required_config_modules.remove(module_name)
+                logger.warning(
+                    "Module in model_index.json has null value, removing from required_config_modules"
+                )
+                if module_name in self.required_config_modules:
+                    self.required_config_modules.remove(module_name)
                continue
            if module_name not in required_modules:
                logger.info("Skipping module %s", module_name)
@@ -217,3 +217,11 @@ class LoRAPipeline(ComposedPipelineBase):
                layer.disable_lora = True
        logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank,
                    lora_path, adapted_count)
+
+    def merge_lora_weights(self) -> None:
+        for name, layer in self.lora_layers.items():
+            layer.merge_lora_weights()
+
+    def unmerge_lora_weights(self) -> None:
+        for name, layer in self.lora_layers.items():
+            layer.unmerge_lora_weights()
@@ -17,10 +17,47 @@ import torch
 if TYPE_CHECKING:
    from torchcodec.decoders import VideoDecoder

+import time
+from collections import OrderedDict
+
 from fastvideo.attention import AttentionMetadata
 from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams


+class PipelineLoggingInfo:
+    """Simple approach using OrderedDict to track stage metrics."""
+
+    def __init__(self):
+        # OrderedDict preserves insertion order and allows easy access
+        self.stages: OrderedDict[str, dict[str, Any]] = OrderedDict()
+
+    def add_stage_execution_time(self, stage_name: str, execution_time: float):
+        """Add execution time for a stage."""
+        if stage_name not in self.stages:
+            self.stages[stage_name] = {}
+        self.stages[stage_name]['execution_time'] = execution_time
+        self.stages[stage_name]['timestamp'] = time.time()
+
+    def add_stage_metric(self, stage_name: str, metric_name: str, value: Any):
+        """Add any metric for a stage."""
+        if stage_name not in self.stages:
+            self.stages[stage_name] = {}
+        self.stages[stage_name][metric_name] = value
+
+    def get_stage_info(self, stage_name: str) -> dict[str, Any]:
+        """Get all info for a specific stage."""
+        return self.stages.get(stage_name, {})
+
+    def get_execution_order(self) -> list[str]:
+        """Get stages in execution order."""
+        return list(self.stages.keys())
+
+    def get_total_execution_time(self) -> float:
+        """Get total pipeline execution time."""
+        return sum(
+            stage.get('execution_time', 0) for stage in self.stages.values())
+
+
@dataclass
 class ForwardBatch:
    """
@@ -40,7 +77,7 @@ class ForwardBatch:
    # Image inputs
    image_path: str | None = None
    image_embeds: list[torch.Tensor] = field(default_factory=list)
-    pil_image: PIL.Image.Image | None = None
+    pil_image: torch.Tensor | PIL.Image.Image | None = None
    preprocessed_image: torch.Tensor | None = None

    # Text inputs
@@ -132,6 +169,10 @@ class ForwardBatch:
    # VSA parameters
    VSA_sparsity: float = 0.0

+    # Logging info
+    logging_info: PipelineLoggingInfo = field(
+        default_factory=PipelineLoggingInfo)
+
    def __post_init__(self):
        """Initialize dependent fields after dataclass initialization."""

@@ -201,4 +242,4 @@ class TrainingBatch:
@dataclass
 class PreprocessBatch(ForwardBatch):
    video_loader: list["VideoDecoder"] = field(default_factory=list)
-    video_file_name: list[str] = field(default_factory=list)
+    video_file_name: list[str] = field(default_factory=list)
@@ -21,10 +21,16 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
    "WanPipeline": "wan",
    "WanDMDPipeline": "wan",
    "WanImageToVideoPipeline": "wan",
+    "WanCausalDMDPipeline": "wan",
    "StepVideoPipeline": "stepvideo",
    "HunyuanVideoPipeline": "hunyuan",
 }

+_PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {
+    WorkloadType.I2V: "PreprocessPipelineI2V",
+    WorkloadType.T2V: "PreprocessPipelineT2V",
+}
+

 class PipelineType(str, Enum):
    """
@@ -68,12 +74,8 @@ class _PipelineRegistry:
    def _load_preprocess_pipeline_cls(
            self, workload_type: WorkloadType,
            arch: str) -> type[ComposedPipelineBase] | None:
-        if workload_type == WorkloadType.I2V:
-            pipeline_name = "I2VPreprocessPipeline"
-        elif workload_type == WorkloadType.T2V:
-            pipeline_name = "T2VPreprocessPipeline"
-        else:
-            raise ValueError(f"Invalid workload type: {workload_type.value}")
+        pipeline_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[
+            workload_type]

        return self.pipelines[
            PipelineType.PREPROCESS.value][arch][pipeline_name]
@@ -1,3 +1,4 @@
+import random
 from collections.abc import Callable
 from typing import cast

@@ -8,7 +9,7 @@ from torchvision import transforms

 from fastvideo.dataset.transform import (CenterCropResizeVideo,
                                         TemporalRandomCrop)
-from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
 from fastvideo.pipelines.pipeline_batch_info import (ForwardBatch,
                                                     PreprocessBatch)
 from fastvideo.pipelines.stages.base import PipelineStage
@@ -68,6 +69,39 @@ class VideoTransformStage(PipelineStage):
        video_pixel_values = rearrange(video_pixel_values,
                                       "b t c h w -> b c t h w")
        video_pixel_values = video_pixel_values.to(torch.uint8)
+
+        if fastvideo_args.workload_type == WorkloadType.I2V:
+            batch.pil_image = video_pixel_values[:, :, 0, :, :]
+
        video_pixel_values = video_pixel_values.float() / 255.0
        batch.latents = video_pixel_values
+        batch.num_frames = [video_pixel_values.shape[2]] * len(
+            batch.video_loader)
+        batch.height = [video_pixel_values.shape[3]] * len(batch.video_loader)
+        batch.width = [video_pixel_values.shape[4]] * len(batch.video_loader)
+        return cast(ForwardBatch, batch)
+
+
+class TextTransformStage(PipelineStage):
+    """
+    Process text data according to the cfg rate.
+    """
+
+    def __init__(self, cfg_uncondition_drop_rate: float, seed: int) -> None:
+        self.cfg_rate = cfg_uncondition_drop_rate
+        self.rng = random.Random(seed)
+
+    def forward(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        batch = cast(PreprocessBatch, batch)
+
+        prompts = []
+        for prompt in batch.prompt:
+            if not isinstance(prompt, list):
+                prompt = [prompt]
+            prompt = self.rng.choice(prompt)
+            prompt = prompt if self.rng.random() > self.cfg_rate else ""
+            prompts.append(prompt)
+
+        batch.prompt = prompts
        return cast(ForwardBatch, batch)
@@ -1,21 +1,17 @@
-import os
-
 from fastvideo.distributed import (
    maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.utils import FlexibleArgumentParser
-from fastvideo.workflow.preprocess.preprocess_workflow_t2v import (
-    PreprocessWorkflowT2V)
+from fastvideo.workflow.workflow_base import WorkflowBase

 logger = init_logger(__name__)


 def main(fastvideo_args: FastVideoArgs) -> None:
    maybe_init_distributed_environment_and_model_parallel(1, 1)
-    num_gpus = int(os.environ["WORLD_SIZE"])
-    assert num_gpus == 1, "Only support 1 GPU"
-    preprocess_workflow = PreprocessWorkflowT2V(fastvideo_args)
+    preprocess_workflow_cls = WorkflowBase.get_workflow_cls(fastvideo_args)
+    preprocess_workflow = preprocess_workflow_cls(fastvideo_args)
    preprocess_workflow.run()


@@ -1,36 +1,66 @@
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.pipelines.preprocess.preprocess_stages import VideoTransformStage
+from fastvideo.pipelines.preprocess.preprocess_stages import (
+    TextTransformStage, VideoTransformStage)
 from fastvideo.pipelines.stages import (EncodingStage, ImageEncodingStage,
                                        TextEncodingStage)
+from fastvideo.pipelines.stages.image_encoding import ImageVAEEncodingStage


-class I2VPreprocessPipeline(ComposedPipelineBase):
+class PreprocessPipelineI2V(ComposedPipelineBase):
    _required_config_modules = [
        "image_encoder", "image_processor", "text_encoder", "tokenizer", "vae"
    ]

    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        assert fastvideo_args.preprocess_config is not None
+        self.add_stage(stage_name="text_transform_stage",
+                       stage=TextTransformStage(
+                           cfg_uncondition_drop_rate=fastvideo_args.
+                           preprocess_config.training_cfg_rate,
+                           seed=fastvideo_args.preprocess_config.seed,
+                       ))
        self.add_stage(stage_name="prompt_encoding_stage",
                       stage=TextEncodingStage(
                           text_encoders=[self.get_module("text_encoder")],
                           tokenizers=[self.get_module("tokenizer")],
                       ))
-        self.add_stage(stage_name="image_encoding_stage",
-                       stage=ImageEncodingStage(
-                           image_encoder=self.get_module("image_encoder"),
-                           image_processor=self.get_module("image_processor"),
-                       ))
+        self.add_stage(
+            stage_name="video_transform_stage",
+            stage=VideoTransformStage(
+                train_fps=fastvideo_args.preprocess_config.train_fps,
+                num_frames=fastvideo_args.preprocess_config.num_frames,
+                max_height=fastvideo_args.preprocess_config.max_height,
+                max_width=fastvideo_args.preprocess_config.max_width,
+                do_temporal_sample=fastvideo_args.preprocess_config.
+                do_temporal_sample,
+            ))
+        if (self.get_module("image_encoder") is not None
+                and self.get_module("image_processor") is not None):
+            self.add_stage(
+                stage_name="image_encoding_stage",
+                stage=ImageEncodingStage(
+                    image_encoder=self.get_module("image_encoder"),
+                    image_processor=self.get_module("image_processor"),
+                ))
+        self.add_stage(stage_name="image_vae_encoding_stage",
+                       stage=ImageVAEEncodingStage(
+                           vae=self.get_module("vae"), ))
        self.add_stage(stage_name="video_encoding_stage",
                       stage=EncodingStage(vae=self.get_module("vae"), ))


-class T2VPreprocessPipeline(ComposedPipelineBase):
+class PreprocessPipelineT2V(ComposedPipelineBase):
    _required_config_modules = ["text_encoder", "tokenizer", "vae"]

    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
        assert fastvideo_args.preprocess_config is not None
-
+        self.add_stage(stage_name="text_transform_stage",
+                       stage=TextTransformStage(
+                           cfg_uncondition_drop_rate=fastvideo_args.
+                           preprocess_config.training_cfg_rate,
+                           seed=fastvideo_args.preprocess_config.seed,
+                       ))
        self.add_stage(stage_name="prompt_encoding_stage",
                       stage=TextEncodingStage(
                           text_encoders=[self.get_module("text_encoder")],
@@ -50,4 +80,4 @@ class T2VPreprocessPipeline(ComposedPipelineBase):
                       stage=EncodingStage(vae=self.get_module("vae"), ))


-EntryClass = [I2VPreprocessPipeline, T2VPreprocessPipeline]
+EntryClass = [PreprocessPipelineI2V, PreprocessPipelineT2V]
@@ -7,6 +7,7 @@ complete diffusion pipelines.
 """

 from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.causal_denoising import CausalDMDDenosingStage
 from fastvideo.pipelines.stages.conditioning import ConditioningStage
 from fastvideo.pipelines.stages.decoding import DecodingStage
 from fastvideo.pipelines.stages.denoising import (DenoisingStage,
@@ -30,6 +31,7 @@ __all__ = [
    "ConditioningStage",
    "DenoisingStage",
    "DmdDenoisingStage",
+    "CausalDMDDenosingStage",
    "EncodingStage",
    "DecodingStage",
    "ImageEncodingStage",
@@ -155,6 +155,8 @@ class PipelineStage(ABC):
                execution_time = time.perf_counter() - start_time
                logger.info("[%s] Execution completed in %s ms", stage_name,
                            execution_time * 1000)
+                batch.logging_info.add_stage_execution_time(
+                    stage_name, execution_time)
            except Exception as e:
                execution_time = time.perf_counter() - start_time
                logger.error("[%s] Error during execution after %s ms: %s",
@@ -0,0 +1,414 @@
+import torch  # type: ignore
+
+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
+
+try:
+    from fastvideo.attention.backends.sliding_tile_attn import (
+        SlidingTileAttentionBackend)
+    st_attn_available = True
+except ImportError:
+    st_attn_available = False
+    SlidingTileAttentionBackend = None  # type: ignore
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (
+        VideoSparseAttentionBackend)
+    vsa_available = True
+except ImportError:
+    vsa_available = False
+    VideoSparseAttentionBackend = None  # type: ignore
+
+logger = init_logger(__name__)
+
+
+class CausalDMDDenosingStage(DenoisingStage):
+    """
+    Denoising stage for causal diffusion.
+    """
+
+    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
+        # Model-dependent constants (aligned with causal_inference.py assumptions)
+        self.num_transformer_blocks = self.transformer.config.arch_config.num_layers
+        self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block
+        self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames
+
+        try:
+            self.local_attn_size = getattr(self.transformer.model,
+                                           "local_attn_size",
+                                           -1)  # type: ignore
+        except Exception:
+            self.local_attn_size = -1
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
+        patch_ratio = self.transformer.config.arch_config.patch_size[
+            -1] * self.transformer.config.arch_config.patch_size[-2]
+        self.frame_seq_length = latent_seq_length // patch_ratio
+        # TODO(will): make this a parameter once we add i2v support
+        independent_first_frame = self.transformer.independent_first_frame
+
+        # Timesteps for DMD
+        timesteps = torch.tensor(
+            fastvideo_args.pipeline_config.dmd_denoising_steps,
+            dtype=torch.long,
+            device=get_local_torch_device())
+
+        # Image kwargs (kept empty unless caller provides compatible args)
+        image_kwargs: dict = {}
+
+        pos_cond_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                # "encoder_hidden_states_2": batch.clip_embedding_pos,
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        # STA
+        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
+            self.prepare_sta_param(batch, fastvideo_args)
+
+        # Latents and prompts
+        assert batch.latents is not None, "latents must be provided"
+        latents = batch.latents  # [B, C, T, H, W]
+        b, c, t, h, w = latents.shape
+        prompt_embeds = batch.prompt_embeds
+        assert torch.isnan(prompt_embeds[0]).sum() == 0
+
+        # Initialize or reset caches
+        if self.kv_cache1 is None:
+            self._initialize_kv_cache(batch_size=latents.shape[0],
+                                      dtype=target_dtype,
+                                      device=latents.device)
+            self._initialize_crossattn_cache(
+                batch_size=latents.shape[0],
+                max_text_len=fastvideo_args.pipeline_config.
+                text_encoder_configs[0].arch_config.text_len,
+                dtype=target_dtype,
+                device=latents.device)
+        else:
+            assert self.crossattn_cache is not None
+            # reset cross-attention cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache[block_index][
+                    "is_init"] = False  # type: ignore
+            # reset kv cache pointers
+            for block_index in range(len(self.kv_cache1)):
+                self.kv_cache1[block_index][
+                    "global_end_index"] = torch.tensor(  # type: ignore
+                        [0],
+                        dtype=torch.long,
+                        device=latents.device)
+                self.kv_cache1[block_index][
+                    "local_end_index"] = torch.tensor(  # type: ignore
+                        [0],
+                        dtype=torch.long,
+                        device=latents.device)
+
+        # Optional: cache context features from provided image latents prior to generation
+        current_start_frame = 0
+        if getattr(batch, "image_latent", None) is not None:
+            image_latent = batch.image_latent
+            assert image_latent is not None
+            input_frames = image_latent.shape[2]
+            # timestep zero (or configured context noise) for cache warm-up
+            t_zero = torch.zeros([latents.shape[0]],
+                                 device=latents.device,
+                                 dtype=torch.long)
+            if independent_first_frame and input_frames >= 1:
+                # warm-up with the very first frame independently
+                image_first_btchw = image_latent[:, :, :1, :, :].to(
+                    target_dtype).permute(0, 2, 1, 3, 4)
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+                    _ = self.transformer(
+                        image_first_btchw,
+                        prompt_embeds,
+                        t_zero,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame *
+                        self.frame_seq_length,
+                        **image_kwargs,
+                        **pos_cond_kwargs,
+                    )
+                current_start_frame += 1
+                remaining_frames = input_frames - 1
+            else:
+                remaining_frames = input_frames
+
+            # process remaining input frames in blocks of num_frame_per_block
+            while remaining_frames > 0:
+                block = min(self.num_frames_per_block, remaining_frames)
+                ref_btchw = image_latent[:, :, current_start_frame:
+                                         current_start_frame +
+                                         block, :, :].to(target_dtype).permute(
+                                             0, 2, 1, 3, 4)
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+                    _ = self.transformer(
+                        ref_btchw,
+                        prompt_embeds,
+                        t_zero,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame *
+                        self.frame_seq_length,
+                        **image_kwargs,
+                        **pos_cond_kwargs,
+                    )
+                current_start_frame += block
+                remaining_frames -= block
+
+        # Base position offset from any cache warm-up
+        pos_start_base = current_start_frame
+
+        # Determine block sizes
+        if not independent_first_frame or (independent_first_frame
+                                           and batch.image_latent is not None):
+            if t % self.num_frames_per_block != 0:
+                raise ValueError(
+                    "num_frames must be divisible by num_frames_per_block for causal DMD denoising"
+                )
+            num_blocks = t // self.num_frames_per_block
+            block_sizes = [self.num_frames_per_block] * num_blocks
+            start_index = 0
+        else:
+            if (t - 1) % self.num_frames_per_block != 0:
+                raise ValueError(
+                    "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True"
+                )
+            num_blocks = (t - 1) // self.num_frames_per_block
+            block_sizes = [1] + [self.num_frames_per_block] * num_blocks
+            start_index = 0
+
+        # DMD loop in causal blocks
+        with self.progress_bar(total=len(block_sizes) *
+                               len(timesteps)) as progress_bar:
+            for current_num_frames in block_sizes:
+                current_latents = latents[:, :, start_index:start_index +
+                                          current_num_frames, :, :]
+                # use BTCHW for DMD conversion routines
+                noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)
+                video_raw_latent_shape = noise_latents_btchw.shape
+
+                for i, t_cur in enumerate(timesteps):
+                    # Copy for pred conversion
+                    noise_latents = noise_latents_btchw.clone()
+                    latent_model_input = current_latents.to(target_dtype)
+
+                    if batch.image_latent is not None and independent_first_frame and start_index == 0:
+                        latent_model_input = torch.cat([
+                            latent_model_input,
+                            batch.image_latent.to(target_dtype)
+                        ],
+                                                       dim=2)
+
+                    # Prepare inputs
+                    t_expand = t_cur.expand(latent_model_input.shape[0])
+                    # t_expand = t_cur * torch.ones((latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long)
+                    # t_expand = t_expand.repeat(1, self.sliding_window_num_frames)
+
+
+                    # Attention metadata if needed
+                    if (vsa_available and self.attn_backend
+                            == VideoSparseAttentionBackend):
+                        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(
+                            )
+                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
+                                current_timestep=i,  # type: ignore
+                                raw_latent_shape=(current_num_frames, h,
+                                                  w),  # type: ignore
+                                patch_size=fastvideo_args.pipeline_config.
+                                dit_config.patch_size,  # type: ignore
+                                STA_param=batch.STA_param,  # type: ignore
+                                VSA_sparsity=fastvideo_args.
+                                VSA_sparsity,  # type: ignore
+                                device=get_local_torch_device(),  # type: ignore
+                            )  # type: ignore
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    else:
+                        attn_metadata = None
+
+                    with torch.autocast(device_type="cuda",
+                                        dtype=target_dtype,
+                                        enabled=autocast_enabled), \
+                        set_forward_context(current_timestep=i,
+                                            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,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=(pos_start_base + start_index) *
+                            self.frame_seq_length,
+                            start_frame=start_index,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                        ).permute(0, 2, 1, 3, 4)
+
+                    # Convert pred noise to pred video with FM Euler scheduler utilities
+                    pred_video_btchw = pred_noise_to_pred_video(
+                        pred_noise=pred_noise_btchw.flatten(0, 1),
+                        noise_input_latent=noise_latents.flatten(0, 1),
+                        timestep=t_expand,
+                        scheduler=self.scheduler).unflatten(
+                            0, pred_noise_btchw.shape[:2])
+
+                    if i < len(timesteps) - 1:
+                        next_timestep = timesteps[i + 1] * torch.ones(
+                            [1],
+                            dtype=torch.long,
+                            device=pred_video_btchw.device)
+                        noise = torch.randn(
+                            video_raw_latent_shape,
+                            dtype=pred_video_btchw.dtype,
+                            generator=(batch.generator[0] if isinstance(
+                                batch.generator, list) else
+                                       batch.generator)).to(self.device)
+                        noise_btchw = noise
+                        noise_latents_btchw = self.scheduler.add_noise(
+                            pred_video_btchw.flatten(0, 1),
+                            noise_btchw.flatten(0, 1),
+                            next_timestep).unflatten(0,
+                                                     pred_video_btchw.shape[:2])
+                        current_latents = noise_latents_btchw.permute(
+                            0, 2, 1, 3, 4)
+                    else:
+                        current_latents = pred_video_btchw.permute(
+                            0, 2, 1, 3, 4)
+
+                    if progress_bar is not None:
+                        progress_bar.update()
+
+                # Write back and advance
+                latents[:, :, start_index:start_index +
+                        current_num_frames, :, :] = current_latents
+
+                # Re-run with context timestep to update KV cache using clean context
+                context_noise = getattr(fastvideo_args.pipeline_config,
+                                        "context_noise", 0)
+                t_context = torch.ones([latents.shape[0]],
+                                       device=latents.device,
+                                       dtype=torch.long) * int(context_noise)
+                context_bcthw = current_latents.to(target_dtype)
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled), \
+                    set_forward_context(current_timestep=0,
+                                        attn_metadata=attn_metadata,
+                                        forward_batch=batch):
+                    t_expanded_context = t_context * torch.ones((context_bcthw.shape[0], 1), device=context_bcthw.device, dtype=torch.long)
+                    _ = self.transformer(
+                        context_bcthw,
+                        prompt_embeds,
+                        t_expanded_context,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=(pos_start_base + start_index) *
+                        self.frame_seq_length,
+                        start_frame=start_index,
+                        **image_kwargs,
+                        **pos_cond_kwargs,
+                    )
+                start_index += current_num_frames
+
+        batch.latents = latents
+        return batch
+
+    def _initialize_kv_cache(self, batch_size, dtype, device) -> None:
+        """
+        Initialize a Per-GPU KV cache aligned with the Wan model assumptions.
+        """
+        kv_cache1 = []
+        num_attention_heads = self.transformer.num_attention_heads
+        attention_head_dim = self.transformer.attention_head_dim
+        if self.local_attn_size != -1:
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames
+
+        for _ in range(self.num_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),
+            })
+
+        self.kv_cache1 = kv_cache1
+
+    def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
+                                    device) -> None:
+        """
+        Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions.
+        """
+        crossattn_cache = []
+        num_attention_heads = self.transformer.num_attention_heads
+        attention_head_dim = self.transformer.attention_head_dim
+        for _ in range(self.num_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,
+            })
+        self.crossattn_cache = crossattn_cache
@@ -673,12 +673,8 @@ class DmdDenoisingStage(DenoisingStage):
        # Get latents and embeddings
        assert batch.latents is not None, "latents must be provided"
        latents = batch.latents
-        # TODO(yongqi) hard code prepare latents
-        latents = torch.randn(
-            latents.permute(0, 2, 1, 3, 4).shape,
-            dtype=torch.bfloat16,
-            device="cuda",
-            generator=torch.Generator(device="cuda").manual_seed(42))
+        latents = latents.permute(0, 2, 1, 3, 4)
+
        video_raw_latent_shape = latents.shape
        prompt_embeds = batch.prompt_embeds
        assert torch.isnan(prompt_embeds[0]).sum() == 0
@@ -782,8 +778,7 @@ class DmdDenoisingStage(DenoisingStage):
                            **pos_cond_kwargs,
                        ).permute(0, 2, 1, 3, 4)

-                    from fastvideo.training.training_utils import (
-                        pred_noise_to_pred_video)
+                    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),
@@ -795,8 +790,9 @@ class DmdDenoisingStage(DenoisingStage):
                        next_timestep = timesteps[i + 1] * torch.ones(
                            [1], dtype=torch.long, device=pred_video.device)
                        noise = torch.randn(video_raw_latent_shape,
-                                            device=self.device,
-                                            dtype=pred_video.dtype)
+                                            dtype=pred_video.dtype,
+                                            generator=batch.generator[0]).to(
+                                                self.device)
                        if sp_group:
                            noise = rearrange(noise,
                                              "b (n t) c h w -> b n t c h w",
@@ -9,7 +9,7 @@ import PIL
 import torch

 from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.fastvideo_args import ExecutionMode, FastVideoArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.vaes.common import ParallelTiledVAE
@@ -120,30 +120,47 @@ class ImageVAEEncodingStage(PipelineStage):
        Returns:
            The batch with encoded outputs.
        """
-        assert batch.height is not None and isinstance(batch.height, int)
-        assert batch.width is not None and isinstance(batch.width, int)
-        assert batch.num_frames is not None and isinstance(
-            batch.num_frames, int)
+        assert batch.pil_image is not None
+        if fastvideo_args.mode == ExecutionMode.INFERENCE:
+            assert batch.pil_image is not None and isinstance(
+                batch.pil_image, PIL.Image.Image)
+            assert batch.height is not None and isinstance(batch.height, int)
+            assert batch.width is not None and isinstance(batch.width, int)
+            assert batch.num_frames is not None and isinstance(
+                batch.num_frames, int)
+            height = batch.height
+            width = batch.width
+            num_frames = batch.num_frames
+        elif fastvideo_args.mode == ExecutionMode.PREPROCESS:
+            assert batch.pil_image is not None and isinstance(
+                batch.pil_image, torch.Tensor)
+            assert batch.height is not None and isinstance(batch.height, list)
+            assert batch.width is not None and isinstance(batch.width, list)
+            assert batch.num_frames is not None and isinstance(
+                batch.num_frames, list)
+            num_frames = batch.num_frames[0]
+            height = batch.height[0]
+            width = batch.width[0]

        self.vae = self.vae.to(get_local_torch_device())

-        latent_height = batch.height // self.vae.spatial_compression_ratio
-        latent_width = batch.width // self.vae.spatial_compression_ratio
+        latent_height = height // self.vae.spatial_compression_ratio
+        latent_width = width // self.vae.spatial_compression_ratio

-        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_local_torch_device(), dtype=torch.float32)
+            height=height,
+            width=width).to(get_local_torch_device(), dtype=torch.float32)

+        # (B, C, H, W) -> (B, C, 1, H, W)
        image = image.unsqueeze(2)

        video_condition = torch.cat([
            image,
-            image.new_zeros(image.shape[0], image.shape[1],
-                            batch.num_frames - 1, batch.height, batch.width)
+            image.new_zeros(image.shape[0], image.shape[1], num_frames - 1,
+                            image.shape[3], image.shape[4])
        ],
                                    dim=2)
        video_condition = video_condition.to(device=get_local_torch_device(),
@@ -167,10 +184,13 @@ class ImageVAEEncodingStage(PipelineStage):
                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)
+        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
+            latent_condition = encoder_output.mean
+        else:
+            generator = batch.generator
+            if generator is None:
+                raise ValueError("Generator must be provided")
+            latent_condition = self.retrieve_latents(encoder_output, generator)

        # Apply shifting if needed
        if (hasattr(self.vae, "shift_factor")
@@ -187,24 +207,27 @@ class ImageVAEEncodingStage(PipelineStage):
        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)
+        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
+            batch.image_latent = latent_condition
+        else:
+            mask_lat_size = torch.ones(1, 1, num_frames, latent_height,
+                                       latent_width)
+            mask_lat_size[:, :, list(range(1, num_frames))] = 0
+            first_frame_mask = mask_lat_size[:, :, 0:1]
+            first_frame_mask = torch.repeat_interleave(
+                first_frame_mask,
+                dim=2,
+                repeats=self.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)
+            batch.image_latent = torch.concat([mask_lat_size, latent_condition],
+                                              dim=1)

        # Offload models if needed
        if hasattr(self, 'maybe_free_model_hooks'):
@@ -228,21 +251,19 @@ class ImageVAEEncodingStage(PipelineStage):

    def preprocess(
            self,
-            image: PIL.Image.Image,
+            image: torch.Tensor | PIL.Image.Image,
            vae_scale_factor: int,
            height: int | None = None,
            width: int | None = 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
+        if isinstance(image, PIL.Image.Image):
+            height, width = get_default_height_width(image, vae_scale_factor,
+                                                     height, width)
+            image = resize(image, height, width, resize_mode=resize_mode)
+            image = pil_to_numpy(image)  # to np
+            image = numpy_to_pt(image)  # to pt

        do_normalize = True
        if image.min() < 0:
@@ -256,12 +277,16 @@ class ImageVAEEncodingStage(PipelineStage):
                     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)
+        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
+            result.add_check("height", batch.height, V.list_not_empty)
+            result.add_check("width", batch.width, V.list_not_empty)
+            result.add_check("num_frames", batch.num_frames, V.list_not_empty)
+        else:
+            result.add_check("height", batch.height, V.positive_int)
+            result.add_check("width", batch.width, V.positive_int)
+            result.add_check("num_frames", batch.num_frames, V.positive_int)
        return result

    def verify_output(self, batch: ForwardBatch,
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import copy
 import gc
+import json
 import os
 import time
 from abc import abstractmethod
@@ -29,16 +30,18 @@ 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 import (ComposedPipelineBase, ForwardBatch,
                                 TrainingBatch)
 from fastvideo.training.activation_checkpoint import (
    apply_activation_checkpointing)
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
-    clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler,
-    load_distillation_checkpoint, pred_noise_to_pred_video,
-    save_distillation_checkpoint, shift_timestep)
-from fastvideo.utils import is_vsa_available, set_random_seed
+    EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
+    get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint,
+    shift_timestep)
+from fastvideo.utils import (is_vsa_available, maybe_download_model,
+                             set_random_seed, verify_model_config_and_directory)

 import wandb  # isort: skip

@@ -87,9 +90,27 @@ class DistillationPipeline(TrainingPipeline):
        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
            shift=self.timestep_shift)

-        # self.transformer is the generator model
-        self.real_score_transformer = self.get_module("real_score_transformer")
-        self.fake_score_transformer = self.get_module("fake_score_transformer")
+        if training_args.real_score_model_path:
+            logger.info(
+                f"Loading real score transformer from: {training_args.real_score_model_path}"
+            )
+            self.real_score_transformer = self.load_module_from_path(
+                training_args.real_score_model_path, "transformer",
+                training_args)
+        else:
+            self.real_score_transformer = self.get_module(
+                "real_score_transformer")
+
+        if training_args.fake_score_model_path:
+            logger.info(
+                f"Loading fake score transformer from: {training_args.fake_score_model_path}"
+            )
+            self.fake_score_transformer = self.load_module_from_path(
+                training_args.fake_score_model_path, "transformer",
+                training_args)
+        else:
+            self.fake_score_transformer = self.get_module(
+                "fake_score_transformer")

        self.real_score_transformer.requires_grad_(False)
        self.real_score_transformer.eval()
@@ -116,10 +137,13 @@ class DistillationPipeline(TrainingPipeline):
        if fake_score_lr == 0.0:
            fake_score_lr = training_args.learning_rate

+        betas_str = training_args.fake_score_betas
+        betas = tuple(float(x.strip()) for x in betas_str.split(","))
+
        self.fake_score_optimizer = torch.optim.AdamW(
            fake_score_params,
            lr=fake_score_lr,
-            betas=(0.9, 0.999),
+            betas=betas,
            weight_decay=training_args.weight_decay,
            eps=1e-8,
        )
@@ -147,8 +171,19 @@ class DistillationPipeline(TrainingPipeline):
            self.training_args.pipeline_config.dmd_denoising_steps,
            dtype=torch.long,
            device=get_local_torch_device())
-        logger.info("Distillation generator model to %s denoising steps",
-                    len(self.denoising_step_list))
+
+        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
+            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
+                                   torch.tensor([0],
+                                                dtype=torch.float32))).cuda()
+            self.denoising_step_list = timesteps[1000 -
+                                                 self.denoising_step_list]
+            logger.info("Warping denoising_step_list")
+
+        self.denoising_step_list = self.denoising_step_list.to(
+            get_local_torch_device())
+        logger.info("Distillation generator model to %s denoising steps: %s",
+                    len(self.denoising_step_list), self.denoising_step_list)
        self.num_train_timestep = self.noise_scheduler.num_train_timesteps

        self.min_timestep = int(self.training_args.min_timestep_ratio *
@@ -158,6 +193,82 @@ class DistillationPipeline(TrainingPipeline):

        self.real_score_guidance_scale = self.training_args.real_score_guidance_scale

+        self.generator_ema = None
+        if (self.training_args.ema_decay
+                is not None) and (self.training_args.ema_decay > 0.0):
+            self.generator_ema = EMA_FSDP(self.transformer,
+                                          decay=self.training_args.ema_decay)
+            logger.info(
+                f"Initialized generator EMA with decay={self.training_args.ema_decay}"
+            )
+        else:
+            logger.info("Generator EMA disabled (ema_decay <= 0.0)")
+
+    def load_module_from_path(self, model_path: str, module_type: str,
+                              training_args: "TrainingArgs"):
+        """
+        Load a module from a specific path using the same loading logic as the pipeline.
+        
+        Args:
+            model_path: Path to the model
+            module_type: Type of module to load (e.g., "transformer")
+            training_args: Training arguments
+            
+        Returns:
+            The loaded module
+        """
+        logger.info(f"Loading {module_type} from custom path: {model_path}")
+        # Set flag to prevent custom weight loading for teacher/critic models
+        training_args._loading_teacher_critic_model = True
+
+        try:
+            from fastvideo.models.loader.component_loader import (
+                PipelineComponentLoader)
+
+            # Download the model if it's a Hugging Face model ID
+            local_model_path = maybe_download_model(model_path)
+            logger.info(f"Model downloaded/found at: {local_model_path}")
+            config = verify_model_config_and_directory(local_model_path)
+
+            if module_type not in config:
+                if hasattr(self, '_extra_config_module_map'
+                           ) and module_type in self._extra_config_module_map:
+                    extra_module = self._extra_config_module_map[module_type]
+                    if extra_module in config:
+                        module_type = extra_module
+                        logger.info(f"Using {extra_module} for {module_type}")
+                    else:
+                        raise ValueError(
+                            f"Module {module_type} not found in config at {local_model_path}"
+                        )
+                else:
+                    raise ValueError(
+                        f"Module {module_type} not found in config at {local_model_path}"
+                    )
+
+            module_info = config[module_type]
+            if module_info is None:
+                raise ValueError(
+                    f"Module {module_type} has null value in config at {local_model_path}"
+                )
+
+            transformers_or_diffusers, architecture = module_info
+            component_path = os.path.join(local_model_path, module_type)
+            module = PipelineComponentLoader.load_module(
+                module_name=module_type,
+                component_model_path=component_path,
+                transformers_or_diffusers=transformers_or_diffusers,
+                fastvideo_args=training_args,
+            )
+
+            logger.info(
+                f"Successfully loaded {module_type} from {component_path}")
+            return module
+        finally:
+            # Always clean up the flag
+            if hasattr(training_args, '_loading_teacher_critic_model'):
+                delattr(training_args, '_loading_teacher_critic_model')
+
    @abstractmethod
    def initialize_validation_pipeline(self, training_args: TrainingArgs):
        """Initialize validation pipeline - must be implemented by subclasses."""
@@ -174,6 +285,110 @@ class DistillationPipeline(TrainingPipeline):

        return training_batch

+    def apply_ema_to_model(self, model):
+        """Apply EMA weights to the model for validation or inference."""
+        if self.generator_ema is not None:
+            with self.generator_ema.apply_to_model(model):
+                return model
+        return model
+
+    def get_ema_model_copy(self):
+        """Get a copy of the model with EMA weights applied."""
+        if self.generator_ema is not None:
+            ema_model = copy.deepcopy(self.transformer)
+            self.generator_ema.copy_to_unwrapped(ema_model)
+            return ema_model
+        return None
+
+    def is_ema_ready(self, current_step: int = None):
+        """Check if EMA is ready for use (after ema_start_step)."""
+        if current_step is None:
+            current_step = getattr(self, 'current_trainstep', 0)
+        return (self.generator_ema is not None
+                and current_step >= self.training_args.ema_start_step)
+
+    def save_ema_weights(self, output_dir: str, step: int):
+        """Save EMA weights separately for inference purposes."""
+        if self.generator_ema is None:
+            logger.warning("Cannot save EMA weights: EMA not initialized")
+            return
+
+        if not self.is_ema_ready():
+            logger.warning(
+                "Cannot save EMA weights: EMA not ready yet (step < ema_start_step)"
+            )
+            return
+
+        try:
+            ema_model = self.get_ema_model_copy()
+            if ema_model is None:
+                logger.warning("Failed to create EMA model copy")
+                return
+
+            ema_save_dir = os.path.join(output_dir, f"ema_checkpoint-{step}")
+            os.makedirs(ema_save_dir, exist_ok=True)
+
+            # save as diffusers format
+            from safetensors.torch import save_file
+
+            from fastvideo.training.training_utils import (
+                custom_to_hf_state_dict, gather_state_dict_on_cpu_rank0)
+            cpu_state = gather_state_dict_on_cpu_rank0(ema_model, device=None)
+
+            if self.global_rank == 0:
+                weight_path = os.path.join(
+                    ema_save_dir, "diffusion_pytorch_model.safetensors")
+                diffusers_state_dict = custom_to_hf_state_dict(
+                    cpu_state, ema_model.reverse_param_names_mapping)
+                save_file(diffusers_state_dict, weight_path)
+
+                config_dict = ema_model.hf_config
+                if "dtype" in config_dict:
+                    del config_dict["dtype"]
+                config_path = os.path.join(ema_save_dir, "config.json")
+                with open(config_path, "w") as f:
+                    json.dump(config_dict, f, indent=4)
+
+                logger.info(f"EMA weights saved to {weight_path}")
+
+            del ema_model
+
+        except Exception as e:
+            logger.error(f"Failed to save EMA weights: {str(e)}")
+
+    def get_ema_stats(self):
+        """Get EMA statistics for monitoring."""
+        if self.generator_ema is None:
+            return {
+                "ema_enabled": False,
+                "ema_decay": None,
+                "ema_start_step": self.training_args.ema_start_step,
+                "ema_ready": False,
+                "ema_step": self.current_trainstep,
+            }
+
+        return {
+            "ema_enabled": True,
+            "ema_decay": self.training_args.ema_decay,
+            "ema_start_step": self.training_args.ema_start_step,
+            "ema_ready": self.is_ema_ready(),
+            "ema_step": self.current_trainstep,
+        }
+
+    def reset_ema(self):
+        """Reset EMA to current model weights."""
+        if self.generator_ema is not None:
+            logger.info("Resetting EMA to current model weights")
+            self.generator_ema.update(self.transformer)
+            # Force update to current weights by setting decay to 0 temporarily
+            original_decay = self.generator_ema.decay
+            self.generator_ema.decay = 0.0
+            self.generator_ema.update(self.transformer)
+            self.generator_ema.decay = original_decay
+            logger.info("EMA reset completed")
+        else:
+            logger.warning("Cannot reset EMA: EMA not initialized")
+
    def _build_distill_input_kwargs(
            self, noise_input: torch.Tensor, timestep: torch.Tensor,
            text_dict: dict[str, torch.Tensor] | None,
@@ -241,6 +456,7 @@ class DistillationPipeline(TrainingPipeline):
                                            len(self.denoising_step_list), [1],
                                            device=self.device,
                                            dtype=torch.long)
+        target_timestep_idx_int = target_timestep_idx.item()
        target_timestep = self.denoising_step_list[target_timestep_idx]

        # Step 2: Simulate the multi-step inference process up to the target timestep
@@ -255,9 +471,12 @@ class DistillationPipeline(TrainingPipeline):
                n=self.sp_world_size).contiguous()
            current_noise_latents = current_noise_latents[:, self.
                                                          rank_in_sp_group, :, :, :, :]
+        current_noise_latents_copy = current_noise_latents.clone()

        # Only run intermediate steps if target_timestep_idx > 0
        max_target_idx = len(self.denoising_step_list) - 1
+        noise_latents = []
+        noise_latent_index = target_timestep_idx_int - 1
        if max_target_idx > 0:
            # Run student model for all steps before the target timestep
            with torch.no_grad():
@@ -294,11 +513,19 @@ class DistillationPipeline(TrainingPipeline):
                    current_noise_latents = self.noise_scheduler.add_noise(
                        pred_clean.flatten(0, 1), noise.flatten(0, 1),
                        next_timestep_tensor).unflatten(0, pred_clean.shape[:2])
+                    latent_copy = current_noise_latents.clone()
+                    noise_latents.append(latent_copy)

        # Step 3: Use the simulated noisy input for the final training step
        # For timestep index 0, this is pure noise
        # For timestep index k > 0, this is the result after k denoising steps + noise at target level
-        noisy_input = current_noise_latents
+        if noise_latent_index >= 0:
+            assert noise_latent_index < len(
+                self.denoising_step_list
+            ) - 1, "noise_latent_index is out of bounds"
+            noisy_input = noise_latents[noise_latent_index]
+        else:
+            noisy_input = current_noise_latents_copy

        # Step 4: Final student prediction (this is what we train on)
        training_batch = self._build_distill_input_kwargs(
@@ -575,6 +802,10 @@ class DistillationPipeline(TrainingPipeline):
            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:
+                self.generator_ema.update(self.transformer)
+
            avg_dmd_loss = torch.tensor(total_dmd_loss /
                                        gradient_accumulation_steps,
                                        device=self.device)
@@ -625,7 +856,8 @@ class DistillationPipeline(TrainingPipeline):
            self.transformer, self.fake_score_transformer, self.global_rank,
            self.training_args.resume_from_checkpoint, self.optimizer,
            self.fake_score_optimizer, self.train_dataloader, self.lr_scheduler,
-            self.fake_score_lr_scheduler, self.noise_random_generator)
+            self.fake_score_lr_scheduler, self.noise_random_generator,
+            self.generator_ema)

        if resumed_step > 0:
            self.init_steps = resumed_step
@@ -656,6 +888,14 @@ class DistillationPipeline(TrainingPipeline):
            sum(p.numel()
                for p in self.fake_score_transformer.parameters()) / 1e9)

+        if self.generator_ema is not None:
+            logger.info("  Generator EMA enabled with decay: %s",
+                        self.training_args.ema_decay)
+            logger.info("  Generator EMA start step: %s",
+                        self.training_args.ema_start_step)
+        else:
+            logger.info("  Generator EMA disabled")
+
    @torch.no_grad()
    def _log_validation(self, transformer, training_args, global_step) -> None:
        training_args.inference_mode = True
@@ -687,6 +927,18 @@ class DistillationPipeline(TrainingPipeline):

        transformer.eval()

+        # Optionally use EMA model for validation if available and ready
+        use_ema_for_validation = (self.training_args.use_ema
+                                  and self.is_ema_ready(global_step))
+        if use_ema_for_validation:
+            logger.info("Using EMA model for validation")
+            validation_transformer = self.transformer
+            ema_context = self.generator_ema.apply_to_model(
+                validation_transformer)
+        else:
+            validation_transformer = transformer
+            ema_context = None
+
        validation_steps = training_args.validation_sampling_steps.split(",")
        validation_steps = [int(step) for step in validation_steps]
        validation_steps = [step for step in validation_steps if step > 0]
@@ -702,50 +954,98 @@ class DistillationPipeline(TrainingPipeline):
            step_videos: list[np.ndarray] = []
            step_captions: list[str] = []

-            for validation_batch in validation_dataloader:
-                batch = self._prepare_validation_batch(sampling_param,
-                                                       training_args,
-                                                       validation_batch,
-                                                       num_inference_steps)
+            if ema_context is not None:
+                with ema_context:
+                    for validation_batch in validation_dataloader:
+                        batch = self._prepare_validation_batch(
+                            sampling_param, training_args, validation_batch,
+                            num_inference_steps)

-                negative_prompt = batch.negative_prompt
-                batch_negative = ForwardBatch(
-                    data_type="video",
-                    prompt=negative_prompt,
-                    prompt_embeds=[],
-                    prompt_attention_mask=[],
-                )
-                result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
-                    batch_negative, training_args)
-                self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
-                    0], result_batch.prompt_attention_mask[0]
+                        negative_prompt = batch.negative_prompt
+                        batch_negative = ForwardBatch(
+                            data_type="video",
+                            prompt=negative_prompt,
+                            prompt_embeds=[],
+                            prompt_attention_mask=[],
+                        )
+                        result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                            batch_negative, training_args)
+                        self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                            0], result_batch.prompt_attention_mask[0]

-                logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                        logger.info(
+                            "rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
                            self.global_rank,
                            self.rank_in_sp_group,
                            batch.prompt,
                            local_main_process_only=False)

-                assert batch.prompt is not None and isinstance(
-                    batch.prompt, str)
-                step_captions.append(batch.prompt)
+                        assert batch.prompt is not None and isinstance(
+                            batch.prompt, str)
+                        step_captions.append(batch.prompt)

-                # Run validation inference
-                with torch.no_grad():
-                    output_batch = self.validation_pipeline.forward(
-                        batch, training_args)
-                samples = output_batch.output
-                if self.rank_in_sp_group != 0:
-                    continue
+                        # Run validation inference
+                        with torch.no_grad():
+                            output_batch = self.validation_pipeline.forward(
+                                batch, training_args)
+                        samples = output_batch.output
+                        if self.rank_in_sp_group != 0:
+                            continue

-                # Process outputs
-                video = rearrange(samples, "b c t h w -> t b c h w")
-                frames = []
-                for x in video:
-                    x = torchvision.utils.make_grid(x, nrow=6)
-                    x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-                    frames.append((x * 255).numpy().astype(np.uint8))
-                step_videos.append(frames)
+                        # Process outputs
+                        video = rearrange(samples, "b c t h w -> t b c h w")
+                        frames = []
+                        for x in video:
+                            x = torchvision.utils.make_grid(x, nrow=6)
+                            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                            frames.append((x * 255).numpy().astype(np.uint8))
+                        step_videos.append(frames)
+            else:
+                # Use original transformer without EMA
+                for validation_batch in validation_dataloader:
+                    batch = self._prepare_validation_batch(
+                        sampling_param, training_args, validation_batch,
+                        num_inference_steps)
+
+                    negative_prompt = batch.negative_prompt
+                    batch_negative = ForwardBatch(
+                        data_type="video",
+                        prompt=negative_prompt,
+                        prompt_embeds=[],
+                        prompt_attention_mask=[],
+                    )
+                    result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                        batch_negative, training_args)
+                    self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                        0], result_batch.prompt_attention_mask[0]
+
+                    logger.info(
+                        "rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                        self.global_rank,
+                        self.rank_in_sp_group,
+                        batch.prompt,
+                        local_main_process_only=False)
+
+                    assert batch.prompt is not None and isinstance(
+                        batch.prompt, str)
+                    step_captions.append(batch.prompt)
+
+                    # Run validation inference
+                    with torch.no_grad():
+                        output_batch = self.validation_pipeline.forward(
+                            batch, training_args)
+                    samples = output_batch.output
+                    if self.rank_in_sp_group != 0:
+                        continue
+
+                    # Process outputs
+                    video = rearrange(samples, "b c t h w -> t b c h w")
+                    frames = []
+                    for x in video:
+                        x = torchvision.utils.make_grid(x, nrow=6)
+                        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                        frames.append((x * 255).numpy().astype(np.uint8))
+                    step_videos.append(frames)

            # Log validation results for this step
            world_group = get_world_group()
@@ -802,7 +1102,7 @@ class DistillationPipeline(TrainingPipeline):
        dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
        fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
        fake_score_log_keys = ['generator_pred_video']
-        dmd_log_keys = ['faker_score_pred_video']
+        dmd_log_keys = ['faker_score_pred_video', 'real_score_pred_video']

        for latent_key in fake_score_log_keys:
            latents = fake_score_latents_vis_dict[latent_key]
@@ -822,16 +1122,16 @@ class DistillationPipeline(TrainingPipeline):
                                                        latents.dtype)
                else:
                    latents += self.vae.shift_factor
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                video = self.vae.decode(latents)
-            video = (video / 2 + 0.5).clamp(0, 1)
-            video = video.cpu().float()
-            video = video.permute(0, 2, 1, 3, 4)
-            video = (video * 255).numpy().astype(np.uint8)
-            wandb_loss_dict[latent_key] = wandb.Video(
-                video, fps=24, format="mp4")  # change to 16 for Wan2.1
-            # Clean up references
-            del video, latents
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    video = self.vae.decode(latents)
+                video = (video / 2 + 0.5).clamp(0, 1)
+                video = video.cpu().float()
+                video = video.permute(0, 2, 1, 3, 4)
+                video = (video * 255).numpy().astype(np.uint8)
+                wandb_loss_dict[latent_key] = wandb.Video(
+                    video, fps=24, format="mp4")  # change to 16 for Wan2.1
+                # Clean up references
+                del video, latents

        # Process DMD training data if available - use decode_stage instead of self.vae.decode
        if 'generator_pred_video' in dmd_latents_vis_dict:
@@ -892,6 +1192,10 @@ class DistillationPipeline(TrainingPipeline):
            device="cpu").manual_seed(self.seed)
        logger.info("Initialized random seeds with seed: %s", seed)

+        # Initialize current_trainstep for EMA ready checks
+        #TODO: check if needed
+        self.current_trainstep = self.init_steps
+
        # Resume from checkpoint if specified (this will restore random states)
        if self.training_args.resume_from_checkpoint:
            self._resume_from_checkpoint()
@@ -935,6 +1239,14 @@ class DistillationPipeline(TrainingPipeline):
            self.current_trainstep = step
            training_batch.current_vsa_sparsity = current_vsa_sparsity

+            if (step >= self.training_args.ema_start_step) and \
+                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+                self.generator_ema = EMA_FSDP(
+                    self.transformer, decay=self.training_args.ema_decay)
+                logger.info(
+                    f"Created generator EMA at step {step} with decay={self.training_args.ema_decay}"
+                )
+
            with torch.autocast("cuda", dtype=torch.bfloat16):
                training_batch = self.train_one_step(training_batch)

@@ -948,11 +1260,19 @@ class DistillationPipeline(TrainingPipeline):
            avg_step_time = sum(step_times) / len(step_times)

            progress_bar.set_postfix({
-                "total_loss": f"{total_loss:.4f}",
-                "generator_loss": f"{generator_loss:.4f}",
-                "fake_score_loss": f"{fake_score_loss:.4f}",
-                "step_time": f"{step_time:.2f}s",
-                "grad_norm": grad_norm,
+                "total_loss":
+                f"{total_loss:.4f}",
+                "generator_loss":
+                f"{generator_loss:.4f}",
+                "fake_score_loss":
+                f"{fake_score_loss:.4f}",
+                "step_time":
+                f"{step_time:.2f}s",
+                "grad_norm":
+                grad_norm,
+                "ema":
+                "✓" if (self.generator_ema is not None and self.is_ema_ready())
+                else "✗",
            })
            progress_bar.update(1)

@@ -980,6 +1300,15 @@ class DistillationPipeline(TrainingPipeline):
                if use_vsa:
                    log_data["VSA_train_sparsity"] = current_vsa_sparsity

+                if self.generator_ema is not None:
+                    log_data["ema_enabled"] = True
+                    log_data["ema_decay"] = self.training_args.ema_decay
+                else:
+                    log_data["ema_enabled"] = False
+
+                ema_stats = self.get_ema_stats()
+                log_data.update(ema_stats)
+
                if training_batch.dmd_latent_vis_dict:
                    dmd_additional_logs = {
                        "generator_timestep":
@@ -1011,7 +1340,8 @@ class DistillationPipeline(TrainingPipeline):
                    self.global_rank, self.training_args.output_dir, step,
                    self.optimizer, self.fake_score_optimizer,
                    self.train_dataloader, self.lr_scheduler,
-                    self.fake_score_lr_scheduler, self.noise_random_generator)
+                    self.fake_score_lr_scheduler, self.noise_random_generator,
+                    self.generator_ema)

                if self.transformer:
                    self.transformer.train()
@@ -1028,7 +1358,11 @@ class DistillationPipeline(TrainingPipeline):
                                             self.global_rank,
                                             self.training_args.output_dir,
                                             f"{step}_weight_only",
-                                             only_save_generator_weight=True)
+                                             only_save_generator_weight=True,
+                                             generator_ema=self.generator_ema)
+
+                if self.training_args.use_ema and self.is_ema_ready():
+                    self.save_ema_weights(self.training_args.output_dir, step)

            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
                if self.training_args.log_visualization:
@@ -1048,7 +1382,11 @@ class DistillationPipeline(TrainingPipeline):
            self.training_args.output_dir, self.training_args.max_train_steps,
            self.optimizer, self.fake_score_optimizer, self.train_dataloader,
            self.lr_scheduler, self.fake_score_lr_scheduler,
-            self.noise_random_generator)
+            self.noise_random_generator, self.generator_ema)
+
+        if self.training_args.use_ema and self.is_ema_ready():
+            self.save_ema_weights(self.training_args.output_dir,
+                                  self.training_args.max_train_steps)

        if get_sp_group():
            cleanup_dist_env_and_memory()
@@ -0,0 +1,996 @@
+# SPDX-License-Identifier: Apache-2.0
+import copy
+import gc
+import logging
+from typing import Any
+from collections import deque
+import time
+
+import torch
+import torch.nn.functional as F
+import wandb
+from tqdm.auto import tqdm
+
+from fastvideo.fastvideo_args import TrainingArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import TrainingBatch
+from fastvideo.training.distillation_pipeline import DistillationPipeline
+from fastvideo.training.training_utils import (
+    clip_grad_norm_while_handling_failing_dtensor_cases,
+    EMA_FSDP,
+    save_distillation_checkpoint,
+)
+from fastvideo.models.utils import pred_noise_to_pred_video
+from fastvideo.distributed import get_world_group
+import torch.distributed as dist
+import numpy as np
+from fastvideo.utils import set_random_seed, is_vsa_available
+import fastvideo.envs as envs
+from einops import rearrange
+
+logger = init_logger(__name__)
+
+vsa_available = is_vsa_available()
+
+
+class SelfForcingDistillationPipeline(DistillationPipeline):
+    """
+    A self-forcing distillation pipeline that alternates between training
+    the generator and critic based on the self-forcing methodology.
+    
+    This implementation follows the self-forcing approach where:
+    1. Generator and critic are trained in alternating steps
+    2. Generator loss uses DMD-style loss with the critic as fake score
+    3. Critic loss trains the fake score model to distinguish real vs fake
+    """
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        """Initialize the self-forcing training pipeline."""
+        logger.info("Initializing self-forcing distillation pipeline...")
+
+        super().initialize_training_pipeline(training_args)
+        self.dfake_gen_update_ratio = getattr(training_args, 'dfake_gen_update_ratio', 5)
+
+        self.num_frame_per_block = getattr(training_args, 'num_frame_per_block', 3)
+        self.independent_first_frame = getattr(training_args, 'independent_first_frame', False)
+        self.same_step_across_blocks = getattr(training_args, 'same_step_across_blocks', False)
+        self.last_step_only = getattr(training_args, 'last_step_only', False)
+        self.context_noise = getattr(training_args, 'context_noise', 0)
+        
+        self.frame_seq_length = 1560  # TODO: Calculate this dynamically based on patch size
+        
+        self.kv_cache1 = None
+        self.crossattn_cache = None
+
+        logger.info(f"Self-forcing generator update ratio: {self.dfake_gen_update_ratio}")
+
+    def generate_and_sync_list(self, num_blocks, num_denoising_steps, device):
+        """Generate and synchronize random exit flags across distributed processes."""
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        if rank == 0:
+            # Generate random indices
+            indices = torch.randint(
+                low=0,
+                high=num_denoising_steps,
+                size=(num_blocks,),
+                device=device
+            )
+            if self.last_step_only:
+                indices = torch.ones_like(indices) * (num_denoising_steps - 1)
+        else:
+            indices = torch.empty(num_blocks, dtype=torch.long, device=device)
+
+        if dist.is_initialized():
+            dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
+        return indices.tolist()
+
+    def generator_loss(self, training_batch: TrainingBatch) -> tuple[torch.Tensor, dict[str, Any]]:
+        """
+        Compute generator loss using DMD-style approach.
+        The generator tries to fool the critic (fake_score_transformer).
+        """
+        with set_forward_context(
+                current_timestep=training_batch.timesteps,
+                attn_metadata=training_batch.attn_metadata_vsa):
+            if self.training_args.simulate_generator_forward:
+                generator_pred_video = self._generator_multi_step_simulation_forward(training_batch)
+            else:
+                generator_pred_video = self._generator_forward(training_batch)
+
+        with set_forward_context(
+                current_timestep=training_batch.timesteps,
+                attn_metadata=training_batch.attn_metadata):
+            dmd_loss = self._dmd_forward(
+                generator_pred_video=generator_pred_video,
+                training_batch=training_batch
+            )
+
+        log_dict = {
+            "dmdtrain_gradient_norm": torch.tensor(0.0, device=self.device)
+        }
+        
+        return dmd_loss, log_dict
+
+    def critic_loss(self, training_batch: TrainingBatch) -> tuple[torch.Tensor, dict[str, Any]]:
+        """
+        Compute critic loss using flow matching between noise and generator output.
+        The critic learns to predict the flow from noise to the generator's output.
+        """
+        updated_batch, flow_matching_loss = self.faker_score_forward(training_batch)
+        training_batch.fake_score_latent_vis_dict = updated_batch.fake_score_latent_vis_dict
+        log_dict = {}
+        
+        return flow_matching_loss, log_dict
+
+    def _generator_forward(self, training_batch: TrainingBatch) -> torch.Tensor:
+        """Forward pass through generator with KV cache support for causal generation."""
+        latents = training_batch.latents
+        dtype = latents.dtype
+        batch_size = latents.shape[0]
+        
+        # Step 1: Sample a timestep from denoising_step_list
+        index = torch.randint(0,
+                              len(self.denoising_step_list), [1],
+                              device=self.device,
+                              dtype=torch.long)
+        timestep = self.denoising_step_list[index]
+        training_batch.dmd_latent_vis_dict["generator_timestep"] = timestep
+
+        # Step 2: Initialize KV cache and cross-attention cache for causal generation
+        kv_cache, crossattn_cache = self._initialize_simulation_caches(batch_size, dtype, self.device)
+        
+        if getattr(self.training_args, 'validate_cache_structure', False):
+            self._validate_cache_structure(kv_cache, crossattn_cache, batch_size)
+
+        # Step 3: Add noise to latents
+        noise = torch.randn(self.video_latent_shape,
+                            device=self.device,
+                            dtype=dtype)
+        if self.sp_world_size > 1:
+            noise = rearrange(noise,
+                              "b (n t) c h w -> b n t c h w",
+                              n=self.sp_world_size).contiguous()
+            noise = noise[:, self.rank_in_sp_group, :, :, :, :]
+        noisy_latent = self.noise_scheduler.add_noise(latents.flatten(0, 1),
+                                                      noise.flatten(0, 1),
+                                                      timestep).unflatten(
+                                                          0,
+                                                          (1, latents.shape[1]))
+
+        # Step 4: Build input kwargs with KV cache support
+        training_batch = self._build_distill_input_kwargs(
+            noisy_latent, timestep, training_batch.conditional_dict,
+            training_batch)
+        
+        # Step 5: Forward pass with KV cache if available
+        if hasattr(self.transformer, '_forward_inference'):
+            # Use causal inference forward with KV cache
+            pred_noise = self.transformer(
+                hidden_states=training_batch.input_kwargs['hidden_states'],
+                encoder_hidden_states=training_batch.input_kwargs['encoder_hidden_states'],
+                timestep=training_batch.input_kwargs['timestep'],
+                encoder_hidden_states_image=training_batch.input_kwargs.get('encoder_hidden_states_image'),
+                kv_cache=kv_cache,
+                crossattn_cache=crossattn_cache,
+                current_start=0,  # Start from beginning for single-step
+                cache_start=0
+            ).permute(0, 2, 1, 3, 4)
+        else:
+            # Fallback to regular forward
+            pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+                0, 2, 1, 3, 4)
+        
+        # Step 6: Convert noise prediction to video prediction
+        pred_video = pred_noise_to_pred_video(
+            pred_noise=pred_noise.flatten(0, 1),
+            noise_input_latent=noisy_latent.flatten(0, 1),
+            timestep=timestep,
+            scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2])
+
+        self._reset_simulation_caches(kv_cache, crossattn_cache)
+
+        return pred_video
+
+    def _generator_multi_step_simulation_forward(
+            self, training_batch: TrainingBatch, return_sim_steps: bool = False) -> torch.Tensor:
+        """Forward pass through student transformer matching inference procedure with KV cache management.
+        
+        This function is adapted from the reference self-forcing implementation's inference_with_trajectory
+        and includes gradient masking logic for dynamic frame generation.
+        """
+        latents = training_batch.latents
+        dtype = latents.dtype
+        batch_size = latents.shape[0]
+        initial_latent = getattr(training_batch, 'image_latent', None)
+        
+        # Dynamic frame generation logic (adapted from _run_generator)
+        num_training_frames = getattr(self.training_args, 'num_frames', 21)
+        
+        # During training, the number of generated frames should be uniformly sampled from
+        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block
+        min_num_frames = 20 if self.independent_first_frame else 21
+        max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames
+        assert max_num_frames % self.num_frame_per_block == 0
+        assert min_num_frames % self.num_frame_per_block == 0
+        max_num_blocks = max_num_frames // self.num_frame_per_block
+        min_num_blocks = min_num_frames // self.num_frame_per_block
+        
+        # Sample number of blocks and sync across processes
+        num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1,), device=self.device)
+        if dist.is_initialized():
+            dist.broadcast(num_generated_blocks, src=0)
+        num_generated_blocks = num_generated_blocks.item()
+        num_generated_frames = num_generated_blocks * self.num_frame_per_block
+        if self.independent_first_frame and initial_latent is None:
+            num_generated_frames += 1
+            min_num_frames += 1
+        
+        # Create noise with dynamic shape
+        if initial_latent is not None:
+            noise_shape = [batch_size, num_generated_frames - 1, *self.video_latent_shape[2:]]
+        else:
+            noise_shape = [batch_size, num_generated_frames, *self.video_latent_shape[2:]]
+        
+        noise = torch.randn(noise_shape, device=self.device, dtype=dtype)
+        if self.sp_world_size > 1:
+            noise = rearrange(noise, "b (n t) c h w -> b n t c h w", n=self.sp_world_size).contiguous()
+            noise = noise[:, self.rank_in_sp_group, :, :, :, :]
+        
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        
+        # Block size calculation
+        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+            
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames
+        output = torch.zeros([batch_size, num_output_frames, num_channels, height, width],
+                           device=noise.device, dtype=noise.dtype)
+
+        # Step 1: Initialize KV cache to all zeros
+        self.kv_cache1, self.crossattn_cache = self._initialize_simulation_caches(batch_size, dtype, self.device)
+        
+        # Validate cache structure (can be disabled in production)
+        if getattr(self.training_args, 'validate_cache_structure', False):
+            self._validate_cache_structure(self.kv_cache1, self.crossattn_cache, batch_size)
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            output[:, :1] = initial_latent
+            with torch.no_grad():
+                # Build input kwargs for initial latent
+                training_batch_temp = self._build_distill_input_kwargs(
+                    initial_latent, timestep * 0, training_batch.conditional_dict, training_batch)
+            
+                self.transformer(
+                    hidden_states=training_batch_temp.input_kwargs['hidden_states'],
+                    encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
+                    timestep=training_batch_temp.input_kwargs['timestep'],
+                    encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length
+                )
+            current_start_frame += 1
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        num_denoising_steps = len(self.denoising_step_list)
+        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
+        start_gradient_frame_index = num_output_frames - 21
+
+        for block_index, current_num_frames in enumerate(all_num_frames):
+            noisy_input = noise[:, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+            # Step 3.1: Spatial denoising loop
+            for index, current_timestep in enumerate(self.denoising_step_list):
+                if self.same_step_across_blocks:
+                    exit_flag = (index == exit_flags[0])
+                else:
+                    exit_flag = (index == exit_flags[block_index])
+                    
+                timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep
+                logger.info("timestep shape at initalization: %s", timestep.shape)
+
+                if not exit_flag:
+                    with torch.no_grad():
+                        # Build input kwargs
+                        training_batch_temp = self._build_distill_input_kwargs(
+                            noisy_input, timestep, training_batch.conditional_dict, training_batch)
+                        
+                        logger.info("timestep shape in generator: %s", timestep.shape)
+                        pred_flow = self.transformer(
+                            hidden_states=training_batch_temp.input_kwargs['hidden_states'],
+                            encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
+                            timestep=training_batch_temp.input_kwargs['timestep'],
+                            encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=current_start_frame * self.frame_seq_length
+                        ).permute(0, 2, 1, 3, 4)
+                        
+                        denoised_pred = pred_noise_to_pred_video(
+                            pred_noise=pred_flow.flatten(0, 1),
+                            noise_input_latent=noisy_input.flatten(0, 1),
+                            timestep=timestep.flatten(),
+                            scheduler=self.noise_scheduler).unflatten(0, pred_flow.shape[:2])
+                        
+                        next_timestep = self.denoising_step_list[index + 1]
+                        logger.info("denoised_pred shape before: %s", denoised_pred.shape)
+                        noisy_input = self.noise_scheduler.add_noise(
+                            denoised_pred.flatten(0, 1),
+                            torch.randn_like(denoised_pred.flatten(0, 1)),
+                            next_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+                        ).unflatten(0, denoised_pred.shape[:2])
+                        logger.info("denoised_pred shape after: %s", denoised_pred.shape)
+                else:
+                    # Final prediction with gradient control
+                    if current_start_frame < start_gradient_frame_index:
+                        logger.info("timestep shape in generator: %s", timestep.shape)
+                        with torch.no_grad():
+                            training_batch_temp = self._build_distill_input_kwargs(
+                                noisy_input, timestep, training_batch.conditional_dict, training_batch)
+                            
+                            pred_flow = self.transformer(
+                                hidden_states=training_batch_temp.input_kwargs['hidden_states'],
+                                encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
+                                timestep=training_batch_temp.input_kwargs['timestep'],
+                                encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
+                                kv_cache=self.kv_cache1,
+                                crossattn_cache=self.crossattn_cache,
+                                current_start=current_start_frame * self.frame_seq_length
+                            ).permute(0, 2, 1, 3, 4)
+                    else:
+                        logger.info("timestep shape in generator: %s", timestep.shape)
+                        training_batch_temp = self._build_distill_input_kwargs(
+                            noisy_input, timestep, training_batch.conditional_dict, training_batch)
+    
+                        pred_flow = self.transformer(
+                            hidden_states=training_batch_temp.input_kwargs['hidden_states'],
+                            encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
+                            timestep=training_batch_temp.input_kwargs['timestep'],
+                            encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=current_start_frame * self.frame_seq_length
+                        ).permute(0, 2, 1, 3, 4)
+                    
+                    denoised_pred = pred_noise_to_pred_video(
+                        pred_noise=pred_flow.flatten(0, 1),
+                        noise_input_latent=noisy_input.flatten(0, 1),
+                        timestep=timestep.flatten(),
+                        scheduler=self.noise_scheduler).unflatten(0, pred_flow.shape[:2])
+                    break
+
+            # Step 3.2: record the model's output
+            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Step 3.3: rerun with timestep zero to update the cache
+            context_timestep = torch.ones_like(timestep) * self.context_noise
+            logger.info("denoised_pred shape before: %s", denoised_pred.shape)
+            denoised_pred = self.noise_scheduler.add_noise(
+                denoised_pred.flatten(0, 1),
+                torch.randn_like(denoised_pred.flatten(0, 1)),
+                context_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+            ).unflatten(0, denoised_pred.shape[:2])
+            logger.info("denoised_pred shape after: %s", denoised_pred.shape)
+            
+            with torch.no_grad():
+                training_batch_temp = self._build_distill_input_kwargs(
+                    denoised_pred, context_timestep, training_batch.conditional_dict, training_batch)
+                
+                self.transformer(
+                    hidden_states=training_batch_temp.input_kwargs['hidden_states'],
+                    encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
+                    timestep=training_batch_temp.input_kwargs['timestep'],
+                    encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length
+                )
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        # Handle last 21 frames logic 
+        pred_image_or_video = output
+        if num_input_frames > 0:
+            pred_image_or_video = output[:, num_input_frames:]
+        
+        # Slice last 21 frames if we generated more
+        gradient_mask = None
+        if pred_image_or_video.shape[1] > 21:
+            with torch.no_grad():
+                # Reencode to get image latent
+                latent_to_decode = pred_image_or_video[:, :-20, ...]
+                # Decode to video
+                latent_to_decode = latent_to_decode.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
+                
+                # Apply VAE scaling and shift factors
+                if isinstance(self.vae.scaling_factor, torch.Tensor):
+                    latent_to_decode = latent_to_decode / self.vae.scaling_factor.to(latent_to_decode.device, latent_to_decode.dtype)
+                else:
+                    latent_to_decode = latent_to_decode / self.vae.scaling_factor
+                
+                if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None:
+                    if isinstance(self.vae.shift_factor, torch.Tensor):
+                        latent_to_decode += self.vae.shift_factor.to(latent_to_decode.device, latent_to_decode.dtype)
+                    else:
+                        latent_to_decode += self.vae.shift_factor
+                
+                # Decode to pixels
+                pixels = self.vae.decode(latent_to_decode)
+                frame = pixels[:, :, -1:, :, :].to(dtype)  # Last frame [B, C, 1, H, W]
+                
+                # Encode frame back to get image latent
+                image_latent = self.vae.encode(frame).mean.to(dtype)
+                image_latent = image_latent.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]
+            
+            pred_image_or_video_last_21 = torch.cat([image_latent, pred_image_or_video[:, -20:, ...]], dim=1)
+        else:
+            pred_image_or_video_last_21 = pred_image_or_video
+        
+        # Set up gradient mask if we generated more than minimum frames
+        if num_generated_frames != min_num_frames:
+            # Currently, we do not use gradient for the first chunk, since it contains image latents
+            gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool)
+            if self.independent_first_frame:
+                gradient_mask[:, :1] = False
+            else:
+                gradient_mask[:, :self.num_frame_per_block] = False
+        
+        # Apply gradient masking if needed
+        final_output = pred_image_or_video_last_21.to(dtype)
+        if gradient_mask is not None:
+            # Apply gradient masking: detach frames that shouldn't contribute gradients
+            final_output = torch.where(
+                gradient_mask,
+                pred_image_or_video_last_21,  # Keep original values where gradient_mask is True
+                pred_image_or_video_last_21.detach()  # Detach where gradient_mask is False
+            )
+
+        # Store visualization data
+        training_batch.dmd_latent_vis_dict["generator_timestep"] = torch.tensor(
+            self.denoising_step_list[exit_flags[0]], dtype=torch.float32, device=self.device)
+        
+        # Store gradient mask information for debugging
+        if gradient_mask is not None:
+            training_batch.dmd_latent_vis_dict["gradient_mask"] = gradient_mask.float()
+            training_batch.dmd_latent_vis_dict["num_generated_frames"] = torch.tensor(
+                num_generated_frames, dtype=torch.float32, device=self.device)
+            training_batch.dmd_latent_vis_dict["min_num_frames"] = torch.tensor(
+                min_num_frames, dtype=torch.float32, device=self.device)
+        
+        self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache)
+
+        return final_output if gradient_mask is not None else pred_image_or_video
+
+    def _initialize_simulation_caches(self, batch_size: int, dtype: torch.dtype, device: torch.device):
+        """Initialize KV cache and cross-attention cache for multi-step simulation."""
+        num_transformer_blocks = len(self.transformer.blocks)
+        
+        # Calculate frame sequence length based on input dimensions and patch size
+        # From the training batch, we can get the actual latent dimensions
+        latent_shape = self.video_latent_shape_sp  # This is set in _prepare_dit_inputs
+        batch_size_actual, num_frames, num_channels, height, width = latent_shape
+        
+        # Get patch size from transformer config
+        p_t, p_h, p_w = self.transformer.patch_size
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+        
+        # Frame sequence length is the spatial sequence length per frame
+        frame_seq_length = post_patch_height * post_patch_width
+        
+        # Get local attention size from transformer config
+        local_attn_size = getattr(self.transformer, 'local_attn_size', -1)
+        
+        # Get model configuration parameters - handle FSDP wrapping
+        if hasattr(self.transformer, 'config'):
+            config = self.transformer.config
+            num_attention_heads = config.num_attention_heads
+            attention_head_dim = config.attention_head_dim
+            text_len = config.text_len
+        else:
+            # Fallback to direct attribute access for non-FSDP models
+            num_attention_heads = getattr(self.transformer, 'num_attention_heads', 40)
+            attention_head_dim = getattr(self.transformer, 'attention_head_dim', 128)
+            text_len = getattr(self.transformer, 'text_len', 512)
+        
+        num_max_frames = getattr(self.training_args, "num_frames", num_frames)
+        kv_cache_size = num_max_frames * frame_seq_length
+
+        kv_cache = []
+        for _ in range(num_transformer_blocks):
+            kv_cache.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)
+            })
+        
+        # Initialize cross-attention cache
+        crossattn_cache = []
+        for _ in range(num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, text_len, num_attention_heads, attention_head_dim], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, text_len, num_attention_heads, attention_head_dim], dtype=dtype, device=device),
+                "is_init": False
+            })
+        
+        return kv_cache, crossattn_cache
+
+    def _reset_simulation_caches(self, kv_cache, crossattn_cache):
+        """Reset KV cache and cross-attention cache to clean state."""
+        if kv_cache is not None:
+            for cache_dict in kv_cache:
+                cache_dict["global_end_index"].fill_(0)
+                cache_dict["local_end_index"].fill_(0)
+                cache_dict["k"].zero_()
+                cache_dict["v"].zero_()
+        
+        if crossattn_cache is not None:
+            for cache_dict in crossattn_cache:
+                cache_dict["is_init"] = False
+                cache_dict["k"].zero_()
+                cache_dict["v"].zero_()
+
+    def _validate_cache_structure(self, kv_cache, crossattn_cache, batch_size: int):
+        """Validate that cache structures are correctly initialized."""
+        num_transformer_blocks = len(self.transformer.blocks)
+        
+        # Get model configuration parameters - handle FSDP wrapping
+        if hasattr(self.transformer, 'config'):
+            config = self.transformer.config
+            num_attention_heads = config.num_attention_heads
+            attention_head_dim = config.attention_head_dim
+            text_len = config.text_len
+        else:
+            # Fallback to direct attribute access for non-FSDP models
+            num_attention_heads = getattr(self.transformer, 'num_attention_heads', 40)
+            attention_head_dim = getattr(self.transformer, 'attention_head_dim', 128)
+            text_len = getattr(self.transformer, 'text_len', 512)
+        
+        if kv_cache is not None:
+            assert len(kv_cache) == num_transformer_blocks, f"Expected {num_transformer_blocks} transformer blocks, got {len(kv_cache)}"
+            for i, cache_dict in enumerate(kv_cache):
+                assert "k" in cache_dict and "v" in cache_dict, f"Missing k/v in kv_cache block {i}"
+                assert "global_end_index" in cache_dict and "local_end_index" in cache_dict, f"Missing indices in kv_cache block {i}"
+                assert cache_dict["k"].shape[0] == batch_size, f"Batch size mismatch in kv_cache block {i}"
+                assert cache_dict["v"].shape[0] == batch_size, f"Batch size mismatch in kv_cache block {i}"
+                assert cache_dict["k"].shape[2] == num_attention_heads, f"Attention heads mismatch in kv_cache block {i}"
+                assert cache_dict["k"].shape[3] == attention_head_dim, f"Attention head dim mismatch in kv_cache block {i}"
+        
+        if crossattn_cache is not None:
+            assert len(crossattn_cache) == num_transformer_blocks, f"Expected {num_transformer_blocks} transformer blocks, got {len(crossattn_cache)}"
+            for i, cache_dict in enumerate(crossattn_cache):
+                assert "k" in cache_dict and "v" in cache_dict, f"Missing k/v in crossattn_cache block {i}"
+                assert "is_init" in cache_dict, f"Missing is_init in crossattn_cache block {i}"
+                assert cache_dict["k"].shape[0] == batch_size, f"Batch size mismatch in crossattn_cache block {i}"
+                assert cache_dict["v"].shape[0] == batch_size, f"Batch size mismatch in crossattn_cache block {i}"
+                assert cache_dict["k"].shape[1] == text_len, f"Text length mismatch in crossattn_cache block {i}"
+                assert cache_dict["k"].shape[2] == num_attention_heads, f"Attention heads mismatch in crossattn_cache block {i}"
+                assert cache_dict["k"].shape[3] == attention_head_dim, f"Attention head dim mismatch in crossattn_cache block {i}"
+
+    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
+        """
+        Self-forcing training step that alternates between generator and critic training.
+        """
+        gradient_accumulation_steps = getattr(self.training_args, 'gradient_accumulation_steps', 1)
+        train_generator = (self.current_trainstep % self.dfake_gen_update_ratio == 0)
+        
+        batches = []
+        for _ in range(gradient_accumulation_steps):
+            batch = self._prepare_distillation(training_batch)
+            batch = self._get_next_batch(batch)
+            batch = self._normalize_dit_input(batch)
+            batch = self._prepare_dit_inputs(batch)
+            batch = self._build_attention_metadata(batch)
+            batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata)
+            if batch.attn_metadata is not None:
+                batch.attn_metadata.VSA_sparsity = 0.0
+            batches.append(batch)
+
+        training_batch.dmd_latent_vis_dict = {}
+        training_batch.fake_score_latent_vis_dict = {}
+
+        if train_generator:
+            logger.debug(f"Training generator at step {self.current_trainstep}")
+            self.optimizer.zero_grad()
+            total_generator_loss = 0.0
+            generator_log_dict = {}
+            
+            for batch in batches:
+                # Create a new batch with detached tensors
+                batch_gen = TrainingBatch()
+                for key, value in batch.__dict__.items():
+                    if isinstance(value, torch.Tensor):
+                        setattr(batch_gen, key, value.detach().clone())
+                    elif isinstance(value, dict):
+                        setattr(batch_gen, key, {k: v.detach().clone() if isinstance(v, torch.Tensor) else copy.deepcopy(v) for k, v in value.items()})
+                    else:
+                        setattr(batch_gen, key, copy.deepcopy(value))
+                
+                generator_loss, gen_log_dict = self.generator_loss(batch_gen)
+                with set_forward_context(
+                        current_timestep=batch_gen.timesteps,
+                        attn_metadata=batch_gen.attn_metadata):
+                    (generator_loss / gradient_accumulation_steps).backward()
+                total_generator_loss += generator_loss.detach().item()
+                generator_log_dict.update(gen_log_dict)
+                # Store visualization data from generator training
+                if hasattr(batch_gen, 'dmd_latent_vis_dict'):
+                    training_batch.dmd_latent_vis_dict.update(batch_gen.dmd_latent_vis_dict)
+            
+            self._clip_model_grad_norm_(batch_gen, self.transformer)
+            self.optimizer.step()
+            self.lr_scheduler.step()
+            
+            if self.generator_ema is not None:
+                self.generator_ema.update(self.transformer)
+            
+            avg_generator_loss = torch.tensor(
+                total_generator_loss / gradient_accumulation_steps,
+                device=self.device
+            )
+            world_group = get_world_group()
+            world_group.all_reduce(avg_generator_loss, op=torch.distributed.ReduceOp.AVG)
+            training_batch.generator_loss = avg_generator_loss.item()
+        else:
+            training_batch.generator_loss = 0.0
+
+        logger.debug(f"Training critic at step {self.current_trainstep}")
+        self.fake_score_optimizer.zero_grad()
+        total_critic_loss = 0.0
+        critic_log_dict = {}
+        
+        for batch in batches:
+            # Create a new batch with detached tensors
+            batch_critic = TrainingBatch()
+            for key, value in batch.__dict__.items():
+                if isinstance(value, torch.Tensor):
+                    setattr(batch_critic, key, value.detach().clone())
+                elif isinstance(value, dict):
+                    setattr(batch_critic, key, {k: v.detach().clone() if isinstance(v, torch.Tensor) else copy.deepcopy(v) for k, v in value.items()})
+                else:
+                    setattr(batch_critic, key, copy.deepcopy(value))
+            
+            critic_loss, crit_log_dict = self.critic_loss(batch_critic)
+            with set_forward_context(
+                    current_timestep=batch_critic.timesteps,
+                    attn_metadata=batch_critic.attn_metadata):
+                (critic_loss / gradient_accumulation_steps).backward()
+            total_critic_loss += critic_loss.detach().item()
+            critic_log_dict.update(crit_log_dict)
+            # Store visualization data from critic training
+            if hasattr(batch_critic, 'fake_score_latent_vis_dict'):
+                training_batch.fake_score_latent_vis_dict.update(batch_critic.fake_score_latent_vis_dict)
+        
+        self._clip_model_grad_norm_(batch_critic, self.fake_score_transformer)
+        self.fake_score_optimizer.step()
+        self.fake_score_lr_scheduler.step()
+        
+        avg_critic_loss = torch.tensor(
+            total_critic_loss / gradient_accumulation_steps,
+            device=self.device
+        )
+        world_group = get_world_group()
+        world_group.all_reduce(avg_critic_loss, op=torch.distributed.ReduceOp.AVG)
+        training_batch.fake_score_loss = avg_critic_loss.item()
+
+        training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss
+        
+        return training_batch
+
+    def _log_training_info(self) -> None:
+        """Log self-forcing specific training information."""
+        super()._log_training_info()
+        logger.info("Self-forcing specific settings:")
+        logger.info("  Generator update ratio: %s", self.dfake_gen_update_ratio) 
+
+    def visualize_intermediate_latents(self, training_batch: TrainingBatch,
+                                     training_args: TrainingArgs, step: int):
+        """Add visualization data to wandb logging and save frames to disk."""
+        wandb_loss_dict = {}
+        
+        # Debug logging
+        logger.info(f"Step {step}: Starting visualization")
+        if hasattr(training_batch, 'dmd_latent_vis_dict'):
+            logger.info(f"DMD latent keys: {list(training_batch.dmd_latent_vis_dict.keys())}")
+        if hasattr(training_batch, 'fake_score_latent_vis_dict'):
+            logger.info(f"Fake score latent keys: {list(training_batch.fake_score_latent_vis_dict.keys())}")
+        
+        # Process generator predictions if available
+        if hasattr(training_batch, 'dmd_latent_vis_dict') and training_batch.dmd_latent_vis_dict:
+            dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
+            dmd_log_keys = ['generator_pred_video', 'real_score_pred_video', 'faker_score_pred_video']
+            
+            for latent_key in dmd_log_keys:
+                if latent_key in dmd_latents_vis_dict:
+                    logger.info(f"Processing DMD latent: {latent_key}")
+                    latents = dmd_latents_vis_dict[latent_key]
+                    if not isinstance(latents, torch.Tensor):
+                        logger.warning(f"Expected tensor for {latent_key}, got {type(latents)}")
+                        continue
+                        
+                    latents = latents.detach()
+                    latents = latents.permute(0, 2, 1, 3, 4)
+
+                    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
+
+                    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
+
+                    try:
+                        with torch.autocast("cuda", dtype=torch.bfloat16):
+                            video = self.vae.decode(latents)
+                        video = (video / 2 + 0.5).clamp(0, 1)
+                        video = video.cpu().float()
+                        video = video.permute(0, 2, 1, 3, 4)
+                        video = (video * 255).numpy().astype(np.uint8)
+                        wandb_loss_dict[f"dmd_{latent_key}"] = wandb.Video(video, fps=24, format="mp4")
+                        logger.info(f"Successfully processed DMD latent: {latent_key}")
+                    except Exception as e:
+                        logger.error(f"Error processing DMD latent {latent_key}: {str(e)}")
+                    del video, latents
+
+        # Process critic predictions
+        if hasattr(training_batch, 'fake_score_latent_vis_dict') and training_batch.fake_score_latent_vis_dict:
+            fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
+            fake_score_log_keys = ['generator_pred_video']
+            
+            for latent_key in fake_score_log_keys:
+                if latent_key in fake_score_latents_vis_dict:
+                    logger.info(f"Processing critic latent: {latent_key}")
+                    latents = fake_score_latents_vis_dict[latent_key]
+                    if not isinstance(latents, torch.Tensor):
+                        logger.warning(f"Expected tensor for {latent_key}, got {type(latents)}")
+                        continue
+                        
+                    latents = latents.detach()
+                    latents = latents.permute(0, 2, 1, 3, 4)
+
+                    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
+
+                    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
+
+                    try:
+                        with torch.autocast("cuda", dtype=torch.bfloat16):
+                            video = self.vae.decode(latents)
+                        video = (video / 2 + 0.5).clamp(0, 1)
+                        video = video.cpu().float()
+                        video = video.permute(0, 2, 1, 3, 4)
+                        video = (video * 255).numpy().astype(np.uint8)
+                        wandb_loss_dict[f"critic_{latent_key}"] = wandb.Video(video, fps=24, format="mp4")
+                        logger.info(f"Successfully processed critic latent: {latent_key}")
+                    except Exception as e:
+                        logger.error(f"Error processing critic latent {latent_key}: {str(e)}")
+                    del video, latents
+
+        # Log metadata
+        if hasattr(training_batch, 'dmd_latent_vis_dict') and training_batch.dmd_latent_vis_dict:
+            if "generator_timestep" in training_batch.dmd_latent_vis_dict:
+                wandb_loss_dict["generator_timestep"] = training_batch.dmd_latent_vis_dict["generator_timestep"].item()
+            if "dmd_timestep" in training_batch.dmd_latent_vis_dict:
+                wandb_loss_dict["dmd_timestep"] = training_batch.dmd_latent_vis_dict["dmd_timestep"].item()
+
+        if hasattr(training_batch, 'fake_score_latent_vis_dict') and training_batch.fake_score_latent_vis_dict:
+            if "fake_score_timestep" in training_batch.fake_score_latent_vis_dict:
+                wandb_loss_dict["fake_score_timestep"] = training_batch.fake_score_latent_vis_dict["fake_score_timestep"].item()
+
+        # Log final dict contents
+        logger.info(f"Final wandb_loss_dict keys: {list(wandb_loss_dict.keys())}")
+
+        if self.global_rank == 0:
+            wandb.log(wandb_loss_dict, step=step)
+
+    def train(self) -> None:
+        """Main training loop with self-forcing specific logging."""
+        assert self.training_args.seed is not None, "seed must be set"
+        seed = self.training_args.seed
+
+        # Set the same seed within each SP group to ensure reproducibility
+        if self.sp_world_size > 1:
+            # Use the same seed for all processes within the same SP group
+            sp_group_seed = seed + (self.global_rank // self.sp_world_size)
+            set_random_seed(sp_group_seed)
+            logger.info("Rank %s: Using SP group seed %s", self.global_rank,
+                        sp_group_seed)
+        else:
+            set_random_seed(seed + self.global_rank)
+
+        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", seed)
+        
+        self.current_trainstep = self.init_steps
+
+        if self.training_args.resume_from_checkpoint:
+            self._resume_from_checkpoint()
+            logger.info("Resumed from checkpoint, random states restored")
+        else:
+            logger.info("Starting training from scratch")
+
+        self.train_loader_iter = iter(self.train_dataloader)
+
+        step_times = deque(maxlen=100)
+
+        self._log_training_info()
+        self._log_validation(self.transformer, self.training_args,
+                             self.init_steps)
+
+        progress_bar = tqdm(
+            range(0, self.training_args.max_train_steps),
+            initial=self.init_steps,
+            desc="Steps",
+            disable=self.local_rank > 0,
+        )
+
+        use_vsa = vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN"
+        for step in range(self.init_steps + 1,
+                          self.training_args.max_train_steps + 1):
+            start_time = time.perf_counter()
+            if use_vsa:
+                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
+                if vsa_decay_interval_steps > 1:
+                    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 = vsa_sparsity
+            else:
+                current_vsa_sparsity = 0.0
+
+            training_batch = TrainingBatch()
+            self.current_trainstep = step
+            training_batch.current_vsa_sparsity = current_vsa_sparsity
+
+            if (step >= self.training_args.ema_start_step) and \
+                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+                self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
+                logger.info(f"Created generator EMA at step {step} with decay={self.training_args.ema_decay}")
+
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                training_batch = self.train_one_step(training_batch)
+
+            total_loss = training_batch.total_loss
+            generator_loss = training_batch.generator_loss
+            fake_score_loss = training_batch.fake_score_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({
+                "total_loss": f"{total_loss:.4f}",
+                "generator_loss": f"{generator_loss:.4f}",
+                "fake_score_loss": f"{fake_score_loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+                "ema": "✓" if (self.generator_ema is not None and self.is_ema_ready()) else "✗",
+            })
+            progress_bar.update(1)
+
+            if self.global_rank == 0:
+                log_data = {
+                    "train_total_loss": total_loss,
+                    "train_fake_score_loss": fake_score_loss,
+                    "learning_rate": self.lr_scheduler.get_last_lr()[0],
+                    "fake_score_learning_rate": self.fake_score_lr_scheduler.get_last_lr()[0],
+                    "step_time": step_time,
+                    "avg_step_time": avg_step_time,
+                    "grad_norm": grad_norm,
+                }
+                if (step % self.dfake_gen_update_ratio == 0):
+                    log_data["train_generator_loss"] = generator_loss
+                if use_vsa:
+                    log_data["VSA_train_sparsity"] = current_vsa_sparsity
+
+                if self.generator_ema is not None:
+                    log_data["ema_enabled"] = True
+                    log_data["ema_decay"] = self.training_args.ema_decay
+                else:
+                    log_data["ema_enabled"] = False
+
+                ema_stats = self.get_ema_stats()
+                log_data.update(ema_stats)
+
+                if training_batch.dmd_latent_vis_dict:
+                    dmd_additional_logs = {
+                        "generator_timestep": training_batch.dmd_latent_vis_dict["generator_timestep"].item(),
+                        "dmd_timestep": training_batch.dmd_latent_vis_dict["dmd_timestep"].item(),
+                    }
+                    log_data.update(dmd_additional_logs)
+
+                faker_score_additional_logs = {
+                    "fake_score_timestep": training_batch.fake_score_latent_vis_dict["fake_score_timestep"].item(),
+                }
+                log_data.update(faker_score_additional_logs)
+
+                wandb.log(log_data, step=step)
+
+                if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
+                    if self.training_args.log_visualization:
+                        self.visualize_intermediate_latents(training_batch, self.training_args, step)
+
+            if (self.training_args.training_state_checkpointing_steps > 0
+                    and step % self.training_args.training_state_checkpointing_steps == 0):
+                print("rank", self.global_rank,
+                      "save training state checkpoint at step", step)
+                save_distillation_checkpoint(
+                    self.transformer, self.fake_score_transformer,
+                    self.global_rank, self.training_args.output_dir, step,
+                    self.optimizer, self.fake_score_optimizer,
+                    self.train_dataloader, self.lr_scheduler,
+                    self.fake_score_lr_scheduler, self.noise_random_generator,
+                    self.generator_ema)
+
+                if self.transformer:
+                    self.transformer.train()
+                self.sp_group.barrier()
+
+            if (self.training_args.weight_only_checkpointing_steps > 0
+                    and step % self.training_args.weight_only_checkpointing_steps == 0):
+                print("rank", self.global_rank,
+                      "save weight-only checkpoint at step", step)
+                save_distillation_checkpoint(self.transformer,
+                                             self.fake_score_transformer,
+                                             self.global_rank,
+                                             self.training_args.output_dir,
+                                             f"{step}_weight_only",
+                                             only_save_generator_weight=True,
+                                             generator_ema=self.generator_ema)
+                
+                if self.training_args.use_ema and self.is_ema_ready():
+                    self.save_ema_weights(self.training_args.output_dir, step)
+
+            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
+                self._log_validation(self.transformer, self.training_args, step)
+
+        wandb.finish()
+
+        print("rank", self.global_rank,
+              "save final training state checkpoint at step",
+              self.training_args.max_train_steps)
+        save_distillation_checkpoint(
+            self.transformer, self.fake_score_transformer, self.global_rank,
+            self.training_args.output_dir, self.training_args.max_train_steps,
+            self.optimizer, self.fake_score_optimizer, self.train_dataloader,
+            self.lr_scheduler, self.fake_score_lr_scheduler,
+            self.noise_random_generator, self.generator_ema)
+
+        if self.training_args.use_ema and self.is_ema_ready():
+            self.save_ema_weights(self.training_args.output_dir, self.training_args.max_train_steps)
+
+        if get_sp_group():
+            cleanup_dist_env_and_memory() 
@@ -117,10 +117,14 @@ class TrainingPipeline(LoRAPipeline, ABC):
        params_to_optimize = self.transformer.parameters()
        params_to_optimize = list(
            filter(lambda p: p.requires_grad, params_to_optimize))
+        # Parse betas from string format "beta1,beta2"
+        betas_str = training_args.betas
+        betas = tuple(float(x.strip()) for x in betas_str.split(","))
+        
        self.optimizer = torch.optim.AdamW(
            params_to_optimize,
            lr=training_args.learning_rate,
-            betas=(0.9, 0.999),
+            betas=betas,
            weight_decay=training_args.weight_decay,
            eps=1e-8,
        )
@@ -202,6 +202,7 @@ def save_distillation_checkpoint(generator_transformer,
                                 generator_scheduler=None,
                                 fake_score_scheduler=None,
                                 noise_generator=None,
+                                 generator_ema=None,
                                 only_save_generator_weight=False) -> None:
    """
    Save distillation checkpoint with both generator and fake_score models.
@@ -233,6 +234,8 @@ def save_distillation_checkpoint(generator_transformer,
        if generator_scheduler is not None:
            generator_states["scheduler"] = SchedulerWrapper(
                generator_scheduler)
+        if generator_ema is not None:
+            generator_states["ema"] = generator_ema.state_dict()

        generator_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
                                         "generator")
@@ -402,7 +405,8 @@ def load_distillation_checkpoint(generator_transformer,
                                 dataloader=None,
                                 generator_scheduler=None,
                                 fake_score_scheduler=None,
-                                 noise_generator=None) -> int:
+                                 noise_generator=None,
+                                 generator_ema=None) -> int:
    """
    Load distillation checkpoint with both generator and fake_score models.
    Returns the step number from which training should resume.
@@ -456,6 +460,18 @@ def load_distillation_checkpoint(generator_transformer,
        end_time - begin_time,
        local_main_process_only=False)

+    # Load EMA state if available and generator_ema is provided
+    if generator_ema is not None:
+        try:
+            ema_state = generator_states.get("ema")
+            if ema_state is not None:
+                generator_ema.load_state_dict(ema_state)
+                logger.info("rank: %s, generator EMA state loaded successfully", rank)
+            else:
+                logger.info("rank: %s, no EMA state found in checkpoint", rank)
+        except Exception as e:
+            logger.warning("rank: %s, failed to load EMA state: %s", rank, str(e))
+
    # Load critic distributed checkpoint
    critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
                                  "critic")
@@ -826,27 +842,6 @@ def custom_to_hf_state_dict(
    return new_state_dict


-def pred_noise_to_pred_video(pred_noise: torch.Tensor,
-                             noise_input_latent: torch.Tensor,
-                             timestep: torch.Tensor,
-                             scheduler: Any) -> torch.Tensor:
-    """
-    Convert predicted noise to clean latent.
-    """
-    timestep = timestep.expand(noise_input_latent.shape[0])
-    dtype = pred_noise.dtype
-    device = pred_noise.device
-    pred_noise = pred_noise.float().to(device)
-    noise_input_latent = noise_input_latent.float().to(device)
-    sigmas = scheduler.sigmas.float().to(device)
-    timesteps = scheduler.timesteps.float().to(device)
-    timestep_id = torch.argmin(
-        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
-    pred_video = noise_input_latent - sigma_t * pred_noise
-    return pred_video.to(dtype)
-
-
 def shift_timestep(timestep: torch.Tensor, shift: float,
                   num_train_timestep: float) -> torch.Tensor:
    if shift == 1:
@@ -1299,3 +1294,156 @@ def get_scheduler(
                         num_warmup_steps=num_warmup_steps,
                         num_training_steps=num_training_steps,
                         last_epoch=last_epoch)
+
+
+class EMA_FSDP:
+    """
+    FSDP2-friendly EMA with two modes:
+      - mode="local_shard" (default): maintain float32 CPU EMA of local parameter shards on every rank.
+        Provides a context manager to temporarily swap EMA weights into the live model for teacher forward.
+      - mode="rank0_full": maintain a consolidated float32 CPU EMA of full parameters on rank 0 only
+        using gather_state_dict_on_cpu_rank0(). Useful for checkpoint export; not for teacher forward.
+
+    Usage (local_shard for CM teacher):
+      ema = EMA_FSDP(model, decay=0.999, mode="local_shard")
+      for step in ...:
+          ema.update(model)
+      with ema.apply_to_model(model):
+          with torch.no_grad():
+              y_teacher = model(...)
+
+    Usage (rank0_full for export):
+      ema = EMA_FSDP(model, decay=0.999, mode="rank0_full")
+      ema.update(model)
+      ema.state_dict()  # on rank 0
+    """
+    def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"):
+        self.decay = float(decay)
+        self.mode = mode
+        self.shadow: dict[str, torch.Tensor] = {}
+        self.rank = dist.get_rank() if dist.is_initialized() else 0
+        if self.mode not in {"local_shard", "rank0_full"}:
+            raise ValueError(f"Unsupported EMA_FSDP mode: {self.mode}")
+        self._init_shadow(module)
+
+    @staticmethod
+    def _to_local_tensor(t: torch.Tensor) -> torch.Tensor:
+        # DTensor-aware to_local fetch; fall back to raw tensor
+        try:
+            from torch.distributed.tensor import DTensor  # type: ignore
+            if isinstance(t, DTensor):
+                return t.to_local()
+        except Exception:
+            pass
+        return t
+
+    @torch.no_grad()
+    def _init_shadow(self, module):
+        if self.mode == "rank0_full":
+            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
+            if self.rank == 0:
+                self.shadow = {k: v.detach().clone().float().cpu() for k, v in cpu_state.items()}
+            else:
+                self.shadow = {}
+            return
+
+        # local_shard: maintain EMA of local shards for requires_grad params
+        self.shadow = {}
+        for name, p in module.named_parameters():
+            if not p.requires_grad:
+                continue
+            local = self._to_local_tensor(p.detach())
+            self.shadow[name] = local.clone().float().cpu()
+
+    @torch.no_grad()
+    def update(self, module):
+        d = self.decay
+        if self.mode == "rank0_full":
+            if self.rank != 0:
+                return
+            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
+            for n, v in cpu_state.items():
+                v_cpu = v.detach().float().cpu()
+                if n not in self.shadow:
+                    self.shadow[n] = v_cpu.clone()
+                else:
+                    self.shadow[n].mul_(d).add_(v_cpu, alpha=1.0 - d)
+            return
+
+        # local_shard: update local shard EMA on every rank
+        for name, p in module.named_parameters():
+            if not p.requires_grad:
+                continue
+            local = self._to_local_tensor(p.detach())
+            v_cpu = local.float().cpu()
+            if name not in self.shadow:
+                self.shadow[name] = v_cpu.clone()
+            else:
+                self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
+
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        if self.mode == "rank0_full":
+            return {k: v.clone() for k, v in self.shadow.items()} if self.rank == 0 else {}
+        return {k: v.clone() for k, v in self.shadow.items()}
+
+    def load_state_dict(self, sd: dict[str, torch.Tensor]):
+        self.shadow = {k: v.clone() for k, v in sd.items()}
+
+    @torch.no_grad()
+    def copy_to_unwrapped(self, module) -> None:
+        """
+        Copy EMA weights into a non-sharded (unwrapped) module. Intended for export/eval.
+        For mode="rank0_full", only rank 0 has the full EMA state.
+        """
+        if self.mode == "rank0_full" and self.rank != 0:
+            return
+        name_to_param = dict(module.named_parameters())
+        for n, w in self.shadow.items():
+            if n in name_to_param:
+                p = name_to_param[n]
+                p.data.copy_(w.to(dtype=p.dtype, device=p.device))
+
+    class _ApplyEMACtx:
+        def __init__(self, ema: "EMA_FSDP", module):
+            self.ema = ema
+            self.module = module
+            self.saved: dict[str, torch.Tensor] = {}
+
+        def __enter__(self):
+            if self.ema.mode != "local_shard":
+                raise RuntimeError("EMA apply_to_model is only supported for mode='local_shard'")
+            with torch.no_grad():
+                for name, p in self.module.named_parameters():
+                    if not p.requires_grad:
+                        continue
+                    # Save local shard
+                    p_local = EMA_FSDP._to_local_tensor(p.detach())
+                    if p_local.numel() == 0:
+                        # Nothing to swap on this rank for this param
+                        continue
+                    self.saved[name] = p_local.clone().to(device=p_local.device, dtype=p_local.dtype)
+                    if name in self.ema.shadow:
+                        ema_cpu = self.ema.shadow[name]
+                        if ema_cpu.numel() != p_local.numel():
+                            # Shard shape mismatch (e.g., empty shard here), skip
+                            continue
+                        # Copy EMA shard into local param shard
+                        p_local.copy_(ema_cpu.to(dtype=p_local.dtype, device=p_local.device))
+            return self.module
+
+        def __exit__(self, exc_type, exc, tb):
+            with torch.no_grad():
+                for name, p in self.module.named_parameters():
+                    if name in self.saved:
+                        p_local = EMA_FSDP._to_local_tensor(p.detach())
+                        if p_local.numel() == 0:
+                            continue
+                        saved_local = self.saved[name]
+                        if saved_local.numel() != p_local.numel():
+                            continue
+                        p_local.copy_(saved_local)
+            self.saved.clear()
+            return False
+
+    def apply_to_model(self, module):
+        return EMA_FSDP._ApplyEMACtx(self, module)
@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import WanCausalDMDPipeline
+from fastvideo.training.self_forcing_distillation_pipeline import SelfForcingDistillationPipeline
+from fastvideo.utils import is_vsa_available
+
+vsa_available = is_vsa_available()
+
+logger = init_logger(__name__)
+
+
+class WanSelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
+    """
+    A self-forcing distillation pipeline for Wan that uses the self-forcing methodology
+    with DMD for video generation.
+    """
+    _required_config_modules = [
+        "scheduler", "transformer", "vae", "real_score_transformer",
+        "fake_score_transformer"
+    ]
+
+    def 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.
+        """
+        pass
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+
+        args_copy.inference_mode = True
+        validation_pipeline = WanCausalDMDPipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,  # type: ignore
+            inference_mode=True,
+            loaded_modules={"transformer": self.get_module("transformer")},
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True)
+
+        self.validation_pipeline = validation_pipeline
+
+
+def main(args) -> None:
+    logger.info("Starting Wan self-forcing distillation pipeline...")
+
+    pipeline = WanSelfForcingDistillationPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, args=args)
+
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Wan self-forcing distillation pipeline completed")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    main(args) 
@@ -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:
@@ -57,6 +57,20 @@ class Executor(ABC):
        """
        raise NotImplementedError

+    @abstractmethod
+    def unmerge_lora_weights(self) -> None:
+        """
+        Unmerge the LoRA weights for the workers.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def merge_lora_weights(self) -> None:
+        """
+        Merge the LoRA weights for the workers.
+        """
+        raise NotImplementedError
+
    @abstractmethod
    def collective_rpc(self,
                       method: str | Callable[..., _R],
@@ -11,13 +11,14 @@ from typing import Any, TextIO, cast
 import psutil
 import torch

+import fastvideo.envs as envs
 from fastvideo.distributed import (
    cleanup_dist_env_and_memory,
    maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.distributed.parallel_state import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
-from fastvideo.pipelines import ForwardBatch, build_pipeline
+from fastvideo.pipelines import ForwardBatch, LoRAPipeline, build_pipeline
 from fastvideo.platforms import current_platform
 from fastvideo.utils import (get_exception_traceback,
                             kill_itself_when_parent_died)
@@ -116,6 +117,14 @@ class Worker:
                    local_main_process_only=False)
        return {"status": "shutdown_complete"}

+    def unmerge_lora_weights(self) -> None:
+        if isinstance(self.pipeline, LoRAPipeline):
+            self.pipeline.unmerge_lora_weights()
+
+    def merge_lora_weights(self) -> None:
+        if isinstance(self.pipeline, LoRAPipeline):
+            self.pipeline.merge_lora_weights()
+
    def event_loop(self) -> None:
        """Event loop for the worker."""
        logger.info("Worker %d starting event loop...",
@@ -139,7 +148,13 @@ class Worker:
                    fastvideo_args = recv_rpc['kwargs']['fastvideo_args']
                    output_batch = self.execute_forward(forward_batch,
                                                        fastvideo_args)
-                    self.pipe.send({"output_batch": output_batch.output.cpu()})
+                    logging_info = None
+                    if envs.FASTVIDEO_STAGE_LOGGING:
+                        logging_info = output_batch.logging_info
+                    self.pipe.send({
+                        "output_batch": output_batch.output.cpu(),
+                        "logging_info": logging_info
+                    })
                elif method_name == 'set_lora_adapter':
                    lora_nickname = recv_rpc['kwargs']['lora_nickname']
                    lora_path = recv_rpc['kwargs']['lora_path']
@@ -147,6 +162,14 @@ class Worker:
                    logger.info("Worker %d set LoRA adapter %s with path %s",
                                self.rank, lora_nickname, lora_path)
                    self.pipe.send({"status": "lora_adapter_set"})
+                elif method_name == 'unmerge_lora_weights':
+                    self.unmerge_lora_weights()
+                    logger.info("Worker %d unmerged LoRA weights", self.rank)
+                    self.pipe.send({"status": "lora_adapter_unmerged"})
+                elif method_name == 'merge_lora_weights':
+                    self.merge_lora_weights()
+                    logger.info("Worker %d merged LoRA weights", self.rank)
+                    self.pipe.send({"status": "lora_adapter_merged"})
                else:
                    # Handle other methods dynamically if needed
                    args = recv_rpc.get('args', ())
@@ -8,8 +8,9 @@ import socket
 import time
 from collections.abc import Callable
 from multiprocessing.process import BaseProcess
-from typing import Any, cast
+from typing import Any

+import fastvideo.envs as envs
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
@@ -40,7 +41,8 @@ class MultiprocExecutor(Executor):
            logger.info("Using provided master port: %s", self.master_port)
        else:
            # Auto-find available port
-            for port in range(29503, 65535):
+            import random
+            for port in range(29503 + random.randint(0, 10000), 65535):
                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    if s.connect_ex(('localhost', port)) != 0:
                        self.master_port = port
@@ -80,7 +82,19 @@ class MultiprocExecutor(Executor):
                                            "forward_batch": forward_batch,
                                            "fastvideo_args": fastvideo_args
                                        })
-        return cast(ForwardBatch, responses[0]["output_batch"])
+        output = responses[0]["output_batch"]
+
+        logging_info = None
+        if envs.FASTVIDEO_STAGE_LOGGING:
+            logging_info = responses[0]["logging_info"]
+        else:
+            logging_info = None
+
+        result_batch = ForwardBatch(data_type=forward_batch.data_type,
+                                    output=output,
+                                    logging_info=logging_info)
+
+        return result_batch

    def set_lora_adapter(self,
                         lora_nickname: str,
@@ -95,6 +109,18 @@ class MultiprocExecutor(Executor):
                raise RuntimeError(
                    f"Worker {i} failed to set LoRA adapter to {lora_path}")

+    def unmerge_lora_weights(self) -> None:
+        responses = self.collective_rpc("unmerge_lora_weights", kwargs={})
+        for i, response in enumerate(responses):
+            if response["status"] != "lora_adapter_unmerged":
+                raise RuntimeError(f"Worker {i} failed to unmerge LoRA weights")
+
+    def merge_lora_weights(self) -> None:
+        responses = self.collective_rpc("merge_lora_weights", kwargs={})
+        for i, response in enumerate(responses):
+            if response["status"] != "lora_adapter_merged":
+                raise RuntimeError(f"Worker {i} failed to merge LoRA weights")
+
    def collective_rpc(self,
                       method: str | Callable,
                       timeout: float | None = None,
@@ -1,3 +1,4 @@
+import dataclasses
 import gc
 import multiprocessing
 import os
@@ -10,7 +11,10 @@ import numpy as np
 import pyarrow as pa
 import pyarrow.parquet as pq
 import torch
+from datasets import Dataset, Video, load_dataset

+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

@@ -130,6 +134,9 @@ class PreprocessingDataValidator:

 class VideoForwardBatchBuilder:

+    def __init__(self, seed: int):
+        self.seed = seed
+
    def __call__(self, batch: list) -> PreprocessBatch:
        forward_batch = PreprocessBatch(
            video_loader=[item["video"] for item in batch],
@@ -138,78 +145,14 @@ class VideoForwardBatchBuilder:
            width=[item["resolution"]["width"] for item in batch],
            fps=[item["fps"] for item in batch],
            num_frames=[item["num_frames"] for item in batch],
-            prompt=[
-                item["caption"]
-                if isinstance(item["caption"], str) else item["caption"][0]
-                for item in batch
-            ],
+            prompt=[item["caption"] for item in batch],
            prompt_attention_mask=[],
            data_type="video",
+            generator=torch.Generator("cpu").manual_seed(self.seed),
        )
        return forward_batch


-def basic_record_creator(
-        video_name: str,
-        vae_latent: np.ndarray,
-        text_embedding: np.ndarray,
-        batch: PreprocessBatch,
-        idx: int,
-        extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
-    """Create a record for the Parquet dataset from PreprocessBatch."""
-    # For batch processing, we need to handle the case where some fields might be single values
-    # or lists depending on the batch size
-
-    # Get caption - handle both single string and list cases
-    caption = ""
-    if batch.prompt:
-        if isinstance(batch.prompt, list):
-            caption = batch.prompt[idx] if idx < len(
-                batch.prompt) else batch.prompt[0]
-        else:
-            caption = batch.prompt
-
-    # Get dimensions - these are typically single values in PreprocessBatch
-    assert isinstance(batch.width, list)
-    assert isinstance(batch.height, list)
-    assert isinstance(batch.fps, list)
-    assert isinstance(batch.num_frames, list)
-
-    width = batch.width[idx] if batch.width is not None else 0
-    height = batch.height[idx] if batch.height is not None else 0
-
-    # Get FPS - single value in PreprocessBatch
-    fps_val = float(batch.fps[idx]) if batch.fps is not None else 0.0
-
-    # For duration, we need to calculate it or use a default since it's not in PreprocessBatch
-    # duration = num_frames / fps if available
-    duration_val = 0.0
-    if batch.num_frames[idx] and batch.fps[idx] and batch.fps[idx] > 0:
-        duration_val = float(batch.num_frames[idx]) / float(batch.fps[idx])
-
-    record = {
-        "id": video_name,
-        "vae_latent_bytes": vae_latent.tobytes(),
-        "vae_latent_shape": list(vae_latent.shape),
-        "vae_latent_dtype": str(vae_latent.dtype),
-        "text_embedding_bytes": text_embedding.tobytes(),
-        "text_embedding_shape": list(text_embedding.shape),
-        "text_embedding_dtype": str(text_embedding.dtype),
-        "file_name": video_name,
-        "caption": caption,
-        "media_type": "video",
-        "width": int(width),
-        "height": int(height),
-        "num_frames": vae_latent.shape[1] if len(vae_latent.shape) > 1 else 0,
-        "duration_sec": duration_val,
-        "fps": fps_val,
-    }
-
-    if extra_features:
-        record.update(extra_features)
-    return record
-
-
 class ParquetDatasetSaver:
    """Component for saving and writing Parquet datasets"""

@@ -217,20 +160,20 @@ class ParquetDatasetSaver:
                 flush_frequency: int,
                 samples_per_file: int,
                 schema_fields: list[str],
-                 record_creator: Callable[..., dict[str, Any]] | None = None,
+                 record_creator: Callable[..., list[dict[str, Any]]],
                 file_writer_fn: Callable | None = None):
        """
        Initialize ParquetDatasetSaver
        
        Args:
-            schema_fields_provider: Function that returns schema fields list
+            schema_fields: schema fields list
            record_creator: Function for creating records
            file_writer_fn: Function for writing records to files, uses default implementation if None
        """
        self.flush_frequency = flush_frequency
        self.samples_per_file = samples_per_file
        self.schema_fields = schema_fields
-        self.create_record = record_creator or basic_record_creator
+        self.create_records_from_batch = record_creator
        self.file_writer_fn: Callable[
            [tuple], int] = file_writer_fn or self._default_file_writer_fn
        self.all_tables: list[pa.Table] = []
@@ -257,39 +200,29 @@ class ParquetDatasetSaver:
        assert isinstance(batch.prompt_embeds, list)
        assert isinstance(batch.prompt_attention_mask, list)

-        prompt_embeds: list[torch.Tensor] = []
        # Process non-padded embeddings (if needed)
        if batch.prompt_attention_mask is not None:
-            prompt_embeds = self._process_non_padded_embeddings(
+            batch.prompt_embeds = self._process_non_padded_embeddings(
                batch.prompt_embeds[0], batch.prompt_attention_mask[0])
        else:
            raise ValueError("prompt_attention_mask is None")

        # Prepare batch data for Parquet dataset
-        batch_data = []
+        batch_data: list[dict[str, Any]] = []

-        for idx, video_name in enumerate(batch.video_file_name):
-            # Get the corresponding latent and info using video name
-            vae_latent = batch.latents[idx].cpu().numpy()
-            text_embedding = prompt_embeds[idx].cpu().numpy()
+        for key in dataclasses.fields(batch):
+            value = getattr(batch, key.name)
+            if isinstance(value, list):
+                for idx in range(len(value)):
+                    if isinstance(value[idx], torch.Tensor):
+                        value[idx] = value[idx].cpu().numpy()
+            elif isinstance(value, torch.Tensor):
+                value = value.cpu().numpy()
+                setattr(batch, key.name, value)

-            # Get extra features for this sample if needed
-            sample_extra_features = {}
-            if extra_features:
-                for key, value in extra_features.items():
-                    if isinstance(value, torch.Tensor):
-                        sample_extra_features[key] = value[idx].cpu().numpy()
-                    else:
-                        sample_extra_features[key] = value[idx]
-
-            # Create record for Parquet dataset
-            record = self.create_record(video_name=video_name,
-                                        vae_latent=vae_latent,
-                                        text_embedding=text_embedding,
-                                        batch=batch,
-                                        idx=idx,
-                                        extra_features=sample_extra_features)
-            batch_data.append(record)
+        # Create record for Parquet dataset
+        records = self.create_records_from_batch(batch)
+        batch_data.extend(records)

        if batch_data:
            self.num_processed_samples += len(batch_data)
@@ -465,3 +398,42 @@ class ParquetDatasetSaver:
            written_count += len(chunk_table)

        return written_count
+
+
+def build_dataset(preprocess_config: PreprocessConfig, split: str,
+                  validator: Callable[[dict[str, Any]], bool]) -> Dataset:
+    if preprocess_config.dataset_type == DatasetType.HF:
+        dataset = load_dataset(preprocess_config.dataset_path, split=split)
+        dataset = dataset.filter(validator)
+        dataset = dataset.shard(num_shards=get_world_size(),
+                                index=get_world_rank())
+    elif preprocess_config.dataset_type == DatasetType.MERGED:
+        metadata_json_path = os.path.join(preprocess_config.dataset_path,
+                                          "videos2caption.json")
+        video_folder = os.path.join(preprocess_config.dataset_path, "videos")
+        dataset = load_dataset("json",
+                               data_files=metadata_json_path,
+                               split=split)
+        column_names = dataset.column_names
+        # rename columns to match the schema
+        if "cap" in column_names:
+            dataset = dataset.rename_column("cap", "caption")
+        if "path" in column_names:
+            dataset = dataset.rename_column("path", "name")
+
+        dataset = dataset.filter(validator)
+        dataset = dataset.shard(num_shards=get_world_size(),
+                                index=get_world_rank())
+
+        # add video column
+        def add_video_column(item: dict[str, Any]) -> dict[str, Any]:
+            item["video"] = os.path.join(video_folder, item["name"])
+            return item
+
+        dataset = dataset.map(add_video_column)
+        dataset = dataset.cast_column("video", Video())
+    else:
+        raise ValueError(
+            f"Invalid dataset type: {preprocess_config.dataset_type}")
+
+    return dataset
@@ -1,16 +1,24 @@
 import os
 from typing import cast

-from datasets import load_dataset
 from torch.utils.data import DataLoader

 from fastvideo.configs.configs import PreprocessConfig
+from fastvideo.dataset.dataloader.schema import (pyarrow_schema_i2v,
+                                                 pyarrow_schema_t2v)
+from fastvideo.distributed.parallel_state import get_world_rank
 from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
+from fastvideo.logger import init_logger
 from fastvideo.pipelines.pipeline_registry import PipelineType
 from fastvideo.workflow.preprocess.components import (
-    PreprocessingDataValidator, VideoForwardBatchBuilder)
+    ParquetDatasetSaver, PreprocessingDataValidator, VideoForwardBatchBuilder,
+    build_dataset)
+from fastvideo.workflow.preprocess.record_schema import (
+    basic_t2v_record_creator, i2v_record_creator)
 from fastvideo.workflow.workflow_base import WorkflowBase

+logger = init_logger(__name__)
+

 class PreprocessWorkflow(WorkflowBase):

@@ -36,9 +44,9 @@ class PreprocessWorkflow(WorkflowBase):
        self.add_component("raw_data_validator", raw_data_validator)

        # training dataset
-        training_dataset = load_dataset(preprocess_config.dataset_path,
-                                        split="train")
-        training_dataset = training_dataset.filter(raw_data_validator)
+        training_dataset = build_dataset(preprocess_config,
+                                         split="train",
+                                         validator=raw_data_validator)
        # we do not use collate_fn here because we use iterable-style Dataset
        # and want to keep the original type of the dataset
        training_dataloader = DataLoader(
@@ -49,35 +57,62 @@ class PreprocessWorkflow(WorkflowBase):
        )
        self.add_component("training_dataloader", training_dataloader)

-        # validation dataset
-        validation_dataset = load_dataset(preprocess_config.dataset_path,
-                                          split="validation")
-        validation_dataset = validation_dataset.filter(raw_data_validator)
-        validation_dataloader = DataLoader(
-            validation_dataset,
-            batch_size=preprocess_config.preprocess_video_batch_size,
-            num_workers=preprocess_config.dataloader_num_workers,
-            collate_fn=lambda x: x,
-        )
+        # try to load validation dataset if it exists
+        try:
+            validation_dataset = build_dataset(preprocess_config,
+                                               split="validation",
+                                               validator=raw_data_validator)
+            validation_dataloader = DataLoader(
+                validation_dataset,
+                batch_size=preprocess_config.preprocess_video_batch_size,
+                num_workers=preprocess_config.dataloader_num_workers,
+                collate_fn=lambda x: x,
+            )
+        except ValueError:
+            logger.warning(
+                "Validation dataset not found, skipping validation dataset preprocessing."
+            )
+            validation_dataloader = None
+
        self.add_component("validation_dataloader", validation_dataloader)

        # forward batch builder
-        video_forward_batch_builder = VideoForwardBatchBuilder()
+        video_forward_batch_builder = VideoForwardBatchBuilder(
+            seed=self.fastvideo_args.preprocess_config.seed)
        self.add_component("video_forward_batch_builder",
                           video_forward_batch_builder)

+        # record creator
+        if self.fastvideo_args.workload_type == WorkloadType.I2V:
+            record_creator = i2v_record_creator
+            schema_fields = [f.name for f in pyarrow_schema_i2v]
+        else:
+            record_creator = basic_t2v_record_creator
+            schema_fields = [f.name for f in pyarrow_schema_t2v]
+        processed_dataset_saver = ParquetDatasetSaver(
+            flush_frequency=self.fastvideo_args.preprocess_config.
+            flush_frequency,
+            samples_per_file=self.fastvideo_args.preprocess_config.
+            samples_per_file,
+            schema_fields=schema_fields,
+            record_creator=record_creator,
+        )
+        self.add_component("processed_dataset_saver", processed_dataset_saver)
+
    def prepare_system_environment(self) -> None:
        assert self.fastvideo_args.preprocess_config is not None
        dataset_output_dir = self.fastvideo_args.preprocess_config.dataset_output_dir
        os.makedirs(dataset_output_dir, exist_ok=True)

-        validation_dataset_output_dir = os.path.join(dataset_output_dir,
-                                                     "validation_dataset")
+        validation_dataset_output_dir = os.path.join(
+            dataset_output_dir, "validation_dataset",
+            f"worker_{get_world_rank()}")
        os.makedirs(validation_dataset_output_dir, exist_ok=True)
        self.validation_dataset_output_dir = validation_dataset_output_dir

-        training_dataset_output_dir = os.path.join(dataset_output_dir,
-                                                   "training_dataset")
+        training_dataset_output_dir = os.path.join(
+            dataset_output_dir, "training_dataset",
+            f"worker_{get_world_rank()}")
        os.makedirs(training_dataset_output_dir, exist_ok=True)
        self.training_dataset_output_dir = training_dataset_output_dir

@@ -1,14 +1,56 @@
-from fastvideo.fastvideo_args import FastVideoArgs
+from typing import TYPE_CHECKING
+
+from tqdm import tqdm
+
+from fastvideo.dataset.preprocessing_datasets import PreprocessBatch
+from fastvideo.workflow.preprocess.components import ParquetDatasetSaver
 from fastvideo.workflow.preprocess.preprocess_workflow import PreprocessWorkflow

+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+    from fastvideo.workflow.preprocess.components import (
+        VideoForwardBatchBuilder)
+

 class PreprocessWorkflowI2V(PreprocessWorkflow):
+    training_dataloader: "DataLoader"
+    validation_dataloader: "DataLoader"
+    preprocess_pipeline: "ComposedPipelineBase"
+    processed_dataset_saver: "ParquetDatasetSaver"
+    video_forward_batch_builder: "VideoForwardBatchBuilder"

-    def __init__(self, fastvideo_args: FastVideoArgs):
-        super().__init__(fastvideo_args)
+    def run(self) -> None:
+        # Training dataset preprocessing
+        for batch in tqdm(self.training_dataloader,
+                          desc="Preprocessing training dataset",
+                          unit="batch"):
+            forward_batch: PreprocessBatch = self.video_forward_batch_builder(
+                batch)

-    def register_pipelines(self) -> None:
-        pass
+            forward_batch = self.preprocess_pipeline.forward(
+                forward_batch, self.fastvideo_args)

-    def register_components(self) -> None:
-        pass
+            self.processed_dataset_saver.save_and_write_parquet_batch(
+                forward_batch, self.training_dataset_output_dir)
+
+        self.processed_dataset_saver.flush_tables(
+            self.training_dataset_output_dir)
+        self.processed_dataset_saver.clean_up()
+
+        # Validation dataset preprocessing
+        if self.validation_dataloader is not None:
+            for batch in tqdm(self.validation_dataloader,
+                              desc="Preprocessing validation dataset",
+                              unit="batch"):
+                forward_batch = self.video_forward_batch_builder(batch)
+
+                forward_batch = self.preprocess_pipeline.forward(
+                    forward_batch, self.fastvideo_args)
+
+                self.processed_dataset_saver.save_and_write_parquet_batch(
+                    forward_batch, self.validation_dataset_output_dir)
+            self.processed_dataset_saver.flush_tables(
+                self.validation_dataset_output_dir)
+            self.processed_dataset_saver.clean_up()
@@ -1,8 +1,7 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional

 from tqdm import tqdm

-from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
 from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
 from fastvideo.workflow.preprocess.components import ParquetDatasetSaver
 from fastvideo.workflow.preprocess.preprocess_workflow import PreprocessWorkflow
@@ -17,24 +16,11 @@ if TYPE_CHECKING:

 class PreprocessWorkflowT2V(PreprocessWorkflow):
    training_dataloader: "DataLoader"
-    validation_dataloader: "DataLoader"
+    validation_dataloader: Optional["DataLoader"]
    preprocess_pipeline: "ComposedPipelineBase"
    processed_dataset_saver: "ParquetDatasetSaver"
    video_forward_batch_builder: "VideoForwardBatchBuilder"

-    def register_components(self) -> None:
-        assert self.fastvideo_args.preprocess_config is not None
-        super().register_components()
-        self.add_component(
-            "processed_dataset_saver",
-            ParquetDatasetSaver(
-                flush_frequency=self.fastvideo_args.preprocess_config.
-                flush_frequency,
-                samples_per_file=self.fastvideo_args.preprocess_config.
-                samples_per_file,
-                schema_fields=[f.name for f in pyarrow_schema_t2v],
-            ))
-
    def run(self) -> None:
        # Training dataset preprocessing
        for batch in tqdm(self.training_dataloader,
@@ -54,16 +40,17 @@ class PreprocessWorkflowT2V(PreprocessWorkflow):
        self.processed_dataset_saver.clean_up()

        # Validation dataset preprocessing
-        for batch in tqdm(self.validation_dataloader,
-                          desc="Preprocessing validation dataset",
-                          unit="batch"):
-            forward_batch = self.video_forward_batch_builder(batch)
+        if self.validation_dataloader is not None:
+            for batch in tqdm(self.validation_dataloader,
+                              desc="Preprocessing validation dataset",
+                              unit="batch"):
+                forward_batch = self.video_forward_batch_builder(batch)

-            forward_batch = self.preprocess_pipeline.forward(
-                forward_batch, self.fastvideo_args)
+                forward_batch = self.preprocess_pipeline.forward(
+                    forward_batch, self.fastvideo_args)

-            self.processed_dataset_saver.save_and_write_parquet_batch(
-                forward_batch, self.validation_dataset_output_dir)
-        self.processed_dataset_saver.flush_tables(
-            self.validation_dataset_output_dir)
-        self.processed_dataset_saver.clean_up()
+                self.processed_dataset_saver.save_and_write_parquet_batch(
+                    forward_batch, self.validation_dataset_output_dir)
+            self.processed_dataset_saver.flush_tables(
+                self.validation_dataset_output_dir)
+            self.processed_dataset_saver.clean_up()
@@ -0,0 +1,122 @@
+from typing import Any
+
+from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
+
+
+def basic_t2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
+    """Create a record for the Parquet dataset from PreprocessBatch."""
+    # For batch processing, we need to handle the case where some fields might be single values
+    # or lists depending on the batch size
+
+    assert isinstance(batch.prompt, list)
+    assert isinstance(batch.width, list)
+    assert isinstance(batch.height, list)
+    assert isinstance(batch.fps, list)
+    assert isinstance(batch.num_frames, list)
+
+    records = []
+    for idx, video_name in enumerate(batch.video_file_name):
+        width = batch.width[idx] if batch.width is not None else 0
+        height = batch.height[idx] if batch.height is not None else 0
+
+        # Get FPS - single value in PreprocessBatch
+        fps_val = float(batch.fps[idx]) if batch.fps is not None else 0.0
+
+        # For duration, we need to calculate it or use a default since it's not in PreprocessBatch
+        # duration = num_frames / fps if available
+        duration_val = 0.0
+        if batch.num_frames[idx] and batch.fps[idx] and batch.fps[idx] > 0:
+            duration_val = float(batch.num_frames[idx]) / float(batch.fps[idx])
+
+        record = {
+            "id":
+            video_name,
+            "vae_latent_bytes":
+            batch.latents[idx].tobytes(),
+            "vae_latent_shape":
+            list(batch.latents[idx].shape),
+            "vae_latent_dtype":
+            str(batch.latents[idx].dtype),
+            "text_embedding_bytes":
+            batch.prompt_embeds[idx].tobytes(),
+            "text_embedding_shape":
+            list(batch.prompt_embeds[idx].shape),
+            "text_embedding_dtype":
+            str(batch.prompt_embeds[idx].dtype),
+            "file_name":
+            video_name,
+            "caption":
+            batch.prompt[idx],
+            "media_type":
+            "video",
+            "width":
+            int(width),
+            "height":
+            int(height),
+            "num_frames":
+            batch.latents[idx].shape[1]
+            if len(batch.latents[idx].shape) > 1 else 0,
+            "duration_sec":
+            duration_val,
+            "fps":
+            fps_val,
+        }
+        records.append(record)
+
+    return records
+
+
+def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
+    """Create a record for the Parquet dataset with CLIP features."""
+    records = basic_t2v_record_creator(batch)
+
+    assert len(
+        batch.image_embeds) == 1, "image embedding should be a single tensor"
+    image_embeds = batch.image_embeds[0]
+    image_latent = batch.image_latent
+    pil_image = batch.pil_image
+
+    for idx, record in enumerate(records):
+        if image_embeds is not None:
+            record.update({
+                "clip_feature_bytes": image_embeds[idx].tobytes(),
+                "clip_feature_shape": list(image_embeds[idx].shape),
+                "clip_feature_dtype": str(image_embeds[idx].dtype),
+            })
+        else:
+            record.update({
+                "clip_feature_bytes": b"",
+                "clip_feature_shape": [],
+                "clip_feature_dtype": "",
+            })
+
+        if image_latent is not None:
+            record.update({
+                "first_frame_latent_bytes":
+                image_latent[idx].tobytes(),
+                "first_frame_latent_shape":
+                list(image_latent[idx].shape),
+                "first_frame_latent_dtype":
+                str(image_latent[idx].dtype),
+            })
+        else:
+            record.update({
+                "first_frame_latent_bytes": b"",
+                "first_frame_latent_shape": [],
+                "first_frame_latent_dtype": "",
+            })
+
+        if pil_image is not None:
+            record.update({
+                "pil_image_bytes": pil_image[idx].tobytes(),
+                "pil_image_shape": list(pil_image[idx].shape),
+                "pil_image_dtype": str(pil_image[idx].dtype),
+            })
+        else:
+            record.update({
+                "pil_image_bytes": b"",
+                "pil_image_shape": [],
+                "pil_image_dtype": "",
+            })
+
+    return records
@@ -180,4 +180,7 @@ class WorkflowBase(ABC):
            from fastvideo.workflow.preprocess.preprocess_workflow import (
                PreprocessWorkflow)
            return PreprocessWorkflow.get_workflow_cls(fastvideo_args)
-        return None
+        else:
+            raise ValueError(
+                f"Execution mode: {fastvideo_args.mode} is not supported in workflow."
+            )
@@ -49,7 +49,7 @@ dependencies = [
    "av",

    # Preprocessing Dependencies
-    "torchcodec"
+    "torchcodec==0.5.0"
 ]

 [tool.uv]
@@ -171,4 +171,4 @@ Repository = "https://github.com/hao-ai-lab/FastVideo"
 [tool.comfy]
 PublisherId = "fastvideo"
 DisplayName = "FastVideo"
-Icon = "https://raw.githubusercontent.com/hao-ai-lab/FastVideo/main/comfyui/assets/logo.png"
+Icon = "https://raw.githubusercontent.com/hao-ai-lab/FastVideo/main/comfyui/assets/icon_simple.svg"
@@ -0,0 +1,161 @@
+from huggingface_hub import save_torch_state_dict, load_state_dict_from_file
+# from safetensors import safetensors
+from safetensors.torch import save_file
+import torch
+import re
+from collections import OrderedDict
+
+_param_names_mapping: dict = {
+    r"^text_embedding\.0\.(.*)$":
+    r"condition_embedder.text_embedder.linear_1.\1",
+    r"^text_embedding\.2\.(.*)$":
+    r"condition_embedder.text_embedder.linear_2.\1",
+    r"^time_embedding\.0\.(.*)$":
+    r"condition_embedder.time_embedder.linear_1.\1",
+    r"^time_embedding\.2\.(.*)$":
+    r"condition_embedder.time_embedder.linear_2.\1",
+    r"^time_projection\.1\.(.*)$":
+    r"condition_embedder.time_proj.\1",
+    r"^img_emb\.proj\.0\.(.*)$":
+    r"condition_embedder.image_embedder.norm1.\1",
+    r"^img_emb\.proj\.1\.(.*)$":
+    r"condition_embedder.image_embedder.ff.net.0.proj.\1",
+    r"^img_emb\.proj\.3\.(.*)$":
+    r"condition_embedder.image_embedder.ff.net.2.\1",
+    r"^img_emb\.proj\.4\.(.*)$":
+    r"condition_embedder.image_embedder.norm2.\1",
+    r"^head\.modulation":
+    r"scale_shift_table",
+    r"^head\.head\.(.*)$":
+    r"proj_out.\1",
+    r"^blocks\.(\d+)\.self_attn\.q\.(.*)$":
+    r"blocks.\1.attn1.to_q.\2",
+    r"^blocks\.(\d+)\.self_attn\.k\.(.*)$":
+    r"blocks.\1.attn1.to_k.\2",
+    r"^blocks\.(\d+)\.self_attn\.v\.(.*)$":
+    r"blocks.\1.attn1.to_v.\2",
+    r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
+    r"blocks.\1.attn1.to_out.0.\2",
+    r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
+    r"blocks.\1.attn1.norm_q.\2",
+    r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
+    r"blocks.\1.attn1.norm_k.\2",
+    r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$":
+    r"blocks.\1.attn2.to_q.\2",
+    r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$":
+    r"blocks.\1.attn2.to_k.\2",
+    r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
+    r"blocks.\1.attn2.add_k_proj.\2",
+    r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$":
+    r"blocks.\1.attn2.to_v.\2",
+    r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
+    r"blocks.\1.attn2.add_v_proj.\2",
+    r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
+    r"blocks.\1.attn2.to_out.0.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
+    r"blocks.\1.attn2.norm_q.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
+    r"blocks.\1.attn2.norm_k.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
+    r"blocks.\1.attn2.norm_added_k.\2",
+    r"^blocks\.(\d+)\.ffn\.0\.(.*)$":
+    r"blocks.\1.ffn.net.0.proj.\2",
+    r"^blocks\.(\d+)\.ffn\.2\.(.*)$":
+    r"blocks.\1.ffn.net.2.\2",
+    r"^blocks\.(\d+)\.modulation":
+    r"blocks.\1.scale_shift_table",
+    r"^blocks\.(\d+)\.norm3\.(.*)$":
+    r"blocks.\1.norm2.\2",
+}
+
+# The following mapping has an extra 'patch_embedding' field and also contains
+# the 'model' prefixes
+_self_forcing_to_diffusers_param_names_mapping: dict = {
+    r"^model.patch_embedding\.(.*)$":
+    r"patch_embedding.\1",
+    r"^model.text_embedding\.0\.(.*)$":
+    r"condition_embedder.text_embedder.linear_1.\1",
+    r"^model.text_embedding\.2\.(.*)$":
+    r"condition_embedder.text_embedder.linear_2.\1",
+    r"^model.time_embedding\.0\.(.*)$":
+    r"condition_embedder.time_embedder.linear_1.\1",
+    r"^model.time_embedding\.2\.(.*)$":
+    r"condition_embedder.time_embedder.linear_2.\1",
+    r"^model.time_projection\.1\.(.*)$":
+    r"condition_embedder.time_proj.\1",
+    r"^model.img_emb\.proj\.0\.(.*)$":
+    r"condition_embedder.image_embedder.norm1.\1",
+    r"^model.img_emb\.proj\.1\.(.*)$":
+    r"condition_embedder.image_embedder.ff.net.0.proj.\1",
+    r"^model.img_emb\.proj\.3\.(.*)$":
+    r"condition_embedder.image_embedder.ff.net.2.\1",
+    r"^model.img_emb\.proj\.4\.(.*)$":
+    r"condition_embedder.image_embedder.norm2.\1",
+    r"^model.head\.modulation":
+    r"scale_shift_table",
+    r"^model.head\.head\.(.*)$":
+    r"proj_out.\1",
+    r"^model.blocks\.(\d+)\.self_attn\.q\.(.*)$":
+    r"blocks.\1.attn1.to_q.\2",
+    r"^model.blocks\.(\d+)\.self_attn\.k\.(.*)$":
+    r"blocks.\1.attn1.to_k.\2",
+    r"^model.blocks\.(\d+)\.self_attn\.v\.(.*)$":
+    r"blocks.\1.attn1.to_v.\2",
+    r"^model.blocks\.(\d+)\.self_attn\.o\.(.*)$":
+    r"blocks.\1.attn1.to_out.0.\2",
+    r"^model.blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
+    r"blocks.\1.attn1.norm_q.\2",
+    r"^model.blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
+    r"blocks.\1.attn1.norm_k.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.q\.(.*)$":
+    r"blocks.\1.attn2.to_q.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.k\.(.*)$":
+    r"blocks.\1.attn2.to_k.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
+    r"blocks.\1.attn2.add_k_proj.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.v\.(.*)$":
+    r"blocks.\1.attn2.to_v.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
+    r"blocks.\1.attn2.add_v_proj.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.o\.(.*)$":
+    r"blocks.\1.attn2.to_out.0.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
+    r"blocks.\1.attn2.norm_q.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
+    r"blocks.\1.attn2.norm_k.\2",
+    r"^model.blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
+    r"blocks.\1.attn2.norm_added_k.\2",
+    r"^model.blocks\.(\d+)\.ffn\.0\.(.*)$":
+    r"blocks.\1.ffn.net.0.proj.\2",
+    r"^model.blocks\.(\d+)\.ffn\.2\.(.*)$":
+    r"blocks.\1.ffn.net.2.\2",
+    r"^model.blocks\.(\d+)\.modulation":
+    r"blocks.\1.scale_shift_table",
+    r"^model.blocks\.(\d+)\.norm3\.(.*)$":
+    r"blocks.\1.norm2.\2",
+}
+
+state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt")
+state_dict = state_dict["generator_ema"]
+new_state_dict = OrderedDict()
+for k, v in state_dict.items():
+    new_key = k
+    for pattern, replacement in _self_forcing_to_diffusers_param_names_mapping.items():
+        if re.match(pattern, k):
+            new_key = re.sub(pattern, replacement, k)
+            break  # Stop at the first match
+    else:
+        # print(f"No match found for {k}")
+        raise ValueError(f"No match found for {k}")
+    new_state_dict[new_key] = v
+    if "norm_added_k" in new_key:
+        dummy_key = new_key.replace("norm_added_k", "norm_added_q")
+        dummy_value = torch.zeros_like(v)
+        new_state_dict[dummy_key] = dummy_value
+del state_dict
+
+save_torch_state_dict(
+    new_state_dict,
+    "new2/",
+    max_shard_size="10GB"
+)
@@ -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",
 )
Author	SHA1	Message	Date
Matthew Noto	4477dbea67	training refactor done	2025-09-08 03:48:01 +00:00
Matthew Noto	04685a3ecd	kv cache debugging	2025-09-08 02:27:19 +00:00
Matthew Noto	b8392e9b2a	revert	2025-09-07 23:38:02 +00:00
Matthew Noto	60b71c5053	new branch	2025-09-07 23:36:43 +00:00
JerryZhou54	2b7bf88a4d	new branch	2025-09-07 08:49:36 +00:00
RandNMR73	328eb611c4	training refactor	2025-09-07 02:40:19 +00:00
RandNMR73	3408e20d7a	Merge branch 'matthew/causal' of github.com:hao-ai-lab/FastVideo into matthew/causal	2025-09-07 00:17:20 +00:00
SolitaryThinker	4328fe1ebf	fix warping	2025-09-05 10:20:10 +00:00
SolitaryThinker	6d0eba5789	warp timestep	2025-09-05 09:54:35 +00:00
RandNMR73	b3e76ae7dd	Merge branch 'matthew/causal' of github.com:hao-ai-lab/FastVideo into matthew/causal	2025-09-02 00:57:20 +00:00
Matthew Noto	1ffd80ee51	training experiments in progress	2025-09-02 00:03:58 +00:00
Matthew Noto	e84fdaedde	add single example preprocessing	2025-09-01 03:49:33 +00:00
Matthew Noto	dd0fe401c9	fix kv-cache + training loop	2025-09-01 01:52:11 +00:00
Matthew Noto	60eac9f18b	clean-up + kv cache debugging	2025-08-31 08:31:19 +00:00
Matthew Noto	a75bccb75d	debugging kv cache	2025-08-31 05:30:21 +00:00
RandNMR73	c52ff91747	Merge branch 'matthew/causal' of github.com:hao-ai-lab/FastVideo into matthew/causal	2025-08-28 12:21:18 +00:00
RandNMR73	36bb0935a8	val data path fix	2025-08-28 12:18:41 +00:00
RandNMR73	d1e26abd63	Merge branch 'matthew/causal' of github.com:hao-ai-lab/FastVideo into matthew/causal	2025-08-28 11:53:46 +00:00
RandNMR73	b6f187f338	working training script	2025-08-28 11:32:53 +00:00
RandNMR73	de4938c3a7	working training script	2025-08-28 11:27:18 +00:00
RandNMR73	a9f7407228	debugging training	2025-08-28 08:06:31 +00:00
RandNMR73	258d1da0d3	training script testing	2025-08-28 05:03:12 +00:00
RandNMR73	17f6dff632	training branch	2025-08-28 04:49:51 +00:00
RandNMR73	e66f16057f	training script testing	2025-08-28 04:42:51 +00:00
JerryZhou54	8acc7c3655	small fix	2025-08-28 04:10:51 +00:00
JerryZhou54	8dd0b6536d	Fix pre-commit tests	2025-08-28 04:07:11 +00:00
JerryZhou54	1a79b30ea4	Fix inter-block issues	2025-08-28 03:56:37 +00:00
JerryZhou54	772ead0d34	Wan2.1 causal, few-step inference runnable	2025-08-28 03:56:37 +00:00
SolitaryThinker	1575102965	debugging model	2025-08-28 03:56:37 +00:00
SolitaryThinker	b198ba5607	checkpoint	2025-08-28 03:56:34 +00:00
William Lin	cf1942fd47	[dev] Will/causal infer (#764 )	2025-08-28 03:55:16 +00:00
Wei Zhou	b5519f1f91	Delete fastvideo/configs/models/dits/causal_wanvideo.py	2025-08-28 03:55:16 +00:00
JerryZhou54	a464f96b95	Fix file dir	2025-08-28 03:55:16 +00:00
JerryZhou54	363cf0d173	First version of causalwan	2025-08-28 03:55:15 +00:00
JerryZhou54	36371c5689	checkpoint	2025-08-28 03:55:15 +00:00
Zhang Peiyuan	7c554e5da8	Update Community Link (#765 )	2025-08-27 16:12:47 -07:00
SolitaryThinker	8fea7c02b5	debugging model	2025-08-27 02:56:31 -07:00
SolitaryThinker	e2b6f49879	checkpoint	2025-08-27 02:12:43 -07:00
William Lin	9f24aef7cf	[dev] Will/causal infer (#764 )	2025-08-26 23:13:20 -07:00
William Lin	663ea33ff1	[bugfix] Fix wrong HF model string for FastWan2.2 5B (#763 )	2025-08-26 22:05:40 -07:00
Wei Zhou	7a489da74d	Delete fastvideo/configs/models/dits/causal_wanvideo.py	2025-08-26 16:33:25 -07:00
JerryZhou54	cf230dcccd	Fix file dir	2025-08-26 23:32:30 +00:00
JerryZhou54	c4521e8953	First version of causalwan	2025-08-26 23:31:46 +00:00
JerryZhou54	026ee8d9f4	checkpoint	2025-08-26 19:49:57 +00:00
William Lin	3ef04f1654	[misc] [docs] Various fixes for logging and docs (#758 )	2025-08-23 21:13:50 -07:00
Jinzhe Pan	0eced76a41	[Feat][Preprocess] support multi-gpus (#753 )	2025-08-23 11:34:42 +08:00
Jinzhe Pan	3ab6470d1a	[Feat][Preprocess] support merged dataset (#752 )	2025-08-22 15:29:33 -07:00
Wenxuan Tan	989a03532c	Optionally use unmerged weights for inference (#745 )	2025-08-22 15:20:31 -07:00
William Lin	fa15369a02	[bugfix] Check that model_index.json module is in required_modules list before removing (#756 )	2025-08-22 14:36:44 -07:00
Zhang Peiyuan	78a9cb88d8	[Fix] fix seed in dmd denoising loop (#736 )	2025-08-21 18:06:16 -07:00
Peng Xiaoand肖鹏	a0bff12746	[bugfix] [dmd] Align backward simulation with dmd2 sample back (#744 ) Co-authored-by: 肖鹏 <xiaopeng1@aishi.ai>	2025-08-20 22:25:33 -07:00
William Lin	98f2af94e5	[bugfix] Missing Docker file for cuda12.9 (#750 )	2025-08-20 15:34:31 -07:00
William Lin	46f7b6d574	[Docker] add 12.9 docker image and also fix py3.10 and py3.11 dockerfile (#749 )	2025-08-20 15:31:15 -07:00
Jinzhe Pan	911a6a6a35	[Feat][Preprocessing] i2v preprocessing workflow (#737 )	2025-08-14 20:47:25 -07:00
Zhang Peiyuan	38c7949d5c	Update WeChat group link (#739 )	2025-08-14 15:03:35 -07:00
Jinzhe Pan	7e7a0dba9d	feat: preprocess validation dataset only when exist (#734 )	2025-08-12 02:16:31 -07:00
Zhang Peiyuan	f62e210ae6	Fix vsa backward gQ (#735 )	2025-08-11 21:43:13 -07:00
William Lin	6ceb4942a0	[bugfix] [dmd] Fix backward simulation and also naming in wan_i2v_dmd_pipeline (#731 )	2025-08-10 21:13:30 -07:00
William LinandRandNMR73	8cae5e4708	[feature] add Gradio live serving demo code (#727 ) Co-authored-by: RandNMR73 <notomatthew31@gmail.com>	2025-08-10 15:34:03 -07:00