chekcpoint

Co-authored-by: 肖鹏 <xiaopeng1@aishi.ai>

.github/workflows/build-image.yml

@@ -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

.pre-commit-config.yaml

@@ -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
   )

README.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/wZPZTLKg" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">

assets/logos/icon_simple.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>

assets/logos/logo.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>

comfyui/assets/icon_simple.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>

csrc/attn/tests/test_vsa.py

@@ -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.")

csrc/attn/vsa/block_sparse_h100.cu

@@ -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();
-}
+}

docker/Dockerfile.python3.10

@@ -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

docker/Dockerfile.python3.11

@@ -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

docker/Dockerfile.python3.12

@@ -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 . .
 

docker/Dockerfile.python3.12.cuda12.9.1

@@ -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

docs/source/conf.py

@@ -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/',

docs/source/contributing/developer_env/docker.md

@@ -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:

docs/source/contributing/developer_env/runpod.md

@@ -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)
 

docs/source/contributing/overview.md

@@ -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
+```

docs/source/getting_started/installation/gpu.md

@@ -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

docs/source/getting_started/installation/mps.md

@@ -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:

docs/source/index.md

@@ -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}

docs/source/inference/inference_quick_start.md

@@ -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

docs/source/inference/support_matrix.md

@@ -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
 

docs/source/video_sparse_attention/demo.md

@@ -1,3 +0,0 @@
-(vsa-demo)=
-
-# 🎬 Demo

docs/source/video_sparse_attention/installation.md

@@ -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
 ```

examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/README.md

@@ -1,16 +0,0 @@
-# Wan2.1-I2V-1.3B-InP Crush-Smol Example
-These are e2e example scripts for finetuning Wan2.1 T2V 1.3B InP on the crush-smol dataset.
-
-## Execute the following commands from `FastVideo/` to run training:
-
-### Download crush-smol dataset:
-
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
-
-### Preprocess the videos and captions into latents:
-
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
-
-### Edit the following file and run finetuning:
-
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`

examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/distill_i2v.sh

@@ -1,111 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
-DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp_single/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
-NUM_GPUS=8
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_i2v_distill"
-  --output_dir "checkpoints/wan_i2v_distill"
-  --wandb_run_name "wan_i2v_distill"
-  --max_train_steps 2000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 2
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  # --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 8
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  # --log_visualization
-  --log_visualization
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
-  --validation_sampling_steps "4"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 4e-6
-  --lr_scheduler "constant"
-  --fake_score_learning_rate 8e-7
-  --fake_score_lr_scheduler "constant"
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 2000
-  --weight_only_checkpointing_steps 2000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 8
-  --seed 1000
-  # --enable_gradient_checkpointing_type "full"
-)
-
-dmd_args=(
-  --dmd_denoising_steps '1000,750,500,250'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --generator_update_interval 5
-  --simulate_generator_forward
-  --real_score_guidance_scale 5
-  --VSA_sparsity 0.8
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_i2v_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"

examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/download_dataset.sh

@@ -1,3 +0,0 @@
-#!/bin/bash
-
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh

@@ -1,111 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
-NUM_GPUS=8
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-export HF_HUB_ENABLE_HF_TRANSFER=1
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_i2v_distill"
-  --output_dir "checkpoints/wan_i2v_distill"
-  --wandb_run_name "wan_i2v_distill"
-  --max_train_steps 2000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 2
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  # --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 8
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  # --log_visualization
-  --log_visualization
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 4e-6
-  --lr_scheduler "constant"
-  --fake_score_learning_rate 8e-7
-  --fake_score_lr_scheduler "constant"
-  --mixed_precision "bf16"
-  --checkpointing_steps 2000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 8
-  --seed 1000
-  # --enable_gradient_checkpointing_type "full"
-)
-
-dmd_args=(
-  --dmd_denoising_steps '1000,757,522'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --generator_update_interval 5
-  --simulate_generator_forward
-  --real_score_guidance_scale 3.5
-  --VSA_sparsity 0.8
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_i2v_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"

examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm

@@ -1,129 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=i2v
-#SBATCH --partition=main
-#SBATCH --nodes=4
-#SBATCH --ntasks=4
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=i2v_output/i2v_%j.out
-#SBATCH --error=i2v_output/i2v_%j.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-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=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-
-MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
-DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_i2v_finetune
-  --output_dir "checkpoints/wan_i2v_finetune"
-  --max_train_steps 2000
-  --train_batch_size 2
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size $NUM_GPUS
-  --tp_size $NUM_GPUS
-  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
-  --hsdp_shard_dim $NUM_GPUS
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 10
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 100
-  --validation_sampling_steps "40"
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-  --enable_gradient_checkpointing_type "full"
-)
-
-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_i2v_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"

examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/preprocess_wan_data_i2v.sh

@@ -1,24 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
-MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_i2v_1_3b_inp/"
-
-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 77 \
-    --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 "i2v" 

examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json

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

examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_old.json

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

examples/distill/Wan2.1-I2V/crush_smol/README.md

@@ -1,16 +0,0 @@
-# Wan2.1-I2V-1.3B-InP Crush-Smol Example
-These are e2e example scripts for finetuning Wan2.1 T2V 1.3B InP on the crush-smol dataset.
-
-## Execute the following commands from `FastVideo/` to run training:
-
-### Download crush-smol dataset:
-
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
-
-### Preprocess the videos and captions into latents:
-
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
-
-### Edit the following file and run finetuning:
-
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`

examples/distill/Wan2.1-I2V/crush_smol/distill_i2v_14B.sh

@@ -1,122 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-# MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
-DATA_DIR="data/crush-smol_processed_wan21_i2v_14b/combined_parquet_dataset/"
-# DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/distill/Wan2.1-I2V/crush_smol/validation_orig.json"
-NUM_GPUS=8
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_i2v_distill"
-  --output_dir "checkpoints/wan_i2v_distill"
-  --wandb_run_name "14b_i2v_cm"
-  --max_train_steps 1500
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 8
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  # --log_visualization
-  --log_visualization
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 20
-  --validation_sampling_steps "3"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 4e-6
-  --lr_scheduler "constant"
-  # --min_lr_ratio 0.5
-  # --lr_warmup_steps 50
-  --fake_score_learning_rate 7e-7
-  --fake_score_lr_scheduler "constant"
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 2000
-  --weight_only_checkpointing_steps 2000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 8
-  --seed 1000
-  # --enable_gradient_checkpointing_type "full"
-)
-
-dmd_args=(
-  --dmd_denoising_steps '1000,750,500,250'
-  --warp_denoising_step True
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.96
-  --generator_update_interval 5
-  --simulate_generator_forward
-  --simulate_forward_interval 1
-  --real_score_guidance_scale 5
-  --VSA_sparsity 0.8
-  --regression_loss_weight 0.01
-  --use_regression_loss False
-  --cm_loss_weight 1
-  --ema_decay 0.999
-  --cm_use_ema_teacher
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_i2v_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"

examples/distill/Wan2.1-I2V/crush_smol/download_dataset.sh

@@ -1,3 +0,0 @@
-#!/bin/bash
-
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/distill/Wan2.1-I2V/crush_smol/finetune_i2v.sh

@@ -1,111 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-DATA_DIR="data/crush-smol_processed_wan21_i2v_14b/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/distill/Wan2.1-I2V/crush_smol/validation.json"
-NUM_GPUS=8
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-export HF_HUB_ENABLE_HF_TRANSFER=1
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_i2v_distill"
-  --output_dir "checkpoints/wan_i2v_distill"
-  --wandb_run_name "wan21_i2v_14b_ode_50steps"
-  --max_train_steps 2000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 2
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  # --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 8
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 8
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  # --log_visualization
-  --log_visualization
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 4e-6
-  --lr_scheduler "constant"
-  --fake_score_learning_rate 8e-7
-  --fake_score_lr_scheduler "constant"
-  --mixed_precision "bf16"
-  --checkpointing_steps 2000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 8
-  --seed 1000
-  # --enable_gradient_checkpointing_type "full"
-)
-
-dmd_args=(
-  --dmd_denoising_steps '1000,757,522'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --generator_update_interval 5
-  --simulate_generator_forward
-  --real_score_guidance_scale 3.5
-  --VSA_sparsity 0.8
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_i2v_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"

examples/distill/Wan2.1-I2V/crush_smol/finetune_i2v.slurm

@@ -1,129 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=i2v
-#SBATCH --partition=main
-#SBATCH --nodes=4
-#SBATCH --ntasks=4
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=i2v_output/i2v_%j.out
-#SBATCH --error=i2v_output/i2v_%j.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-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=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-
-MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
-DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_i2v_finetune
-  --output_dir "checkpoints/wan_i2v_finetune"
-  --max_train_steps 2000
-  --train_batch_size 2
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size $NUM_GPUS
-  --tp_size $NUM_GPUS
-  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
-  --hsdp_shard_dim $NUM_GPUS
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 10
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 100
-  --validation_sampling_steps "40"
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-  --enable_gradient_checkpointing_type "full"
-)
-
-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_i2v_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"

examples/distill/Wan2.1-I2V/crush_smol/preprocess_wan_data_i2v.sh

@@ -1,26 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
-MODEL_TYPE="wan"
-# DATA_MERGE_PATH="data/single/merge.txt"
-DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_wan21_t2v_14b/"
-# OUTPUT_DIR="data/single_processed_i2v_14b/"
-
-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 77 \
-    --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" 

examples/distill/Wan2.1-I2V/crush_smol/validation.json

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

examples/distill/Wan2.1-I2V/crush_smol/validation_orig.json

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

examples/distill/consistency_distillation/Wan2.1-I2V/crush_smol/distill_i2v_14B.sh

@@ -1,118 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
-# MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
-DATA_DIR="data/crush-smol_processed_wan21_t2v_14b/combined_parquet_dataset/"
-# DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/distill/Wan2.1-I2V/crush_smol/validation_orig.json"
-NUM_GPUS=8
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_i2v_distill"
-  --output_dir "checkpoints/wan_i2v_distill"
-  --wandb_run_name "cm_t2v"
-  --max_train_steps 1500
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 8
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  # --log_visualization
-  --log_visualization
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 20
-  --validation_sampling_steps "3"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 4e-6
-  --lr_scheduler "constant"
-  # --min_lr_ratio 0.5
-  # --lr_warmup_steps 50
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 2000
-  --weight_only_checkpointing_steps 2000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 8
-  --seed 1000
-  # --enable_gradient_checkpointing_type "full"
-)
-
-dmd_args=(
-  --dmd_denoising_steps '1000,750,500,250'
-  --warp_denoising_step True
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.96
-  --real_score_guidance_scale 5
-  --VSA_sparsity 0.8
-  --regression_loss_weight 0.01
-  --use_regression_loss False
-  --cm_loss_weight 1
-  --ema_decay 0.999
-  --cm_weighing_function "constant"
-  # --cm_use_ema_teacher
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_cm_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"

examples/inference/gradio/README.md

@@ -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

examples/inference/gradio/gradio_demo.py

@@ -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)

examples/inference/gradio/gradio_frontend.py

@@ -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() 

examples/inference/gradio/ray_serve_backend.py

@@ -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) 

examples/inference/gradio/start.sh

@@ -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"

examples/inference/gradio/start_ray_serve_app.py

@@ -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()

examples/inference/lora/wan_lora_inference_from_ckpt.py

@@ -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,

examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh

@@ -9,19 +9,17 @@ MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
 DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
 NUM_GPUS=8
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # Training arguments
 training_args=(
   --tracker_project_name "wan_i2v_finetune"
-  --output_dir "outputs/wan_i2v_finetune"
-  --wandb_run_name "wan_i2v_finetune"
+  --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
   --max_train_steps 2000
-  --train_batch_size 1
+  --train_batch_size 4
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 2
+  --gradient_accumulation_steps 1
   --num_latent_t 8
   --num_height 480
   --num_width 832
@@ -32,10 +30,10 @@ training_args=(
 # Parallel arguments
 parallel_args=(
   --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 8
+  --sp_size 4
+  --tp_size 4
+  --hsdp_replicate_dim 2
+  --hsdp_shard_dim 4
 )
 
 # Model arguments
@@ -61,7 +59,7 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 2e-5
   --mixed_precision "bf16"
   --checkpointing_steps 2000
   --weight_decay 1e-4

examples/training/finetune/wan_i2v_14B_480p/crush_smol/preprocess_wan_data_i2v_new.sh

@@ -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

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

@@ -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 \

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

@@ -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 \

fastvideo/configs/configs.py

@@ -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):

fastvideo/configs/pipelines/registry.py

@@ -34,7 +34,6 @@ 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,
-    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
     # Add other specific weight variants
 }
 

fastvideo/configs/sample/registry.py

@@ -17,7 +17,6 @@ from fastvideo.configs.sample.wan import (
     WanI2V_14B_720P_SamplingParam,
     WanT2V_1_3B_SamplingParam,
     WanT2V_14B_SamplingParam,
-    Wan2_1_Fun_1_3B_InP_SamplingParam,
 )
 # isort: on
 from fastvideo.logger import init_logger
@@ -39,8 +38,6 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "FastVideo/FastWan2.2-TI2V-5B-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,
-    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
-    Wan2_1_Fun_1_3B_InP_SamplingParam,
     # Add other specific weight variants
 }
 

fastvideo/configs/sample/wan.py

@@ -107,28 +107,13 @@ class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
     fps: int = 16
 
 
-# =============================================
-# ============= Wan2.1 Fun Models =============
-# =============================================
-@dataclass
-class Wan2_1_Fun_1_3B_InP_SamplingParam(SamplingParam):
-    """Sampling parameters for Wan2.1 Fun 1.3B InP model."""
-    height: int = 480
-    width: int = 832
-    num_frames: int = 81
-    fps: int = 16
-    negative_prompt: str = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
-    guidance_scale: float = 6.0
-    num_inference_steps: int = 50
-
-
 # =============================================
 # ============= Wan2.2 TI2V Models =============
 # =============================================
 @dataclass
 class Wan2_2_Base_SamplingParam(SamplingParam):
     """Sampling parameters for Wan2.2 TI2V 5B model."""
-    negative_prompt: str = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
 
 
 @dataclass

fastvideo/entrypoints/video_generator.py

@@ -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.

fastvideo/fastvideo_args.py

@@ -161,9 +161,6 @@ class FastVideoArgs:
     # MoE parameters used by Wan2.2
     boundary_ratio: float | None = None
 
-    # scheduler parameters
-    warp_denoising_step: bool = False
-
     @property
     def training_mode(self) -> bool:
         return not self.inference_mode
@@ -385,15 +382,6 @@ class FastVideoArgs:
             default=FastVideoArgs.enable_stage_verification,
             help="Enable input/output verification for pipeline stages",
         )
-
-        # scheduler parameters
-        parser.add_argument(
-            "--warp-denoising-step",
-            action=StoreBoolean,
-            default=FastVideoArgs.warp_denoising_step,
-            help="Warp denoising step for scheduler",
-        )
-
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 
@@ -604,6 +592,8 @@ class TrainingArgs(FastVideoArgs):
     dit_model_name_or_path: str = ""
 
     # diffusion setting
+    ema_decay: float = 0.0
+    ema_start_step: int = 0
     training_cfg_rate: float = 0.0
     precondition_outputs: bool = False
 
@@ -684,16 +674,6 @@ class TrainingArgs(FastVideoArgs):
     log_visualization: bool = False
     # simulate generator forward to match inference
     simulate_generator_forward: bool = False
-    simulate_forward_interval: int = 1
-    regression_loss_weight: float = 0.0
-    use_regression_loss: bool = False
-
-    # consistency model parameters
-    cm_loss_weight: float = 0.0
-    cm_use_ema_teacher: bool = False
-    ema_decay: float = 0.999
-    ema_start_step: int = 0
-    cm_weighing_function: str = "constant"
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -796,6 +776,14 @@ class TrainingArgs(FastVideoArgs):
                             help="Directory to cache models")
 
         # Diffusion settings
+        parser.add_argument("--ema-decay",
+                            type=float,
+                            default=0.999,
+                            help="EMA decay rate")
+        parser.add_argument("--ema-start-step",
+                            type=int,
+                            default=0,
+                            help="Step to start EMA")
         parser.add_argument("--training-cfg-rate",
                             type=float,
                             help="Classifier-free guidance scale")
@@ -1030,40 +1018,6 @@ class TrainingArgs(FastVideoArgs):
             "--simulate-generator-forward",
             action=StoreBoolean,
             help="Whether to simulate generator forward to match inference")
-        parser.add_argument("--simulate-forward-interval",
-                            type=int,
-                            default=TrainingArgs.simulate_forward_interval,
-                            help="Ratio of steps to simulate generator forward")
-        parser.add_argument("--regression-loss-weight",
-                            type=float,
-                            default=TrainingArgs.regression_loss_weight,
-                            help="Weight for regression loss")
-        parser.add_argument("--use-regression-loss",
-                            action=StoreBoolean,
-                            help="Whether to use regression loss")
-
-        # Consistency model arguments
-        parser.add_argument("--ema-decay",
-                            type=float,
-                            default=0.999,
-                            help="EMA decay rate")
-        parser.add_argument("--ema-start-step",
-                            type=int,
-                            default=0,
-                            help="Step to start EMA")
-        parser.add_argument("--cm-loss-weight",
-                            type=float,
-                            default=TrainingArgs.cm_loss_weight,
-                            help="Weight for consistency loss")
-        parser.add_argument(
-            "--cm-use-ema-teacher",
-            action=StoreBoolean,
-            help="Whether to use EMA teacher for consistency loss")
-        parser.add_argument("--cm-weighing-function",
-                            type=str,
-                            choices=["constant", "sigma_sqrt"],
-                            default=TrainingArgs.cm_weighing_function,
-                            help="Weighting function for consistency loss")
 
         return parser
 

fastvideo/layers/lora/linear.py

@@ -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)))

fastvideo/logger.py

@@ -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,
         )

fastvideo/pipelines/composed_pipeline_base.py

@@ -282,7 +282,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)

fastvideo/pipelines/lora_pipeline.py

@@ -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()

fastvideo/pipelines/pipeline_batch_info.py

@@ -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
@@ -118,8 +155,6 @@ class ForwardBatch:
     # Misc
     save_video: bool = True
     return_frames: bool = False
-    return_trajectory_latents: bool = False
-    trajectory_latents: list[torch.Tensor] = field(default_factory=list)
 
     # TeaCache parameters
     enable_teacache: bool = False
@@ -134,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."""
 
@@ -195,8 +234,6 @@ class TrainingBatch:
     # Distillation losses
     generator_loss: float = 0.0
     fake_score_loss: float = 0.0
-    regression_loss: float = 0.0
-    consistency_loss: float = 0.0
 
     dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
     fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)

fastvideo/pipelines/pipeline_registry.py

@@ -25,6 +25,11 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "HunyuanVideoPipeline": "hunyuan",
 }
 
+_PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {
+    WorkloadType.I2V: "PreprocessPipelineI2V",
+    WorkloadType.T2V: "PreprocessPipelineT2V",
+}
+
 
 class PipelineType(str, Enum):
     """
@@ -68,12 +73,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]

fastvideo/pipelines/preprocess/preprocess_stages.py

@@ -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)

fastvideo/pipelines/preprocess/v1_preprocessing_new.py

@@ -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()
 
 

fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py

@@ -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]

fastvideo/pipelines/stages/base.py

@@ -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",

fastvideo/pipelines/stages/denoising.py

@@ -197,7 +197,6 @@ class DenoisingStage(PipelineStage):
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
-                logger.info('timesteps: %s', timesteps)
                 # Skip if interrupted
                 if hasattr(self, 'interrupt') and self.interrupt:
                     continue
@@ -381,8 +380,7 @@ class DenoisingStage(PipelineStage):
 
     def progress_bar(self,
                      iterable: Iterable | None = None,
-                     total: int | None = None,
-                     disable: bool = False) -> tqdm:
+                     total: int | None = None) -> tqdm:
         """
         Create a progress bar for the denoising process.
         
@@ -395,7 +393,7 @@ class DenoisingStage(PipelineStage):
         """
         local_rank = get_world_group().local_rank
         if local_rank == 0:
-            return tqdm(iterable=iterable, total=total, disable=disable)
+            return tqdm(iterable=iterable, total=total)
         else:
             return tqdm(iterable=iterable, total=total, disable=True)
 
@@ -675,25 +673,11 @@ 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,
-            # latents.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
-        timesteps_50 = [
-            999, 993, 986, 979, 971, 964, 956, 948, 940, 931, 922, 913, 904,
-            895, 885, 874, 864, 853, 841, 830, 818, 805, 792, 778, 764, 749,
-            734, 718, 702, 684, 666, 647, 627, 607, 585, 562, 538, 513, 486,
-            458, 428, 396, 363, 328, 290, 249, 206, 160, 111, 57
-        ]
-        # timesteps = torch.tensor(timesteps_50,
-        #                          dtype=torch.long,
-        #                          device=get_local_torch_device())
         timesteps = torch.tensor(
             fastvideo_args.pipeline_config.dmd_denoising_steps,
             dtype=torch.long,
@@ -717,9 +701,7 @@ class DmdDenoisingStage(DenoisingStage):
                 batch.image_latent = image_latent
 
         # Run denoising loop
-        with self.progress_bar(
-                total=len(timesteps),
-                disable=batch.return_trajectory_latents) as progress_bar:
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
             for i, t in enumerate(timesteps):
                 # Skip if interrupted
                 if hasattr(self, 'interrupt') and self.interrupt:
@@ -809,8 +791,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",
@@ -819,10 +802,6 @@ class DmdDenoisingStage(DenoisingStage):
                         latents = self.scheduler.add_noise(
                             pred_video.flatten(0, 1), noise.flatten(0, 1),
                             next_timestep).unflatten(0, pred_video.shape[:2])
-
-                        if batch.return_trajectory_latents:
-                            batch.trajectory_latents.append(latents.clone())
-
                     else:
                         latents = pred_video
 

fastvideo/pipelines/stages/image_encoding.py

@@ -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,

fastvideo/tests/training/distill/test_ema.py

@@ -1,207 +0,0 @@
-import types
-
-import pytest  # type: ignore
-import torch  # type: ignore
-import torch.nn as nn  # type: ignore
-
-from fastvideo.training.training_utils import EMA_FSDP
-
-
-class TinyNet(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.lin = nn.Linear(4, 3, bias=True)
-        # Non-trainable param should be ignored by EMA
-        self.register_parameter("frozen", nn.Parameter(torch.randn(2), requires_grad=False))
-        # Buffer should be ignored by EMA
-        self.register_buffer("buf", torch.ones(1))
-
-
-def named_params_to_cpu_float_dict(module: nn.Module):
-    return {n: p.detach().clone().float().cpu() for n, p in module.named_parameters() if p.requires_grad}
-
-
-def test_ema_init_local_shard_copies_params_cpu_float():
-    torch.manual_seed(0)
-    model = TinyNet()
-    ema = EMA_FSDP(model, decay=0.9, mode="local_shard")
-
-    expected = named_params_to_cpu_float_dict(model)
-    assert set(ema.shadow.keys()) == set(expected.keys())
-    for n, v in expected.items():
-        assert torch.equal(ema.shadow[n], v)
-
-
-def test_ema_update_local_shard_matches_formula():
-    torch.manual_seed(0)
-    model = TinyNet()
-    decay = 0.9
-    ema = EMA_FSDP(model, decay=decay, mode="local_shard")
-
-    # Save initial snapshot
-    start = {n: p.detach().clone().float().cpu() for n, p in model.named_parameters() if p.requires_grad}
-
-    # Mutate model parameters to new values
-    with torch.no_grad():
-        for _, p in model.named_parameters():
-            if p.requires_grad:
-                p.add_(torch.randn_like(p))
-
-    after = {n: p.detach().clone().float().cpu() for n, p in model.named_parameters() if p.requires_grad}
-
-    ema.update(model)
-
-    for n in start.keys():
-        expected = start[n] * decay + after[n] * (1.0 - decay)
-        assert torch.allclose(ema.shadow[n], expected, rtol=1e-6, atol=1e-8)
-
-
-def test_apply_to_model_swaps_and_restores():
-    torch.manual_seed(0)
-    model = TinyNet()
-    ema = EMA_FSDP(model, decay=0.0, mode="local_shard")  # decay 0 -> shadow becomes last param on update
-
-    # Change model and update EMA so shadow != current
-    with torch.no_grad():
-        for _, p in model.named_parameters():
-            if p.requires_grad:
-                p.add_(1.2345)
-    ema.update(model)
-
-    # Change model again so current != shadow
-    with torch.no_grad():
-        for _, p in model.named_parameters():
-            if p.requires_grad:
-                p.add_(2.0)
-
-    # Snapshot of values at context entry
-    entry_values = {n: p.detach().clone() for n, p in model.named_parameters() if p.requires_grad}
-
-    # Inside context, params should equal EMA shadow; outside, restored
-    with ema.apply_to_model(model):
-        for n, p in model.named_parameters():
-            if p.requires_grad:
-                assert torch.allclose(p.detach().cpu().float(), ema.shadow[n])
-
-    for n, p in model.named_parameters():
-        if p.requires_grad:
-            assert torch.equal(p.detach(), entry_values[n])
-
-
-def test_state_dict_roundtrip():
-    torch.manual_seed(0)
-    model = TinyNet()
-    ema = EMA_FSDP(model, decay=0.8, mode="local_shard")
-
-    with torch.no_grad():
-        for _, p in model.named_parameters():
-            if p.requires_grad:
-                p.add_(torch.randn_like(p))
-    ema.update(model)
-
-    sd = ema.state_dict()
-    ema2 = EMA_FSDP(model, decay=0.8, mode="local_shard")
-    ema2.load_state_dict(sd)
-
-    for k in sd.keys():
-        assert torch.equal(ema2.shadow[k], sd[k])
-
-
-def test_copy_to_unwrapped_and_rank0_full_guard():
-    torch.manual_seed(0)
-    src = TinyNet()
-    ema = EMA_FSDP(src, decay=0.7, mode="local_shard")
-
-    # Make EMA shadow distinct from a fresh target
-    with torch.no_grad():
-        for _, p in src.named_parameters():
-            if p.requires_grad:
-                p.mul_(3.14)
-    ema.update(src)
-
-    tgt = TinyNet()
-    # copy in local_shard mode always applies
-    ema.copy_to_unwrapped(tgt)
-    for n, p in tgt.named_parameters():
-        if p.requires_grad:
-            assert torch.allclose(p.detach().cpu().float(), ema.shadow[n])
-
-    # If mode is rank0_full but rank != 0, copy is a no-op
-    ema.mode = "rank0_full"
-    ema.rank = 1
-    before = {n: p.detach().clone() for n, p in tgt.named_parameters()}
-    ema.copy_to_unwrapped(tgt)
-    for n, p in tgt.named_parameters():
-        assert torch.equal(p.detach(), before[n])
-
-
-def test_rank0_full_init_and_update_with_stubbed_gather(monkeypatch):
-    torch.manual_seed(0)
-    model = TinyNet()
-
-    # Stub gather_state_dict_on_cpu_rank0 to avoid requiring initialized dist
-    def fake_gather(mod, device=None):
-        return {n: p.detach().clone() for n, p in mod.named_parameters() if p.requires_grad}
-
-    # Patch in the module where EMA_FSDP is defined
-    import fastvideo.training.training_utils as tu
-    monkeypatch.setattr(tu, "gather_state_dict_on_cpu_rank0", fake_gather, raising=True)
-
-    ema = EMA_FSDP(model, decay=0.5, mode="rank0_full")
-    assert len(ema.shadow) > 0
-
-    start = {n: t.clone().float().cpu() for n, t in ema.shadow.items()}
-
-    with torch.no_grad():
-        for _, p in model.named_parameters():
-            if p.requires_grad:
-                p.add_(1.0)
-
-    ema.update(model)
-
-    cpu_state = {n: p.detach().clone().float().cpu() for n, p in model.named_parameters() if p.requires_grad}
-    for n in start.keys():
-        expected = start[n] * 0.5 + cpu_state[n] * 0.5
-        assert torch.allclose(ema.shadow[n], expected)
-
-    # state_dict should be empty if rank != 0 in rank0_full mode
-    ema.rank = 1
-    assert ema.state_dict() == {}
-
-
-def test_apply_to_model_raises_when_rank0_full():
-    model = TinyNet()
-    ema = EMA_FSDP(model, decay=0.9, mode="rank0_full")
-    with pytest.raises(RuntimeError):
-        with ema.apply_to_model(model):
-            pass
-
-
-def test_copy_to_unwrapped_rank0_full_rank0(monkeypatch):
-    torch.manual_seed(0)
-    model = TinyNet()
-
-    # Stub gather to produce keys matching named_parameters
-    def fake_gather(mod, device=None):
-        return {n: p.detach().clone() for n, p in mod.named_parameters() if p.requires_grad}
-
-    import fastvideo.training.training_utils as tu
-    monkeypatch.setattr(tu, "gather_state_dict_on_cpu_rank0", fake_gather, raising=True)
-
-    ema = EMA_FSDP(model, decay=0.0, mode="rank0_full")
-    ema.rank = 0
-
-    # Modify model and update so EMA has distinct weights
-    with torch.no_grad():
-        for _, p in model.named_parameters():
-            if p.requires_grad:
-                p.add_(1.0)
-    ema.update(model)
-
-    tgt = TinyNet()
-    ema.copy_to_unwrapped(tgt)
-    for n, p in tgt.named_parameters():
-        if p.requires_grad:
-            assert torch.allclose(p.detach().cpu().float(), ema.shadow[n])
-
-

fastvideo/training/cm_distillation_pipeline.py

@@ -1,846 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import copy
-import gc
-import os
-import time
-from abc import abstractmethod
-from collections import deque
-from collections.abc import Iterator
-from typing import Any
-
-import imageio
-import numpy as np
-import torch
-import torch.nn.functional as F
-import torchvision
-from einops import rearrange
-from torch.utils.data import DataLoader
-from torchdata.stateful_dataloader import StatefulDataLoader
-from torchmetrics.image.lpip import (
-    LearnedPerceptualImagePatchSimilarity as LPIPSimilarity)
-from tqdm.auto import tqdm
-
-import fastvideo.envs as envs
-from fastvideo.configs.sample import SamplingParam
-from fastvideo.dataset.validation_dataset import ValidationDataset
-from fastvideo.distributed import (cleanup_dist_env_and_memory,
-                                   get_local_torch_device, get_sp_group,
-                                   get_world_group)
-from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler)
-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 (
-    EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
-    get_scheduler, load_checkpoint, pred_noise_to_pred_video, save_checkpoint,
-    shift_timestep)
-from fastvideo.utils import is_vsa_available, set_random_seed
-
-import wandb  # isort: skip
-
-vsa_available = is_vsa_available()
-
-logger = init_logger(__name__)
-
-
-class CMDistillationPipeline(TrainingPipeline):
-    """
-    A distillation pipeline for training a 3 step model.
-    Inherits from TrainingPipeline to reuse training infrastructure.
-    """
-    _required_config_modules = [
-        "scheduler", "transformer", "vae"
-    ]
-    validation_pipeline: ComposedPipelineBase
-    train_dataloader: StatefulDataLoader
-    train_loader_iter: Iterator[dict[str, Any]]
-    current_epoch: int = 0
-    init_steps: int
-    current_trainstep: int
-    num_generator_updates: int = 0
-    video_latent_shape: tuple[int, ...]
-    video_latent_shape_sp: tuple[int, ...]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        raise RuntimeError(
-            "create_pipeline_stages should not be called for training pipeline")
-
-    def initialize_training_pipeline(self, training_args: TrainingArgs):
-        """Initialize the distillation training pipeline with multiple models."""
-        logger.info("Initializing distillation pipeline...")
-
-        super().initialize_training_pipeline(training_args)
-
-        self.noise_scheduler = self.get_module("scheduler")
-        self.vae = self.get_module("vae")
-        self.vae.requires_grad_(False)
-
-        self.timestep_shift = self.training_args.pipeline_config.flow_shift
-        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
-            shift=self.timestep_shift)
-
-        if self.training_args.warp_denoising_step:
-            # timesteps = self.noise_scheduler.timesteps.cpu()
-            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
-                                   torch.tensor([0], dtype=torch.float32)))
-            self.denoising_step_list = torch.tensor(
-                self.training_args.pipeline_config.dmd_denoising_steps,
-                dtype=torch.long,
-                device=torch.device("cpu"))
-            self.denoising_step_list = timesteps[1000 -
-                                                 self.denoising_step_list].to(
-                                                     get_local_torch_device())
-            logger.info(
-                "Warp denoising step is enabled, using %s denoising steps",
-                self.denoising_step_list)
-        else:
-            self.denoising_step_list = torch.tensor(
-                self.training_args.pipeline_config.dmd_denoising_steps,
-                dtype=torch.long,
-                device=get_local_torch_device())
-            logger.info(
-                "Warp denoising step is disabled, using %s denoising steps",
-                self.denoising_step_list)
-        logger.info("Distillation generator model to %s denoising steps",
-                    len(self.denoising_step_list))
-        self.num_train_timestep = self.noise_scheduler.num_train_timesteps
-
-        self.min_timestep = int(self.training_args.min_timestep_ratio *
-                                self.num_train_timestep)
-        self.max_timestep = int(self.training_args.max_timestep_ratio *
-                                self.num_train_timestep)
-
-        # self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
-
-        # Initialize EMA teacher for CM if enabled
-        self.use_cm_with_ema = getattr(self.training_args, "cm_use_ema_teacher",
-                                       False)
-        if self.use_cm_with_ema:
-            # Use local_shard mode for teacher forward compatibility with FSDP2
-            self.ema_teacher = EMA_FSDP(self.transformer,
-                                        decay=self.training_args.ema_decay,
-                                        mode="local_shard")
-        else:
-            self.ema_teacher = None
-        self.lpips = LPIPSimilarity().to(self.device)
-
-    def _compute_lpips_loss(self, pred_video: torch.Tensor,
-                            latents: torch.Tensor) -> torch.Tensor:
-        """
-        Helper method to compute LPIPS loss between predicted video and ground truth latents.
-        This method handles VAE scaling, shifting, decoding, and frame processing consistently.
-        """
-        print("Computing LPIPS loss...")
-        with torch.autocast("cuda", dtype=torch.bfloat16), torch.no_grad():
-            # Apply VAE scaling factor and shift factor before decoding (same as visualize_intermediate_latents)
-            pred_video = pred_video.permute(0, 2, 1, 3, 4)
-            latents = latents.permute(0, 2, 1, 3, 4)
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                pred_video_scaled = pred_video / self.vae.scaling_factor.to(
-                    pred_video.device, pred_video.dtype)
-                latents_scaled = latents / self.vae.scaling_factor.to(
-                    latents.device, latents.dtype)
-            else:
-                pred_video_scaled = pred_video / self.vae.scaling_factor
-                latents_scaled = latents / self.vae.scaling_factor
-
-            # Apply shifting if needed
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    pred_video_scaled += self.vae.shift_factor.to(
-                        pred_video.device, pred_video.dtype)
-                    latents_scaled += self.vae.shift_factor.to(
-                        latents.device, latents.dtype)
-                else:
-                    pred_video_scaled += self.vae.shift_factor
-                    latents_scaled += self.vae.shift_factor
-
-            # Permute from [batch, channels, frames, height, width] to [batch, frames, channels, height, width] for VAE decode
-            # pred_video_scaled = pred_video_scaled.permute(0, 2, 1, 3, 4)
-            # latents_scaled = latents_scaled.permute(0, 2, 1, 3, 4)
-
-            # print(f"pred_video_scaled shape after permute: {pred_video_scaled.shape}")
-            # print(f"latents_scaled shape after permute: {latents_scaled.shape}")
-            pred_video_frames = self.vae.decode(pred_video_scaled)
-            latents_frames = self.vae.decode(latents_scaled)
-
-        # VAE output is already in [-1, 1] range, keep it for LPIPS
-        # Just permute back to [B, T, C, H, W] format for frame processing
-        pred_video_frames = pred_video_frames.permute(0, 2, 1, 3, 4)
-        latents_frames = latents_frames.permute(0, 2, 1, 3, 4)
-
-        pred_video_frames = rearrange(pred_video_frames,
-                                      "b n c h w -> (b n) c h w")
-        latents_frames = rearrange(latents_frames, "b n c h w -> (b n) c h w")
-
-        return self.lpips(pred_video_frames, latents_frames)
-
-    @abstractmethod
-    def initialize_validation_pipeline(self, training_args: TrainingArgs):
-        """Initialize validation pipeline - must be implemented by subclasses."""
-        raise NotImplementedError(
-            "Distillation pipelines must implement this method")
-
-    def _prepare_distillation(self,
-                              training_batch: TrainingBatch) -> TrainingBatch:
-        """Prepare training environment for distillation."""
-        self.transformer.requires_grad_(True)
-        self.transformer.train()
-
-        return training_batch
-
-    def _build_distill_input_kwargs(
-            self, noise_input: torch.Tensor, timestep: torch.Tensor,
-            text_dict: dict[str, torch.Tensor] | None,
-            training_batch: TrainingBatch) -> TrainingBatch:
-        if text_dict is None:
-            raise ValueError(
-                "text_dict cannot be None for distillation pipeline")
-
-        training_batch.input_kwargs = {
-            "hidden_states": noise_input.permute(0, 2, 1, 3, 4),
-            "encoder_hidden_states": text_dict["encoder_hidden_states"],
-            "encoder_attention_mask": text_dict["encoder_attention_mask"],
-            "timestep": timestep,
-            "return_dict": False,
-        }
-
-        return training_batch
-
-    def _generator_forward(self, training_batch: TrainingBatch) -> torch.Tensor:
-        """
-        Forward pass through student transformer for a single randomly sampled
-        denoising step; returns predicted clean video latents and records
-        auxiliary info/losses on training_batch.
-        """
-        latents = training_batch.latents
-        dtype = latents.dtype
-        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
-
-        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]))
-
-        training_batch = self._build_distill_input_kwargs(noisy_latent, timestep,
-                                                          training_batch.conditional_dict,
-                                                          training_batch)
-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
-        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])
-
-        # Compute regression loss (LPIPS) for optional auxiliary loss/logging
-        # regression_loss = self._compute_lpips_loss(pred_video, latents)
-        # training_batch.regression_loss = regression_loss
-        training_batch.dmd_latent_vis_dict.update({
-            "generator_pred_video": pred_video.detach().clone(),
-        })
-        return pred_video
-
-    def _calculate_consistency_loss(
-            self, pred_video: torch.Tensor,
-            training_batch: TrainingBatch) -> torch.Tensor:
-        """
-        Consistency loss (CM-like): two variants
-          - Without EMA: teacher = model at higher noise with stop-grad
-          - With EMA: teacher = EMA(model) at lower noise or higher noise depending on schedule (use higher noise as teacher)
-        """
-        # Need at least two steps to form a pair
-        if len(self.denoising_step_list) < 2:
-            return torch.tensor(0.0, device=self.device, dtype=pred_video.dtype)
-
-        # Choose a random index k in [1, N-1] so that (k-1, k) is valid
-        idx_high = torch.randint(1,
-                                 len(self.denoising_step_list), [1],
-                                 device=self.device).item()
-        t_high = self.denoising_step_list[idx_high]
-        t_low = self.denoising_step_list[idx_high - 1]
-        t_low = t_low * torch.ones(1, device=self.device, dtype=torch.long)
-        t_high = t_high * torch.ones(1, device=self.device, dtype=torch.long)
-        logger.info("t_high: %s, t_low: %s", t_high, t_low)
-
-        # Use student output as base; construct noisy sample at higher timestep
-        base_student = pred_video
-        base_noise = torch.randn(self.video_latent_shape, device=self.device, dtype=base_student.dtype)
-        if self.sp_world_size > 1:
-            base_noise = rearrange(base_noise, "b (n t) c h w -> b n t c h w", n=self.sp_world_size).contiguous()
-            base_noise = base_noise[:, self.rank_in_sp_group, :, :, :, :]
-
-        noisy_high = self.noise_scheduler.add_noise(base_student.flatten(0, 1),
-                                                    base_noise.flatten(0, 1),
-                                                    t_high).unflatten(0, (1, base_student.shape[1]))
-
-        # Compute ODE-adjacent lower-t sample via a single Euler step using teacher output at t_high
-        # Temporarily switch to eval and no-grad for teacher forward(s)
-        was_training = self.transformer.training
-        self.transformer.eval()
-        with torch.no_grad():
-            tb_high_teacher = self._build_distill_input_kwargs(noisy_high, t_high,
-                                                               training_batch.conditional_dict, training_batch)
-            pred_flow_high_teacher = self.transformer(**tb_high_teacher.input_kwargs).permute(0, 2, 1, 3, 4)
-
-            # Prepare flat tensors for scheduler stepping
-            sample_flat = noisy_high.flatten(0, 1).to(dtype=torch.float32)
-            flow_flat = pred_flow_high_teacher.flatten(0, 1).to(dtype=torch.float32)
-
-            # Reset step index so stepping starts at the provided timestep
-            prev_step_index = getattr(self.noise_scheduler, "_step_index", None)
-            self.noise_scheduler._step_index = None
-            stepped = self.noise_scheduler.step(model_output=flow_flat,
-                                                timestep=t_high,
-                                                sample=sample_flat,
-                                                return_dict=True)
-            # Restore step index state
-            self.noise_scheduler._step_index = prev_step_index
-
-            noisy_low_ode = stepped.prev_sample.unflatten(0, (1, noisy_high.shape[1])).to(noisy_high.dtype)
-
-            # Teacher target at lower t (clean latent)
-            tb_low_teacher = self._build_distill_input_kwargs(noisy_low_ode, t_low,
-                                                              training_batch.conditional_dict, training_batch)
-            pred_noise_low_teacher = self.transformer(**tb_low_teacher.input_kwargs).permute(0, 2, 1, 3, 4)
-            y_low_teacher = pred_noise_to_pred_video(
-                pred_noise=pred_noise_low_teacher.flatten(0, 1),
-                noise_input_latent=noisy_low_ode.flatten(0, 1),
-                timestep=t_low,
-                scheduler=self.noise_scheduler,
-            ).unflatten(0, pred_noise_low_teacher.shape[:2])
-        if was_training:
-            self.transformer.train()
-
-        # Student prediction at lower t on the ODE-stepped sample
-        tb_low_student = self._build_distill_input_kwargs(noisy_low_ode, t_low,
-                                                          training_batch.conditional_dict, training_batch)
-        pred_noise_low_student = self.transformer(**tb_low_student.input_kwargs).permute(0, 2, 1, 3, 4)
-        y_low_student = pred_noise_to_pred_video(
-            pred_noise=pred_noise_low_student.flatten(0, 1),
-            noise_input_latent=noisy_low_ode.flatten(0, 1),
-            timestep=t_low,
-            scheduler=self.noise_scheduler,
-        ).unflatten(0, pred_noise_low_student.shape[:2])
-
-        cm_loss_weight = self.training_args.cm_loss_weight
-        weighing_fn = getattr(self.training_args, "cm_weighing_function", "constant")
-
-        if weighing_fn == "constant":
-            cm_weighing_function = lambda x: cm_loss_weight
-        elif weighing_fn == "sigma_sqrt":
-            cm_weighing_function = lambda x: cm_loss_weight * (x**0.5)
-        else:
-            raise ValueError(f"Invalid cm_weighing_function: {weighing_fn}")
-
-        # Match student at lower t to teacher at lower t (adjacent PF ODE point)
-        cm_loss = F.mse_loss(y_low_student, y_low_teacher.detach())
-        cm_loss = cm_loss * cm_weighing_function(t_low)
-
-        # Optionally record for logging
-        training_batch.dmd_latent_vis_dict.update({
-            "cm_timestep_high": t_high,
-            "cm_timestep_low": t_low,
-        })
-
-        return cm_loss
-
-    def _clip_model_grad_norm_(self, training_batch: TrainingBatch,
-                               transformer) -> TrainingBatch:
-
-        max_grad_norm = self.training_args.max_grad_norm
-
-        if max_grad_norm is not None:
-            model_parts = [transformer]
-            grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
-                [p for m in model_parts for p in m.parameters()],
-                max_grad_norm,
-                foreach=None,
-            )
-            assert grad_norm is not float('nan') or grad_norm is not float(
-                'inf')
-            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
-        else:
-            grad_norm = 0.0
-        training_batch.grad_norm = grad_norm
-        return training_batch
-
-    def _prepare_dit_inputs(self,
-                            training_batch: TrainingBatch) -> TrainingBatch:
-        super()._prepare_dit_inputs(training_batch)
-        conditional_dict = {
-            "encoder_hidden_states": training_batch.encoder_hidden_states,
-            "encoder_attention_mask": training_batch.encoder_attention_mask,
-        }
-        unconditional_dict = {
-            "encoder_hidden_states": self.negative_prompt_embeds,
-            "encoder_attention_mask": self.negative_prompt_attention_mask,
-        }
-
-        training_batch.dmd_latent_vis_dict = {}
-        training_batch.fake_score_latent_vis_dict = {}
-
-        training_batch.conditional_dict = conditional_dict
-        training_batch.unconditional_dict = unconditional_dict
-        training_batch.raw_latent_shape = training_batch.latents.shape
-        training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4)
-        self.video_latent_shape = training_batch.latents.shape
-
-        if self.sp_world_size > 1:
-            training_batch.latents = rearrange(
-                training_batch.latents,
-                "b (n t) c h w -> b n t c h w",
-                n=self.sp_world_size).contiguous()
-            training_batch.latents = training_batch.latents[:, self.
-                                                            rank_in_sp_group, :, :, :, :]
-
-        self.video_latent_shape_sp = training_batch.latents.shape
-
-        return training_batch
-
-    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
-        gradient_accumulation_steps = getattr(self.training_args, 'gradient_accumulation_steps', 1)
-        batches: list[TrainingBatch] = []
-        # Collect N batches for gradient accumulation
-        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  # type: ignore
-            batches.append(batch)
-
-        self.optimizer.zero_grad()
-        total_loss_value = 0.0
-        dmd_latent_vis_dict: dict[str, torch.Tensor] = {}
-        batch_gen = None
-
-        for batch in batches:
-            batch_gen = copy.deepcopy(batch)
-            # Forward student once to obtain a clean prediction to anchor CM pair
-            with set_forward_context(current_timestep=batch_gen.timesteps,
-                                     attn_metadata=batch_gen.attn_metadata_vsa):
-                generator_pred_video = self._generator_forward(batch_gen)
-
-            # Consistency loss
-            with set_forward_context(current_timestep=batch_gen.timesteps,
-                                     attn_metadata=batch_gen.attn_metadata):
-                cm_loss = self._calculate_consistency_loss(
-                    pred_video=generator_pred_video, training_batch=batch_gen)
-
-                # Optional auxiliary regression loss
-                if getattr(self.training_args, "use_regression_loss", False):
-                    cm_loss = cm_loss + batch_gen.regression_loss * self.training_args.regression_loss_weight
-
-            with set_forward_context(current_timestep=batch_gen.timesteps,
-                                     attn_metadata=batch_gen.attn_metadata_vsa):
-                (cm_loss / gradient_accumulation_steps).backward()
-            total_loss_value += cm_loss.detach().item()
-            dmd_latent_vis_dict.update(batch_gen.dmd_latent_vis_dict)
-
-        # Clip and step
-        assert batch_gen is not None
-        self._clip_model_grad_norm_(batch_gen, self.transformer)
-        self.optimizer.step()
-        self.lr_scheduler.step()
-
-        # Update EMA teacher after parameter update
-        if self.use_cm_with_ema:
-            assert self.ema_teacher is not None
-            self.ema_teacher.update(self.transformer)
-
-        self.optimizer.zero_grad(set_to_none=True)
-        avg_loss = torch.tensor(total_loss_value / max(1, gradient_accumulation_steps), device=self.device)
-        world_group = get_world_group()
-        world_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
-
-        training_batch.total_loss = avg_loss.item()
-        training_batch.dmd_latent_vis_dict = dmd_latent_vis_dict
-        training_batch.grad_norm = getattr(training_batch, "grad_norm", 0.0)
-        return training_batch
-
-    def _resume_from_checkpoint(self) -> None:
-        """Resume training from checkpoint for the generator model only."""
-
-        logger.info("Loading checkpoint from %s",
-                    self.training_args.resume_from_checkpoint)
-
-        resumed_step = load_checkpoint(self.transformer, self.global_rank,
-                                       self.training_args.resume_from_checkpoint,
-                                       self.optimizer, self.train_dataloader,
-                                       self.lr_scheduler, self.noise_random_generator)
-
-        if resumed_step > 0:
-            self.init_steps = resumed_step
-            logger.info("Successfully resumed from step %s", resumed_step)
-        else:
-            logger.warning("Failed to load checkpoint, starting from step 0")
-            self.init_steps = -1
-
-    def _log_training_info(self) -> None:
-        """Log distillation-specific training information."""
-        # First call parent class method to get basic training info
-        super()._log_training_info()
-
-        # Then add distillation-specific information
-        logger.info("Distillation-specific settings:")
-        assert isinstance(self.training_args, TrainingArgs)
-        logger.info("  Max gradient norm: %s", self.training_args.max_grad_norm)
-
-
-    @torch.no_grad()
-    def _log_validation(self, transformer, training_args, global_step) -> None:
-        training_args.inference_mode = True
-        training_args.dit_cpu_offload = True
-        if not training_args.log_validation:
-            return
-        if self.validation_pipeline is None:
-            raise ValueError("Validation pipeline is not set")
-
-        logger.info("Starting validation")
-
-        # Create sampling parameters if not provided
-        sampling_param = SamplingParam.from_pretrained(training_args.model_path)
-
-        # Set deterministic seed for validation
-
-        logger.info("Using validation seed: %s", self.seed)
-
-        # Prepare validation prompts
-        logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s',
-                    self.global_rank,
-                    training_args.validation_dataset_file,
-                    local_main_process_only=False)
-        validation_dataset = ValidationDataset(
-            training_args.validation_dataset_file)
-        validation_dataloader = DataLoader(validation_dataset,
-                                           batch_size=None,
-                                           num_workers=0)
-
-        transformer.eval()
-
-        validation_steps = training_args.validation_sampling_steps.split(",")
-        validation_steps = [int(step) for step in validation_steps]
-        validation_steps = [step for step in validation_steps if step > 0]
-        # Log validation results for this step
-        world_group = get_world_group()
-        num_sp_groups = world_group.world_size // self.sp_group.world_size
-        # Process each validation prompt for each validation step
-        for num_inference_steps in validation_steps:
-            logger.info("rank: %s: num_inference_steps: %s",
-                        self.global_rank,
-                        num_inference_steps,
-                        local_main_process_only=False)
-            step_videos: list[np.ndarray] = []
-            step_captions: list[str] = []
-
-            for validation_batch in validation_dataloader:
-                batch = self._prepare_validation_batch(sampling_param,
-                                                       training_args,
-                                                       validation_batch,
-                                                       num_inference_steps)
-
-                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()
-            num_sp_groups = world_group.world_size // self.sp_group.world_size
-
-            # Only sp_group leaders (rank_in_sp_group == 0) need to send their
-            # results to global rank 0
-            if self.rank_in_sp_group == 0:
-                if self.global_rank == 0:
-                    # Global rank 0 collects results from all sp_group leaders
-                    all_videos = step_videos  # Start with own results
-                    all_captions = step_captions
-
-                    # Receive from other sp_group leaders
-                    for sp_group_idx in range(1, num_sp_groups):
-                        src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
-                        recv_videos = world_group.recv_object(src=src_rank)
-                        recv_captions = world_group.recv_object(src=src_rank)
-                        all_videos.extend(recv_videos)
-                        all_captions.extend(recv_captions)
-
-                    video_filenames = []
-                    for i, (video, caption) in enumerate(
-                            zip(all_videos, all_captions, strict=True)):
-                        os.makedirs(training_args.output_dir, exist_ok=True)
-                        filename = os.path.join(
-                            training_args.output_dir,
-                            f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4"
-                        )
-                        imageio.mimsave(filename, video, fps=sampling_param.fps)
-                        video_filenames.append(filename)
-
-                    logs = {
-                        f"validation_videos_{num_inference_steps}_steps": [
-                            wandb.Video(filename, caption=caption)
-                            for filename, caption in zip(
-                                video_filenames, all_captions, strict=True)
-                        ]
-                    }
-                    wandb.log(logs, step=global_step)
-                else:
-                    # Other sp_group leaders send their results to global rank 0
-                    world_group.send_object(step_videos, dst=0)
-                    world_group.send_object(step_captions, dst=0)
-
-        # Re-enable gradients for training
-        transformer.train()
-        gc.collect()
-
-    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 = {}
-        dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
-
-        # Only log generator output for CM pipeline
-        if 'generator_pred_video' in dmd_latents_vis_dict:
-            latents = dmd_latents_vis_dict['generator_pred_video']
-            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
-
-            # Apply shifting if needed
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    latents += self.vae.shift_factor.to(latents.device,
-                                                        latents.dtype)
-                else:
-                    latents += self.vae.shift_factor
-            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['generator_pred_video'] = wandb.Video(
-                video, fps=24, format="mp4")
-            # Clean up references
-            del video, latents
-
-        # Log to wandb
-        if self.global_rank == 0:
-            wandb.log(wandb_loss_dict, step=step)
-
-    def train(self) -> None:
-        """Main training loop with distillation-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)
-
-        # Set random seeds for deterministic training
-        self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
-            self.seed)
-        self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed(
-            self.seed)
-        self.validation_random_generator = torch.Generator(
-            device="cpu").manual_seed(self.seed)
-        logger.info("Initialized random seeds with seed: %s", seed)
-
-        # Resume from checkpoint if specified (this will restore random states)
-        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[float] = 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
-
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                training_batch = self.train_one_step(training_batch)
-
-            total_loss = training_batch.total_loss
-            grad_norm = training_batch.grad_norm
-
-            step_time = time.perf_counter() - start_time
-            step_times.append(step_time)
-            avg_step_time = sum(step_times) / len(step_times)
-
-            progress_bar.set_postfix({
-                "total_loss":
-                f"{total_loss:.4f}",
-                "lpips_loss":
-                f"{training_batch.regression_loss:.4f}"
-                if hasattr(training_batch, 'regression_loss')
-                and training_batch.regression_loss is not None else "N/A",
-                "step_time":
-                f"{step_time:.2f}s",
-                "grad_norm":
-                grad_norm,
-            })
-            progress_bar.update(1)
-
-            if self.global_rank == 0:
-                # Prepare logging data
-                log_data = {
-                    "train_total_loss": total_loss,
-                    "learning_rate": self.lr_scheduler.get_last_lr()[0],
-                    "step_time": step_time,
-                    "avg_step_time": avg_step_time,
-                    "grad_norm": grad_norm,
-                }
-                if use_vsa:
-                    log_data["VSA_train_sparsity"] = current_vsa_sparsity
-
-                if training_batch.dmd_latent_vis_dict:
-                    dmd_additional_logs = {
-                        "generator_timestep": training_batch.dmd_latent_vis_dict["generator_timestep"].item(),
-                        "regression_loss": training_batch.regression_loss,
-                    }
-                    log_data.update(dmd_additional_logs)
-
-                wandb.log(log_data, step=step)
-
-            # Save training state checkpoint (for resuming training)
-            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_checkpoint(self.transformer, self.global_rank,
-                                self.training_args.output_dir, step,
-                                self.optimizer, self.train_dataloader,
-                                self.lr_scheduler, self.noise_random_generator)
-
-                if self.transformer:
-                    self.transformer.train()
-                self.sp_group.barrier()
-
-            # Save weight-only checkpoint (export consolidated generator weights)
-            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_checkpoint(self.transformer, self.global_rank,
-                                self.training_args.output_dir,
-                                f"{step}", self.optimizer, self.train_dataloader,
-                                self.lr_scheduler, self.noise_random_generator)
-
-            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)
-                self._log_validation(self.transformer, self.training_args, step)
-
-        wandb.finish()
-
-        # Save final training state checkpoint
-        print("rank", self.global_rank,
-              "save final training state checkpoint at step",
-              self.training_args.max_train_steps)
-        save_checkpoint(self.transformer, self.global_rank,
-                        self.training_args.output_dir,
-                        self.training_args.max_train_steps, self.optimizer,
-                        self.train_dataloader, self.lr_scheduler,
-                        self.noise_random_generator)
-
-        if get_sp_group():
-            cleanup_dist_env_and_memory()

fastvideo/training/distillation_pipeline.py

@@ -16,8 +16,6 @@ import torchvision
 from einops import rearrange
 from torch.utils.data import DataLoader
 from torchdata.stateful_dataloader import StatefulDataLoader
-from torchmetrics.image.lpip import (
-    LearnedPerceptualImagePatchSimilarity as LPIPSimilarity)
 from tqdm.auto import tqdm
 
 import fastvideo.envs as envs
@@ -37,8 +35,8 @@ from fastvideo.training.activation_checkpoint import (
     apply_activation_checkpointing)
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
-    EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
-    get_scheduler, load_distillation_checkpoint, pred_noise_to_pred_video,
+    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
 
@@ -68,7 +66,6 @@ class DistillationPipeline(TrainingPipeline):
     current_epoch: int = 0
     init_steps: int
     current_trainstep: int
-    num_generator_updates: int = 0
     video_latent_shape: tuple[int, ...]
     video_latent_shape_sp: tuple[int, ...]
 
@@ -146,28 +143,10 @@ class DistillationPipeline(TrainingPipeline):
             "Distillation pipeline initialized with generator_update_interval=%s",
             self.generator_update_interval)
 
-        if self.training_args.warp_denoising_step:
-            # timesteps = self.noise_scheduler.timesteps.cpu()
-            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
-                                   torch.tensor([0], dtype=torch.float32)))
-            self.denoising_step_list = torch.tensor(
-                self.training_args.pipeline_config.dmd_denoising_steps,
-                dtype=torch.long,
-                device=torch.device("cpu"))
-            self.denoising_step_list = timesteps[1000 -
-                                                 self.denoising_step_list].to(
-                                                     get_local_torch_device())
-            logger.info(
-                "Warp denoising step is enabled, using %s denoising steps",
-                self.denoising_step_list)
-        else:
-            self.denoising_step_list = torch.tensor(
-                self.training_args.pipeline_config.dmd_denoising_steps,
-                dtype=torch.long,
-                device=get_local_torch_device())
-            logger.info(
-                "Warp denoising step is disabled, using %s denoising steps",
-                self.denoising_step_list)
+        self.denoising_step_list = torch.tensor(
+            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))
         self.num_train_timestep = self.noise_scheduler.num_train_timesteps
@@ -179,70 +158,6 @@ class DistillationPipeline(TrainingPipeline):
 
         self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
 
-        # Initialize EMA teacher for CM if enabled
-        self.use_cm_with_ema = getattr(self.training_args, "cm_use_ema_teacher",
-                                       False)
-        if self.use_cm_with_ema:
-            # Use local_shard mode for teacher forward compatibility with FSDP2
-            self.ema_teacher = EMA_FSDP(self.transformer,
-                                        decay=self.training_args.ema_decay,
-                                        mode="local_shard")
-        else:
-            self.ema_teacher = None
-        self.lpips = LPIPSimilarity().to(self.device)
-
-    def _compute_lpips_loss(self, pred_video: torch.Tensor,
-                            latents: torch.Tensor) -> torch.Tensor:
-        """
-        Helper method to compute LPIPS loss between predicted video and ground truth latents.
-        This method handles VAE scaling, shifting, decoding, and frame processing consistently.
-        """
-        print("Computing LPIPS loss...")
-        with torch.autocast("cuda", dtype=torch.bfloat16), torch.no_grad():
-            # Apply VAE scaling factor and shift factor before decoding (same as visualize_intermediate_latents)
-            pred_video = pred_video.permute(0, 2, 1, 3, 4)
-            latents = latents.permute(0, 2, 1, 3, 4)
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                pred_video_scaled = pred_video / self.vae.scaling_factor.to(
-                    pred_video.device, pred_video.dtype)
-                latents_scaled = latents / self.vae.scaling_factor.to(
-                    latents.device, latents.dtype)
-            else:
-                pred_video_scaled = pred_video / self.vae.scaling_factor
-                latents_scaled = latents / self.vae.scaling_factor
-
-            # Apply shifting if needed
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    pred_video_scaled += self.vae.shift_factor.to(
-                        pred_video.device, pred_video.dtype)
-                    latents_scaled += self.vae.shift_factor.to(
-                        latents.device, latents.dtype)
-                else:
-                    pred_video_scaled += self.vae.shift_factor
-                    latents_scaled += self.vae.shift_factor
-
-            # Permute from [batch, channels, frames, height, width] to [batch, frames, channels, height, width] for VAE decode
-            # pred_video_scaled = pred_video_scaled.permute(0, 2, 1, 3, 4)
-            # latents_scaled = latents_scaled.permute(0, 2, 1, 3, 4)
-
-            # print(f"pred_video_scaled shape after permute: {pred_video_scaled.shape}")
-            # print(f"latents_scaled shape after permute: {latents_scaled.shape}")
-            pred_video_frames = self.vae.decode(pred_video_scaled)
-            latents_frames = self.vae.decode(latents_scaled)
-
-        # VAE output is already in [-1, 1] range, keep it for LPIPS
-        # Just permute back to [B, T, C, H, W] format for frame processing
-        pred_video_frames = pred_video_frames.permute(0, 2, 1, 3, 4)
-        latents_frames = latents_frames.permute(0, 2, 1, 3, 4)
-
-        pred_video_frames = rearrange(pred_video_frames,
-                                      "b n c h w -> (b n) c h w")
-        latents_frames = rearrange(latents_frames, "b n c h w -> (b n) c h w")
-
-        return self.lpips(pred_video_frames, latents_frames)
-
     @abstractmethod
     def initialize_validation_pipeline(self, training_args: TrainingArgs):
         """Initialize validation pipeline - must be implemented by subclasses."""
@@ -313,18 +228,7 @@ class DistillationPipeline(TrainingPipeline):
             timestep=timestep,
             scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2])
 
-        # regression_loss = F.mse_loss(pred_video, latents)
-
-        # LPIPS loss
-        regression_loss = self._compute_lpips_loss(pred_video, latents)
-        pred_video = pred_video.type_as(noisy_latent)
-        training_batch.regression_loss = regression_loss
-        # regression_log_dict = {
-        #     "regression_loss": regression_Loss.detach(),
-        #     "regression_clean_latents": latents,
-        # }
         return pred_video
-        # return pred_video
 
     def _generator_multi_step_simulation_forward(
             self, training_batch: TrainingBatch) -> torch.Tensor:
@@ -337,23 +241,22 @@ class DistillationPipeline(TrainingPipeline):
                                             len(self.denoising_step_list), [1],
                                             device=self.device,
                                             dtype=torch.long)
-        # target_timestep_idx = torch.tensor([2],
-        #                                    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
         # Start from pure noise like in inference
-        noise_latent = torch.randn(self.video_latent_shape,
-                                   device=self.device,
-                                   dtype=dtype)
+        current_noise_latents = torch.randn(self.video_latent_shape,
+                                            device=self.device,
+                                            dtype=dtype)
         if self.sp_world_size > 1:
-            noise_latents = rearrange(noise_latent,
-                                      "b (n t) c h w -> b n t c h w",
-                                      n=self.sp_world_size).contiguous()
-            noise_latent = noise_latent[:, self.rank_in_sp_group, :, :, :, :]
-        current_noise_latents = noise_latent.clone()
+            current_noise_latents = rearrange(
+                current_noise_latents,
+                "b (n t) c h w -> b n t c h w",
+                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
@@ -407,7 +310,7 @@ class DistillationPipeline(TrainingPipeline):
             ) - 1, "noise_latent_index is out of bounds"
             noisy_input = noise_latents[noise_latent_index]
         else:
-            noisy_input = noise_latent
+            noisy_input = current_noise_latents_copy
 
         # Step 4: Final student prediction (this is what we train on)
         training_batch = self._build_distill_input_kwargs(
@@ -420,122 +323,10 @@ class DistillationPipeline(TrainingPipeline):
             noise_input_latent=noisy_input.flatten(0, 1),
             timestep=target_timestep,
             scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2])
-
-        regression_loss = F.mse_loss(pred_video, latents)
-        # LPIPS loss - decode latents to actual frames first
-        regression_loss = self._compute_lpips_loss(pred_video, latents)
-        training_batch.regression_loss = regression_loss
         training_batch.dmd_latent_vis_dict[
             "generator_timestep"] = target_timestep.float().detach()
-
         return pred_video
 
-    def _calculate_consistency_loss(
-            self, pred_video: torch.Tensor,
-            training_batch: TrainingBatch) -> torch.Tensor:
-        """
-        Consistency loss (CM-like): two variants
-          - Without EMA: teacher = model at higher noise with stop-grad
-          - With EMA: teacher = EMA(model) at lower noise or higher noise depending on schedule (use higher noise as teacher)
-        """
-        # Need at least two steps to form a pair
-        if len(self.denoising_step_list) < 2:
-            return torch.tensor(0.0, device=self.device, dtype=pred_video.dtype)
-
-        # Choose a random index k in [1, N-1] so that (k-1, k) is valid
-        idx_high = torch.randint(1,
-                                 len(self.denoising_step_list), [1],
-                                 device=self.device).item()
-        t_high = self.denoising_step_list[idx_high]
-        t_low = self.denoising_step_list[idx_high - 1]
-        t_low = t_low * torch.ones(1, device=self.device, dtype=torch.long)
-        t_high = t_high * torch.ones(1, device=self.device, dtype=torch.long)
-        logger.info("t_high: %s, t_low: %s", t_high, t_low)
-
-        # Use student output as base for CM pairing (noisy inputs derived from student's prediction)
-        base_student = pred_video
-        # base_student = training_batch.latents
-        # Shared base noise for both timesteps
-        base_noise = torch.randn(self.video_latent_shape,
-                                 device=self.device,
-                                 dtype=base_student.dtype)
-        if self.sp_world_size > 1:
-            base_noise = rearrange(base_noise,
-                                   "b (n t) c h w -> b n t c h w",
-                                   n=self.sp_world_size).contiguous()
-            base_noise = base_noise[:, self.rank_in_sp_group, :, :, :, :]
-
-        # Build noisy inputs for the selected timesteps using the same student base and same noise
-        noisy_high = self.noise_scheduler.add_noise(
-            base_student.flatten(0, 1), base_noise.flatten(0, 1),
-            t_high).unflatten(0, (1, base_student.shape[1]))
-
-        # Teacher ODE step: t_high → t_low and teacher target at t_low
-        was_training = self.transformer.training
-        self.transformer.eval()
-        with torch.no_grad():
-            tb_high_teacher = self._build_distill_input_kwargs(
-                noisy_high, t_high, training_batch.conditional_dict, training_batch)
-            pred_flow_high_teacher = self.transformer(
-                **tb_high_teacher.input_kwargs).permute(0, 2, 1, 3, 4)
-
-            sample_flat = noisy_high.flatten(0, 1).float()
-            flow_flat = pred_flow_high_teacher.flatten(0, 1).float()
-
-            prev_step_index = getattr(self.noise_scheduler, "_step_index", None)
-            self.noise_scheduler._step_index = None
-            stepped = self.noise_scheduler.step(
-                model_output=flow_flat, timestep=t_high, sample=sample_flat,
-                return_dict=True)
-            self.noise_scheduler._step_index = prev_step_index
-
-            noisy_low_ode = stepped.prev_sample.unflatten(0, (1, noisy_high.shape[1])).to(noisy_high.dtype)
-
-            tb_low_teacher = self._build_distill_input_kwargs(
-                noisy_low_ode, t_low, training_batch.conditional_dict, training_batch)
-            pred_noise_low_teacher = self.transformer(
-                **tb_low_teacher.input_kwargs).permute(0, 2, 1, 3, 4)
-            y_low_teacher = pred_noise_to_pred_video(
-                pred_noise=pred_noise_low_teacher.flatten(0, 1),
-                noise_input_latent=noisy_low_ode.flatten(0, 1),
-                timestep=t_low,
-                scheduler=self.noise_scheduler,
-            ).unflatten(0, pred_noise_low_teacher.shape[:2])
-        if was_training:
-            self.transformer.train()
-
-        # Student prediction at lower t on the ODE-stepped sample
-        tb_low_student = self._build_distill_input_kwargs(
-            noisy_low_ode, t_low, training_batch.conditional_dict, training_batch)
-        pred_noise_low_student = self.transformer(
-            **tb_low_student.input_kwargs).permute(0, 2, 1, 3, 4)
-        y_low_student = pred_noise_to_pred_video(
-            pred_noise=pred_noise_low_student.flatten(0, 1),
-            noise_input_latent=noisy_low_ode.flatten(0, 1),
-            timestep=t_low,
-            scheduler=self.noise_scheduler,
-        ).unflatten(0, pred_noise_low_student.shape[:2])
-
-        cm_loss_weight = self.training_args.cm_loss_weight
-        weighing_fn = getattr(self.training_args, "cm_weighing_function", "constant")
-        if weighing_fn == "constant":
-            cm_weighing_function = lambda x: cm_loss_weight
-        elif weighing_fn == "sigma_sqrt":
-            cm_weighing_function = lambda x: cm_loss_weight * (x**0.5)
-        else:
-            raise ValueError(f"Invalid cm_weighing_function: {weighing_fn}")
-
-        cm_loss = F.mse_loss(y_low_student, y_low_teacher.detach())
-        cm_loss = cm_loss * cm_weighing_function(t_low)
-
-        # Optionally record for logging
-        training_batch.dmd_latent_vis_dict.update({
-            "cm_timestep_high": t_high,
-            "cm_timestep_low": t_low,
-        })
-
-        return cm_loss
-
     def _dmd_forward(self, generator_pred_video: torch.Tensor,
                      training_batch: TrainingBatch) -> torch.Tensor:
         """Compute DMD (Diffusion Model Distillation) loss."""
@@ -619,9 +410,6 @@ class DistillationPipeline(TrainingPipeline):
             generator_pred_video.float(),
             (generator_pred_video.float() - grad.float()).detach())
 
-        if self.training_args.use_regression_loss:
-            dmd_loss = dmd_loss + training_batch.regression_loss * self.training_args.regression_loss_weight
-
         training_batch.dmd_latent_vis_dict.update({
             "training_batch_dmd_fwd_clean_latent":
             training_batch.latents,
@@ -633,8 +421,6 @@ class DistillationPipeline(TrainingPipeline):
             faker_score_pred_video,
             "dmd_timestep":
             timestep,
-            "regression_loss":
-            training_batch.regression_loss,
         })
 
         return dmd_loss
@@ -646,12 +432,8 @@ class DistillationPipeline(TrainingPipeline):
                 current_timestep=training_batch.timesteps,
                 attn_metadata=training_batch.attn_metadata_vsa):
             if self.training_args.simulate_generator_forward:
-                if False:
-                    generator_pred_video = self._generator_forward_with_validation_pipeline(
-                        training_batch)
-                else:
-                    generator_pred_video = self._generator_multi_step_simulation_forward(
-                        training_batch)
+                generator_pred_video = self._generator_multi_step_simulation_forward(
+                    training_batch)
             else:
                 generator_pred_video = self._generator_forward(training_batch)
 
@@ -781,57 +563,29 @@ class DistillationPipeline(TrainingPipeline):
             for batch in batches:
                 batch_gen = copy.deepcopy(batch)
 
-                # Generator forward and DMD loss after CM
                 with set_forward_context(
                         current_timestep=batch_gen.timesteps,
                         attn_metadata=batch_gen.attn_metadata_vsa):
-                    if self.training_args.simulate_generator_forward and self.num_generator_updates % self.training_args.simulate_forward_interval == 0:
-                        logger.info(
-                            "using simulate_generator_forward, num_generator_updates: %s",
-                            self.num_generator_updates)
-                        if False:
-                            generator_pred_video = self._generator_forward_with_validation_pipeline(
-                                batch_gen)
-                        else:
-                            generator_pred_video = self._generator_multi_step_simulation_forward(
-                                batch_gen)
+                    if self.training_args.simulate_generator_forward:
+                        generator_pred_video = self._generator_multi_step_simulation_forward(
+                            batch_gen)
                     else:
                         generator_pred_video = self._generator_forward(
                             batch_gen)
 
-                # Consistency loss first (to avoid in-place param swaps after any grad-tracked forward)
-                cm_weight = self.training_args.cm_loss_weight
-                dmd_loss_extra = 0.0
-                if cm_weight > 0.0:
-                    with set_forward_context(
-                            current_timestep=batch_gen.timesteps,
-                            attn_metadata=batch_gen.attn_metadata):
-                        cm_loss = self._calculate_consistency_loss(
-                            pred_video=generator_pred_video,
-                            training_batch=batch_gen,
-                        )
-                        dmd_loss_extra = cm_loss
-
                 with set_forward_context(current_timestep=batch_gen.timesteps,
                                          attn_metadata=batch_gen.attn_metadata):
                     dmd_loss = self._dmd_forward(
                         generator_pred_video=generator_pred_video,
                         training_batch=batch_gen)
-                    logger.info("dmd_loss: %s", dmd_loss.item())
-                    dmd_loss = dmd_loss + dmd_loss_extra
 
                 with set_forward_context(
                         current_timestep=batch_gen.timesteps,
                         attn_metadata=batch_gen.attn_metadata_vsa):
                     (dmd_loss / gradient_accumulation_steps).backward()
                 total_dmd_loss += dmd_loss.detach().item()
-            self.num_generator_updates += 1
             self._clip_model_grad_norm_(batch_gen, self.transformer)
             self.optimizer.step()
-            # Update EMA teacher after generator parameter update
-            if self.use_cm_with_ema:
-                assert self.ema_teacher is not None
-                self.ema_teacher.update(self.transformer)
             self.optimizer.zero_grad(set_to_none=True)
             avg_dmd_loss = torch.tensor(total_dmd_loss /
                                         gradient_accumulation_steps,
@@ -1206,20 +960,11 @@ 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}",
-                "lpips_loss":
-                f"{training_batch.regression_loss:.4f}"
-                if hasattr(training_batch, 'regression_loss')
-                and training_batch.regression_loss is not None else "N/A",
-                "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,
             })
             progress_bar.update(1)
 
@@ -1255,8 +1000,6 @@ class DistillationPipeline(TrainingPipeline):
                         "dmd_timestep":
                         training_batch.dmd_latent_vis_dict["dmd_timestep"].item(
                         ),
-                        "regression_loss":
-                        training_batch.regression_loss,
                     }
                     log_data.update(dmd_additional_logs)
 

fastvideo/training/training_utils.py

@@ -1299,156 +1299,3 @@ def get_scheduler(
                          num_warmup_steps=num_warmup_steps,
                          num_training_steps=num_training_steps,
                          last_epoch=last_epoch)
-
-
-class EMA_FSDP:
-    """
-    FSDP2-friendly EMA with two modes:
-      - mode="local_shard" (default): maintain float32 CPU EMA of local parameter shards on every rank.
-        Provides a context manager to temporarily swap EMA weights into the live model for teacher forward.
-      - mode="rank0_full": maintain a consolidated float32 CPU EMA of full parameters on rank 0 only
-        using gather_state_dict_on_cpu_rank0(). Useful for checkpoint export; not for teacher forward.
-
-    Usage (local_shard for CM teacher):
-      ema = EMA_FSDP(model, decay=0.999, mode="local_shard")
-      for step in ...:
-          ema.update(model)
-      with ema.apply_to_model(model):
-          with torch.no_grad():
-              y_teacher = model(...)
-
-    Usage (rank0_full for export):
-      ema = EMA_FSDP(model, decay=0.999, mode="rank0_full")
-      ema.update(model)
-      ema.state_dict()  # on rank 0
-    """
-    def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"):
-        self.decay = float(decay)
-        self.mode = mode
-        self.shadow: dict[str, torch.Tensor] = {}
-        self.rank = dist.get_rank() if dist.is_initialized() else 0
-        if self.mode not in {"local_shard", "rank0_full"}:
-            raise ValueError(f"Unsupported EMA_FSDP mode: {self.mode}")
-        self._init_shadow(module)
-
-    @staticmethod
-    def _to_local_tensor(t: torch.Tensor) -> torch.Tensor:
-        # DTensor-aware to_local fetch; fall back to raw tensor
-        try:
-            from torch.distributed.tensor import DTensor  # type: ignore
-            if isinstance(t, DTensor):
-                return t.to_local()
-        except Exception:
-            pass
-        return t
-
-    @torch.no_grad()
-    def _init_shadow(self, module):
-        if self.mode == "rank0_full":
-            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
-            if self.rank == 0:
-                self.shadow = {k: v.detach().clone().float().cpu() for k, v in cpu_state.items()}
-            else:
-                self.shadow = {}
-            return
-
-        # local_shard: maintain EMA of local shards for requires_grad params
-        self.shadow = {}
-        for name, p in module.named_parameters():
-            if not p.requires_grad:
-                continue
-            local = self._to_local_tensor(p.detach())
-            self.shadow[name] = local.clone().float().cpu()
-
-    @torch.no_grad()
-    def update(self, module):
-        d = self.decay
-        if self.mode == "rank0_full":
-            if self.rank != 0:
-                return
-            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
-            for n, v in cpu_state.items():
-                v_cpu = v.detach().float().cpu()
-                if n not in self.shadow:
-                    self.shadow[n] = v_cpu.clone()
-                else:
-                    self.shadow[n].mul_(d).add_(v_cpu, alpha=1.0 - d)
-            return
-
-        # local_shard: update local shard EMA on every rank
-        for name, p in module.named_parameters():
-            if not p.requires_grad:
-                continue
-            local = self._to_local_tensor(p.detach())
-            v_cpu = local.float().cpu()
-            if name not in self.shadow:
-                self.shadow[name] = v_cpu.clone()
-            else:
-                self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
-
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        if self.mode == "rank0_full":
-            return {k: v.clone() for k, v in self.shadow.items()} if self.rank == 0 else {}
-        return {k: v.clone() for k, v in self.shadow.items()}
-
-    def load_state_dict(self, sd: dict[str, torch.Tensor]):
-        self.shadow = {k: v.clone() for k, v in sd.items()}
-
-    @torch.no_grad()
-    def copy_to_unwrapped(self, module) -> None:
-        """
-        Copy EMA weights into a non-sharded (unwrapped) module. Intended for export/eval.
-        For mode="rank0_full", only rank 0 has the full EMA state.
-        """
-        if self.mode == "rank0_full" and self.rank != 0:
-            return
-        name_to_param = dict(module.named_parameters())
-        for n, w in self.shadow.items():
-            if n in name_to_param:
-                p = name_to_param[n]
-                p.data.copy_(w.to(dtype=p.dtype, device=p.device))
-
-    class _ApplyEMACtx:
-        def __init__(self, ema: "EMA_FSDP", module):
-            self.ema = ema
-            self.module = module
-            self.saved: dict[str, torch.Tensor] = {}
-
-        def __enter__(self):
-            if self.ema.mode != "local_shard":
-                raise RuntimeError("EMA apply_to_model is only supported for mode='local_shard'")
-            with torch.no_grad():
-                for name, p in self.module.named_parameters():
-                    if not p.requires_grad:
-                        continue
-                    # Save local shard
-                    p_local = EMA_FSDP._to_local_tensor(p.detach())
-                    if p_local.numel() == 0:
-                        # Nothing to swap on this rank for this param
-                        continue
-                    self.saved[name] = p_local.clone().to(device=p_local.device, dtype=p_local.dtype)
-                    if name in self.ema.shadow:
-                        ema_cpu = self.ema.shadow[name]
-                        if ema_cpu.numel() != p_local.numel():
-                            # Shard shape mismatch (e.g., empty shard here), skip
-                            continue
-                        # Copy EMA shard into local param shard
-                        p_local.copy_(ema_cpu.to(dtype=p_local.dtype, device=p_local.device))
-            return self.module
-
-        def __exit__(self, exc_type, exc, tb):
-            with torch.no_grad():
-                for name, p in self.module.named_parameters():
-                    if name in self.saved:
-                        p_local = EMA_FSDP._to_local_tensor(p.detach())
-                        if p_local.numel() == 0:
-                            continue
-                        saved_local = self.saved[name]
-                        if saved_local.numel() != p_local.numel():
-                            continue
-                        p_local.copy_(saved_local)
-            self.saved.clear()
-            return False
-
-    def apply_to_model(self, module):
-        return EMA_FSDP._ApplyEMACtx(self, module)

fastvideo/training/wan_cm_distillation_pipeline.py

@@ -1,79 +0,0 @@
-# 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_dmd_pipeline import WanDMDPipeline
-from fastvideo.training.cm_distillation_pipeline import CMDistillationPipeline
-from fastvideo.utils import is_vsa_available
-
-vsa_available = is_vsa_available()
-
-logger = init_logger(__name__)
-
-
-class WanCMDistillationPipeline(CMDistillationPipeline):
-    """
-    A distillation pipeline for Wan that uses a single transformer model.
-    The main transformer serves as the student model, and copies are made for teacher and critic.
-    """
-    _required_config_modules = [
-        "scheduler", "transformer", "vae"
-    ]
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        """Initialize Wan-specific scheduler."""
-        shift_value = fastvideo_args.pipeline_config.flow_shift
-        shift = 1.0 if shift_value is None else float(shift_value)
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=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 = WanDMDPipeline.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 distillation pipeline...")
-
-    # Create pipeline with original args
-    pipeline = WanCMDistillationPipeline.from_pretrained(
-        args.pretrained_model_name_or_path, args=args)
-
-    args = pipeline.training_args
-    # Start training
-    pipeline.train()
-    logger.info("Wan 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)