[Docs] Docs update for Training and MPS (#641)

Co-authored-by: root <a1286225768@gmail,com>

.buildkite/pipeline.yml

@@ -13,7 +13,7 @@ steps:
   - label: "Trigger Tests"
     plugins:
       - monorepo-diff#v1.4.0:
-          diff: "git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD"
+          diff: 'git fetch origin "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" && git diff --name-only origin/"$BUILDKITE_PULL_REQUEST_BASE_BRANCH"...HEAD'
           watch:
             - path: 
                 - "fastvideo/v1/models/encoders/**"
@@ -58,8 +58,10 @@ steps:
                   queue: "default"
             - path: 
                 - "fastvideo/v1/**/*.py"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 45m .buildkite/scripts/pr_test.sh"
                 label: "SSIM Tests"
                 env:
                   - TEST_TYPE=ssim
@@ -91,6 +93,17 @@ steps:
                   - TEST_TYPE=training
                 agents:
                   queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "LoRA Training Tests"
+                env:
+                  - TEST_TYPE=training_lora
+                agents:
+                  queue: "default"
             - path:
                 - "fastvideo/v1/**"
                 - "csrc/attn/vsa/**"

.buildkite/scripts/pr_test.sh

@@ -81,6 +81,10 @@ case "$TEST_TYPE" in
         log "Running training tests..."
         MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests"
         ;;
+    "training_lora")
+        log "Running LoRA training tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_lora_tests"
+        ;;
     "training_vsa")
         log "Running training VSA tests..."
         MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests_VSA"

.github/ISSUE_TEMPLATE/3-RFC.yml

@@ -0,0 +1,56 @@
+name: 💬 Request for comments (RFC).
+description: Ask for feedback on major architectural changes or design choices.
+title: "[RFC]: "
+labels: ["RFC"]
+
+body:
+- type: markdown
+  attributes:
+    value: >
+      #### Please take a look at previous [RFCs](https://github.com/hao-ai-lab/FastVideo/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
+- type: textarea
+  attributes:
+    label: Motivation.
+    description: >
+      The motivation of the RFC.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Proposed Change.
+    description: >
+      The proposed change of the RFC.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Feedback Period.
+    description: >
+      The feedback period of the RFC. Usually at least one week.
+  validations:
+    required: false
+- type: textarea
+  attributes:
+    label: CC List.
+    description: >
+      The list of people you want to CC.
+  validations:
+    required: false
+- type: textarea
+  attributes:
+    label: Any Other Things.
+    description: >
+      Any other things you would like to mention.
+  validations:
+    required: false
+- type: markdown
+  attributes:
+    value: >
+      Thanks for contributing 🎉!
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues.
+        required: true

.gitignore

@@ -20,6 +20,7 @@ samples/
 data/
 outputs/
 outputs_video
+checkpoints/
 sbatch.sh
 *.out
 env
@@ -40,6 +41,7 @@ eggs/
 docs/_build/
 docs/source/getting_started/examples/
 docs/source/inference/examples/
+docs/source/training/examples/
 
 # VSCode
 .vscode/

cog.yaml

@@ -1,24 +0,0 @@
-# Configuration for Cog ⚙️
-# Reference: https://cog.run/yaml
-
-build:
-  gpu: true
-  cuda: "12.1"
-  python_version: "3.10"
-  python_packages:
-    - "torch==2.4.0"
-    - "torchvision"
-    - "ninja==1.11.1.3"
-    - "transformers==4.46.1"
-    - "git+https://github.com/huggingface/diffusers.git@bf64b32652a63a1865a0528a73a13652b201698b"
-    - "accelerate==1.0.1"
-    - "safetensors==0.4.5"
-    - "peft==0.13.2"
-    - "packaging==24.2"
-    - "git+https://github.com/hao-ai-lab/FastVideo"
-
-  run:
-    - FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation
-    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/latest/download/pget_$(uname -s)_$(uname -m)" && chmod +x /usr/local/bin/pget
-
-predict: "predict.py:Predictor"

csrc/attn/README.md

@@ -1,11 +1,23 @@
 
 
-# Sliding Tile Atteniton Kernel
+# Attention Kernel Used in FastVideo
 
+## Sliding Tile Attention (STA)
+We only support H100 for STA.
+```bash
+git submodule update --init --recursive
+python setup_sta.py install
+```
 
-## Installation
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only support H100/H200, because ThunderKittens uses TMA but doesn't support Blackwell yet.
-First, install C++20 for ThunderKittens:
+## Video Sparse Attention (VSA)
+We support H100 (via TK) and RTX 4090 (via triton) for VSA.
+```bash
+git submodule update --init --recursive
+python setup_vsa.py install
+```
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
 ```bash
 sudo apt update
 sudo apt install gcc-11 g++-11
@@ -15,27 +27,23 @@ sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave
 sudo apt update
 sudo apt install clang-11
 ```
-
-## Environment Setup
-First, set up your CUDA environment:
+(If you use CUDA12.4)
 ```bash
 export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-git submodule update --init --recursive
 ```
 
-## Install Sliding Tile Attention (STA)
-```bash
-python setup_sta.py install
-```
 
-## Install Video Sparse Attention (VSA)
-```bash
-python setup_vsa.py install
-```
 
 ## Usage
+### STA
+End-2-end inference with FastVideo:
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```
+
+If you want to use sliding tile attention in your custom model:
 ```python
 from st_attn import sliding_tile_attention
 # assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
@@ -47,9 +55,11 @@ from st_attn import sliding_tile_attention
 out = sliding_tile_attention(q, k, v, window_size, text_length)
 # If your attention does not contain text token (StepVideo)
 out = sliding_tile_attention(q, k, v, window_size, 0, False)
-
 ```
 
+### VSA
+We do not officially supoort end-2-end inference with VSA in FastVideo yet. Stay tuned.
+
 
 ## Test
 ```bash

csrc/attn/benchmarks/bench_vsa_hopper.py

@@ -1,7 +1,7 @@
 import torch
 import argparse
 from flash_attn.utils.benchmark import benchmark_forward
-from vsa import block_sparse_attention_fwd, block_sparse_attention_backward
+from vsa import block_sparse_fwd, block_sparse_bwd
 from vsa import BLOCK_M, BLOCK_N
 
 import numpy as np
@@ -130,12 +130,12 @@ def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_
     
     # Forward pass
     # Warm-up run
-    o, l_vec = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
+    o, l_vec = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
     torch.cuda.synchronize()
     
     # Benchmark forward
     _, fwd_time = benchmark_forward(
-        block_sparse_attention_fwd,
+        block_sparse_fwd,
         q, k, v, q2k_block_sparse_index, q2k_block_sparse_num,
         repeats=20,
         verbose=False,
@@ -150,12 +150,12 @@ def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_
     
     # Warm-up runs
     for _ in range(5):
-        block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
+        block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
     torch.cuda.synchronize()
     
     # Benchmark backward
     _, bwd_time = benchmark_forward(
-        block_sparse_attention_backward,
+        block_sparse_bwd,
         q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num,
         repeats=20,
         verbose=False,

csrc/attn/benchmarks/bench_vsa_triton.py

@@ -0,0 +1,217 @@
+import torch
+import argparse
+import triton.testing
+from vsa import block_sparse_attn
+from vsa import BLOCK_M, BLOCK_N
+
+import numpy as np
+import random
+
+def set_seed(seed: int = 42):
+    # Python random module
+    random.seed(seed)
+
+    # NumPy
+    np.random.seed(seed)
+
+    # PyTorch
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
+    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+    parser.add_argument('--num_heads', type=int, default=12, help='Number of heads')
+    parser.add_argument('--head_dim', type=int, default=64, help='Head dimension')
+    parser.add_argument('--topk', type=int, default=None, help='Number of kv blocks each q block attends to')
+    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[49152], help='Sequence lengths to benchmark')
+    return parser.parse_args()
+
+def create_input_tensors(batch, head, seq_len, headdim):
+    """Create random input tensors for attention."""
+    q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    return q, k, v
+
+def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"):
+    """
+    Generate a block sparse pattern where each q block attends to exactly k kv blocks.
+    
+    Args:
+        bs: batch size
+        h: number of heads
+        num_q_blocks: number of query blocks
+        num_kv_blocks: number of key-value blocks
+        k: number of kv blocks each q block attends to
+        device: device to create tensors on
+        
+    Returns:
+        q2k_block_sparse_index: [bs, h, num_q_blocks, k]
+            Contains the indices of kv blocks that each q block attends to.
+        q2k_block_sparse_num: [bs, h, num_q_blocks]
+            Contains the number of kv blocks that each q block attends to (all equal to k).
+        k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks]
+            Contains the indices of q blocks that attend to each kv block.
+        k2q_block_sparse_num: [bs, h, num_kv_blocks]
+            Contains the number of q blocks that attend to each kv block.
+        block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]
+            Binary mask where 1 indicates attention connection.
+    """
+    # Ensure k is not larger than num_kv_blocks
+    k = min(k, num_kv_blocks)
+    
+    # Create random scores for sampling
+    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device)
+    
+    # Get top-k indices for each q block
+    _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1)
+    q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32)
+    
+    # sort q2k_block_sparse_index
+    q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1)
+
+    # All q blocks attend to exactly k kv blocks
+    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device)
+    
+    # Create the corresponding mask
+    block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
+    
+    # Fill in the mask based on the indices
+    for b in range(bs):
+        for head in range(h):
+            for q_idx in range(num_q_blocks):
+                kv_indices = q2k_block_sparse_index[b, head, q_idx]
+                block_sparse_mask[b, head, q_idx, kv_indices] = True
+    
+    # Create the reverse mapping (k2q)
+    # First, initialize lists to collect q indices for each kv block
+    k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)]
+    
+    # Populate the lists based on q2k mapping
+    for b in range(bs):
+        for head in range(h):
+            flat_idx = b * h + head
+            for q_idx in range(num_q_blocks):
+                kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist()
+                for kv_idx in kv_indices:
+                    k2q_indices_list[flat_idx][kv_idx].append(q_idx)
+    
+    # Find the maximum number of q blocks that attend to any kv block
+    max_q_per_kv = 0
+    for flat_idx in range(bs * h):
+        for kv_idx in range(num_kv_blocks):
+            max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx]))
+    
+    # Create tensors for k2q mapping
+    k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, 
+                                        dtype=torch.int32, device=device)
+    k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), 
+                                       dtype=torch.int32, device=device)
+    
+    # Fill the tensors
+    for b in range(bs):
+        for head in range(h):
+            flat_idx = b * h + head
+            for kv_idx in range(num_kv_blocks):
+                q_indices = k2q_indices_list[flat_idx][kv_idx]
+                num_q = len(q_indices)
+                k2q_block_sparse_num[b, head, kv_idx] = num_q
+                if num_q > 0:
+                    k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor(
+                        q_indices, dtype=torch.int32, device=device)
+                
+    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
+
+def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops):
+    """Benchmark block sparse attention forward+backward pass."""
+    print("\n=== BLOCK SPARSE ATTENTION FORWARD+BACKWARD BENCHMARK ===")
+    
+    # Combined forward+backward pass
+    # Warm-up run
+    q_fwd = q.clone().requires_grad_(True)
+    k_fwd = k.clone().requires_grad_(True)
+    v_fwd = v.clone().requires_grad_(True)
+    o = block_sparse_attn(q_fwd, k_fwd, v_fwd, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
+    grad_output = torch.randn_like(o)
+    o.backward(grad_output)
+    torch.cuda.synchronize()
+    
+    # Benchmark forward+backward
+    def forward_backward_fn():
+        q_fwd = q.clone().requires_grad_(True)
+        k_fwd = k.clone().requires_grad_(True)
+        v_fwd = v.clone().requires_grad_(True)
+        o = block_sparse_attn(q_fwd, k_fwd, v_fwd, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
+        grad_output = torch.randn_like(o)
+        o.backward(grad_output)
+    
+    total_time = triton.testing.do_bench(
+        forward_backward_fn,
+        warmup=25,
+        rep=100,
+        return_mode='mean'
+    )
+    
+    # Total flops for forward + backward (forward + 2.5x backward approximation)
+    total_flops = flops + 2.5 * flops  # 3.5x the forward flops
+    sparse_tflops = total_flops / total_time * 1e-12 * 1e3
+    print(f"Block Sparse Forward+Backward - TFLOPS: {sparse_tflops:.2f}")
+    
+    return sparse_tflops
+
+def main():
+    args = parse_arguments()
+
+    set_seed(42)
+    
+    # Extract parameters
+    batch = args.batch_size
+    head = args.num_heads
+    headdim = args.head_dim
+    
+    print(f"Block Sparse Attention Benchmark")
+    print(f"batch: {batch}, head: {head}, headdim: {headdim}")
+    
+    # Test with different sequence lengths
+    for seq_len in args.seq_lengths:
+        # Skip very long sequences if they might cause OOM
+        if seq_len > 16384 and batch > 1:
+            continue
+            
+        print("="*100)
+        print(f"\nSequence length: {seq_len}")
+        
+        # Calculate theoretical FLOPs for attention
+        flops = 4 * batch * head * headdim * seq_len * seq_len
+        
+        # Create input tensors
+        q, k, v = create_input_tensors(batch, head, seq_len, headdim)
+        
+        # Setup block sparse parameters
+        num_q_blocks = seq_len // BLOCK_M
+        num_kv_blocks = seq_len // BLOCK_N
+        
+        # Determine k value (number of kv blocks per q block)
+        topk = args.topk
+        if topk is None:
+            topk = num_kv_blocks // 10  # Default to ~90% sparsity if k is not specified
+        topk = max(1, topk)       
+        print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)")
+        
+        # Generate block sparse pattern
+        q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, _ = generate_block_sparse_pattern(
+            batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda")
+        
+        # Benchmark block sparse attention
+        sparse_fwd = benchmark_block_sparse_attention(
+            q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops
+        )
+        
+        # Print results
+        print("\n=== PERFORMANCE RESULTS ===")
+        print(f"Block Sparse Forward+Backward - TFLOPS: {sparse_fwd:.2f}")
+
+if __name__ == "__main__":
+    main()

csrc/attn/scratch.sh

@@ -0,0 +1,4 @@
+    off_hz = tl.program_id(2)
+    b = off_hz // H
+    h = off_hz % H
+    meta_base = ((b * H + h) * q_tiles + q_blk)

csrc/attn/setup_vsa.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
 
-from csrc.attn.config_vsa import kernels, sources, target
+from config_vsa import kernels, sources, target
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
@@ -51,21 +51,29 @@ for k in kernels:
         source_files.append(sources[k]['source_files'][target])
     cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}')
 
+ext_modules = []
+import torch
+major, minor = torch.cuda.get_device_capability(0)
+if major == 9 and minor == 0:# check if H100
+    ext_modules = [
+        CUDAExtension('vsa_cuda',
+            sources=source_files,
+            extra_compile_args={
+                'cxx': cpp_flags,
+                'nvcc': cuda_flags
+            },
+            libraries=['cuda'])
+    ]
+
+
+
 setup(name=PACKAGE_NAME,
       version=VERSION,
       author=AUTHOR,
       description=DESCRIPTION,
       url=URL,
       packages=find_packages(),
-      ext_modules=[
-          CUDAExtension('vsa_cuda',
-                        sources=source_files,
-                        extra_compile_args={
-                            'cxx': cpp_flags,
-                            'nvcc': cuda_flags
-                        },
-                        libraries=['cuda'])
-      ],
+      ext_modules=ext_modules,
       cmdclass={'build_ext': BuildExtension},
       classifiers=[
           "Programming Language :: Python :: 3",

csrc/attn/tests/test_block_sparse.py

@@ -2,7 +2,7 @@ import torch
 import argparse
 from flash_attn.utils.benchmark import benchmark_forward
 from flash_attn import flash_attn_func
-from vsa import block_sparse_attention_fwd, block_sparse_attention_backward, BlockSparseAttentionFunction
+from vsa import block_sparse_attn
 from vsa import BLOCK_M, BLOCK_N
 
 import numpy as np
@@ -188,7 +188,7 @@ def main(args):
 
 
             # testing forward
-            o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
+            o = block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
             del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded
             grad_o = torch.randn_like(o)
             o.backward(grad_o)
@@ -215,7 +215,7 @@ def main(args):
             forward_metrics['l1'].append(l1)
             forward_metrics['rmse'].append(rmse)
 
-            print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            print(f"block_sparse_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
 
             # test backward
             o_sdpa.backward(grad_o)
@@ -228,7 +228,7 @@ def main(args):
             grad_q_metrics['sim'].append(sim)
             grad_q_metrics['l1'].append(l1)
             grad_q_metrics['rmse'].append(rmse)
-            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
+            print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
             
             sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad)
             assert sim > 0.9999, f"SSIM too low: {sim}"
@@ -237,7 +237,7 @@ def main(args):
             grad_k_metrics['sim'].append(sim)
             grad_k_metrics['l1'].append(l1)
             grad_k_metrics['rmse'].append(rmse)
-            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
             
             sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad)
             assert sim > 0.9999, f"SSIM too low: {sim}"
@@ -246,7 +246,7 @@ def main(args):
             grad_v_metrics['sim'].append(sim)
             grad_v_metrics['l1'].append(l1)
             grad_v_metrics['rmse'].append(rmse)
-            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
             
             del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa
             gc.collect()

csrc/attn/tests/test_block_sparse_triton.py

@@ -0,0 +1,289 @@
+import torch
+import argparse
+from flash_attn.utils.benchmark import benchmark_forward
+from flash_attn import flash_attn_func
+from vsa import triton_attention_sparse
+from vsa import BLOCK_M, BLOCK_N
+
+import numpy as np
+import random
+import gc
+
+def set_seed(seed: int = 42):
+    # Python random module
+    random.seed(seed)
+
+    # NumPy
+    np.random.seed(seed)
+
+    # PyTorch
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
+
+
+@torch.no_grad
+def precision_metric(quant_o, fa2_o): 
+    x, xx = quant_o.float(), fa2_o.float() 
+    sim = torch.nn.functional.cosine_similarity(x.reshape(1, -1), xx.reshape(1, -1)).item()
+    l1 =   ((x - xx).abs().sum() / xx.abs().sum() ).item()
+    rmse = torch.sqrt(torch.mean((x -xx) ** 2)).item()
+
+    return sim, l1, rmse
+
+def create_input_tensors(batch, head, seq_len, headdim):
+    """Create random input tensors for attention."""
+    q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    return q, k, v
+
+def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"):
+    """
+    Generate a block sparse pattern where each q block attends to exactly k kv blocks.
+    
+    Args:
+        bs: batch size
+        h: number of heads
+        num_q_blocks: number of query blocks
+        num_kv_blocks: number of key-value blocks
+        k: number of kv blocks each q block attends to
+        device: device to create tensors on
+        
+    Returns:
+        q2k_block_sparse_index: [bs, h, num_q_blocks, k]
+            Contains the indices of kv blocks that each q block attends to.
+        q2k_block_sparse_num: [bs, h, num_q_blocks]
+            Contains the number of kv blocks that each q block attends to (all equal to k).
+        k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks]
+            Contains the indices of q blocks that attend to each kv block.
+        k2q_block_sparse_num: [bs, h, num_kv_blocks]
+            Contains the number of q blocks that attend to each kv block.
+        block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]
+            Binary mask where 1 indicates attention connection.
+    """
+    # Ensure k is not larger than num_kv_blocks
+    k = min(k, num_kv_blocks)
+    
+    # Create random scores for sampling
+    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device)
+    
+    # Get top-k indices for each q block
+    _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1)
+    q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32)
+
+    # sort q2k_block_sparse_index
+    q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1)
+
+    # All q blocks attend to exactly k kv blocks
+    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device)
+    
+    # Create the corresponding mask
+    block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
+    
+    # Fill in the mask based on the indices
+    for b in range(bs):
+        for head in range(h):
+            for q_idx in range(num_q_blocks):
+                kv_indices = q2k_block_sparse_index[b, head, q_idx]
+                block_sparse_mask[b, head, q_idx, kv_indices] = True
+    
+    # Create the reverse mapping (k2q)
+    # First, initialize lists to collect q indices for each kv block
+    k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)]
+    
+    # Populate the lists based on q2k mapping
+    for b in range(bs):
+        for head in range(h):
+            flat_idx = b * h + head
+            for q_idx in range(num_q_blocks):
+                kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist()
+                for kv_idx in kv_indices:
+                    k2q_indices_list[flat_idx][kv_idx].append(q_idx)
+    
+    # Find the maximum number of q blocks that attend to any kv block
+    max_q_per_kv = 0
+    for flat_idx in range(bs * h):
+        for kv_idx in range(num_kv_blocks):
+            max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx]))
+    
+    # Create tensors for k2q mapping
+    k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, 
+                                        dtype=torch.int32, device=device)
+    k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), 
+                                       dtype=torch.int32, device=device)
+    
+    # Fill the tensors
+    for b in range(bs):
+        for head in range(h):
+            flat_idx = b * h + head
+            for kv_idx in range(num_kv_blocks):
+                q_indices = k2q_indices_list[flat_idx][kv_idx]
+                num_q = len(q_indices)
+                k2q_block_sparse_num[b, head, kv_idx] = num_q
+                if num_q > 0:
+                    k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor(
+                        q_indices, dtype=torch.int32, device=device)
+                
+    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
+
+def main(args):
+    set_seed(42)
+    
+    # Extract parameters
+    batch = args.batch_size
+    head = args.num_heads
+    headdim = args.head_dim
+    num_iterations = args.num_iterations
+    
+    print(f"Block Sparse Attention Benchmark")
+    print(f"batch: {batch}, head: {head}, headdim: {headdim}, iterations: {num_iterations}")
+    
+    # Test with different sequence lengths
+    for seq_len in args.seq_lengths:
+        # Skip very long sequences if they might cause OOM
+        # if seq_len > 16384 and batch > 1:
+        #     continue
+            
+        print("="*100)
+        print(f"\nSequence length: {seq_len}")
+        
+        # Collect metrics across iterations
+        forward_metrics = {'sim': [], 'l1': [], 'rmse': []}
+        grad_q_metrics = {'sim': [], 'l1': [], 'rmse': []}
+        grad_k_metrics = {'sim': [], 'l1': [], 'rmse': []}
+        grad_v_metrics = {'sim': [], 'l1': [], 'rmse': []}
+        
+        for iter_idx in range(num_iterations):
+            if num_iterations > 1:
+                print(f"\nIteration {iter_idx+1}/{num_iterations}")
+            
+            # Create input tensors
+            q, k, v = create_input_tensors(batch, head, seq_len, headdim)
+            
+            # Setup block sparse parameters
+            num_q_blocks = seq_len // BLOCK_M
+            num_kv_blocks = seq_len // BLOCK_N
+            
+            # Determine k value (number of kv blocks per q block)
+            topk = args.topk
+            if topk is None:
+                topk = num_kv_blocks // 10  # Default to ~90% sparsity if k is not specified
+            topk = max(1, topk)
+            if iter_idx == 0:  # Only print this once
+                print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)")
+
+            # Generate block sparse pattern
+            q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask = generate_block_sparse_pattern(
+                batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda")
+
+            # expand block_sparse_mask to full mask
+            block_mask_expanded = block_sparse_mask.unsqueeze(-1).unsqueeze(-2)  # [b, h, num_q_blocks, num_kv_blocks, 1, 1]
+            block_mask_expanded = block_mask_expanded.expand(-1, -1, -1, -1, BLOCK_M, BLOCK_N)  # [b, h, num_q_blocks, num_kv_blocks, BLOCK_M, BLOCK_N]
+            full_mask = block_mask_expanded.permute(0, 1, 2, 4, 3, 5).reshape(batch, head, seq_len, seq_len)
+
+            q.requires_grad = True
+            k.requires_grad = True
+            v.requires_grad = True
+
+
+            # testing forward
+            o = triton_attention_sparse(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
+            del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded
+            grad_o = torch.randn_like(o)
+            o.backward(grad_o)
+            # clear memory
+            q_sdpa = q.detach().clone()
+            k_sdpa = k.detach().clone()
+            v_sdpa = v.detach().clone()
+            q_sdpa.requires_grad = True
+            k_sdpa.requires_grad = True
+            v_sdpa.requires_grad = True
+            q.data = torch.empty(0, device=q.device)
+            k.data = torch.empty(0, device=k.device)
+            v.data = torch.empty(0, device=v.device)
+            torch.cuda.empty_cache()
+
+            o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask)
+
+
+            sim, l1, rmse = precision_metric(o, o_sdpa)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 8e-5, f"l1 too large: {l1}"
+            assert rmse < 5e-5, f"RMSE too large: {rmse}"
+            forward_metrics['sim'].append(sim)
+            forward_metrics['l1'].append(l1)
+            forward_metrics['rmse'].append(rmse)
+
+            print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+
+            # test backward
+            o_sdpa.backward(grad_o)
+
+            sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad)
+            # Error bounds collected on H100
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 5e-4, f"RMSE too large: {rmse}"
+            grad_q_metrics['sim'].append(sim)
+            grad_q_metrics['l1'].append(l1)
+            grad_q_metrics['rmse'].append(rmse)
+            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
+            
+            sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 5e-4, f"RMSE too large: {rmse}"
+            grad_k_metrics['sim'].append(sim)
+            grad_k_metrics['l1'].append(l1)
+            grad_k_metrics['rmse'].append(rmse)
+            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            
+            sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 5e-4, f"RMSE too large: {rmse}"
+            grad_v_metrics['sim'].append(sim)
+            grad_v_metrics['l1'].append(l1)
+            grad_v_metrics['rmse'].append(rmse)
+            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            
+            del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        # Print summary statistics if multiple iterations were run
+        if num_iterations > 1:
+            print("\n" + "="*50)
+            print(f"Summary Statistics (over {num_iterations} iterations):")
+            
+            print("\nForward metrics:")
+            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}")
+            
+            print("\nGradient Q metrics:")
+            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}")
+            
+            print("\nGradient K metrics:")
+            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}")
+            
+            print("\nGradient V metrics:")
+            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
+    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
+    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
+    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
+    parser.add_argument('--topk', type=int, default=4, help='Number of kv blocks each q block attends to')
+    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[4096], help='Sequence lengths to benchmark')
+    parser.add_argument('--num_iterations', type=int, default=10, help='Number of test iterations to run')
+    args = parser.parse_args()
+    main(args)

csrc/attn/tests/test_triton.py

@@ -3,10 +3,15 @@ from tqdm import tqdm
 import matplotlib.pyplot as plt
 import numpy as np
 
+from vsa import triton_attention
+
 def pytorch_test(Q, K, V, dO):
     q_ = Q.to(torch.float64).requires_grad_()
     k_ = K.to(torch.float64).requires_grad_()
     v_ = V.to(torch.float64).requires_grad_()
+    q_.grad = None
+    k_.grad = None
+    v_.grad = None
     dO_ = dO.to(torch.float64)
     
     # manual pytorch implementation of scaled dot product attention
@@ -30,11 +35,33 @@ def fa2_test(Q, K, V, dO):
     Q.requires_grad = True
     K.requires_grad = True
     V.requires_grad = True
+    Q.grad = None
+    K.grad = None
+    V.grad = None
+    
     output = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False)
     output.backward(dO)
     
     return output, Q.grad, K.grad, V.grad
 
+def triton_test(Q, K, V, dO):
+    Q.requires_grad = True
+    K.requires_grad = True
+    V.requires_grad = True
+    Q.grad = None
+    K.grad = None
+    V.grad = None
+
+    output = triton_attention(Q, K, V)
+
+    output.backward(dO)
+    
+    q_grad = Q.grad
+    k_grad = K.grad
+    v_grad = V.grad
+    
+    return output.to(Q.dtype) if output is not None else None, q_grad, k_grad, v_grad
+
 def generate_tensor(shape, mean, std, dtype, device):
     tensor = torch.randn(shape, dtype=dtype, device=device)
 
@@ -45,7 +72,8 @@ def generate_tensor(shape, mean, std, dtype, device):
 
 def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'):
     results = {
-        'FA2 vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0},
+        'FA2 vs PT': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
+        'Triton vs PT': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
     }
 
     for _ in range(num_iterations):
@@ -58,21 +86,31 @@ def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all
         
         pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO)
         fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO)
+        triton_o, triton_qg, triton_kg, triton_vg = triton_test(Q, K, V, dO)
         
         if test_mode == 'forward_only':
             tensors_fa2_pt = [(pt_o, fa2_o)]
+            tensors_triton_pt = [(pt_o, triton_o)]
         else:  # 'forward_backward'
             if error_mode == 'output':
                 tensors_fa2_pt = [(pt_o, fa2_o)]
+                tensors_triton_pt = [(pt_o, triton_o)]
             elif error_mode == 'backward':
                 tensors_fa2_pt = [(pt_qg, fa2_qg),
                                (pt_kg, fa2_kg),
                                (pt_vg, fa2_vg)]
+                tensors_triton_pt = [(pt_qg, triton_qg),
+                                 (pt_kg, triton_kg),
+                                 (pt_vg, triton_vg)]
             else:  # 'all'
                 tensors_fa2_pt = [(pt_o, fa2_o),
                                (pt_qg, fa2_qg),
                                (pt_kg, fa2_kg),
                                (pt_vg, fa2_vg)]
+                tensors_triton_pt = [(pt_o, triton_o),
+                                 (pt_qg, triton_qg),
+                                 (pt_kg, triton_kg),
+                                 (pt_vg, triton_vg)]
                 
         for pt, fa2 in tensors_fa2_pt:
             diff = pt - fa2
@@ -80,6 +118,13 @@ def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all
             results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item()
             results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
             results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item())
+        
+        for pt, triton in tensors_triton_pt:
+            diff = pt - triton
+            abs_diff = torch.abs(diff)
+            results['Triton vs PT']['sum_diff'] += torch.sum(abs_diff).item()
+            results['Triton vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
+            results['Triton vs PT']['max_diff'] = max(results['Triton vs PT']['max_diff'], torch.max(abs_diff).item())
                 
         torch.cuda.empty_cache()
 
@@ -99,38 +144,40 @@ def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all
 def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'):
     seq_lengths = [768 * (2**i) for i in range(1)]
 
-    print(f"\n{'='*80}")
+    print(f"\n{'='*100}")
     print(f"ATTENTION ERROR COMPARISON TABLE (b={b}, h={h}, d={d}, mean={mean}, std={std})")
     print(f"Mode: {error_mode}, Test: {test_mode}")
-    print(f"{'='*80}")
+    print(f"{'='*100}")
     
     # Print header
-    print(f"{'Seq Length':<12} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15}")
-    print(f"{'-'*12} | {'-'*15} | {'-'*15}")
+    print(f"{'Seq Length':<12} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15} | {'Triton vs PT Avg':<15} | {'Triton vs PT Max':<15}")
+    print(f"{'-'*12} | {'-'*15} | {'-'*15} | {'-'*15} | {'-'*15}")
     
     for n in seq_lengths:
         results = check_correctness(b, h, n, d, mean, std, error_mode=error_mode, test_mode=test_mode)
         
         fa2_pt_avg = results['FA2 vs PT']['avg_diff']
         fa2_pt_max = results['FA2 vs PT']['max_diff']
+        triton_pt_avg = results['Triton vs PT']['avg_diff']
+        triton_pt_max = results['Triton vs PT']['max_diff']
         
-        # Print row
-        print(f"{n:<12} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e}")
+        # Print row with both comparisons
+        print(f"{n:<12} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e} | {triton_pt_avg:<15.6e} | {triton_pt_max:<15.6e}")
     
-    print(f"{'='*80}\n")
+    print(f"{'='*100}\n")
 
 # fix random seed
 torch.manual_seed(0)
 
-# Example usage
-b, h, d = 2, 2, 64
 mean = 1e-1
 std = 10
+configs = [
+    (4, 1, 128),  # Larger batch, single head, larger dim
+    (2, 8, 64),   # Medium batch, many heads, medium dim
+]
 
-# Test forward only
-generate_error_tables(b, h, d, mean, std, error_mode='output', test_mode='forward_only')
-
-# Test forward and backward
-generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward')
+for b, h, d in configs:
+    print(f"\nConfiguration: batch={b}, heads={h}, dim={d}")
+    generate_error_tables(b, h, d, mean, std, error_mode='backward', test_mode='forward_backward')
 
 print("Attention error comparison completed.")

csrc/attn/vsa/__init__.py

@@ -1,18 +1,27 @@
-import math
-
 import torch
-from torch.utils.checkpoint import detach_variable
 from typing import Tuple
+from vsa.vsa import block_sparse_attn
 try:
     from vsa_cuda import block_sparse_fwd, block_sparse_bwd
 except ImportError:
     block_sparse_fwd = None
     block_sparse_bwd = None
-
-
+from vsa.block_sparse_attn_triton import attention as triton_attention, attention_sparse as triton_attention_sparse
 BLOCK_M = 64
 BLOCK_N = 64
 
+
+def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
+    QK = torch.matmul(q, k.transpose(-2, -1))
+    QK /= (q.size(-1)**0.5)
+
+    # Causal mask removed since causal is always false
+
+    QK = torch.nn.functional.softmax(QK, dim=-1)
+    output = torch.matmul(QK, v)
+    return output, QK
+
+
 def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None):
     """
     q: [batch_size, num_heads, seq_len, head_dim]
@@ -67,15 +76,6 @@ def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None):
         final_output = output_compress + output_select
     return final_output
 
-def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
-    QK = torch.matmul(q, k.transpose(-2, -1))
-    QK /= (q.size(-1)**0.5)
-
-    # Causal mask removed since causal is always false
-
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v)
-    return output, QK
 
 def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,
                                        topk: int):
@@ -123,107 +123,11 @@ def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,
 
     return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
 
-@torch._dynamo.disable
-def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
-    """
-    Differentiable block sparse attention function.
-    
-    Args:
-        q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
-        k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
-        v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
-        q2k_block_sparse_index: Indices for query-to-key sparse blocks
-        q2k_block_sparse_num: Number of sparse blocks for each query block
-        k2q_block_sparse_index: Indices for key-to-query sparse blocks (for backward pass)
-        k2q_block_sparse_num: Number of sparse blocks for each key block (for backward pass)
-    
-    Returns:
-        output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
-    """
-    return BlockSparseAttentionFunction.apply(
-        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
-    )
-
-def block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num):
-    """
-    block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]. 
-        [*, *, i, j] = 1 means the i-th q block should attend to the j-th kv block.
-    """
-    # assert all elements in q2k_block_sparse_num can be devisible by 2
-    o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-    return o, lse
-
-def block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num):
-    grad_output = grad_output.contiguous()
-    grad_q, grad_k, grad_v = block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
-    return grad_q, grad_k, grad_v
 
 ## pytorch sdpa version of block sparse ##
 import triton
 import triton.language as tl
 
-@triton.jit
-def index_to_mask_kernel(
-    q2k_block_sparse_index_ptr,
-    q2k_block_sparse_num_ptr,
-    mask_ptr,
-    batch_size: tl.constexpr,
-    num_heads: tl.constexpr,
-    num_q_blocks: tl.constexpr,
-    num_k_blocks: tl.constexpr,
-    max_kv_blocks: tl.constexpr,
-    BLOCK_Q: tl.constexpr,
-    BLOCK_K: tl.constexpr,
-):
-    bh, q, id = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64)
-    b = bh // num_heads
-    h = bh % num_heads
-
-    num_valid_blocks = tl.load(q2k_block_sparse_num_ptr + b * num_heads * num_q_blocks + h * num_q_blocks + q)
-
-    if num_valid_blocks <= id:
-        return
-    k = tl.load(q2k_block_sparse_index_ptr + b * num_heads * num_q_blocks * max_kv_blocks + h * num_q_blocks * max_kv_blocks + q * max_kv_blocks + id)
-
-    full_mask = (tl.arange(0, BLOCK_Q)[:, None] < BLOCK_Q) & (tl.arange(0, BLOCK_K)[None, :] < BLOCK_K)
-
-    q_lengths = num_q_blocks * BLOCK_Q
-    k_lengths = num_k_blocks * BLOCK_K
-    mask_ptr_base = mask_ptr + b * num_heads * q_lengths * k_lengths + h * q_lengths * k_lengths + q * BLOCK_Q * k_lengths + k * BLOCK_K
-
-    tl.store(mask_ptr_base + tl.arange(0, BLOCK_Q)[:, None] * k_lengths + tl.arange(0, BLOCK_K)[None, :], full_mask)
-
-def index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, BLOCK_Q, BLOCK_K, num_k_blocks):
-    """
-    Convert block sparse indices to a mask.
-    
-    Args:
-        q2k_block_sparse_index: Indices for query-to-key sparse blocks
-        q2k_block_sparse_num: Number of sparse blocks for each query block
-    
-    Returns:
-        mask: Block sparse mask tensor
-    """
-    batch_size, num_heads, num_q_blocks, max_kv_blocks = q2k_block_sparse_index.shape
-    assert q2k_block_sparse_num.shape == (batch_size, num_heads, num_q_blocks)
-
-    mask = torch.zeros((batch_size, num_heads, num_q_blocks * BLOCK_Q, num_k_blocks * BLOCK_K), dtype=torch.bool, device=q2k_block_sparse_index.device)
-
-    grid = (batch_size * num_heads, num_q_blocks, max_kv_blocks)
-    index_to_mask_kernel[grid](
-        q2k_block_sparse_index,
-        q2k_block_sparse_num,
-        mask,
-        batch_size,
-        num_heads,
-        num_q_blocks,
-        num_k_blocks,
-        max_kv_blocks,
-        BLOCK_Q=BLOCK_Q,
-        BLOCK_K=BLOCK_K,
-    )
-
-    return mask
 
 @triton.jit
 def topk_index_to_map_kernel(
@@ -372,99 +276,3 @@ def map_to_index(block_map: torch.Tensor):
 
     return index, index_num
 
-class BlockSparseAttentionFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
-        o, lse = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-        ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num)
-        return o
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors
-        grad_q, grad_k, grad_v = block_sparse_attention_backward(
-            q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num
-        )
-        return grad_q, grad_k, grad_v, None, None, None, None
-
-
-class DummyOperator(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x):
-        return x
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        return grad_output
-
-class CheckpointSDPA(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, obj, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k):
-        """Forward pass."""
-        with torch.no_grad():
-            mask = index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k, k.shape[2] // block_k)
-            outputs = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-        ctx.save_for_backward(*detach_variable((q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)))
-        ctx.block_q = block_q
-        ctx.block_k = block_k
-        # the obj is passed in, then it can access the saved input
-        # tensors later for recomputation
-        obj.ctx = ctx
-        return outputs
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        """Backward pass."""
-        inputs = ctx.saved_tensors
-        output = ctx.output
-        torch.autograd.backward(output, grad_output)
-        ctx.output = None
-        grads = tuple(inp.grad for inp in inputs)
-        return (None, ) + grads + (None, None)
-
-
-class BlockSparseAttnTorch:
-    def __init__(self):
-        self.ctx = None
-    
-    def recompute_mask(self, _):
-        recomputed_mask = index_to_mask(self.q2k_block_sparse_index, self.q2k_block_sparse_num, self.block_q, self.block_k, self.num_kv_blocks)
-        mask_size = recomputed_mask.untyped_storage().size()
-        self.mask.untyped_storage().resize_(mask_size)
-        self.mask.untyped_storage().copy_(recomputed_mask.untyped_storage()) 
-    
-    def recompute(self, _):
-        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num = self.ctx.saved_tensors
-        block_q = self.ctx.block_q
-        block_k = self.ctx.block_k
-        mask = index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k, k.shape[2] // block_k)
-        with torch.enable_grad():
-            output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-        self.ctx.output = output
-        self.ctx = None
-
-    @torch._dynamo.disable
-    def forward(self, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k):
-        """
-        Differentiable block sparse attention function using PyTorch.
-        
-        Args:
-            q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
-            k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
-            v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
-            q2k_block_sparse_index: Indices for query-to-key sparse blocks
-            q2k_block_sparse_num: Number of sparse blocks for each query block
-            block_q: Block size for query
-            block_k: Block size for key-value
-
-        Returns:
-            output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
-        """
-
-        output = CheckpointSDPA.apply(
-            self, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k
-        )
-
-        o = DummyOperator.apply(output)
-        o.register_hook(self.recompute)
-        return o

csrc/attn/vsa/block_sparse_attn_triton.py

@@ -0,0 +1,707 @@
+"""
+Fused Attention
+===============
+
+This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
+(https://tridao.me/publications/flash2/flash2.pdf)
+
+Credits: OpenAI kernel team
+"""
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+import math  # small utility needed by the sparse wrapper
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+@triton.jit
+def _attn_fwd_inner(acc, l_i, m_i, q,  #
+                    K_block_ptr, V_block_ptr,  #
+                    start_m, qk_scale,  #
+                    BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr,  #
+                    STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr,  #
+                    N_CTX: tl.constexpr, fp8_v: tl.constexpr):
+
+    # loop over k, v and update accumulator
+    for start_n in range(0, N_CTX, BLOCK_N):
+        # -- compute qk ----
+        k = tl.load(K_block_ptr)
+        qk = tl.dot(q, k)
+
+        m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
+        qk = qk * qk_scale - m_ij[:, None]
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        # -- update m_i and l_i
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        # -- update output accumulator --
+        acc = acc * alpha[:, None]
+        # update acc
+        v = tl.load(V_block_ptr)
+        if fp8_v:
+            p = p.to(tl.float8e5)
+        else:
+            p = p.to(tl.bfloat16)
+        acc = tl.dot(p, v, acc)
+        # update m_i and l_i
+        m_i = m_ij
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+    return acc, l_i, m_i
+
+
+# We don't run auto-tuning every time to keep the tutorial fast. Keeping
+# the code below and commenting out the equivalent parameters is convenient for
+# re-tuning.
+configs = [
+    triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+    for BM in [64]\
+    for BN in [64]\
+    for s in [3, 4, 7]\
+    for w in [4, 8]\
+]
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
+                     q2k_index, q2k_num, max_kv_blks,           #
+                     M, Out,                                    #
+                     stride_qz, stride_qh, stride_qm, stride_qk,
+                     stride_kz, stride_kh, stride_kn, stride_kk,
+                     stride_vz, stride_vh, stride_vk, stride_vn,
+                     stride_oz, stride_oh, stride_om, stride_on,
+                     Z, H, N_CTX,                               #
+                     HEAD_DIM: tl.constexpr,                    #
+                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
+                     STAGE: tl.constexpr):
+    """
+    64×64 **block-sparse** forward kernel. Back-prop kernels remain dense
+    (32×64 and 64×32) – memory footprint unchanged.
+    """
+
+    # ----- program-id mapping -----
+    q_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(1)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_M
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
+    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
+
+    # ----- base pointers -----
+    qvk_off = (b.to(tl.int64) * stride_qz +
+               h.to(tl.int64) * stride_qh)
+
+    Q_ptr = tl.make_block_ptr(
+        base=Q + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(q_blk * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
+
+    K_base = tl.make_block_ptr(
+        base=K + qvk_off, shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N), order=(0, 1))
+
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
+    V_base = tl.make_block_ptr(
+        base=V + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM), order=v_order)
+
+    O_ptr = tl.make_block_ptr(
+        base=Out + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_on),
+        offsets=(q_blk * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
+
+    # ----- accumulators -----
+    offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    qk_scale = sm_scale * 1.44269504  # 1/ln2
+    q = tl.load(Q_ptr)
+
+    # ----- sparse loop over valid K/V tiles -----
+    for i in range(0, kv_blocks):
+        kv_idx = tl.load(kv_ptr + i).to(tl.int32)
+
+        K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N))
+        V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0))
+
+        k = tl.load(K_ptr)
+        qk = tl.dot(q, k)
+
+        m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
+        p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
+        l_ij = tl.sum(p, 1)
+
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        acc = acc * alpha[:, None]
+
+        v = tl.load(V_ptr)
+        acc = tl.dot(p.to(tl.bfloat16), v, acc)
+        m_i = m_ij
+
+    # ----- epilogue -----
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    tl.store(M + off_hz * N_CTX + offs_m, m_i)
+    tl.store(O_ptr, acc.to(Out.type.element_ty))
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd(Q, K, V, sm_scale, M, Out,  #
+              stride_qz, stride_qh, stride_qm, stride_qk,  #
+              stride_kz, stride_kh, stride_kn, stride_kk,  #
+              stride_vz, stride_vh, stride_vk, stride_vn,  #
+              stride_oz, stride_oh, stride_om, stride_on,  #
+              Z, H, N_CTX,  #
+              HEAD_DIM: tl.constexpr,  #
+              BLOCK_M: tl.constexpr,  #
+              BLOCK_N: tl.constexpr,  #
+              STAGE: tl.constexpr  #
+              ):
+    tl.static_assert(BLOCK_N <= HEAD_DIM)
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+
+    # block pointers
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
+    V_block_ptr = tl.make_block_ptr(
+        base=V + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=v_order,
+    )
+    K_block_ptr = tl.make_block_ptr(
+        base=K + qvk_offset,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N),
+        order=(0, 1),
+    )
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    # load scales
+    qk_scale = sm_scale
+    qk_scale *= 1.44269504  # 1/log(2)
+    # load q: it will stay in SRAM throughout
+    q = tl.load(Q_block_ptr)
+    acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr,  #
+                                    start_m, qk_scale,  #
+                                    BLOCK_M, HEAD_DIM, BLOCK_N,  #
+                                    3, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5  #
+                                    )
+
+    # epilogue
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    m_ptrs = M + off_hz * N_CTX + offs_m
+    tl.store(m_ptrs, m_i)
+    tl.store(O_block_ptr, acc.to(Out.type.element_ty))
+
+
+@triton.jit
+def _attn_bwd_preprocess(O, DO,  #
+                         Delta,  #
+                         Z, H, N_CTX,  #
+                         BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr  #
+                         ):
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_hz = tl.program_id(1)
+    off_n = tl.arange(0, HEAD_DIM)
+    # load
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hz * N_CTX + off_m, delta)
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _attn_bwd_dkdv(dk, dv,  #
+                   Q, k, v, sm_scale,  #
+                   DO,  #
+                   M, D,  #
+                   k2q_index, k2q_num, max_q_blks,
+                   # shared by Q/K/V/DO.
+                   stride_tok, stride_d,  #
+                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
+                   BLOCK_N1: tl.constexpr,  #
+                   HEAD_DIM: tl.constexpr,  #
+                   # Filled in by the wrapper.
+                   start_n, start_m, num_steps):
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    step_m = BLOCK_M1
+    kv_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(2)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_N1
+    meta_base = ((b * H + h) * q_tiles + kv_blk)
+
+    q_blocks = tl.load(k2q_num  + meta_base)                 # int32
+    q_ptr    = k2q_index + meta_base * max_q_blks           # ptr to list
+    
+        
+    for blk_idx in range(q_blocks*2):
+        block_sparse_offset = (tl.load(q_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_m
+        qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
+        m = tl.load(M + offs_m)
+        qkT = tl.dot(k, qT)
+        pT = tl.math.exp2(qkT - m[None, :])
+        do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
+        # Compute dV.
+        ppT = pT
+        ppT = ppT.to(tl.bfloat16)
+        dv += tl.dot(ppT, do)
+        # D (= delta) is pre-divided by ds_scale.
+        Di = tl.load(D + offs_m)
+        # Compute dP and dS.
+        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
+        dsT = pT * (dpT - Di[None, :])
+        dsT = dsT.to(tl.bfloat16)
+        dk += tl.dot(dsT, tl.trans(qT))
+        # Increment pointers.
+    return dk, dv
+
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq(dq, q, K, V,  #
+                 do, m, D,
+                 # shared by Q/K/V/DO.
+                 q2k_index, q2k_num, max_kv_blks,
+                 stride_tok, stride_d,  #
+                 H, N_CTX,  #
+                 BLOCK_M2: tl.constexpr,  #
+                 BLOCK_N2: tl.constexpr,  #
+                 HEAD_DIM: tl.constexpr,
+                 # Filled in by the wrapper.
+                 start_m, start_n, num_steps):
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+    offs_n = start_n + tl.arange(0, BLOCK_N2)
+    offs_k = tl.arange(0, HEAD_DIM)
+    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    # D (= delta) is pre-divided by ds_scale.
+    Di = tl.load(D + offs_m)
+    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
+    step_n = BLOCK_N2
+    
+    q_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(2)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_M2
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
+    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
+    
+    
+    for blk_idx in range(kv_blocks*2):
+        block_sparse_offset = (tl.load(kv_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_n * stride_tok
+        kT = tl.load(kT_ptrs + block_sparse_offset)
+        vT = tl.load(vT_ptrs + block_sparse_offset)
+        qk = tl.dot(q, kT)
+        p = tl.math.exp2(qk - m)
+        # Compute dP and dS.
+        dp = tl.dot(do, vT).to(tl.float32)
+        ds = p * (dp - Di[:, None])
+        ds = ds.to(tl.bfloat16)
+        # Compute dQ.
+        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        dq += tl.dot(ds, tl.trans(kT))
+        # Increment pointers.
+    return dq
+
+
+
+@triton.jit
+def _attn_bwd(Q, K, V, sm_scale,  #
+              DO,  #
+              DQ, DK, DV,  #
+              M, D,
+              q2k_index, q2k_num, max_kv_blks,
+              k2q_index, k2q_num, max_q_blks,
+              # shared by Q/K/V/DO.
+              stride_z, stride_h, stride_tok, stride_d,  #
+              H, N_CTX,  #
+              BLOCK_M1: tl.constexpr,  #
+              BLOCK_N1: tl.constexpr,  #
+              BLOCK_M2: tl.constexpr,  #
+              BLOCK_N2: tl.constexpr,  #
+              HEAD_DIM: tl.constexpr):
+    LN2 = 0.6931471824645996  # = ln(2)
+
+    bhid = tl.program_id(2)
+    off_chz = (bhid * N_CTX).to(tl.int64)
+    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
+    pid = tl.program_id(0)
+
+    # offset pointers for batch/head
+    Q += adj
+    K += adj
+    V += adj
+    DO += adj
+    DQ += adj
+    DK += adj
+    DV += adj
+    M += off_chz
+    D += off_chz
+
+    # load scales
+    offs_k = tl.arange(0, HEAD_DIM)
+
+    start_n = pid * BLOCK_N1
+    start_m = 0
+
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+
+    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+
+    num_steps = N_CTX // BLOCK_M1
+
+    dk, dv = _attn_bwd_dkdv(  #
+        dk, dv,  #
+        Q, k, v, sm_scale,  #
+        DO,  #
+        M, D,  #
+        k2q_index, k2q_num, max_q_blks,
+        stride_tok, stride_d,  #
+        H, N_CTX,  #
+        BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
+        start_n, start_m, num_steps  #
+    )
+
+    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dv_ptrs, dv)
+
+    # Write back dK.
+    dk *= sm_scale
+    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dk_ptrs, dk)
+
+    # THIS BLOCK DOES DQ:
+    start_m = pid * BLOCK_M2
+    end_n = 0
+
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+
+    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    m = tl.load(M + offs_m)
+    m = m[:, None]
+
+    num_steps = N_CTX // BLOCK_N2
+    dq = _attn_bwd_dq(dq, q, K, V,  #
+                      do, m, D,  #
+                      q2k_index, q2k_num, max_kv_blks,
+                      stride_tok, stride_d,  #
+                      H, N_CTX,  #
+                      BLOCK_M2, BLOCK_N2, HEAD_DIM,  #
+                      start_m, end_n, num_steps  #
+                      )
+    # Write back dQ.
+    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    dq *= LN2
+    tl.store(dq_ptrs, dq)
+    
+    
+
+class _attention(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, k, v):
+        # shape constraints
+        sm_scale = 1.0 / (q.shape[-1] ** 0.5)
+        HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
+        # when v is in float8_e5m2 it is transposed.
+        HEAD_DIM_V = v.shape[-1]
+        assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
+        assert HEAD_DIM_K in {16, 32, 64, 128, 256}
+        o = torch.empty_like(q)
+        stage = 1
+        extra_kern_args = {}
+
+
+        grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
+        M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+        _attn_fwd[grid](
+            q, k, v, sm_scale, M, o,  #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            k.stride(0), k.stride(1), k.stride(2), k.stride(3),  #
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),  #
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),  #
+            q.shape[0], q.shape[1],  #
+            N_CTX=q.shape[2],  #
+            HEAD_DIM=HEAD_DIM_K,  #
+            STAGE=stage,  #
+            **extra_kern_args)
+
+        ctx.save_for_backward(q, k, v, o, M)
+        ctx.grid = grid
+        ctx.sm_scale = sm_scale
+        ctx.HEAD_DIM = HEAD_DIM_K
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, M = ctx.saved_tensors
+        assert do.is_contiguous()
+        assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        BATCH, N_HEAD, N_CTX = q.shape[:3]
+        PRE_BLOCK = 128
+        BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
+        RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+        arg_k = k
+        arg_k = arg_k * (ctx.sm_scale * RCP_LN2)
+        PRE_BLOCK = 128
+        assert N_CTX % PRE_BLOCK == 0
+        pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
+        delta = torch.empty_like(M)
+        _attn_bwd_preprocess[pre_grid](
+            o, do,  #
+            delta,  #
+            BATCH, N_HEAD, N_CTX,  #
+            BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM  #
+        )
+        grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+        _attn_bwd[grid](
+            q, arg_k, v, ctx.sm_scale, do, dq, dk, dv,  #
+            M, delta,  #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            N_HEAD, N_CTX,  #
+            BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1,  #
+            BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2,  #
+            HEAD_DIM=ctx.HEAD_DIM  #
+        )
+
+        return dq, dk, dv, None, None
+
+
+attention = _attention.apply
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+class _attention_sparse(torch.autograd.Function):
+    """
+    Thin autograd wrapper that uses the sparse forward kernel above and the
+    standard dense backward kernels defined earlier (no extra memory use).
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, v, q2k_index, q2k_num, k2q_index, k2q_num):
+        B, H, T, D = q.shape
+        sm_scale = 1.0 / math.sqrt(D)
+        max_kv_blks = q2k_index.shape[-1]
+        assert T % 64 == 0, f"T must be a multiple of 64, but got {T}"
+        assert T // 64 == q2k_num.shape[-1], f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
+        o = torch.empty_like(q)
+        M = torch.empty((B, H, T), dtype=torch.float32, device=q.device)
+
+        grid = lambda _: (triton.cdiv(T, 64), B * H, 1)
+        _attn_fwd_sparse[grid](
+            q, k, v, sm_scale,
+            q2k_index, q2k_num, max_kv_blks,
+            M, o,
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+            k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+            B, H, T,
+            HEAD_DIM=D, STAGE=3
+        )
+
+        ctx.save_for_backward(q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num)
+        ctx.grid = None
+        ctx.sm_scale = sm_scale
+        ctx.HEAD_DIM = D
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num = ctx.saved_tensors
+        assert do.is_contiguous()
+        assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        BATCH, N_HEAD, N_CTX = q.shape[:3]
+        PRE_BLOCK = 128
+        BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
+        RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+        arg_k = k
+        arg_k = arg_k * (ctx.sm_scale * RCP_LN2)
+        PRE_BLOCK = 128
+        assert N_CTX % PRE_BLOCK == 0
+        pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
+        delta = torch.empty_like(M)
+        _attn_bwd_preprocess[pre_grid](
+            o, do,  #
+            delta,  #
+            BATCH, N_HEAD, N_CTX,  #
+            BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM  #
+        )
+        
+        
+        max_q_blks = k2q_index.shape[-1]
+        max_kv_blks = q2k_index.shape[-1]
+        
+        
+        grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+        _attn_bwd[grid](
+            q, arg_k, v, ctx.sm_scale, do, dq, dk, dv,  #
+            M, delta,  #
+            q2k_index, q2k_num, max_kv_blks,
+            k2q_index, k2q_num, max_q_blks,
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            N_HEAD, N_CTX,  #
+            BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1,  #
+            BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2,  #
+            HEAD_DIM=ctx.HEAD_DIM  #
+        )
+
+        return dq, dk, dv, None, None, None, None
+
+attention_sparse = _attention_sparse.apply
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+try:
+    from flash_attn.flash_attn_interface import \
+        flash_attn_qkvpacked_func as flash_attn_func
+    HAS_FLASH = True
+except BaseException:
+    HAS_FLASH = False
+
+TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2')
+BATCH, N_HEADS, HEAD_DIM = 4, 32, 128
+# vary seq length for fixed head and batch=4
+configs = []
+for mode in ["fwd", "bwd"]:
+        configs.append(
+            triton.testing.Benchmark(
+                x_names=["N_CTX"],
+                x_vals=[2**i for i in range(10, 15)],
+                line_arg="provider",
+                line_vals=["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else []) +
+                (["flash"] if HAS_FLASH else []),
+                line_names=["Triton [FP16]"] + (["Triton [FP8]"] if TORCH_HAS_FP8 else []) +
+                (["Flash-2"] if HAS_FLASH else []),
+                styles=[("red", "-"), ("blue", "-"), ("green", "-")],
+                ylabel="ms",
+                plot_name=f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}",
+                args={
+                    "H": N_HEADS,
+                    "BATCH": BATCH,
+                    "HEAD_DIM": HEAD_DIM,
+                    "mode": mode,
+                },
+            ))
+
+
+@triton.testing.perf_report(configs)
+def bench_flash_attention(BATCH, H, N_CTX, HEAD_DIM, mode, provider, device="cuda"):
+    assert mode in ["fwd", "bwd"]
+    warmup = 25
+    rep = 100
+    dtype = torch.bfloat16
+    if "triton" in provider:
+        q = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        k = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        v = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        if mode == "fwd" and "fp8" in provider:
+            q = q.to(torch.float8_e5m2)
+            k = k.to(torch.float8_e5m2)
+            v = v.permute(0, 1, 3, 2).contiguous()
+            v = v.permute(0, 1, 3, 2)
+            v = v.to(torch.float8_e5m2)
+        fn = lambda: attention(q, k, v)
+        if mode == "bwd":
+            o = fn()
+            do = torch.randn_like(o)
+            fn = lambda: o.backward(do, retain_graph=True)
+        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
+    if provider == "flash":
+        qkv = torch.randn((BATCH, N_CTX, 3, H, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        fn = lambda: flash_attn_func(qkv, causal=False)
+        if mode == "bwd":
+            o = fn()
+            do = torch.randn_like(o)
+            fn = lambda: o.backward(do, retain_graph=True)
+        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
+    flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * HEAD_DIM
+    total_flops = 2 * flops_per_matmul
+    if mode == "bwd":
+        total_flops *= 2.5  # 2.0(bwd) + 0.5(recompute)
+    return total_flops / ms * 1e-9
+
+
+if __name__ == "__main__":
+    # only works on post-Ampere GPUs right now
+    bench_flash_attention.run(save_path=".", print_data=True)

csrc/attn/vsa/vsa.py

@@ -0,0 +1,47 @@
+import torch
+try:
+    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
+except ImportError:
+    block_sparse_fwd = None
+    block_sparse_bwd = None
+from .block_sparse_attn_triton import attention_sparse as block_sparse_attn_triton
+
+class BlockSparseAttentionFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
+        o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
+        ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num)
+        return o
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors
+        grad_q, grad_k, grad_v = block_sparse_bwd(
+            q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num
+        )
+        return grad_q, grad_k, grad_v, None, None, None, None
+
+
+@torch._dynamo.disable
+def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
+    """
+    Differentiable block sparse attention function.
+    
+    Args:
+        q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
+        k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
+        v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
+        q2k_block_sparse_index: Indices for query-to-key sparse blocks
+        q2k_block_sparse_num: Number of sparse blocks for each query block
+        k2q_block_sparse_index: Indices for key-to-query sparse blocks (for backward pass)
+        k2q_block_sparse_num: Number of sparse blocks for each key block (for backward pass)
+    
+    Returns:
+        output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
+    """
+    if block_sparse_fwd is not None:
+        return BlockSparseAttentionFunction.apply(
+            q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
+        )
+    else:
+        return block_sparse_attn_triton(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)

docs/Makefile

@@ -23,3 +23,4 @@ clean:
 	@$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 	rm -rf "$(SOURCEDIR)/getting_started/examples"
 	rm -rf "$(SOURCEDIR)/inference/examples"
+	rm -rf "$(SOURCEDIR)/training/examples"

docs/requirements-docs.txt

@@ -11,5 +11,5 @@ commonmark # Required by sphinx-argparse when using :markdownhelp:
 
 # packages to install to build the documentation
 cachetools
--f https://download.pytorch.org/whl/cpu
+# -f https://download.pytorch.org/whl/cpu
 torch

docs/source/generate_examples.py

@@ -27,6 +27,15 @@ def fix_case(text: str) -> str:
         "openai": "OpenAI",
         "multilora": "MultiLoRA",
         "mlpspeculator": "MLPSpeculator",
+        "finetune": "Finetune",
+        "distillation": "Distillation",
+        "wan": "Wan",
+        "i2v": "I2V",
+        "t2v": "T2V",
+        "1.3b": "1.3B",
+        "14b": "14B",
+        "480p": "480P",
+        "720p": "720P",
         r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
         r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
     }
@@ -161,31 +170,35 @@ class Example:
         return content
 
 
-def generate_examples(generate_main_index=False):
-    """
-    Generate example documentation.
-    
-    Args:
-        generate_main_index (bool): Whether to generate the main examples index.
-            If False, only category-specific indices will be generated.
-    """
-    # Create empty indices with dynamic paths
+@dataclass
+class NestedStructure:
+    """Helper class to manage nested documentation structures for training/distillation."""
+    category: str
+    method: str
+    model: str
+    dataset: str
+    example: Example
+
+    @property
+    def filename(self) -> str:
+        return f"{self.model}_{self.dataset}"
+
+    @property
+    def title(self) -> str:
+        return fix_case(self.dataset.replace('_', ' '))
+
+    @property
+    def description(self) -> str:
+        category_name = self.category.title()
+        return f"{category_name} example using the {self.dataset} dataset with the {self.model} model."
+
+
+def create_category_indices() -> dict[str, Index]:
+    """Create category indices with their respective configurations."""
     main_index_dir = ROOT_DIR / "docs/source/examples"
     if not main_index_dir.exists():
         main_index_dir.mkdir(parents=True)
 
-    # Create the main examples index only if requested
-    examples_index = None
-    if generate_main_index:
-        examples_index = Index(
-            path=main_index_dir / "examples_index.md",
-            title="💡 Examples",
-            description=
-            "A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",  # noqa: E501
-            caption="Examples",
-            maxdepth=2)
-
-    # Category indices with dynamic paths based on category names
     category_indices = {
         "inference":
         Index(
@@ -193,28 +206,54 @@ def generate_examples(generate_main_index=False):
             "docs/source/inference/examples/examples_inference_index.md",
             title="🚀 Examples",
             description=
-            "Inference examples demonstrate how to use FastVideo in an offline setting, where the model is queried for predictions in batches. We recommend starting with <project:basic.md>.",  # noqa: E501
+            "Inference examples demonstrate how to use FastVideo inference. We recommend starting with <project:basic.md>.",
             caption="Examples",
+            maxdepth=1,
+        ),
+        "training":
+        Index(
+            path=ROOT_DIR /
+            "docs/source/training/examples/examples_training_index.md",
+            title="🚀 Examples",
+            description=
+            "Training examples demonstrate how to use FastVideo training.",
+            caption="Examples",
+            maxdepth=3,
+        ),
+        "distillation":
+        Index(
+            path=ROOT_DIR /
+            "docs/source/distillation/examples/examples_distillation_index.md",
+            title="🚀 Examples",
+            description=
+            "Distillation examples demonstrate how to use FastVideo distillation.",
+            caption="Examples",
+            maxdepth=3,
         ),
     }
 
     # Ensure all category doc directories exist
-    for category, index in category_indices.items():
-        category_dir = index.path.parent
-        if not category_dir.exists():
-            category_dir.mkdir(parents=True)
+    for index in category_indices.values():
+        if not index.path.parent.exists():
+            index.path.parent.mkdir(parents=True)
 
+    return category_indices
+
+
+def find_examples(category_indices: dict[str, Index],
+                  generate_main_index: bool) -> list[Example]:
+    """Find all examples from the examples directory."""
     examples = []
     glob_patterns = ["*.py", "*.md", "*.sh"]
+
     # Find categorised examples
     for category in category_indices:
-        print(category)
         category_dir = EXAMPLE_DIR / category
         globs = [category_dir.glob(pattern) for pattern in glob_patterns]
         for path in itertools.chain(*globs):
             examples.append(Example(path, category))
-        # Find examples in subdirectories
-        for path in category_dir.glob("*/*.md"):
+        # Find examples in subdirectories (recursively)
+        for path in category_dir.glob("**/*.md"):
             examples.append(Example(path.parent, category))
 
     # Find uncategorised examples only if we're generating a main index
@@ -229,36 +268,173 @@ def generate_examples(generate_main_index=False):
                 continue
             examples.append(Example(path.parent))
 
-    # Create document directories for each category based on category name and generate files
-    for example in sorted(examples, key=lambda e: e.path.stem):
-        print(example)
+    return examples
+
+
+def create_nested_structures(
+    examples: list[Example]
+) -> dict[str, dict[str, dict[str, dict[str, NestedStructure]]]]:
+    """Create nested structures for training and distillation categories."""
+    nested_structures: dict[str, dict[str, dict[str,
+                                                dict[str,
+                                                     NestedStructure]]]] = {}
+
+    for example in examples:
+        if example.category not in ["training", "distillation"]:
+            continue
+
+        category_dir = EXAMPLE_DIR / example.category
+        relative_path = example.path.relative_to(category_dir)
+        path_parts = relative_path.parts
+
+        # For nested examples like finetune/wan_i2v_14b_480p/crush_smol
+        if len(path_parts) >= 3:
+            method = path_parts[0]  # e.g., "finetune"
+            model = path_parts[1]  # e.g., "wan_i2v_14b_480p"
+            dataset = path_parts[2]  # e.g., "crush_smol"
+
+            # Initialize nested structure
+            if example.category not in nested_structures:
+                nested_structures[example.category] = {}
+            if method not in nested_structures[example.category]:
+                nested_structures[example.category][method] = {}
+            if model not in nested_structures[example.category][method]:
+                nested_structures[example.category][method][model] = {}
+
+            # Store the nested structure
+            nested_structures[
+                example.category][method][model][dataset] = NestedStructure(
+                    category=example.category,
+                    method=method,
+                    model=model,
+                    dataset=dataset,
+                    example=example)
+
+    return nested_structures
+
+
+def generate_flat_examples(examples: list[Example],
+                           category_indices: dict[str, Index],
+                           examples_index: Index | None,
+                           generate_main_index: bool) -> None:
+    """Generate documentation for flat structure examples (inference, etc.)."""
+    for example in examples:
+        if example.category in ["training", "distillation"]:
+            continue  # Skip nested structure examples
 
         # Determine which index to use for this example
         if example.category is not None and example.category in category_indices:
             index = category_indices[example.category]
         elif generate_main_index:
             assert examples_index is not None
-            index = examples_index  # Default to main index if available
+            index = examples_index
         else:
-            # Skip examples without a category if no main index
-            print(f"Skipping {example.path} (no category and no main index)")
             continue
 
-        # Place generated example markdown in the same directory as its index
+        # Generate the example documentation
         doc_path = index.path.parent / f"{example.path.stem}.md"
         with open(doc_path, "w+") as f:
             f.write(example.generate())
-        # Add the example to the index
         index.documents.append(example.path.stem)
 
+
+def generate_nested_examples(nested_structures: dict[str, dict[str, dict[
+    str, dict[str, NestedStructure]]]], category_indices: dict[str,
+                                                               Index]) -> None:
+    """Generate documentation for nested structure examples (training, distillation)."""
+    for category_name in ["training", "distillation"]:
+        if category_name not in category_indices or category_name not in nested_structures:
+            continue
+
+        category_index = category_indices[category_name]
+        category_base_dir = category_index.path.parent
+
+        for method, models in nested_structures[category_name].items():
+            # Create method-level index
+            method_index = Index(path=category_base_dir / f"{method}.md",
+                                 title=fix_case(method),
+                                 description=f"Examples using {method}.",
+                                 caption=f"{fix_case(method)} Examples",
+                                 maxdepth=2)
+
+            for model, datasets in models.items():
+                # Generate dataset examples using the Example class
+                for dataset, nested_struct in datasets.items():
+                    doc_path = category_base_dir / f"{nested_struct.filename}.md"
+                    with open(doc_path, "w+") as f:
+                        f.write(nested_struct.example.generate())
+
+                # Create model-level index
+                model_index = Index(
+                    path=category_base_dir / f"{model}.md",
+                    title=fix_case(model.replace('_', ' ')),
+                    description=f"Examples for the {model} model.",
+                    caption=f"{fix_case(model.replace('_', ' '))} Datasets",
+                    maxdepth=1)
+
+                # Add dataset indices to model index
+                for dataset, nested_struct in datasets.items():
+                    model_index.documents.append(nested_struct.filename)
+
+                # Write model index
+                with open(model_index.path, "w+") as f:
+                    f.write(model_index.generate())
+
+                # Add model to method index
+                method_index.documents.append(model)
+
+            # Write method index
+            with open(method_index.path, "w+") as f:
+                f.write(method_index.generate())
+
+            # Add method to main category index
+            category_index.documents.append(method)
+
+
+def generate_examples(generate_main_index=False):
+    """
+    Generate example documentation.
+    
+    Args:
+        generate_main_index (bool): Whether to generate the main examples index.
+            If False, only category-specific indices will be generated.
+    """
+    # Create category indices
+    category_indices = create_category_indices()
+
+    # Create the main examples index only if requested
+    examples_index = None
+    if generate_main_index:
+        main_index_dir = ROOT_DIR / "docs/source/examples"
+        examples_index = Index(
+            path=main_index_dir / "examples_index.md",
+            title="💡 Examples",
+            description=
+            "A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",
+            caption="Examples",
+            maxdepth=2)
+
+    # Find all examples
+    examples = find_examples(category_indices, generate_main_index)
+
+    # Create nested structures for training and distillation
+    nested_structures = create_nested_structures(examples)
+
+    # Generate flat structure examples (inference, etc.)
+    generate_flat_examples(examples, category_indices, examples_index,
+                           generate_main_index)
+
+    # Generate nested structure examples (training, distillation)
+    generate_nested_examples(nested_structures, category_indices)
+
     # Generate the index files for categories
     for category_index in category_indices.values():
         if category_index.documents:
             # Add to main index if it exists
-            if generate_main_index:
+            if generate_main_index and examples_index:
+                main_index_dir = examples_index.path.parent
                 rel_path = category_index.path.relative_to(
                     main_index_dir.parent)
-                assert examples_index is not None
                 examples_index.documents.insert(
                     0,
                     str(rel_path).replace(".md", ""))

docs/source/getting_started/installation.md

@@ -1,120 +1,18 @@
-(fastvideo-installation)=
+(installation-index)=
 
 # 🔧 Installation
 
-FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
+FastVideo supports the following hardware platforms:
 
-## Requirements
+:::{toctree}
+:maxdepth: 1
+:hidden:
 
-- **OS: Linux**
-- **Python: 3.10-3.12**
-- **CUDA 12.4**
-- **At least 1 NVIDIA GPU**
-
-## Set up using Python
-### Create a new Python environment
-
-#### Conda
-You can create a new python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html)
-##### 1. Install Miniconda (if not already installed)
-
-```bash
-wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
-bash Miniconda3-latest-Linux-x86_64.sh
-source ~/.bashrc
-```
-
-##### 2. Create and activate a Conda environment for FastVideo
-
-```bash
-# (Recommended) Create a new conda environment.
-conda create -n fastvideo python=3.12 -y
-conda activate fastvideo
-```
-
-:::{note}
-[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
+installation/gpu
+installation/mps
 :::
 
-#### uv
-
-:::{tip}
-We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
-:::
-
-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:
-
-```console
-# (Recommended) Create a new uv environment. Use `--seed` to install `pip` and `setuptools` in the environment.
-uv venv --python 3.12 --seed
-source .venv/bin/activate
-```
-
-### Installation
-
-```bash
-pip install fastvideo
-
-# or if you are using uv
-uv pip install fastvideo
-```
-
-Also optionally install flash-attn:
-
-```bash
-pip install flash-attn==2.7.4.post1 --no-build-isolation
-```
-
-### Installation from Source
-
-#### 1. Clone the FastVideo repository
-
-```bash
-git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
-```
-
-#### 2. Install FastVideo
-
-Basic installation:
-
-```bash
-pip install -e .
-
-# or if you are using uv
-uv pip install -e .
-```
-
-### Optional Dependencies
-
-#### Flash Attention
-
-```bash
-pip install flash-attn==2.7.4.post1 --no-build-isolation
-```
-
-## Set up using Docker
-We also have prebuilt docker images with FastVideo dependencies pre-installed:
-[Docker Images](#docker)
-
-## Development Environment Setup
-
-If you're planning to contribute to FastVideo please see the following page:
-[Contributor Guide](#developer-overview)
-
-## Hardware Requirements
-
-### For Basic Inference
-- NVIDIA GPU with CUDA 12.4 support
-
-### For Lora Finetuning
-- 40GB GPU memory each for 2 GPUs with lora
-- 30GB GPU memory each for 2 GPUs with CPU offload and lora
-
-### For Full Finetuning/Distillation
-- Multiple high-memory GPUs recommended (e.g., H100)
-
-## Troubleshooting
-
-If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
-
-You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.
+- <project:installation/gpu.md>
+  - NVIDIA CUDA
+- <project:installation/mps.md>
+  - Apple silicon

docs/source/getting_started/installation/gpu.md

@@ -0,0 +1,118 @@
+# NVIDIA GPU
+
+Instructions to install FastVideo for NVIDIA CUDA GPUs.
+
+## Requirements
+
+- **OS: Linux or Windows WSL**
+- **Python: 3.10-3.12**
+- **CUDA 12.4**
+- **At least 1 NVIDIA GPU**
+
+## Set up using Python
+### Create a new Python environment
+
+#### Conda
+You can create a new python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html)
+##### 1. Install Miniconda (if not already installed)
+
+```bash
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+bash Miniconda3-latest-Linux-x86_64.sh
+source ~/.bashrc
+```
+
+##### 2. Create and activate a Conda environment for FastVideo
+
+```bash
+# (Recommended) Create a new conda environment.
+conda create -n fastvideo python=3.12 -y
+conda activate fastvideo
+```
+
+:::{note}
+[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
+:::
+
+#### uv
+
+:::{tip}
+We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
+:::
+
+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:
+
+```console
+# (Recommended) Create a new uv environment. Use `--seed` to install `pip` and `setuptools` in the environment.
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+### Installation
+
+```bash
+pip install fastvideo
+
+# or if you are using uv
+uv pip install fastvideo
+```
+
+Also optionally install flash-attn:
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+### Installation from Source
+
+#### 1. Clone the FastVideo repository
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+```
+
+#### 2. Install FastVideo
+
+Basic installation:
+
+```bash
+pip install -e .
+
+# or if you are using uv
+uv pip install -e .
+```
+
+### Optional Dependencies
+
+#### Flash Attention
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+## Set up using Docker
+We also have prebuilt docker images with FastVideo dependencies pre-installed:
+[Docker Images](#docker)
+
+## Development Environment Setup
+
+If you're planning to contribute to FastVideo please see the following page:
+[Contributor Guide](#developer-overview)
+
+## Hardware Requirements
+
+### For Basic Inference
+- NVIDIA GPU with CUDA 12.4 support
+
+### For Lora Finetuning
+- 40GB GPU memory each for 2 GPUs with lora
+- 30GB GPU memory each for 2 GPUs with CPU offload and lora
+
+### For Full Finetuning/Distillation
+- Multiple high-memory GPUs recommended (e.g., H100)
+
+## Troubleshooting
+
+If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
+
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.

docs/source/getting_started/installation/mps.md

@@ -0,0 +1,101 @@
+# MPS (Apple Silicon)
+
+Instructions to install FastVideo for Apple Silicon.
+
+## Requirements
+
+- **OS: MacOS**
+- **Python: 3.12.4**
+
+## Set up using Python
+
+### Create a new Python environment
+
+#### Conda
+
+You can create a new python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html)
+
+##### 1. Install Miniconda (if not already installed)
+
+```bash
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh
+bash Miniconda3-latest-MacOSX-arm64.sh
+source ~/.zshrc
+```
+
+##### 2. Create and activate a Conda environment for FastVideo
+
+```bash
+# (Recommended) Create a new conda environment.
+conda create -n fastvideo python=3.12.4 -y
+conda activate fastvideo
+```
+
+:::{note}
+[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
+:::
+
+#### uv
+
+:::{tip}
+We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
+:::
+
+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:
+
+```console
+# (Recommended) Create a new uv environment. Use `--seed` to install `pip` and `setuptools` in the environment.
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+### Dependencies
+
+```
+brew install ffmpeg
+```
+
+### Installation
+
+```bash
+pip install fastvideo
+
+# or if you are using uv
+uv pip install fastvideo
+```
+
+### Installation from Source
+
+#### 1. Clone the FastVideo repository
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+```
+
+#### 2. Install FastVideo
+
+Basic installation:
+
+```bash
+pip install -e .
+
+# or if you are using uv
+uv pip install -e .
+```
+
+## Development Environment Setup
+
+If you're planning to contribute to FastVideo please see the following page:
+[Contributor Guide](#developer-overview)
+
+## Hardware Requirements
+
+### For Basic Inference
+
+- Mac M1, M2, M3, or M4 (at least 32 GB RAM is preferable for high quality video generation)
+
+## Troubleshooting
+
+If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
+
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.

docs/source/index.md

@@ -63,22 +63,33 @@ getting_started/installation
 :maxdepth: 1
 
 inference/inference_quick_start
+inference/examples/examples_inference_index
 inference/configuration
 inference/optimizations
+inference/comfyui
 inference/support_matrix
-inference/examples/examples_inference_index
 inference/cli
 inference/add_pipeline
-inference/v0_inference
 :::
 
 :::{toctree}
 :caption: Training
 :maxdepth: 1
 
+training/examples/examples_training_index
 training/data_preprocess
 training/distillation
-training/finetune
+<!-- training/finetune -->
+:::
+
+<!-- :::{toctree}
+:caption: Distillation
+:maxdepth: 1
+
+distillation/examples/examples_distillation_index
+distillation/data_preprocess
+distillation/dmd -->
+<!-- training/finetune -->
 :::
 
 % What is STA Kernel?

docs/source/inference/comfyui.md

@@ -0,0 +1,5 @@
+# FastVideo + ComfyUI
+
+FastVideo provides a custom node suite for ComfyUI.
+
+See this [README](https://github.com/hao-ai-lab/FastVideo/tree/main/comfyui) for instructions.

docs/source/inference/v0_inference.md

@@ -1,74 +0,0 @@
-(v0-inference)=
-
-# [Deprecated] V0 Inference
-The following commands and APIs are deprecated but still supported until V1's API can completely replace all the features in this page.
-
-## Inference StepVideo with Sliding Tile Attention
-First, download the model:
-
-```
-python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model
-```
-
-Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204×768×768 (currently, this is the only supported shape).
-
-```bash
-sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA
-sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
-```
-
-## Inference HunyuanVideo with Sliding Tile Attention
-First, download the model:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model
-```
-
-We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
-
-```bash
-sh scripts/inference/inference_hunyuan_STA.sh
-```
-
-## Video Demos using STA + Teacache
-Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.
-
-## Inference FastHunyuan on single RTX4090
-We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan_hf_quantization.sh
-```
-
-For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
-
-| Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
-|--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
-| BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
-| INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
-| NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
-
-For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
-
-## FastHunyuan
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan.sh
-```
-
-You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).
-
-## FastMochi
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffusers --local_dir=data/FastMochi-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_mochi_sp.sh
-```

docs/source/training/data_preprocess.md

@@ -7,7 +7,7 @@ To save GPU memory, we precompute text embeddings and VAE latents to eliminate t
 We provide a sample dataset to help you get started. Download the source media using the following command:
 
 ```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/mini_i2v_dataset --local_dir=FastVideo/mini_i2v_dataset --repo_type=dataset
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/mini_i2v_dataset --local_dir=data/mini_i2v_dataset --repo_type=dataset
 ```
 
 The folder `crush-smol_raw/` contains raw videos and captions for testing preprocessing, while `crush-smol_preprocessed/` contains latents prepared for testing training.

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

@@ -0,0 +1,3 @@
+# DMD Distillation Wan2.1-I2V-14B-480P Crush-Smol Example
+
+Coming soon!

examples/inference/basic/basic_mps.py

@@ -0,0 +1,40 @@
+from fastvideo import VideoGenerator, PipelineConfig
+from fastvideo.v1.configs.sample import SamplingParam
+
+def main():
+    config = PipelineConfig.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    config.text_encoder_precisions = ["fp16"]
+    
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        pipeline_config=config,
+        use_fsdp_inference=False,      # Disable FSDP for MPS
+        use_cpu_offload=True,          
+        text_encoder_offload=True,    
+        pin_cpu_memory=True,           
+        disable_autocast=False,        
+        num_gpus=1,      
+    )
+
+    # Create sampling parameters with reduced number of frames
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    sampling_param.num_frames = 3  # Reduce from default 81 to 25 frames bc we have to use the SDPA attn backend for mps
+    sampling_param.height = 256
+    sampling_param.width = 256
+
+    prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+             "wide with interest. The playful yet serene atmosphere is complemented by soft "
+             "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")
+    
+    video = generator.generate_video(prompt, sampling_param=sampling_param)
+
+    prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    
+    video2 = generator.generate_video(prompt2, sampling_param=sampling_param)
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/test.py

@@ -0,0 +1,39 @@
+from fastvideo import VideoGenerator, PipelineConfig, SamplingParam
+
+# from fastvideo.v1.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_fp16"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    model = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    pipeline_config = PipelineConfig.from_pretrained(model)
+    pipeline_config.text_encoder_precisions = ("bf16", )
+    generator = VideoGenerator.from_pretrained(
+        model,
+        # if num_gpus > 1, FastVideo will automatically handle distributed setup
+        pipeline_config=pipeline_config,
+        use_fsdp_inference=False,      # Disable FSDP for MPS
+        use_cpu_offload=True,          
+        text_encoder_offload=True,    
+        pin_cpu_memory=True,           
+        disable_autocast=False,        
+        num_gpus=1, 
+    )
+
+    sampling_param = SamplingParam.from_pretrained(model)
+    sampling_param.num_frames = 30
+    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    # Generate videos with the same simple API, regardless of GPU count
+    prompt = "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting."
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+
+
+if __name__ == "__main__":
+    main()

examples/inference/lora/wan_lora_inference.py

@@ -1,7 +1,7 @@
 from fastvideo import VideoGenerator
 from fastvideo.v1.configs.sample import SamplingParam
 
-OUTPUT_PATH = "./lora"
+OUTPUT_PATH = "./lora_out"
 def main():
     # Initialize VideoGenerator with the Wan model
     generator = VideoGenerator.from_pretrained(
@@ -32,7 +32,7 @@ def main():
     )
     del generator
     
-    # Until FSDP resharding bug is fixed, multi-lora requires reloading the model
+    # Until FSDP resharding bug is fixed, multi-lora requires reloading the model or disabling FSDP
     # see https://github.com/pytorch/pytorch/issues/157209
     generator = VideoGenerator.from_pretrained(
         "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",

examples/inference/lora/wan_lora_inference_from_ckpt.py

@@ -0,0 +1,35 @@
+"""
+Inference using a LoRA checkpoint from FastVideo trainer.
+"""
+from fastvideo import VideoGenerator
+from fastvideo.v1.configs.sample import SamplingParam
+
+OUTPUT_PATH = "./lora_out"
+def main():
+    # Initialize VideoGenerator with the Wan model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,
+        lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1250/transformer",
+        lora_nickname="crush_smol"
+    )
+    kwargs = {
+        "height": 480,
+        "width": 832,
+        "num_frames": 77,
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "seed": 42,
+    }
+    # Generate video with LoRA style
+    prompt = "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."
+
+    video = generator.generate_video(
+        prompt,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+        **kwargs
+    )    
+
+if __name__ == "__main__":
+    main()

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

@@ -1,10 +1,16 @@
-This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
+# 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:
+## Execute the following commands from `FastVideo/` to run training:
+
+### Download crush-smol dataset:
 
-- Download crush-smol dataset:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
-- Preprocess the videos and captions into latents:
+
+### 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`
+
+### Edit the following file and run finetuning:
+
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`

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

@@ -5,21 +5,19 @@ 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/"
-VALIDATION_PATH="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
 
 torchrun --nproc_per_node=$GPU_NUM \
     fastvideo/v1/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 \
-    --model_type $MODEL_TYPE \
     --train_fps 16 \
-    --validation_dataset_file $VALIDATION_PATH \
     --samples_per_file 8 \
     --flush_frequency 8 \
     --video_length_tolerance_range 5 \

examples/training/finetune/wan_i2v_14b_480p/crush_smol/README.md

@@ -1,10 +1,16 @@
-This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
+# Wan2.1-I2V-14B-480P Crush-Smol Example
+These are e2e examples scripts for finetuning Wan2.1 I2V 14B 480P on the crush-smol dataset.
 
-Execute the following commands from `FastVideo/` to run training:
+## Execute the following commands from `FastVideo/` to run training:
+
+### Download crush-smol dataset:
 
-- Download crush-smol dataset:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
-- Preprocess the videos and captions into latents:
+
+### 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:
+
+### Edit the following file and run finetuning:
+
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh

@@ -15,7 +15,7 @@ NUM_GPUS=8
 # Training arguments
 training_args=(
   --tracker_project_name "wan_i2v_finetune"
-  --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
+  --output_dir "checkpoints/wan_i2v_finetune"
   --max_train_steps 2000
   --train_batch_size 1
   --train_sp_batch_size 1

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm

@@ -49,7 +49,7 @@ VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/
 # Training arguments
 training_args=(
   --tracker_project_name wan_i2v_finetune
-  --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
+  --output_dir="checkpoints/wan_i2v_finetune"
   --max_train_steps=2000
   --train_batch_size=2
   --train_sp_batch_size 1

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v_lora.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
+VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
+NUM_GPUS=4
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_i2v_finetune"
+  --output_dir "checkpoints/wan_i2v_finetune"
+  --max_train_steps 2000
+  --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
+  --lora_rank 32
+  --lora_training True
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size $NUM_GPUS
+  --tp_size $NUM_GPUS
+  --hsdp_replicate_dim 1
+  --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 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_preprocessed_path "$VALIDATION_DIR"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "1.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
+  --allow_tf32
+  --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
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh

@@ -5,21 +5,19 @@ MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_i2v/"
-VALIDATION_PATH="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
 
 torchrun --nproc_per_node=$GPU_NUM \
     fastvideo/v1/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 \
-    --model_type $MODEL_TYPE \
     --train_fps 16 \
-    --validation_dataset_file $VALIDATION_PATH \
     --samples_per_file 8 \
     --flush_frequency 8 \
     --video_length_tolerance_range 5 \

examples/training/finetune/wan_t2v_1_3b/crush_smol/README.md

@@ -1,10 +1,16 @@
-This directory contain e2e examples scripts for finetuning Wan2.1 T2v. 
+# Wan2.1-T2V-1.3B Crush-Smol Example
+These are e2e example scripts for finetuning Wan2.1 T2V 1.3B on the crush-smol dataset.
 
-Execute the following commands from `FastVideo/` to run training:
+## Execute the following commands from `FastVideo/` to run training:
+
+### Download crush-smol dataset:
 
-- Download crush-smol dataset:
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/download_dataset.sh`
-- Preprocess the videos and captions into latents:
+
+### Preprocess the videos and captions into latents:
+
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh`
-- Edit the following file and run finetuning:
+
+### Edit the following file and run finetuning:
+
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh`

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh

@@ -15,12 +15,12 @@ NUM_GPUS=4
 # Training arguments
 training_args=(
   --tracker_project_name "wan_t2v_finetune"
-  --output_dir "outputs/wan_t2v_finetune"
+  --output_dir "checkpoints/wan_t2v_finetune"
   --max_train_steps 5000
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 8
-  --num_latent_t 8 
+  --num_latent_t 20
   --num_height 480
   --num_width 832
   --num_frames 77
@@ -61,7 +61,7 @@ validation_args=(
 optimizer_args=(
   --learning_rate 5e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 6000
+  --checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm

@@ -46,7 +46,7 @@ VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/vali
 # Training arguments
 training_args=(
   --tracker_project_name wan_t2v_finetune
-  --output_dir="outputs/wan_t2v_finetune"
+  --output_dir="checkpoints/wan_t2v_finetune"
   --max_train_steps=1000
   --train_batch_size=4
   --train_sp_batch_size 1

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v_lora.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
+NUM_GPUS=2
+# export CUDA_VISIBLE_DEVICES=4,5
+
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_t2v_finetune"
+  --output_dir "checkpoints/wan_t2v_finetune_lora"
+  --max_train_steps 5000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 8
+  --num_latent_t 20
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --lora_rank 32
+  --lora_training True
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS 
+  --sp_size $NUM_GPUS 
+  --tp_size $NUM_GPUS
+  --hsdp_replicate_dim 1
+  --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 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation 
+  --validation_dataset_file $VALIDATION_DATASET_FILE
+  --validation_steps 50
+  --validation_sampling_steps "50" 
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 500
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --allow_tf32
+  --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
+)
+
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+  --master_port 29501 \
+    fastvideo/v1/training/wan_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh

@@ -5,21 +5,19 @@ MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_t2v/"
-VALIDATION_PATH="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
 
 torchrun --nproc_per_node=$GPU_NUM \
     fastvideo/v1/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 \
-    --model_type $MODEL_TYPE \
     --train_fps 16 \
-    --validation_dataset_file $VALIDATION_PATH \
     --samples_per_file 8 \
     --flush_frequency 8 \
     --video_length_tolerance_range 5 \

fastvideo/data_preprocess/preprocess_text_embeddings.py

@@ -1,147 +0,0 @@
-import argparse
-import json
-import os
-
-import torch
-import torch.distributed as dist
-from accelerate.logging import get_logger
-from diffusers.utils import export_to_video
-from diffusers.video_processor import VideoProcessor
-from torch.utils.data import DataLoader, Dataset
-from torch.utils.data.distributed import DistributedSampler
-from tqdm import tqdm
-
-from fastvideo.utils.load import load_text_encoder, load_vae
-
-logger = get_logger(__name__)
-
-
-class T5dataset(Dataset):
-
-    def __init__(
-        self,
-        json_path,
-        vae_debug,
-    ):
-        self.json_path = json_path
-        self.vae_debug = vae_debug
-        with open(self.json_path, "r") as f:
-            train_dataset = json.load(f)
-            self.train_dataset = sorted(train_dataset, key=lambda x: x["latent_path"])
-
-    def __getitem__(self, idx):
-        caption = self.train_dataset[idx]["caption"]
-        filename = self.train_dataset[idx]["latent_path"].split(".")[0]
-        length = self.train_dataset[idx]["length"]
-        if self.vae_debug:
-            latents = torch.load(
-                os.path.join(args.output_dir, "latent", self.train_dataset[idx]["latent_path"]),
-                map_location="cpu",
-            )
-        else:
-            latents = []
-
-        return dict(caption=caption, latents=latents, filename=filename, length=length)
-
-    def __len__(self):
-        return len(self.train_dataset)
-
-
-def main(args):
-    local_rank = int(os.getenv("RANK", 0))
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    print("world_size", world_size, "local rank", local_rank)
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    torch.cuda.set_device(local_rank)
-    if not dist.is_initialized():
-        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
-
-    videoprocessor = VideoProcessor(vae_scale_factor=8)
-    os.makedirs(args.output_dir, exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "video"), exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "prompt_embed"), exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True)
-
-    latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
-    train_dataset = T5dataset(latents_json_path, args.vae_debug)
-    text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
-    vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
-    if args.model_type != "wan":
-        vae.enable_tiling()
-    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-    )
-
-    json_data = []
-    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
-        with torch.inference_mode():
-            with torch.autocast("cuda", dtype=autocast_type):
-                prompt_embeds, prompt_attention_mask = text_encoder.encode_prompt(prompt=data["caption"], )
-                if args.vae_debug:
-                    latents = data["latents"]
-                    video = vae.decode(latents.to(device), return_dict=False)[0]
-                    video = videoprocessor.postprocess_video(video)
-                for idx, video_name in enumerate(data["filename"]):
-                    prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt")
-                    video_path = os.path.join(args.output_dir, "video", video_name + ".mp4")
-                    prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask",
-                                                              video_name + ".pt")
-                    # save latent
-                    torch.save(prompt_embeds[idx], prompt_embed_path)
-                    torch.save(prompt_attention_mask[idx], prompt_attention_mask_path)
-                    print(f"sample {video_name} saved")
-                    if args.vae_debug:
-                        export_to_video(video[idx], video_path, fps=fps)
-                    item = {}
-                    item["length"] = int(data["length"][idx])
-                    item["latent_path"] = video_name + ".pt"
-                    item["prompt_embed_path"] = video_name + ".pt"
-                    item["prompt_attention_mask"] = video_name + ".pt"
-                    item["caption"] = data["caption"][idx]
-                    json_data.append(item)
-    dist.barrier()
-    local_data = json_data
-    gathered_data = [None] * world_size
-    dist.all_gather_object(gathered_data, local_data)
-    if local_rank == 0:
-        # os.remove(latents_json_path)
-        all_json_data = [item for sublist in gathered_data for item in sublist]
-        with open(os.path.join(args.output_dir, "videos2caption.json"), "w") as f:
-            json.dump(all_json_data, f, indent=4)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    # dataset & dataloader
-    parser.add_argument("--model_path", type=str, default="data/mochi")
-    parser.add_argument("--model_type", type=str, default="mochi")
-    # text encoder & vae & diffusion model
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=1,
-        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
-    )
-    parser.add_argument(
-        "--train_batch_size",
-        type=int,
-        default=1,
-        help="Batch size (per device) for the training dataloader.",
-    )
-    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
-    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default=None,
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-    parser.add_argument("--vae_debug", action="store_true")
-    args = parser.parse_args()
-    main(args)

fastvideo/data_preprocess/preprocess_vae_latents.py

@@ -1,110 +0,0 @@
-import argparse
-import json
-import os
-
-import torch
-import torch.distributed as dist
-from accelerate.logging import get_logger
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from tqdm import tqdm
-
-from fastvideo.dataset import getdataset
-from fastvideo.utils.load import load_vae
-
-logger = get_logger(__name__)
-
-
-def main(args):
-    local_rank = int(os.getenv("RANK", 0))
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    print("world_size", world_size, "local rank", local_rank)
-    train_dataset = getdataset(args)
-    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-    )
-
-    encoder_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    torch.cuda.set_device(local_rank)
-    if not dist.is_initialized():
-        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
-    vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
-    if args.model_type != "wan":
-        vae.enable_tiling()
-    os.makedirs(args.output_dir, exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
-
-    json_data = []
-    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
-        with torch.inference_mode():
-            with torch.autocast("cuda", dtype=autocast_type):
-                latents = vae.encode(data["pixel_values"].to(encoder_device))["latent_dist"].sample()
-            for idx, video_path in enumerate(data["path"]):
-                video_name = os.path.basename(video_path).split(".")[0]
-                latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt")
-                torch.save(latents[idx].to(torch.bfloat16), latent_path)
-                item = {}
-                item["length"] = latents[idx].shape[1]
-                item["latent_path"] = video_name + ".pt"
-                item["caption"] = data["text"][idx]
-                json_data.append(item)
-                print(f"{video_name} processed")
-    dist.barrier()
-    local_data = json_data
-    gathered_data = [None] * world_size
-    dist.all_gather_object(gathered_data, local_data)
-    if local_rank == 0:
-        all_json_data = [item for sublist in gathered_data for item in sublist]
-        with open(os.path.join(args.output_dir, "videos2caption_temp.json"), "w") as f:
-            json.dump(all_json_data, f, indent=4)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    # dataset & dataloader
-    parser.add_argument("--model_path", type=str, default="data/mochi")
-    parser.add_argument("--model_type", type=str, default="mochi")
-    parser.add_argument("--data_merge_path", type=str, required=True)
-    parser.add_argument("--num_frames", type=int, default=163)
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=1,
-        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
-    )
-    parser.add_argument(
-        "--train_batch_size",
-        type=int,
-        default=16,
-        help="Batch size (per device) for the training dataloader.",
-    )
-    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
-    parser.add_argument("--max_height", type=int, default=480)
-    parser.add_argument("--max_width", type=int, default=848)
-    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
-    parser.add_argument("--group_frame", action="store_true")  # TODO
-    parser.add_argument("--group_resolution", action="store_true")  # TODO
-    parser.add_argument("--dataset", default="t2v")
-    parser.add_argument("--train_fps", type=int, default=30)
-    parser.add_argument("--use_image_num", type=int, default=0)
-    parser.add_argument("--text_max_length", type=int, default=256)
-    parser.add_argument("--speed_factor", type=float, default=1.0)
-    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
-    # text encoder & vae & diffusion model
-    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
-    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-    parser.add_argument("--cfg", type=float, default=0.0)
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default=None,
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-
-
-    args = parser.parse_args()
-    main(args)

fastvideo/data_preprocess/preprocess_validation_text_embeddings.py

@@ -1,67 +0,0 @@
-import argparse
-import os
-
-import torch
-import torch.distributed as dist
-from accelerate.logging import get_logger
-
-from fastvideo.utils.load import load_text_encoder
-
-logger = get_logger(__name__)
-
-
-def main(args):
-    local_rank = int(os.getenv("RANK", 0))
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    print("world_size", world_size, "local rank", local_rank)
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    torch.cuda.set_device(local_rank)
-    if not dist.is_initialized():
-        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
-
-    text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
-    autocast_type = torch.float16 if args.model_type == "hunyuan" else torch.bfloat16
-    # output_dir/validation/prompt_attention_mask
-    # output_dir/validation/prompt_embed
-    os.makedirs(os.path.join(args.output_dir, "validation"), exist_ok=True)
-    os.makedirs(
-        os.path.join(args.output_dir, "validation", "prompt_attention_mask"),
-        exist_ok=True,
-    )
-    os.makedirs(os.path.join(args.output_dir, "validation", "prompt_embed"), exist_ok=True)
-
-    with open(args.validation_prompt_txt, "r", encoding="utf-8") as file:
-        lines = file.readlines()
-    prompts = [line.strip() for line in lines]
-    for prompt in prompts:
-        with torch.inference_mode():
-            with torch.autocast("cuda", dtype=autocast_type):
-                prompt_embeds, prompt_attention_mask = text_encoder.encode_prompt(prompt)
-                file_name = prompt.split(".")[0]
-                prompt_embed_path = os.path.join(args.output_dir, "validation", "prompt_embed", f"{file_name}.pt")
-                prompt_attention_mask_path = os.path.join(
-                    args.output_dir,
-                    "validation",
-                    "prompt_attention_mask",
-                    f"{file_name}.pt",
-                )
-                torch.save(prompt_embeds[0], prompt_embed_path)
-                torch.save(prompt_attention_mask[0], prompt_attention_mask_path)
-                print(f"sample {file_name} saved")
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    # dataset & dataloader
-    parser.add_argument("--model_path", type=str, default="data/mochi")
-    parser.add_argument("--model_type", type=str, default="mochi")
-    parser.add_argument("--validation_prompt_txt", type=str)
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default=None,
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-    args = parser.parse_args()
-    main(args)

fastvideo/dataset/__init__.py

@@ -1,97 +0,0 @@
-from torchvision import transforms
-from torchvision.transforms import Lambda
-from transformers import AutoTokenizer
-
-from fastvideo.dataset.t2v_datasets import T2V_dataset
-from fastvideo.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop
-
-
-def getdataset(args):
-    temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
-    norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
-    resize_topcrop = [
-        CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True),
-    ]
-    resize = [
-        CenterCropResizeVideo((args.max_height, args.max_width)),
-    ]
-    transform = transforms.Compose([
-        # Normalize255(),
-        *resize,
-    ])
-    transform_topcrop = transforms.Compose([
-        Normalize255(),
-        *resize_topcrop,
-        norm_fun,
-    ])
-    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
-    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
-    if args.dataset == "t2v":
-        return T2V_dataset(
-            args,
-            transform=transform,
-            temporal_sample=temporal_sample,
-            tokenizer=tokenizer,
-            transform_topcrop=transform_topcrop,
-        )
-
-    raise NotImplementedError(args.dataset)
-
-
-if __name__ == "__main__":
-    import random
-
-    from accelerate import Accelerator
-    from tqdm import tqdm
-
-    from fastvideo.dataset.t2v_datasets import dataset_prog
-
-    args = type(
-        "args",
-        (),
-        {
-            "ae": "CausalVAEModel_4x8x8",
-            "dataset": "t2v",
-            "attention_mode": "xformers",
-            "use_rope": True,
-            "text_max_length": 300,
-            "max_height": 320,
-            "max_width": 240,
-            "num_frames": 1,
-            "use_image_num": 0,
-            "interpolation_scale_t": 1,
-            "interpolation_scale_h": 1,
-            "interpolation_scale_w": 1,
-            "cache_dir": "../cache_dir",
-            "image_data": "/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt",
-            "video_data": "1",
-            "train_fps": 24,
-            "drop_short_ratio": 1.0,
-            "use_img_from_vid": False,
-            "speed_factor": 1.0,
-            "cfg": 0.1,
-            "text_encoder_name": "google/mt5-xxl",
-            "dataloader_num_workers": 10,
-        },
-    )
-    accelerator = Accelerator()
-    dataset = getdataset(args)
-    num = len(dataset_prog.img_cap_list)
-    zero = 0
-    for idx in tqdm(range(num)):
-        image_data = dataset_prog.img_cap_list[idx]
-        caps = [i["cap"] if isinstance(i["cap"], list) else [i["cap"]] for i in image_data]
-        try:
-            caps = [[random.choice(i)] for i in caps]
-        except Exception as e:
-            print(e)
-            # import ipdb;ipdb.set_trace()
-            print(image_data)
-            zero += 1
-            continue
-        assert caps[0] is not None and len(caps[0]) > 0
-    print(num, zero)
-    import ipdb
-
-    ipdb.set_trace()
-    print("end")

fastvideo/dataset/latent_datasets.py

@@ -1,118 +0,0 @@
-import json
-import os
-import random
-
-import torch
-from torch.utils.data import Dataset
-
-
-class LatentDataset(Dataset):
-
-    def __init__(
-        self,
-        json_path,
-        num_latent_t,
-        cfg_rate,
-    ):
-        # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path
-        self.json_path = json_path
-        self.cfg_rate = cfg_rate
-        self.datase_dir_path = os.path.dirname(json_path)
-        self.video_dir = os.path.join(self.datase_dir_path, "video")
-        self.latent_dir = os.path.join(self.datase_dir_path, "latent")
-        self.prompt_embed_dir = os.path.join(self.datase_dir_path, "prompt_embed")
-        self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask")
-        with open(self.json_path, "r") as f:
-            self.data_anno = json.load(f)
-        # json.load(f) already keeps the order
-        # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path'])
-        self.num_latent_t = num_latent_t
-        # just zero embeddings [256, 4096]
-        self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32)
-        # 256 zeros
-        self.uncond_prompt_mask = torch.zeros(256).bool()
-        self.lengths = [data_item["length"] if "length" in data_item else 1 for data_item in self.data_anno]
-
-    def __getitem__(self, idx):
-        latent_file = self.data_anno[idx]["latent_path"]
-        prompt_embed_file = self.data_anno[idx]["prompt_embed_path"]
-        prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"]
-        # load
-        latent = torch.load(
-            os.path.join(self.latent_dir, latent_file),
-            map_location="cpu",
-            weights_only=True,
-        )
-        latent = latent.squeeze(0)[:, -self.num_latent_t:]
-        if random.random() < self.cfg_rate:
-            prompt_embed = self.uncond_prompt_embed
-            prompt_attention_mask = self.uncond_prompt_mask
-        else:
-            prompt_embed = torch.load(
-                os.path.join(self.prompt_embed_dir, prompt_embed_file),
-                map_location="cpu",
-                weights_only=True,
-            )
-            prompt_attention_mask = torch.load(
-                os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file),
-                map_location="cpu",
-                weights_only=True,
-            )
-        return latent, prompt_embed, prompt_attention_mask
-
-    def __len__(self):
-        return len(self.data_anno)
-
-
-def latent_collate_function(batch):
-    # return latent, prompt, latent_attn_mask, text_attn_mask
-    # latent_attn_mask: # b t h w
-    # text_attn_mask: b 1 l
-    # needs to check if the latent/prompt' size and apply padding & attn mask
-    latents, prompt_embeds, prompt_attention_masks = zip(*batch)
-    # calculate max shape
-    max_t = max([latent.shape[1] for latent in latents])
-    max_h = max([latent.shape[2] for latent in latents])
-    max_w = max([latent.shape[3] for latent in latents])
-
-    # padding
-    latents = [
-        torch.nn.functional.pad(
-            latent,
-            (
-                0,
-                max_t - latent.shape[1],
-                0,
-                max_h - latent.shape[2],
-                0,
-                max_w - latent.shape[3],
-            ),
-        ) for latent in latents
-    ]
-    # attn mask
-    latent_attn_mask = torch.ones(len(latents), max_t, max_h, max_w)
-    # set to 0 if padding
-    for i, latent in enumerate(latents):
-        latent_attn_mask[i, latent.shape[1]:, :, :] = 0
-        latent_attn_mask[i, :, latent.shape[2]:, :] = 0
-        latent_attn_mask[i, :, :, latent.shape[3]:] = 0
-
-    prompt_embeds = torch.stack(prompt_embeds, dim=0)
-    prompt_attention_masks = torch.stack(prompt_attention_masks, dim=0)
-    latents = torch.stack(latents, dim=0)
-    return latents, prompt_embeds, latent_attn_mask, prompt_attention_masks
-
-
-if __name__ == "__main__":
-    dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt", num_latent_t=28)
-    dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function)
-    for latent, prompt_embed, latent_attn_mask, prompt_attention_mask in dataloader:
-        print(
-            latent.shape,
-            prompt_embed.shape,
-            latent_attn_mask.shape,
-            prompt_attention_mask.shape,
-        )
-        import pdb
-
-        pdb.set_trace()

fastvideo/dataset/t2v_datasets.py

@@ -1,324 +0,0 @@
-import json
-import math
-import os
-import random
-from collections import Counter
-from os.path import join as opj
-
-import numpy as np
-import torch
-import torchvision
-from einops import rearrange
-from PIL import Image
-from torch.utils.data import Dataset
-
-from fastvideo.utils.dataset_utils import DecordInit
-from fastvideo.utils.logging_ import main_print
-
-
-class SingletonMeta(type):
-    _instances = {}
-
-    def __call__(cls, *args, **kwargs):
-        if cls not in cls._instances:
-            instance = super().__call__(*args, **kwargs)
-            cls._instances[cls] = instance
-        return cls._instances[cls]
-
-
-class DataSetProg(metaclass=SingletonMeta):
-
-    def __init__(self):
-        self.cap_list = []
-        self.elements = []
-        self.num_workers = 1
-        self.n_elements = 0
-        self.worker_elements = dict()
-        self.n_used_elements = dict()
-
-    def set_cap_list(self, num_workers, cap_list, n_elements):
-        self.num_workers = num_workers
-        self.cap_list = cap_list
-        self.n_elements = n_elements
-        self.elements = list(range(n_elements))
-        random.shuffle(self.elements)
-        print(f"n_elements: {len(self.elements)}", flush=True)
-
-        for i in range(self.num_workers):
-            self.n_used_elements[i] = 0
-            per_worker = int(math.ceil(len(self.elements) / float(self.num_workers)))
-            start = i * per_worker
-            end = min(start + per_worker, len(self.elements))
-            self.worker_elements[i] = self.elements[start:end]
-
-    def get_item(self, work_info):
-        if work_info is None:
-            worker_id = 0
-        else:
-            worker_id = work_info.id
-
-        idx = self.worker_elements[worker_id][self.n_used_elements[worker_id] % len(self.worker_elements[worker_id])]
-        self.n_used_elements[worker_id] += 1
-        return idx
-
-
-dataset_prog = DataSetProg()
-
-
-def filter_resolution(h, w, max_h_div_w_ratio=17 / 16, min_h_div_w_ratio=8 / 16):
-    if h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio:
-        return True
-    return False
-
-
-class T2V_dataset(Dataset):
-
-    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
-        self.data = args.data_merge_path
-        self.num_frames = args.num_frames
-        self.train_fps = args.train_fps
-        self.use_image_num = args.use_image_num
-        self.transform = transform
-        self.transform_topcrop = transform_topcrop
-        self.temporal_sample = temporal_sample
-        self.tokenizer = tokenizer
-        self.text_max_length = args.text_max_length
-        self.cfg = args.cfg
-        self.speed_factor = args.speed_factor
-        self.max_height = args.max_height
-        self.max_width = args.max_width
-        self.drop_short_ratio = args.drop_short_ratio
-        assert self.speed_factor >= 1
-        self.v_decoder = DecordInit()
-        self.video_length_tolerance_range = args.video_length_tolerance_range
-        self.support_Chinese = True
-        if "mt5" not in args.text_encoder_name:
-            self.support_Chinese = False
-
-        cap_list = self.get_cap_list()
-
-        assert len(cap_list) > 0
-        cap_list, self.sample_num_frames = self.define_frame_index(cap_list)
-        self.lengths = self.sample_num_frames
-
-        n_elements = len(cap_list)
-        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list, n_elements)
-
-        print(f"video length: {len(dataset_prog.cap_list)}", flush=True)
-
-    def set_checkpoint(self, n_used_elements):
-        for i in range(len(dataset_prog.n_used_elements)):
-            dataset_prog.n_used_elements[i] = n_used_elements
-
-    def __len__(self):
-        return dataset_prog.n_elements
-
-    def __getitem__(self, idx):
-
-        data = self.get_data(idx)
-        return data
-
-    def get_data(self, idx):
-        path = dataset_prog.cap_list[idx]["path"]
-        if path.endswith(".mp4"):
-            return self.get_video(idx)
-        else:
-            return self.get_image(idx)
-
-    def get_video(self, idx):
-        video_path = dataset_prog.cap_list[idx]["path"]
-        assert os.path.exists(video_path), f"file {video_path} do not exist!"
-        frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
-        torchvision_video, _, metadata = torchvision.io.read_video(video_path, output_format="TCHW")
-        video = torchvision_video[frame_indices]
-        video = self.transform(video)
-        video = rearrange(video, "t c h w -> c t h w")
-        video = video.to(torch.uint8)
-        assert video.dtype == torch.uint8
-
-        h, w = video.shape[-2:]
-        assert (
-            h / w <= 17 / 16 and h / w >= 8 / 16
-        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({video_path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
-
-        video = video.float() / 127.5 - 1.0
-
-        text = dataset_prog.cap_list[idx]["cap"]
-        if not isinstance(text, list):
-            text = [text]
-        text = [random.choice(text)]
-
-        text = text[0] if random.random() > self.cfg else ""
-        text_tokens_and_mask = self.tokenizer(
-            text,
-            max_length=self.text_max_length,
-            padding="max_length",
-            truncation=True,
-            return_attention_mask=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        input_ids = text_tokens_and_mask["input_ids"]
-        cond_mask = text_tokens_and_mask["attention_mask"]
-        return dict(
-            pixel_values=video,
-            text=text,
-            input_ids=input_ids,
-            cond_mask=cond_mask,
-            path=video_path,
-        )
-
-    def get_image(self, idx):
-        image_data = dataset_prog.cap_list[idx]  # [{'path': path, 'cap': cap}, ...]
-
-        image = Image.open(image_data["path"]).convert("RGB")  # [h, w, c]
-        image = torch.from_numpy(np.array(image))  # [h, w, c]
-        image = rearrange(image, "h w c -> c h w").unsqueeze(0)  #  [1 c h w]
-        # for i in image:
-        #     h, w = i.shape[-2:]
-        #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
-
-        image = (self.transform_topcrop(image) if "human_images" in image_data["path"] else self.transform(image)
-                 )  #  [1 C H W] -> num_img [1 C H W]
-        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
-
-        image = image.float() / 127.5 - 1.0
-
-        caps = (image_data["cap"] if isinstance(image_data["cap"], list) else [image_data["cap"]])
-        caps = [random.choice(caps)]
-        text = caps
-        input_ids, cond_mask = [], []
-        text = text[0] if random.random() > self.cfg else ""
-        text_tokens_and_mask = self.tokenizer(
-            text,
-            max_length=self.text_max_length,
-            padding="max_length",
-            truncation=True,
-            return_attention_mask=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        input_ids = text_tokens_and_mask["input_ids"]  # 1, l
-        cond_mask = text_tokens_and_mask["attention_mask"]  # 1, l
-        return dict(
-            pixel_values=image,
-            text=text,
-            input_ids=input_ids,
-            cond_mask=cond_mask,
-            path=image_data["path"],
-        )
-
-    def define_frame_index(self, cap_list):
-        new_cap_list = []
-        sample_num_frames = []
-        cnt_too_long = 0
-        cnt_too_short = 0
-        cnt_no_cap = 0
-        cnt_no_resolution = 0
-        cnt_resolution_mismatch = 0
-        cnt_movie = 0
-        cnt_img = 0
-        for i in cap_list:
-            path = i["path"]
-            cap = i.get("cap", None)
-            # ======no caption=====
-            if cap is None:
-                cnt_no_cap += 1
-                continue
-            if path.endswith(".mp4"):
-                # ======no fps and duration=====
-                duration = i.get("duration", None)
-                fps = i.get("fps", None)
-                if fps is None or duration is None:
-                    continue
-
-                # ======resolution mismatch=====
-                resolution = i.get("resolution", None)
-                if resolution is None:
-                    cnt_no_resolution += 1
-                    continue
-                else:
-                    if (resolution.get("height", None) is None or resolution.get("width", None) is None):
-                        cnt_no_resolution += 1
-                        continue
-                    height, width = i["resolution"]["height"], i["resolution"]["width"]
-                    aspect = self.max_height / self.max_width
-                    hw_aspect_thr = 1.5
-                    is_pick = filter_resolution(
-                        height,
-                        width,
-                        max_h_div_w_ratio=hw_aspect_thr * aspect,
-                        min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
-                    )
-                    if not is_pick:
-                        print("resolution mismatch")
-                        cnt_resolution_mismatch += 1
-                        continue
-
-                # import ipdb;ipdb.set_trace()
-                i["num_frames"] = math.ceil(fps * duration)
-                # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
-                if i["num_frames"] / fps > self.video_length_tolerance_range * (
-                        self.num_frames / self.train_fps *
-                        self.speed_factor):  # too long video is not suitable for this training stage (self.num_frames)
-                    cnt_too_long += 1
-                    continue
-
-                # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
-                frame_interval = fps / self.train_fps
-                start_frame_idx = 0
-                frame_indices = np.arange(start_frame_idx, i["num_frames"], frame_interval).astype(int)
-
-                # comment out it to enable dynamic frames training
-                if (len(frame_indices) < self.num_frames and random.random() < self.drop_short_ratio):
-                    cnt_too_short += 1
-                    continue
-
-                #  too long video will be temporal-crop randomly
-                if len(frame_indices) > self.num_frames:
-                    begin_index, end_index = self.temporal_sample(len(frame_indices))
-                    frame_indices = frame_indices[begin_index:end_index]
-                    # frame_indices = frame_indices[:self.num_frames]  # head crop
-                i["sample_frame_index"] = frame_indices.tolist()
-                new_cap_list.append(i)
-                i["sample_num_frames"] = len(i["sample_frame_index"])  # will use in dataloader(group sampler)
-                sample_num_frames.append(i["sample_num_frames"])
-            elif path.endswith(".jpg"):  # image
-                cnt_img += 1
-                new_cap_list.append(i)
-                i["sample_num_frames"] = 1
-                sample_num_frames.append(i["sample_num_frames"])
-            else:
-                raise NameError(
-                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image")
-        # import ipdb;ipdb.set_trace()
-        main_print(
-            f"no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, "
-            f"no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, "
-            f"Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, "
-            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}")
-        return new_cap_list, sample_num_frames
-
-    def decord_read(self, path, frame_indices):
-        decord_vr = self.v_decoder(path)
-        video_data = decord_vr.get_batch(frame_indices).asnumpy()
-        video_data = torch.from_numpy(video_data)
-        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
-        return video_data
-
-    def read_jsons(self, data):
-        cap_lists = []
-        with open(data, "r") as f:
-            folder_anno = [i.strip().split(",") for i in f.readlines() if len(i.strip()) > 0]
-        print(folder_anno)
-        for folder, anno in folder_anno:
-            with open(anno, "r") as f:
-                sub_list = json.load(f)
-            for i in range(len(sub_list)):
-                sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
-            cap_lists += sub_list
-        return cap_lists
-
-    def get_cap_list(self):
-        cap_lists = self.read_jsons(self.data)
-        return cap_lists

fastvideo/dataset/transform.py

@@ -1,608 +0,0 @@
-import numbers
-import random
-
-import torch
-from PIL import Image
-
-
-def _is_tensor_video_clip(clip):
-    if not torch.is_tensor(clip):
-        raise TypeError("clip should be Tensor. Got %s" % type(clip))
-
-    if not clip.ndimension() == 4:
-        raise ValueError("clip should be 4D. Got %dD" % clip.dim())
-
-    return True
-
-
-def center_crop_arr(pil_image, image_size):
-    """
-    Center cropping implementation from ADM.
-    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
-    """
-    while min(*pil_image.size) >= 2 * image_size:
-        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
-
-    scale = image_size / min(*pil_image.size)
-    pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
-
-    arr = np.array(pil_image)
-    crop_y = (arr.shape[0] - image_size) // 2
-    crop_x = (arr.shape[1] - image_size) // 2
-    return Image.fromarray(arr[crop_y:crop_y + image_size, crop_x:crop_x + image_size])
-
-
-def crop(clip, i, j, h, w):
-    """
-    Args:
-        clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-    """
-    if len(clip.size()) != 4:
-        raise ValueError("clip should be a 4D tensor")
-    return clip[..., i:i + h, j:j + w]
-
-
-def resize(clip, target_size, interpolation_mode):
-    if len(target_size) != 2:
-        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
-    return torch.nn.functional.interpolate(
-        clip,
-        size=target_size,
-        mode=interpolation_mode,
-        align_corners=True,
-        antialias=True,
-    )
-
-
-def resize_scale(clip, target_size, interpolation_mode):
-    if len(target_size) != 2:
-        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
-    H, W = clip.size(-2), clip.size(-1)
-    scale_ = target_size[0] / min(H, W)
-    return torch.nn.functional.interpolate(
-        clip,
-        scale_factor=scale_,
-        mode=interpolation_mode,
-        align_corners=True,
-        antialias=True,
-    )
-
-
-def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"):
-    """
-    Do spatial cropping and resizing to the video clip
-    Args:
-        clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        i (int): i in (i,j) i.e coordinates of the upper left corner.
-        j (int): j in (i,j) i.e coordinates of the upper left corner.
-        h (int): Height of the cropped region.
-        w (int): Width of the cropped region.
-        size (tuple(int, int)): height and width of resized clip
-    Returns:
-        clip (torch.tensor): Resized and cropped clip. Size is (T, C, H, W)
-    """
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    clip = crop(clip, i, j, h, w)
-    clip = resize(clip, size, interpolation_mode)
-    return clip
-
-
-def center_crop(clip, crop_size):
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    h, w = clip.size(-2), clip.size(-1)
-    th, tw = crop_size
-    if h < th or w < tw:
-        raise ValueError("height and width must be no smaller than crop_size")
-
-    i = int(round((h - th) / 2.0))
-    j = int(round((w - tw) / 2.0))
-    return crop(clip, i, j, th, tw)
-
-
-def center_crop_using_short_edge(clip):
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    h, w = clip.size(-2), clip.size(-1)
-    if h < w:
-        th, tw = h, h
-        i = 0
-        j = int(round((w - tw) / 2.0))
-    else:
-        th, tw = w, w
-        i = int(round((h - th) / 2.0))
-        j = 0
-    return crop(clip, i, j, th, tw)
-
-
-def center_crop_th_tw(clip, th, tw, top_crop):
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-
-    # import ipdb;ipdb.set_trace()
-    h, w = clip.size(-2), clip.size(-1)
-    tr = th / tw
-    if h / w > tr:
-        new_h = int(w * tr)
-        new_w = w
-    else:
-        new_h = h
-        new_w = int(h / tr)
-
-    i = 0 if top_crop else int(round((h - new_h) / 2.0))
-    j = int(round((w - new_w) / 2.0))
-    return crop(clip, i, j, new_h, new_w)
-
-
-def random_shift_crop(clip):
-    """
-    Slide along the long edge, with the short edge as crop size
-    """
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    h, w = clip.size(-2), clip.size(-1)
-
-    if h <= w:
-        short_edge = h
-    else:
-        short_edge = w
-
-    th, tw = short_edge, short_edge
-
-    i = torch.randint(0, h - th + 1, size=(1, )).item()
-    j = torch.randint(0, w - tw + 1, size=(1, )).item()
-    return crop(clip, i, j, th, tw)
-
-
-def normalize_video(clip):
-    """
-    Convert tensor data type from uint8 to float, divide value by 255.0 and
-    permute the dimensions of clip tensor
-    Args:
-        clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W)
-    Return:
-        clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W)
-    """
-    _is_tensor_video_clip(clip)
-    if not clip.dtype == torch.uint8:
-        raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype))
-    # return clip.float().permute(3, 0, 1, 2) / 255.0
-    return clip.float() / 255.0
-
-
-def normalize(clip, mean, std, inplace=False):
-    """
-    Args:
-        clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W)
-        mean (tuple): pixel RGB mean. Size is (3)
-        std (tuple): pixel standard deviation. Size is (3)
-    Returns:
-        normalized clip (torch.tensor): Size is (T, C, H, W)
-    """
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    if not inplace:
-        clip = clip.clone()
-    mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device)
-    # print(mean)
-    std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device)
-    clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None])
-    return clip
-
-
-def hflip(clip):
-    """
-    Args:
-        clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W)
-    Returns:
-        flipped clip (torch.tensor): Size is (T, C, H, W)
-    """
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    return clip.flip(-1)
-
-
-class RandomCropVideo:
-
-    def __init__(self, size):
-        if isinstance(size, numbers.Number):
-            self.size = (int(size), int(size))
-        else:
-            self.size = size
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: randomly cropped video clip.
-                size is (T, C, OH, OW)
-        """
-        i, j, h, w = self.get_params(clip)
-        return crop(clip, i, j, h, w)
-
-    def get_params(self, clip):
-        h, w = clip.shape[-2:]
-        th, tw = self.size
-
-        if h < th or w < tw:
-            raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}")
-
-        if w == tw and h == th:
-            return 0, 0, h, w
-
-        i = torch.randint(0, h - th + 1, size=(1, )).item()
-        j = torch.randint(0, w - tw + 1, size=(1, )).item()
-
-        return i, j, th, tw
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size})"
-
-
-class SpatialStrideCropVideo:
-
-    def __init__(self, stride):
-        self.stride = stride
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: cropped video clip by stride.
-                size is (T, C, OH, OW)
-        """
-        i, j, h, w = self.get_params(clip)
-        return crop(clip, i, j, h, w)
-
-    def get_params(self, clip):
-        h, w = clip.shape[-2:]
-
-        th, tw = h // self.stride * self.stride, w // self.stride * self.stride
-
-        return 0, 0, th, tw  # from top-left
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size})"
-
-
-class LongSideResizeVideo:
-    """
-    First use the long side,
-    then resize to the specified size
-    """
-
-    def __init__(
-        self,
-        size,
-        skip_low_resolution=False,
-        interpolation_mode="bilinear",
-    ):
-        self.size = size
-        self.skip_low_resolution = skip_low_resolution
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: scale resized video clip.
-                size is (T, C, 512, *) or (T, C, *, 512)
-        """
-        _, _, h, w = clip.shape
-        if self.skip_low_resolution and max(h, w) <= self.size:
-            return clip
-        if h > w:
-            w = int(w * self.size / h)
-            h = self.size
-        else:
-            h = int(h * self.size / w)
-            w = self.size
-        resize_clip = resize(clip, target_size=(h, w), interpolation_mode=self.interpolation_mode)
-        return resize_clip
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
-
-
-class CenterCropResizeVideo:
-    """
-    First use the short side for cropping length,
-    center crop video, then resize to the specified size
-    """
-
-    def __init__(
-        self,
-        size,
-        top_crop=False,
-        interpolation_mode="bilinear",
-    ):
-        if len(size) != 2:
-            raise ValueError(f"size should be tuple (height, width), instead got {size}")
-        self.size = size
-        self.top_crop = top_crop
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: scale resized / center cropped video clip.
-                size is (T, C, crop_size, crop_size)
-        """
-        # clip_center_crop = center_crop_using_short_edge(clip)
-        clip_center_crop = center_crop_th_tw(clip, self.size[0], self.size[1], top_crop=self.top_crop)
-        # import ipdb;ipdb.set_trace()
-        clip_center_crop_resize = resize(
-            clip_center_crop,
-            target_size=self.size,
-            interpolation_mode=self.interpolation_mode,
-        )
-        return clip_center_crop_resize
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
-
-
-class UCFCenterCropVideo:
-    """
-    First scale to the specified size in equal proportion to the short edge,
-    then center cropping
-    """
-
-    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
-    ):
-        if isinstance(size, tuple):
-            if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
-            self.size = size
-        else:
-            self.size = (size, size)
-
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: scale resized / center cropped video clip.
-                size is (T, C, crop_size, crop_size)
-        """
-        clip_resize = resize_scale(clip=clip, target_size=self.size, interpolation_mode=self.interpolation_mode)
-        clip_center_crop = center_crop(clip_resize, self.size)
-        return clip_center_crop
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
-
-
-class KineticsRandomCropResizeVideo:
-    """
-    Slide along the long edge, with the short edge as crop size. And resie to the desired size.
-    """
-
-    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
-    ):
-        if isinstance(size, tuple):
-            if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
-            self.size = size
-        else:
-            self.size = (size, size)
-
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        clip_random_crop = random_shift_crop(clip)
-        clip_resize = resize(clip_random_crop, self.size, self.interpolation_mode)
-        return clip_resize
-
-
-class CenterCropVideo:
-
-    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
-    ):
-        if isinstance(size, tuple):
-            if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
-            self.size = size
-        else:
-            self.size = (size, size)
-
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: center cropped video clip.
-                size is (T, C, crop_size, crop_size)
-        """
-        clip_center_crop = center_crop(clip, self.size)
-        return clip_center_crop
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
-
-
-class Normalize:
-    """
-    Normalize the video clip by mean subtraction and division by standard deviation
-    Args:
-        mean (3-tuple): pixel RGB mean
-        std (3-tuple): pixel RGB standard deviation
-        inplace (boolean): whether do in-place normalization
-    """
-
-    def __init__(self, mean, std, inplace=False):
-        self.mean = mean
-        self.std = std
-        self.inplace = inplace
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): video clip must be normalized. Size is (C, T, H, W)
-        """
-        return normalize(clip, self.mean, self.std, self.inplace)
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(mean={self.mean}, std={self.std}, inplace={self.inplace})"
-
-
-class Normalize255:
-    """
-    Convert tensor data type from uint8 to float, divide value by 255.0 and
-    """
-
-    def __init__(self):
-        pass
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W)
-        Return:
-            clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W)
-        """
-        return normalize_video(clip)
-
-    def __repr__(self) -> str:
-        return self.__class__.__name__
-
-
-class RandomHorizontalFlipVideo:
-    """
-    Flip the video clip along the horizontal direction with a given probability
-    Args:
-        p (float): probability of the clip being flipped. Default value is 0.5
-    """
-
-    def __init__(self, p=0.5):
-        self.p = p
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Size is (T, C, H, W)
-        Return:
-            clip (torch.tensor): Size is (T, C, H, W)
-        """
-        if random.random() < self.p:
-            clip = hflip(clip)
-        return clip
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(p={self.p})"
-
-
-#  ------------------------------------------------------------
-#  ---------------------  Sampling  ---------------------------
-#  ------------------------------------------------------------
-class TemporalRandomCrop(object):
-    """Temporally crop the given frame indices at a random location.
-
-    Args:
-        size (int): Desired length of frames will be seen in the model.
-    """
-
-    def __init__(self, size):
-        self.size = size
-
-    def __call__(self, total_frames):
-        rand_end = max(0, total_frames - self.size - 1)
-        begin_index = random.randint(0, rand_end)
-        end_index = min(begin_index + self.size, total_frames)
-        return begin_index, end_index
-
-
-class DynamicSampleDuration(object):
-    """Temporally crop the given frame indices at a random location.
-
-    Args:
-        size (int): Desired length of frames will be seen in the model.
-    """
-
-    def __init__(self, t_stride, extra_1):
-        self.t_stride = t_stride
-        self.extra_1 = extra_1
-
-    def __call__(self, t, h, w):
-        if self.extra_1:
-            t = t - 1
-        truncate_t_list = list(range(t + 1))[t // 2:][::self.t_stride]  # need half at least
-        truncate_t = random.choice(truncate_t_list)
-        if self.extra_1:
-            truncate_t = truncate_t + 1
-        return 0, truncate_t
-
-
-if __name__ == "__main__":
-    import os
-
-    import numpy as np
-    import torchvision.io as io
-    from torchvision import transforms
-    from torchvision.utils import save_image
-
-    vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi", pts_unit="sec", output_format="TCHW")
-
-    trans = transforms.Compose([
-        Normalize255(),
-        RandomHorizontalFlipVideo(),
-        UCFCenterCropVideo(512),
-        # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
-    ])
-
-    target_video_len = 32
-    frame_interval = 1
-    total_frames = len(vframes)
-    print(total_frames)
-
-    temporal_sample = TemporalRandomCrop(target_video_len * frame_interval)
-
-    # Sampling video frames
-    start_frame_ind, end_frame_ind = temporal_sample(total_frames)
-    # print(start_frame_ind)
-    # print(end_frame_ind)
-    assert end_frame_ind - start_frame_ind >= target_video_len
-    frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, target_video_len, dtype=int)
-    print(frame_indice)
-
-    select_vframes = vframes[frame_indice]
-    print(select_vframes.shape)
-    print(select_vframes.dtype)
-
-    select_vframes_trans = trans(select_vframes)
-    print(select_vframes_trans.shape)
-    print(select_vframes_trans.dtype)
-
-    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * 255).to(dtype=torch.uint8)
-    print(select_vframes_trans_int.dtype)
-    print(select_vframes_trans_int.permute(0, 2, 3, 1).shape)
-
-    io.write_video("./test.avi", select_vframes_trans_int.permute(0, 2, 3, 1), fps=8)
-
-    for i in range(target_video_len):
-        save_image(
-            select_vframes_trans[i],
-            os.path.join("./test000", "%04d.png" % i),
-            normalize=True,
-            value_range=(-1, 1),
-        )

fastvideo/distill.py

@@ -1,808 +0,0 @@
-# !/bin/python3
-# isort: skip_file
-import argparse
-import math
-import os
-import time
-from collections import deque
-from copy import deepcopy
-
-import torch
-import torch.distributed as dist
-import wandb
-from accelerate.utils import set_seed
-from diffusers import FlowMatchEulerDiscreteScheduler
-from fastvideo.distill.solver import PCMFMScheduler
-from diffusers.optimization import get_scheduler
-from diffusers.utils import check_min_version
-from peft import LoraConfig
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp import ShardingStrategy
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from tqdm.auto import tqdm
-
-from fastvideo.dataset.latent_datasets import (LatentDataset, latent_collate_function)
-from fastvideo.distill.solver import EulerSolver, extract_into_tensor
-from fastvideo.utils.latents_utils import normalize_dit_input
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
-from fastvideo.utils.checkpoint import (resume_lora_optimizer, save_checkpoint, save_lora_checkpoint)
-from fastvideo.utils.communications import (broadcast, sp_parallel_dataloader_wrapper)
-from fastvideo.utils.dataset_utils import LengthGroupedSampler
-from fastvideo.utils.fsdp_util import (apply_fsdp_checkpointing, get_dit_fsdp_kwargs)
-from fastvideo.utils.load import load_transformer
-from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group, get_sequence_parallel_state,
-                                             initialize_sequence_parallel_state)
-from fastvideo.utils.validation import log_validation
-
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.31.0")
-
-
-def main_print(content):
-    if int(os.environ["LOCAL_RANK"]) <= 0:
-        print(content)
-
-
-def reshard_fsdp(model):
-    for m in FSDP.fsdp_modules(model):
-        if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
-            torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
-
-
-def get_norm(model_pred, norms, gradient_accumulation_steps):
-    fro_norm = (
-        torch.linalg.matrix_norm(model_pred, ord="fro") /  # codespell:ignore
-        gradient_accumulation_steps)
-    largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) / gradient_accumulation_steps)
-    absolute_mean = torch.mean(torch.abs(model_pred)) / gradient_accumulation_steps
-    absolute_max = torch.max(torch.abs(model_pred)) / gradient_accumulation_steps
-    dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG)
-    dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG)
-    dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG)
-    norms["fro"] += torch.mean(fro_norm).item()  # codespell:ignore
-    norms["largest singular value"] += torch.mean(largest_singular_value).item()
-    norms["absolute mean"] += absolute_mean.item()
-    norms["absolute max"] += absolute_max.item()
-
-
-def distill_one_step(
-    transformer,
-    model_type,
-    teacher_transformer,
-    ema_transformer,
-    optimizer,
-    lr_scheduler,
-    loader,
-    noise_scheduler,
-    solver,
-    noise_random_generator,
-    gradient_accumulation_steps,
-    sp_size,
-    max_grad_norm,
-    uncond_prompt_embed,
-    uncond_prompt_mask,
-    num_euler_timesteps,
-    multiphase,
-    not_apply_cfg_solver,
-    distill_cfg,
-    ema_decay,
-    pred_decay_weight,
-    pred_decay_type,
-    hunyuan_teacher_disable_cfg,
-):
-    total_loss = 0.0
-    optimizer.zero_grad()
-    model_pred_norm = {
-        "fro": 0.0,  # codespell:ignore
-        "largest singular value": 0.0,
-        "absolute mean": 0.0,
-        "absolute max": 0.0,
-    }
-    for _ in range(gradient_accumulation_steps):
-        (
-            latents,
-            encoder_hidden_states,
-            latents_attention_mask,
-            encoder_attention_mask,
-        ) = next(loader)
-        model_input = normalize_dit_input(model_type, latents)
-        noise = torch.randn_like(model_input)
-        bsz = model_input.shape[0]
-        index = torch.randint(0, num_euler_timesteps, (bsz, ), device=model_input.device).long()
-        if sp_size > 1:
-            broadcast(index)
-        # Add noise according to flow matching.
-        # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
-        sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
-        sigmas_prev = extract_into_tensor(solver.sigmas_prev, index, model_input.shape)
-
-        timesteps = (sigmas * noise_scheduler.config.num_train_timesteps).view(-1)
-        # if squeeze to [], unsqueeze to [1]
-
-        timesteps_prev = (sigmas_prev * noise_scheduler.config.num_train_timesteps).view(-1)
-        noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
-        # Predict the noise residual
-        with torch.autocast("cuda", dtype=torch.bfloat16):
-            if args.model_type == "wan":
-                    teacher_kwargs = {
-                    "hidden_states": noisy_model_input,
-                    "encoder_hidden_states": encoder_hidden_states,
-                    "timestep": timesteps,
-                    "return_dict": True,
-                }
-            else:
-                teacher_kwargs = {
-                    "hidden_states": noisy_model_input,
-                    "encoder_hidden_states": encoder_hidden_states,
-                    "timestep": timesteps,
-                    "encoder_attention_mask": encoder_attention_mask,  # B, L
-                    "return_dict": False,
-                }
-            if hunyuan_teacher_disable_cfg:
-                teacher_kwargs["guidance"] = torch.tensor([1000.0],
-                                                          device=noisy_model_input.device,
-                                                          dtype=torch.bfloat16)
-            model_pred = transformer(**teacher_kwargs)[0]
-
-        # if accelerator.is_main_process:
-        model_pred, end_index = solver.euler_style_multiphase_pred(noisy_model_input, model_pred, index, multiphase)
-        with torch.no_grad():
-            w = distill_cfg
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                if args.model_type == "wan":
-                    cond_teacher_kwargs ={
-                    "hidden_states": noisy_model_input,
-                    "encoder_hidden_states": encoder_hidden_states,
-                    "timestep": timesteps,
-                    "return_dict": True,
-                    }
-                else:
-                    cond_teacher_kwargs = {
-                    "hidden_states": noisy_model_input,
-                    "encoder_hidden_states": encoder_hidden_states,
-                    "timestep": timesteps,
-                    "encoder_attention_mask": encoder_attention_mask,  # B, L
-                    "return_dict": False,
-                    }
-                cond_teacher_output = teacher_transformer(**cond_teacher_kwargs)[0].float()
-            if not_apply_cfg_solver:
-                uncond_teacher_output = cond_teacher_output
-            else:
-                # Get teacher model prediction on noisy_latents and unconditional embedding
-                with torch.autocast("cuda", dtype=torch.bfloat16):
-                    if args.model_type == "wan":
-                        uncond_teacher_kwargs = {
-                        "hidden_states": noisy_model_input,
-                        "encoder_hidden_states":uncond_prompt_embed.unsqueeze(0).expand(bsz, -1, -1),
-                        "timestep": timesteps,
-                        "return_dict": True,
-                        }
-                    else:
-                        uncond_teacher_kwargs = {
-                        "hidden_states": noisy_model_input,
-                        "encoder_hidden_states":uncond_prompt_embed.unsqueeze(0).expand(bsz, -1, -1),
-                        "timestep": timesteps,
-                        "encoder_attention_mask": uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
-                        "return_dict": False,
-                        }
-                    
-                    uncond_teacher_output = teacher_transformer(**uncond_teacher_kwargs)[0].float()
-
-            teacher_output = uncond_teacher_output + w * (cond_teacher_output - uncond_teacher_output)
-            x_prev = solver.euler_step(noisy_model_input, teacher_output, index)
-
-        # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
-        with torch.no_grad():
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                if args.model_type == "wan":
-                    target_pred_kwargs = {
-                        "hidden_states": x_prev.float(),
-                        "encoder_hidden_states": encoder_hidden_states,
-                        "timestep":timesteps_prev,
-                        "return_dict":True,
-                    }
-                else:
-                    target_pred_kwargs = {
-                        "hidden_states": x_prev.float(),
-                        "encoder_hidden_states": encoder_hidden_states,
-                        "timestep":timesteps_prev,
-                        "encoder_attention_mask":encoder_attention_mask,
-                        "return_dict":False,
-                    }
-                if ema_transformer is not None:
-                    target_pred = ema_transformer(**target_pred_kwargs)[0]
-                else:
-                    target_pred = transformer(**target_pred_kwargs)[0]
-
-            target, end_index = solver.euler_style_multiphase_pred(x_prev, target_pred, index, multiphase, True)
-
-        huber_c = 0.001
-        # loss = loss.mean()
-        loss = (torch.mean(torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) - huber_c) /
-                gradient_accumulation_steps)
-        if pred_decay_weight > 0:
-            if pred_decay_type == "l1":
-                pred_decay_loss = (torch.mean(torch.sqrt(model_pred.float()**2)) * pred_decay_weight /
-                                   gradient_accumulation_steps)
-                loss += pred_decay_loss
-            elif pred_decay_type == "l2":
-                # essnetially k2?
-                pred_decay_loss = (torch.mean(model_pred.float()**2) * pred_decay_weight / gradient_accumulation_steps)
-                loss += pred_decay_loss
-            else:
-                assert NotImplementedError("pred_decay_type is not implemented")
-
-        # calculate model_pred norm and mean
-        get_norm(model_pred.detach().float(), model_pred_norm, gradient_accumulation_steps)
-        loss.backward()
-
-        avg_loss = loss.detach().clone()
-        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
-        total_loss += avg_loss.item()
-
-    # update ema
-    if ema_transformer is not None:
-        reshard_fsdp(ema_transformer)
-        for p_averaged, p_model in zip(ema_transformer.parameters(), transformer.parameters()):
-            with torch.no_grad():
-                p_averaged.copy_(torch.lerp(p_averaged.detach(), p_model.detach(), 1 - ema_decay))
-
-    grad_norm = transformer.clip_grad_norm_(max_grad_norm)
-    optimizer.step()
-    lr_scheduler.step()
-
-    return total_loss, grad_norm.item(), model_pred_norm
-
-
-def main(args):
-    torch.backends.cuda.matmul.allow_tf32 = True
-
-    local_rank = int(os.environ["LOCAL_RANK"])
-    rank = int(os.environ["RANK"])
-    world_size = int(os.environ["WORLD_SIZE"])
-    dist.init_process_group("nccl")
-    torch.cuda.set_device(local_rank)
-    device = torch.cuda.current_device()
-    initialize_sequence_parallel_state(args.sp_size)
-
-    # If passed along, set the training seed now. On GPU...
-    if args.seed is not None:
-        # TODO: t within the same seq parallel group should be the same. Noise should be different.
-        set_seed(args.seed + rank)
-    # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch.
-    noise_random_generator = None
-
-    # Handle the repository creation
-    if rank == 0 and args.output_dir is not None:
-        os.makedirs(args.output_dir, exist_ok=True)
-
-    # For mixed precision training we cast all non-trainable weights to half-precision
-    # as these weights are only used for inference, keeping weights in full precision is not required.
-
-    # Create model:
-
-    main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
-
-    transformer = load_transformer(
-        args.model_type,
-        args.dit_model_name_or_path,
-        args.pretrained_model_name_or_path,
-        torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16,
-    )
-
-    teacher_transformer = deepcopy(transformer)
-    if args.use_ema:
-        ema_transformer = deepcopy(transformer)
-    else:
-        ema_transformer = None
-
-    if args.use_lora:
-        assert args.model_type == "mochi", "LoRA is only supported for Mochi model."
-        transformer.requires_grad_(False)
-        transformer_lora_config = LoraConfig(
-            r=args.lora_rank,
-            lora_alpha=args.lora_alpha,
-            init_lora_weights=True,
-            target_modules=["to_k", "to_q", "to_v", "to_out.0"],
-        )
-        transformer.add_adapter(transformer_lora_config)
-
-    main_print(
-        f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M")
-    main_print(f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}")
-    fsdp_kwargs, no_split_modules = get_dit_fsdp_kwargs(
-        transformer,
-        args.fsdp_sharding_startegy,
-        args.use_lora,
-        args.use_cpu_offload,
-        args.master_weight_type,
-    )
-
-    if args.use_lora:
-        transformer.config.lora_rank = args.lora_rank
-        transformer.config.lora_alpha = args.lora_alpha
-        transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
-        transformer._no_split_modules = no_split_modules
-        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)
-
-    transformer = FSDP(
-        transformer,
-        **fsdp_kwargs,
-    )
-    teacher_transformer = FSDP(
-        teacher_transformer,
-        **fsdp_kwargs,
-    )
-    if args.use_ema:
-        ema_transformer = FSDP(
-            ema_transformer,
-            **fsdp_kwargs,
-        )
-    main_print("--> model loaded")
-
-    if args.gradient_checkpointing:
-        apply_fsdp_checkpointing(transformer, no_split_modules, args.selective_checkpointing)
-        apply_fsdp_checkpointing(teacher_transformer, no_split_modules, args.selective_checkpointing)
-        if args.use_ema:
-            apply_fsdp_checkpointing(ema_transformer, no_split_modules, args.selective_checkpointing)
-    # Set model as trainable.
-    transformer.train()
-    teacher_transformer.requires_grad_(False)
-    if args.use_ema:
-        ema_transformer.requires_grad_(False)
-    
-    noise_scheduler = FlowMatchEulerDiscreteScheduler()
-
-    if args.scheduler_type == "pcm_linear_quadratic":
-        linear_steps = int(noise_scheduler.config.num_train_timesteps * args.linear_range)
-        sigmas = linear_quadratic_schedule(
-            noise_scheduler.config.num_train_timesteps,
-            args.linear_quadratic_threshold,
-            linear_steps,
-        )
-        sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
-    else:
-        sigmas = noise_scheduler.sigmas
-    solver = EulerSolver(
-        sigmas.numpy()[::-1],
-        noise_scheduler.config.num_train_timesteps,
-        euler_timesteps=args.num_euler_timesteps,
-    )
-    solver.to(device)
-    params_to_optimize = transformer.parameters()
-    params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize))
-
-    optimizer = torch.optim.AdamW(
-        params_to_optimize,
-        lr=args.learning_rate,
-        betas=(0.9, 0.999),
-        weight_decay=args.weight_decay,
-        eps=1e-8,
-    )
-
-    init_steps = 0
-    if args.resume_from_lora_checkpoint:
-        transformer, optimizer, init_steps = resume_lora_optimizer(transformer, args.resume_from_lora_checkpoint,
-                                                                   optimizer)
-    main_print(f"optimizer: {optimizer}")
-
-    # todo add lr scheduler
-    lr_scheduler = get_scheduler(
-        args.lr_scheduler,
-        optimizer=optimizer,
-        num_warmup_steps=args.lr_warmup_steps * world_size,
-        num_training_steps=args.max_train_steps * world_size,
-        num_cycles=args.lr_num_cycles,
-        power=args.lr_power,
-        last_epoch=init_steps - 1,
-    )
-
-    train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
-    uncond_prompt_embed = train_dataset.uncond_prompt_embed
-    uncond_prompt_mask = train_dataset.uncond_prompt_mask
-    sampler = (LengthGroupedSampler(
-        args.train_batch_size,
-        rank=rank,
-        world_size=world_size,
-        lengths=train_dataset.lengths,
-        group_frame=args.group_frame,
-        group_resolution=args.group_resolution,
-    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
-        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
-
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        collate_fn=latent_collate_function,
-        pin_memory=True,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-        drop_last=True,
-    )
-
-    num_update_steps_per_epoch = math.ceil(
-        len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
-    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
-
-    if rank == 0:
-        project = args.tracker_project_name or "fastvideo"
-        wandb.init(project=project, config=args)
-
-    # Train!
-    total_batch_size = (world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size)
-    main_print("***** Running training *****")
-    main_print(f"  Num examples = {len(train_dataset)}")
-    main_print(f"  Dataloader size = {len(train_dataloader)}")
-    main_print(f"  Num Epochs = {args.num_train_epochs}")
-    main_print(f"  Resume training from step {init_steps}")
-    main_print(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    main_print(f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
-    main_print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
-    main_print(f"  Total optimization steps = {args.max_train_steps}")
-    main_print(
-        f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B"
-    )
-    # print dtype
-    main_print(f"  Master weight dtype: {transformer.parameters().__next__().dtype}")
-
-    # Potentially load in the weights and states from a previous save
-    if args.resume_from_checkpoint:
-        assert NotImplementedError("resume_from_checkpoint is not supported now.")
-        # TODO
-
-    progress_bar = tqdm(
-        range(0, args.max_train_steps),
-        initial=init_steps,
-        desc="Steps",
-        # Only show the progress bar once on each machine.
-        disable=local_rank > 0,
-    )
-
-    loader = sp_parallel_dataloader_wrapper(
-        train_dataloader,
-        device,
-        args.train_batch_size,
-        args.sp_size,
-        args.train_sp_batch_size,
-    )
-
-    step_times = deque(maxlen=100)
-
-    # todo future
-    for i in range(init_steps):
-        next(loader)
-
-    # log_validation(args, transformer, device,
-    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
-    def get_num_phases(multi_phased_distill_schedule, step):
-        # step-phase,step-phase
-        multi_phases = multi_phased_distill_schedule.split(",")
-        phase = multi_phases[-1].split("-")[-1]
-        for step_phases in multi_phases:
-            phase_step, phase = step_phases.split("-")
-            if step <= int(phase_step):
-                return int(phase)
-        return phase
-
-    for step in range(init_steps + 1, args.max_train_steps + 1):
-        start_time = time.perf_counter()
-        assert args.multi_phased_distill_schedule is not None
-        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
-
-        loss, grad_norm, pred_norm = distill_one_step(
-            transformer,
-            args.model_type,
-            teacher_transformer,
-            ema_transformer,
-            optimizer,
-            lr_scheduler,
-            loader,
-            noise_scheduler,
-            solver,
-            noise_random_generator,
-            args.gradient_accumulation_steps,
-            args.sp_size,
-            args.max_grad_norm,
-            uncond_prompt_embed,
-            uncond_prompt_mask,
-            args.num_euler_timesteps,
-            num_phases,
-            args.not_apply_cfg_solver,
-            args.distill_cfg,
-            args.ema_decay,
-            args.pred_decay_weight,
-            args.pred_decay_type,
-            args.hunyuan_teacher_disable_cfg,
-        )
-
-        step_time = time.perf_counter() - start_time
-        step_times.append(step_time)
-        avg_step_time = sum(step_times) / len(step_times)
-
-        progress_bar.set_postfix({
-            "loss": f"{loss:.4f}",
-            "step_time": f"{step_time:.2f}s",
-            "grad_norm": grad_norm,
-            "phases": num_phases,
-        })
-        progress_bar.update(1)
-        if rank == 0:
-            wandb.log(
-                {
-                    "train_loss": loss,
-                    "learning_rate": lr_scheduler.get_last_lr()[0],
-                    "step_time": step_time,
-                    "avg_step_time": avg_step_time,
-                    "grad_norm": grad_norm,
-                    "pred_fro_norm": pred_norm["fro"],  # codespell:ignore
-                    "pred_largest_singular_value": pred_norm["largest singular value"],
-                    "pred_absolute_mean": pred_norm["absolute mean"],
-                    "pred_absolute_max": pred_norm["absolute max"],
-                },
-                step=step,
-            )
-        if step % args.checkpointing_steps == 0:
-            if args.use_lora:
-                # Save LoRA weights
-                save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
-            else:
-                # Your existing checkpoint saving code
-                if args.use_ema:
-                    save_checkpoint(ema_transformer, rank, args.output_dir, step)
-                else:
-                    save_checkpoint(transformer, rank, args.output_dir, step)
-            dist.barrier()
-        if args.log_validation and step % args.validation_steps == 0:
-            log_validation(
-                args,
-                transformer,
-                device,
-                torch.bfloat16,
-                step,
-                scheduler_type=args.scheduler_type,
-                shift=args.shift,
-                num_euler_timesteps=args.num_euler_timesteps,
-                linear_quadratic_threshold=args.linear_quadratic_threshold,
-                linear_range=args.linear_range,
-                ema=False,
-            )
-            if args.use_ema:
-                log_validation(
-                    args,
-                    ema_transformer,
-                    device,
-                    torch.bfloat16,
-                    step,
-                    scheduler_type=args.scheduler_type,
-                    shift=args.shift,
-                    num_euler_timesteps=args.num_euler_timesteps,
-                    linear_quadratic_threshold=args.linear_quadratic_threshold,
-                    linear_range=args.linear_range,
-                    ema=True,
-                )
-
-    if args.use_lora:
-        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
-    else:
-        save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
-
-    if get_sequence_parallel_state():
-        destroy_sequence_parallel_group()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument("--model_type", type=str, default="mochi", help="The type of model to train.")
-
-    # dataset & dataloader
-    parser.add_argument("--data_json_path", type=str, required=True)
-    parser.add_argument("--num_height", type=int, default=480)
-    parser.add_argument("--num_width", type=int, default=848)
-    parser.add_argument("--num_frames", type=int, default=163)
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=10,
-        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
-    )
-    parser.add_argument(
-        "--train_batch_size",
-        type=int,
-        default=16,
-        help="Batch size (per device) for the training dataloader.",
-    )
-    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
-    parser.add_argument("--group_frame", action="store_true")  # TODO
-    parser.add_argument("--group_resolution", action="store_true")  # TODO
-
-    # text encoder & vae & diffusion model
-    parser.add_argument("--pretrained_model_name_or_path", type=str)
-    parser.add_argument("--dit_model_name_or_path", type=str)
-    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-
-    # diffusion setting
-    parser.add_argument("--ema_decay", type=float, default=0.95)
-    parser.add_argument("--ema_start_step", type=int, default=0)
-    parser.add_argument("--cfg", type=float, default=0.1)
-
-    # validation & logs
-    parser.add_argument("--validation_prompt_dir", type=str)
-    parser.add_argument("--validation_sampling_steps", type=str, default="64")
-    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
-
-    parser.add_argument("--validation_steps", type=float, default=64)
-    parser.add_argument("--log_validation", action="store_true")
-    parser.add_argument("--tracker_project_name", type=str, default=None)
-    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default=None,
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-    parser.add_argument(
-        "--checkpoints_total_limit",
-        type=int,
-        default=None,
-        help=("Max number of checkpoints to store."),
-    )
-    parser.add_argument(
-        "--checkpointing_steps",
-        type=int,
-        default=500,
-        help=("Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
-              " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
-              " training using `--resume_from_checkpoint`."),
-    )
-    parser.add_argument("--shift", type=float, default=1.0)
-    parser.add_argument(
-        "--resume_from_checkpoint",
-        type=str,
-        default=None,
-        help=("Whether training should be resumed from a previous checkpoint. Use a path saved by"
-              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
-    )
-    parser.add_argument(
-        "--resume_from_lora_checkpoint",
-        type=str,
-        default=None,
-        help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
-              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
-    )
-
-    # optimizer & scheduler & Training
-    parser.add_argument("--num_train_epochs", type=int, default=100)
-    parser.add_argument(
-        "--max_train_steps",
-        type=int,
-        default=None,
-        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
-    )
-    parser.add_argument(
-        "--gradient_accumulation_steps",
-        type=int,
-        default=1,
-        help="Number of updates steps to accumulate before performing a backward/update pass.",
-    )
-    parser.add_argument(
-        "--learning_rate",
-        type=float,
-        default=1e-4,
-        help="Initial learning rate (after the potential warmup period) to use.",
-    )
-    parser.add_argument(
-        "--scale_lr",
-        action="store_true",
-        default=False,
-        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
-    )
-    parser.add_argument(
-        "--lr_warmup_steps",
-        type=int,
-        default=10,
-        help="Number of steps for the warmup in the lr scheduler.",
-    )
-    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
-    parser.add_argument(
-        "--gradient_checkpointing",
-        action="store_true",
-        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
-    )
-    parser.add_argument("--selective_checkpointing", type=float, default=1.0)
-    parser.add_argument(
-        "--allow_tf32",
-        action="store_true",
-        help=("Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
-              " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"),
-    )
-    parser.add_argument(
-        "--mixed_precision",
-        type=str,
-        default=None,
-        choices=["no", "fp16", "bf16"],
-        help=(
-            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
-            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
-            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."),
-    )
-    parser.add_argument(
-        "--use_cpu_offload",
-        action="store_true",
-        help="Whether to use CPU offload for param & gradient & optimizer states.",
-    )
-
-    parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
-    parser.add_argument(
-        "--train_sp_batch_size",
-        type=int,
-        default=1,
-        help="Batch size for sequence parallel training",
-    )
-
-    parser.add_argument(
-        "--use_lora",
-        action="store_true",
-        default=False,
-        help="Whether to use LoRA for finetuning.",
-    )
-    parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.")
-    parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ")
-    parser.add_argument("--fsdp_sharding_startegy", default="full")
-
-    # lr_scheduler
-    parser.add_argument(
-        "--lr_scheduler",
-        type=str,
-        default="constant",
-        help=('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
-              ' "constant", "constant_with_warmup"]'),
-    )
-    parser.add_argument("--num_euler_timesteps", type=int, default=100)
-    parser.add_argument(
-        "--lr_num_cycles",
-        type=int,
-        default=1,
-        help="Number of cycles in the learning rate scheduler.",
-    )
-    parser.add_argument(
-        "--lr_power",
-        type=float,
-        default=1.0,
-        help="Power factor of the polynomial scheduler.",
-    )
-    parser.add_argument(
-        "--not_apply_cfg_solver",
-        action="store_true",
-        help="Whether to apply the cfg_solver.",
-    )
-    parser.add_argument("--distill_cfg", type=float, default=3.0, help="Distillation coefficient.")
-    # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
-    parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.")
-    parser.add_argument(
-        "--linear_quadratic_threshold",
-        type=float,
-        default=0.025,
-        help="Threshold for linear quadratic scheduler.",
-    )
-    parser.add_argument(
-        "--linear_range",
-        type=float,
-        default=0.5,
-        help="Range for linear quadratic scheduler.",
-    )
-    parser.add_argument("--weight_decay", type=float, default=0.001, help="Weight decay to apply.")
-    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA.")
-    parser.add_argument("--multi_phased_distill_schedule", type=str, default=None)
-    parser.add_argument("--pred_decay_weight", type=float, default=0.0)
-    parser.add_argument("--pred_decay_type", default="l1")
-    parser.add_argument("--hunyuan_teacher_disable_cfg", action="store_true")
-    parser.add_argument(
-        "--master_weight_type",
-        type=str,
-        default="fp32",
-        help="Weight type to use - fp32 or bf16.",
-    )
-    args = parser.parse_args()
-    main(args)

fastvideo/distill/discriminator.py

@@ -1,76 +0,0 @@
-import torch.nn as nn
-from diffusers.utils import logging
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-class DiscriminatorHead(nn.Module):
-
-    def __init__(self, input_channel, output_channel=1):
-        super().__init__()
-        inner_channel = 1024
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(input_channel, inner_channel, 1, 1, 0),
-            nn.GroupNorm(32, inner_channel),
-            nn.LeakyReLU(inplace=True),  # use LeakyReLu instead of GELU shown in the paper to save memory
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(inner_channel, inner_channel, 1, 1, 0),
-            nn.GroupNorm(32, inner_channel),
-            nn.LeakyReLU(inplace=True),  # use LeakyReLu instead of GELU shown in the paper to save memory
-        )
-
-        self.conv_out = nn.Conv2d(inner_channel, output_channel, 1, 1, 0)
-
-    def forward(self, x):
-        b, twh, c = x.shape
-        t = twh // (30 * 53)
-        x = x.view(-1, 30 * 53, c)
-        x = x.permute(0, 2, 1)
-        x = x.view(b * t, c, 30, 53)
-        x = self.conv1(x)
-        x = self.conv2(x) + x
-        x = self.conv_out(x)
-        return x
-
-
-class Discriminator(nn.Module):
-
-    def __init__(
-        self,
-        stride=8,
-        num_h_per_head=1,
-        adapter_channel_dims=[3072],
-        total_layers=48,
-    ):
-        super().__init__()
-        adapter_channel_dims = adapter_channel_dims * (total_layers // stride)
-        self.stride = stride
-        self.num_h_per_head = num_h_per_head
-        self.head_num = len(adapter_channel_dims)
-        self.heads = nn.ModuleList([
-            nn.ModuleList([DiscriminatorHead(adapter_channel) for _ in range(self.num_h_per_head)])
-            for adapter_channel in adapter_channel_dims
-        ])
-
-    def forward(self, features):
-        outputs = []
-
-        def create_custom_forward(module):
-
-            def custom_forward(*inputs):
-                return module(*inputs)
-
-            return custom_forward
-
-        assert len(features) == len(self.heads)
-        for i in range(0, len(features)):
-            for h in self.heads[i]:
-                # out = torch.utils.checkpoint.checkpoint(
-                #     create_custom_forward(h),
-                #     features[i],
-                #     use_reentrant=False
-                # )
-                out = h(features[i])
-                outputs.append(out)
-        return outputs

fastvideo/distill/solver.py

@@ -1,278 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import torch
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.schedulers.scheduling_utils import SchedulerMixin
-from diffusers.utils import BaseOutput, logging
-
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-@dataclass
-class PCMFMSchedulerOutput(BaseOutput):
-    prev_sample: torch.FloatTensor
-
-
-def extract_into_tensor(a, t, x_shape):
-    b, *_ = t.shape
-    out = a.gather(-1, t)
-    return out.reshape(b, *((1, ) * (len(x_shape) - 1)))
-
-
-class PCMFMScheduler(SchedulerMixin, ConfigMixin):
-    _compatibles = []
-    order = 1
-
-    @register_to_config
-    def __init__(
-        self,
-        num_train_timesteps: int = 1000,
-        shift: float = 1.0,
-        pcm_timesteps: int = 50,
-        linear_quadratic=False,
-        linear_quadratic_threshold=0.025,
-        linear_range=0.5,
-    ):
-        if linear_quadratic:
-            linear_steps = int(num_train_timesteps * linear_range)
-            sigmas = linear_quadratic_schedule(num_train_timesteps, linear_quadratic_threshold, linear_steps)
-            sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
-        else:
-            timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
-            timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
-            sigmas = timesteps / num_train_timesteps
-            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
-        self.euler_timesteps = (np.arange(1, pcm_timesteps + 1) *
-                                (num_train_timesteps // pcm_timesteps)).round().astype(np.int64) - 1
-        self.sigmas = sigmas.numpy()[::-1][self.euler_timesteps]
-        self.sigmas = torch.from_numpy((self.sigmas[::-1].copy()))
-        self.timesteps = self.sigmas * num_train_timesteps
-        self._step_index = None
-        self._begin_index = None
-        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
-        self.sigma_min = self.sigmas[-1].item()
-        self.sigma_max = self.sigmas[0].item()
-
-    @property
-    def step_index(self):
-        """
-        The index counter for current timestep. It will increase 1 after each scheduler step.
-        """
-        return self._step_index
-
-    @property
-    def begin_index(self):
-        """
-        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
-        """
-        return self._begin_index
-
-    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
-    def set_begin_index(self, begin_index: int = 0):
-        """
-        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
-
-        Args:
-            begin_index (`int`):
-                The begin index for the scheduler.
-        """
-        self._begin_index = begin_index
-
-    def scale_noise(
-        self,
-        sample: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
-        noise: Optional[torch.FloatTensor] = None,
-    ) -> torch.FloatTensor:
-        """
-        Forward process in flow-matching
-
-        Args:
-            sample (`torch.FloatTensor`):
-                The input sample.
-            timestep (`int`, *optional*):
-                The current timestep in the diffusion chain.
-
-        Returns:
-            `torch.FloatTensor`:
-                A scaled input sample.
-        """
-        if self.step_index is None:
-            self._init_step_index(timestep)
-
-        sigma = self.sigmas[self.step_index]
-        sample = sigma * noise + (1.0 - sigma) * sample
-
-        return sample
-
-    def _sigma_to_t(self, sigma):
-        return sigma * self.config.num_train_timesteps
-
-    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
-        """
-        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
-
-        Args:
-            num_inference_steps (`int`):
-                The number of diffusion steps used when generating samples with a pre-trained model.
-            device (`str` or `torch.device`, *optional*):
-                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        """
-        self.num_inference_steps = num_inference_steps
-        inference_indices = np.linspace(0, self.config.pcm_timesteps, num=num_inference_steps, endpoint=False)
-        inference_indices = np.floor(inference_indices).astype(np.int64)
-        inference_indices = torch.from_numpy(inference_indices).long()
-
-        self.sigmas_ = self.sigmas[inference_indices]
-        timesteps = self.sigmas_ * self.config.num_train_timesteps
-        self.timesteps = timesteps.to(device=device)
-        self.sigmas_ = torch.cat([self.sigmas_, torch.zeros(1, device=self.sigmas_.device)])
-        self._step_index = None
-        self._begin_index = None
-
-    def index_for_timestep(self, timestep, schedule_timesteps=None):
-        if schedule_timesteps is None:
-            schedule_timesteps = self.timesteps
-
-        indices = (schedule_timesteps == timestep).nonzero()
-
-        # The sigma index that is taken for the **very** first `step`
-        # is always the second index (or the last index if there is only 1)
-        # This way we can ensure we don't accidentally skip a sigma in
-        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
-        pos = 1 if len(indices) > 1 else 0
-
-        return indices[pos].item()
-
-    def _init_step_index(self, timestep):
-        if self.begin_index is None:
-            if isinstance(timestep, torch.Tensor):
-                timestep = timestep.to(self.timesteps.device)
-            self._step_index = self.index_for_timestep(timestep)
-        else:
-            self._step_index = self._begin_index
-
-    def step(
-        self,
-        model_output: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
-        sample: torch.FloatTensor,
-        generator: Optional[torch.Generator] = None,
-        return_dict: bool = True,
-    ) -> Union[PCMFMSchedulerOutput, Tuple]:
-        """
-        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
-        process from the learned model outputs (most often the predicted noise).
-
-        Args:
-            model_output (`torch.FloatTensor`):
-                The direct output from learned diffusion model.
-            timestep (`float`):
-                The current discrete timestep in the diffusion chain.
-            sample (`torch.FloatTensor`):
-                A current instance of a sample created by the diffusion process.
-            s_churn (`float`):
-            s_tmin  (`float`):
-            s_tmax  (`float`):
-            s_noise (`float`, defaults to 1.0):
-                Scaling factor for noise added to the sample.
-            generator (`torch.Generator`, *optional*):
-                A random number generator.
-            return_dict (`bool`):
-                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
-                tuple.
-
-        Returns:
-            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
-                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
-                returned, otherwise a tuple is returned where the first element is the sample tensor.
-        """
-
-        if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor)
-                or isinstance(timestep, torch.LongTensor)):
-            raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-                              " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-                              " one of the `scheduler.timesteps` as a timestep."), )
-
-        if self.step_index is None:
-            self._init_step_index(timestep)
-
-        sample = sample.to(torch.float32)
-
-        sigma = self.sigmas_[self.step_index]
-
-        denoised = sample - model_output * sigma
-        derivative = (sample - denoised) / sigma
-
-        dt = self.sigmas_[self.step_index + 1] - sigma
-        prev_sample = sample + derivative * dt
-        prev_sample = prev_sample.to(model_output.dtype)
-        self._step_index += 1
-
-        if not return_dict:
-            return (prev_sample, )
-
-        return PCMFMSchedulerOutput(prev_sample=prev_sample)
-
-    def __len__(self):
-        return self.config.num_train_timesteps
-
-
-class EulerSolver:
-
-    def __init__(self, sigmas, timesteps=1000, euler_timesteps=50):
-        self.step_ratio = timesteps // euler_timesteps
-        self.euler_timesteps = (np.arange(1, euler_timesteps + 1) * self.step_ratio).round().astype(np.int64) - 1
-        self.euler_timesteps_prev = np.asarray([0] + self.euler_timesteps[:-1].tolist())
-        self.sigmas = sigmas[self.euler_timesteps]
-        self.sigmas_prev = np.asarray([sigmas[0]] +
-                                      sigmas[self.euler_timesteps[:-1]].tolist())  # either use sigma0 or 0
-
-        self.euler_timesteps = torch.from_numpy(self.euler_timesteps).long()
-        self.euler_timesteps_prev = torch.from_numpy(self.euler_timesteps_prev).long()
-        self.sigmas = torch.from_numpy(self.sigmas)
-        self.sigmas_prev = torch.from_numpy(self.sigmas_prev)
-
-    def to(self, device):
-        self.euler_timesteps = self.euler_timesteps.to(device)
-        self.euler_timesteps_prev = self.euler_timesteps_prev.to(device)
-
-        self.sigmas = self.sigmas.to(device)
-        self.sigmas_prev = self.sigmas_prev.to(device)
-        return self
-
-    def euler_step(self, sample, model_pred, timestep_index):
-        sigma = extract_into_tensor(self.sigmas, timestep_index, model_pred.shape)
-        sigma_prev = extract_into_tensor(self.sigmas_prev, timestep_index, model_pred.shape)
-        x_prev = sample + (sigma_prev - sigma) * model_pred
-        return x_prev
-
-    def euler_style_multiphase_pred(
-        self,
-        sample,
-        model_pred,
-        timestep_index,
-        multiphase,
-        is_target=False,
-    ):
-        inference_indices = np.linspace(0, len(self.euler_timesteps), num=multiphase, endpoint=False)
-        inference_indices = np.floor(inference_indices).astype(np.int64)
-        inference_indices = (torch.from_numpy(inference_indices).long().to(self.euler_timesteps.device))
-        expanded_timestep_index = timestep_index.unsqueeze(1).expand(-1, inference_indices.size(0))
-        valid_indices_mask = expanded_timestep_index >= inference_indices
-        last_valid_index = valid_indices_mask.flip(dims=[1]).long().argmax(dim=1)
-        last_valid_index = inference_indices.size(0) - 1 - last_valid_index
-        timestep_index_end = inference_indices[last_valid_index]
-
-        if is_target:
-            sigma = extract_into_tensor(self.sigmas_prev, timestep_index, sample.shape)
-        else:
-            sigma = extract_into_tensor(self.sigmas, timestep_index, sample.shape)
-        sigma_prev = extract_into_tensor(self.sigmas_prev, timestep_index_end, sample.shape)
-        x_prev = sample + (sigma_prev - sigma) * model_pred
-
-        return x_prev, timestep_index_end

fastvideo/distill_adv.py

@@ -1,844 +0,0 @@
-# !/bin/python3
-# isort: skip_file
-import argparse
-import math
-import os
-import time
-from collections import deque
-from copy import deepcopy
-
-import torch
-import torch.distributed as dist
-import wandb
-from accelerate.utils import set_seed
-from diffusers import FlowMatchEulerDiscreteScheduler
-from diffusers.optimization import get_scheduler
-from diffusers.utils import check_min_version
-from peft import LoraConfig
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from tqdm.auto import tqdm
-
-from fastvideo.dataset.latent_datasets import (LatentDataset, latent_collate_function)
-from fastvideo.distill.discriminator import Discriminator
-from fastvideo.distill.solver import EulerSolver, extract_into_tensor
-from fastvideo.utils.latents_utils import normalize_dit_input
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
-from fastvideo.utils.checkpoint import (resume_lora_optimizer, resume_training_generator_discriminator, save_checkpoint,
-                                        save_lora_checkpoint)
-from fastvideo.utils.communications import (broadcast, sp_parallel_dataloader_wrapper)
-from fastvideo.utils.dataset_utils import LengthGroupedSampler
-from fastvideo.utils.fsdp_util import (apply_fsdp_checkpointing, get_discriminator_fsdp_kwargs, get_dit_fsdp_kwargs)
-from fastvideo.utils.load import load_transformer
-from fastvideo.utils.logging_ import main_print
-from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group, get_sequence_parallel_state,
-                                             initialize_sequence_parallel_state)
-from fastvideo.utils.validation import log_validation
-
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.31.0")
-
-
-def gan_d_loss(
-    discriminator,
-    teacher_transformer,
-    sample_fake,
-    sample_real,
-    timestep,
-    encoder_hidden_states,
-    encoder_attention_mask,
-    weight,
-    discriminator_head_stride,
-):
-    loss = 0.0
-    # collate sample_fake and sample_real
-    with torch.no_grad():
-        fake_features = teacher_transformer(
-            sample_fake,
-            encoder_hidden_states,
-            timestep,
-            encoder_attention_mask,
-            output_features=True,
-            output_features_stride=discriminator_head_stride,
-            return_dict=False,
-        )[1]
-        real_features = teacher_transformer(
-            sample_real,
-            encoder_hidden_states,
-            timestep,
-            encoder_attention_mask,
-            output_features=True,
-            output_features_stride=discriminator_head_stride,
-            return_dict=False,
-        )[1]
-
-    fake_outputs = discriminator(fake_features)
-    real_outputs = discriminator(real_features)
-    for fake_output, real_output in zip(fake_outputs, real_outputs):
-        loss += (torch.mean(weight * torch.relu(fake_output.float() + 1)) + torch.mean(
-            weight * torch.relu(1 - real_output.float()))) / (discriminator.head_num * discriminator.num_h_per_head)
-    return loss
-
-
-def gan_g_loss(
-    discriminator,
-    teacher_transformer,
-    sample_fake,
-    timestep,
-    encoder_hidden_states,
-    encoder_attention_mask,
-    weight,
-    discriminator_head_stride,
-):
-    loss = 0.0
-    features = teacher_transformer(
-        sample_fake,
-        encoder_hidden_states,
-        timestep,
-        encoder_attention_mask,
-        output_features=True,
-        output_features_stride=discriminator_head_stride,
-        return_dict=False,
-    )[1]
-    fake_outputs = discriminator(features, )
-    for fake_output in fake_outputs:
-        loss += torch.mean(
-            weight * torch.relu(1 - fake_output.float())) / (discriminator.head_num * discriminator.num_h_per_head)
-    return loss
-
-
-def distill_one_step_adv(
-    transformer,
-    model_type,
-    teacher_transformer,
-    optimizer,
-    discriminator,
-    discriminator_optimizer,
-    lr_scheduler,
-    loader,
-    noise_scheduler,
-    solver,
-    noise_random_generator,
-    sp_size,
-    max_grad_norm,
-    uncond_prompt_embed,
-    uncond_prompt_mask,
-    num_euler_timesteps,
-    multiphase,
-    not_apply_cfg_solver,
-    distill_cfg,
-    adv_weight,
-    discriminator_head_stride,
-):
-    optimizer.zero_grad()
-    discriminator_optimizer.zero_grad()
-
-    (
-        latents,
-        encoder_hidden_states,
-        latents_attention_mask,
-        encoder_attention_mask,
-    ) = next(loader)
-    model_input = normalize_dit_input(model_type, latents)
-    noise = torch.randn_like(model_input)
-    bsz = model_input.shape[0]
-    index = torch.randint(0, num_euler_timesteps, (bsz, ), device=model_input.device).long()
-    if sp_size > 1:
-        broadcast(index)
-    # Add noise according to flow matching.
-    # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
-    sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
-    sigmas_prev = extract_into_tensor(solver.sigmas_prev, index, model_input.shape)
-
-    timesteps = (sigmas * noise_scheduler.config.num_train_timesteps).view(-1)
-    # if squeeze to [], unsqueeze to [1]
-
-    timesteps_prev = (sigmas_prev * noise_scheduler.config.num_train_timesteps).view(-1)
-    noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
-
-    # Predict the noise residual
-    with torch.autocast("cuda", dtype=torch.bfloat16):
-        model_pred = transformer(
-            noisy_model_input,
-            encoder_hidden_states,
-            timesteps,
-            encoder_attention_mask,  # B, L
-            return_dict=False,
-        )[0]
-
-    # if accelerator.is_main_process:
-    model_pred, end_index = solver.euler_style_multiphase_pred(noisy_model_input, model_pred, index, multiphase)
-
-    # # simplified flow matching aka 0-rectified flow matching loss
-    # # target = model_input - noise
-    # target = model_input
-    adv_index = torch.empty_like(end_index)
-    for i in range(end_index.size(0)):
-        adv_index[i] = torch.randint(
-            end_index[i].item(),
-            end_index[i].item() + num_euler_timesteps // multiphase,
-            (1, ),
-            dtype=end_index.dtype,
-            device=end_index.device,
-        )
-
-    sigmas_end = extract_into_tensor(solver.sigmas_prev, end_index, model_input.shape)
-    sigmas_adv = extract_into_tensor(solver.sigmas_prev, adv_index, model_input.shape)
-    timesteps_adv = (sigmas_adv * noise_scheduler.config.num_train_timesteps).view(-1)
-
-    with torch.no_grad():
-        w = distill_cfg
-        with torch.autocast("cuda", dtype=torch.bfloat16):
-            cond_teacher_output = teacher_transformer(
-                noisy_model_input,
-                encoder_hidden_states,
-                timesteps,
-                encoder_attention_mask,  # B, L
-                return_dict=False,
-            )[0].float()
-        if not_apply_cfg_solver:
-            uncond_teacher_output = cond_teacher_output
-        else:
-            # Get teacher model prediction on noisy_latents and unconditional embedding
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                uncond_teacher_output = teacher_transformer(
-                    noisy_model_input,
-                    uncond_prompt_embed.unsqueeze(0).expand(bsz, -1, -1),
-                    timesteps,
-                    uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
-                    return_dict=False,
-                )[0].float()
-        teacher_output = cond_teacher_output + w * (cond_teacher_output - uncond_teacher_output)
-        x_prev = solver.euler_step(noisy_model_input, teacher_output, index)
-
-    # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
-    with torch.no_grad():
-        with torch.autocast("cuda", dtype=torch.bfloat16):
-            target_pred = transformer(
-                x_prev.float(),
-                encoder_hidden_states,
-                timesteps_prev,
-                encoder_attention_mask,  # B, L
-                return_dict=False,
-            )[0]
-
-        target, end_index = solver.euler_style_multiphase_pred(x_prev, target_pred, index, multiphase, True)
-
-    real_adv = ((1 - sigmas_adv) * target + (sigmas_adv - sigmas_end) * torch.randn_like(target)) / (1 - sigmas_end)
-    fake_adv = ((1 - sigmas_adv) * model_pred +
-                (sigmas_adv - sigmas_end) * torch.randn_like(model_pred)) / (1 - sigmas_end)
-
-    huber_c = 0.001
-    g_loss = torch.mean(torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) - huber_c)
-    discriminator.requires_grad_(False)
-    with torch.autocast("cuda", dtype=torch.bfloat16):
-        g_gan_loss = adv_weight * gan_g_loss(
-            discriminator,
-            teacher_transformer,
-            fake_adv.float(),
-            timesteps_adv,
-            encoder_hidden_states.float(),
-            encoder_attention_mask,
-            1.0,
-            discriminator_head_stride,
-        )
-    g_loss += g_gan_loss
-    g_loss.backward()
-
-    g_loss = g_loss.detach().clone()
-    dist.all_reduce(g_loss, op=dist.ReduceOp.AVG)
-
-    g_grad_norm = transformer.clip_grad_norm_(max_grad_norm).item()
-    optimizer.step()
-    lr_scheduler.step()
-    optimizer.zero_grad()
-
-    discriminator_optimizer.zero_grad()
-    discriminator.requires_grad_(True)
-    with torch.autocast("cuda", dtype=torch.bfloat16):
-        d_loss = gan_d_loss(
-            discriminator,
-            teacher_transformer,
-            fake_adv.detach(),
-            real_adv.detach(),
-            timesteps_adv,
-            encoder_hidden_states,
-            encoder_attention_mask,
-            1.0,
-            discriminator_head_stride,
-        )
-
-    d_loss.backward()
-    d_grad_norm = discriminator.clip_grad_norm_(max_grad_norm).item()
-    discriminator_optimizer.step()
-    discriminator_optimizer.zero_grad()
-
-    return g_loss, g_grad_norm, d_loss, d_grad_norm
-
-
-def main(args):
-    torch.backends.cuda.matmul.allow_tf32 = True
-
-    local_rank = int(os.environ["LOCAL_RANK"])
-    rank = int(os.environ["RANK"])
-    world_size = int(os.environ["WORLD_SIZE"])
-    dist.init_process_group("nccl")
-    torch.cuda.set_device(local_rank)
-    device = torch.cuda.current_device()
-    initialize_sequence_parallel_state(args.sp_size)
-
-    # If passed along, set the training seed now. On GPU...
-    if args.seed is not None:
-        # TODO: t within the same seq parallel group should be the same. Noise should be different.
-        set_seed(args.seed + rank)
-    # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch.
-    noise_random_generator = None
-
-    # Handle the repository creation
-    if rank == 0 and args.output_dir is not None:
-        os.makedirs(args.output_dir, exist_ok=True)
-
-    # For mixed precision training we cast all non-trainable weights to half-precision
-    # as these weights are only used for inference, keeping weights in full precision is not required.
-
-    # Create model:
-
-    main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
-    # keep the master weight to float32
-    transformer = load_transformer(
-        args.model_type,
-        args.dit_model_name_or_path,
-        args.pretrained_model_name_or_path,
-        torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16,
-    )
-    teacher_transformer = deepcopy(transformer)
-    discriminator = Discriminator(
-        args.discriminator_head_stride,
-        total_layers=48 if args.model_type == "mochi" else 40,
-    )
-
-    if args.use_lora:
-        transformer.requires_grad_(False)
-        transformer_lora_config = LoraConfig(
-            r=args.lora_rank,
-            lora_alpha=args.lora_alpha,
-            init_lora_weights=True,
-            target_modules=["to_k", "to_q", "to_v", "to_out.0"],
-        )
-        transformer.add_adapter(transformer_lora_config)
-
-    main_print(
-        f"  Total transformer parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M"
-    )
-    # discriminator
-    main_print(
-        f"  Total discriminator parameters = {sum(p.numel() for p in discriminator.parameters() if p.requires_grad) / 1e6} M"
-    )
-    main_print(f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}")
-    fsdp_kwargs, no_split_modules = get_dit_fsdp_kwargs(
-        transformer,
-        args.fsdp_sharding_startegy,
-        args.use_lora,
-        args.use_cpu_offload,
-        args.master_weight_type,
-    )
-    discriminator_fsdp_kwargs = get_discriminator_fsdp_kwargs(args.master_weight_type)
-    if args.use_lora:
-        assert args.model_type == "mochi", "LoRA is only supported for Mochi model."
-        transformer.config.lora_rank = args.lora_rank
-        transformer.config.lora_alpha = args.lora_alpha
-        transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
-        transformer._no_split_modules = no_split_modules
-        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)
-
-    transformer = FSDP(
-        transformer,
-        **fsdp_kwargs,
-    )
-    teacher_transformer = FSDP(
-        teacher_transformer,
-        **fsdp_kwargs,
-    )
-    discriminator = FSDP(
-        discriminator,
-        **discriminator_fsdp_kwargs,
-    )
-    main_print("--> model loaded")
-
-    if args.gradient_checkpointing:
-        apply_fsdp_checkpointing(transformer, no_split_modules, args.selective_checkpointing)
-        apply_fsdp_checkpointing(teacher_transformer, no_split_modules, args.selective_checkpointing)
-    # Set model as trainable.
-    transformer.train()
-    teacher_transformer.requires_grad_(False)
-    noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=args.shift)
-    if args.scheduler_type == "pcm_linear_quadratic":
-        sigmas = linear_quadratic_schedule(noise_scheduler.config.num_train_timesteps, args.linear_quadratic_threshold)
-        sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
-    else:
-        sigmas = noise_scheduler.sigmas
-    solver = EulerSolver(
-        sigmas.numpy()[::-1],
-        noise_scheduler.config.num_train_timesteps,
-        euler_timesteps=args.num_euler_timesteps,
-    )
-    solver.to(device)
-    params_to_optimize = transformer.parameters()
-    params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize))
-
-    optimizer = torch.optim.AdamW(
-        params_to_optimize,
-        lr=args.learning_rate,
-        betas=(0.9, 0.999),
-        weight_decay=args.weight_decay,
-        eps=1e-8,
-    )
-
-    discriminator_optimizer = torch.optim.AdamW(
-        discriminator.parameters(),
-        lr=args.discriminator_learning_rate,
-        betas=(0, 0.999),
-        weight_decay=args.weight_decay,
-        eps=1e-8,
-    )
-
-    init_steps = 0
-    if args.resume_from_lora_checkpoint:
-        transformer, optimizer, init_steps = resume_lora_optimizer(transformer, args.resume_from_lora_checkpoint,
-                                                                   optimizer)
-    elif args.resume_from_checkpoint:
-        (
-            transformer,
-            optimizer,
-            discriminator,
-            discriminator_optimizer,
-            init_steps,
-        ) = resume_training_generator_discriminator(
-            transformer,
-            optimizer,
-            discriminator,
-            discriminator_optimizer,
-            args.resume_from_checkpoint,
-            rank,
-        )
-
-    main_print(f"optimizer: {optimizer}")
-
-    lr_scheduler = get_scheduler(
-        args.lr_scheduler,
-        optimizer=optimizer,
-        num_warmup_steps=args.lr_warmup_steps * world_size,
-        num_training_steps=args.max_train_steps * world_size,
-        num_cycles=args.lr_num_cycles,
-        power=args.lr_power,
-        last_epoch=init_steps - 1,
-    )
-
-    train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
-    uncond_prompt_embed = train_dataset.uncond_prompt_embed
-    uncond_prompt_mask = train_dataset.uncond_prompt_mask
-    sampler = (LengthGroupedSampler(
-        args.train_batch_size,
-        rank=rank,
-        world_size=world_size,
-        lengths=train_dataset.lengths,
-        group_frame=args.group_frame,
-        group_resolution=args.group_resolution,
-    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
-        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
-
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        collate_fn=latent_collate_function,
-        pin_memory=True,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-        drop_last=True,
-    )
-    assert args.gradient_accumulation_steps == 1
-    num_update_steps_per_epoch = math.ceil(
-        len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
-    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
-
-    if rank == 0:
-        project = args.tracker_project_name or "fastvideo"
-        wandb.init(project=project, config=args)
-
-    # Train!
-    total_batch_size = (world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size)
-    main_print("***** Running training *****")
-    main_print(f"  Num examples = {len(train_dataset)}")
-    main_print(f"  Dataloader size = {len(train_dataloader)}")
-    main_print(f"  Num Epochs = {args.num_train_epochs}")
-    main_print(f"  Resume training from step {init_steps}")
-    main_print(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    main_print(f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
-    main_print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
-    main_print(f"  Total optimization steps = {args.max_train_steps}")
-    main_print(
-        f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B"
-    )
-    # print dtype
-    main_print(f"  Master weight dtype: {transformer.parameters().__next__().dtype}")
-
-    progress_bar = tqdm(
-        range(0, args.max_train_steps),
-        initial=init_steps,
-        desc="Steps",
-        # Only show the progress bar once on each machine.
-        disable=local_rank > 0,
-    )
-    loader = sp_parallel_dataloader_wrapper(
-        train_dataloader,
-        device,
-        args.train_batch_size,
-        args.sp_size,
-        args.train_sp_batch_size,
-    )
-
-    step_times = deque(maxlen=100)
-
-    # log_validation(args, transformer, device,
-    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
-    def get_num_phases(multi_phased_distill_schedule, step):
-        # step-phase,step-phase
-        multi_phases = multi_phased_distill_schedule.split(",")
-        phase = multi_phases[-1].split("-")[-1]
-        for step_phases in multi_phases:
-            phase_step, phase = step_phases.split("-")
-            if step <= int(phase_step):
-                return int(phase)
-        return phase
-
-    for i in range(init_steps):
-        _ = next(loader)
-    for step in range(init_steps + 1, args.max_train_steps + 1):
-        assert args.multi_phased_distill_schedule is not None
-        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
-        start_time = time.perf_counter()
-        (
-            generator_loss,
-            generator_grad_norm,
-            discriminator_loss,
-            discriminator_grad_norm,
-        ) = distill_one_step_adv(
-            transformer,
-            args.model_type,
-            teacher_transformer,
-            optimizer,
-            discriminator,
-            discriminator_optimizer,
-            lr_scheduler,
-            loader,
-            noise_scheduler,
-            solver,
-            noise_random_generator,
-            args.sp_size,
-            args.max_grad_norm,
-            uncond_prompt_embed,
-            uncond_prompt_mask,
-            args.num_euler_timesteps,
-            num_phases,
-            args.not_apply_cfg_solver,
-            args.distill_cfg,
-            args.adv_weight,
-            args.discriminator_head_stride,
-        )
-
-        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({
-            "g_loss": f"{generator_loss:.4f}",
-            "d_loss": f"{discriminator_loss:.4f}",
-            "g_grad_norm": generator_grad_norm,
-            "d_grad_norm": discriminator_grad_norm,
-            "step_time": f"{step_time:.2f}s",
-        })
-        progress_bar.update(1)
-        if rank == 0:
-            wandb.log(
-                {
-                    "generator_loss": generator_loss,
-                    "discriminator_loss": discriminator_loss,
-                    "generator_grad_norm": generator_grad_norm,
-                    "discriminator_grad_norm": discriminator_grad_norm,
-                    "learning_rate": lr_scheduler.get_last_lr()[0],
-                    "step_time": step_time,
-                    "avg_step_time": avg_step_time,
-                },
-                step=step,
-            )
-        if step % args.checkpointing_steps == 0:
-            main_print(f"--> saving checkpoint at step {step}")
-            if args.use_lora:
-                # Save LoRA weights
-                save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
-            else:
-                # Your existing checkpoint saving code
-                # TODO
-                # save_checkpoint_generator_discriminator(
-                #     transformer,
-                #     optimizer,
-                #     discriminator,
-                #     discriminator_optimizer,
-                #     rank,
-                #     args.output_dir,
-                #     step,
-                # )
-                save_checkpoint(transformer, rank, args.output_dir, step)
-            main_print(f"--> checkpoint saved at step {step}")
-            dist.barrier()
-        if args.log_validation and step % args.validation_steps == 0:
-            log_validation(
-                args,
-                transformer,
-                device,
-                torch.bfloat16,
-                step,
-                scheduler_type=args.scheduler_type,
-                shift=args.shift,
-                num_euler_timesteps=args.num_euler_timesteps,
-                linear_quadratic_threshold=args.linear_quadratic_threshold,
-                linear_range=args.linear_range,
-                ema=False,
-            )
-
-    if args.use_lora:
-        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
-    else:
-        save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
-
-    if get_sequence_parallel_state():
-        destroy_sequence_parallel_group()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument("--model_type", type=str, default="mochi", help="The type of model to train.")
-    # dataset & dataloader
-    parser.add_argument("--data_json_path", type=str, required=True)
-    parser.add_argument("--num_height", type=int, default=480)
-    parser.add_argument("--num_width", type=int, default=848)
-    parser.add_argument("--num_frames", type=int, default=163)
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=10,
-        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
-    )
-    parser.add_argument(
-        "--train_batch_size",
-        type=int,
-        default=16,
-        help="Batch size (per device) for the training dataloader.",
-    )
-    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
-    parser.add_argument("--group_frame", action="store_true")  # TODO
-    parser.add_argument("--group_resolution", action="store_true")  # TODO
-
-    # text encoder & vae & diffusion model
-    parser.add_argument("--pretrained_model_name_or_path", type=str)
-    parser.add_argument("--dit_model_name_or_path", type=str)
-    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-
-    # diffusion setting
-    parser.add_argument("--ema_decay", type=float, default=0.999)
-    parser.add_argument("--ema_start_step", type=int, default=0)
-    parser.add_argument("--cfg", type=float, default=0.1)
-    # validation & logs
-    parser.add_argument("--validation_sampling_steps", type=str, default="64")
-    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
-    parser.add_argument("--validation_steps", type=float, default=64)
-    parser.add_argument("--log_validation", action="store_true")
-    parser.add_argument("--tracker_project_name", type=str, default=None)
-    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default=None,
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-    parser.add_argument(
-        "--checkpoints_total_limit",
-        type=int,
-        default=None,
-        help=("Max number of checkpoints to store."),
-    )
-    parser.add_argument(
-        "--checkpointing_steps",
-        type=int,
-        default=500,
-        help=("Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
-              " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
-              " training using `--resume_from_checkpoint`."),
-    )
-    parser.add_argument("--validation_prompt_dir", type=str)
-    parser.add_argument("--shift", type=float, default=1.0)
-    parser.add_argument(
-        "--resume_from_checkpoint",
-        type=str,
-        default=None,
-        help=("Whether training should be resumed from a previous checkpoint. Use a path saved by"
-              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
-    )
-    parser.add_argument(
-        "--resume_from_lora_checkpoint",
-        type=str,
-        default=None,
-        help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
-              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
-    )
-
-    # optimizer & scheduler & Training
-    parser.add_argument("--num_train_epochs", type=int, default=100)
-    parser.add_argument(
-        "--max_train_steps",
-        type=int,
-        default=None,
-        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
-    )
-    parser.add_argument(
-        "--learning_rate",
-        type=float,
-        default=1e-4,
-        help="Initial learning rate (after the potential warmup period) to use.",
-    )
-    parser.add_argument(
-        "--discriminator_learning_rate",
-        type=float,
-        default=1e-5,
-        help="Initial learning rate (after the potential warmup period) to use.",
-    )
-    parser.add_argument(
-        "--scale_lr",
-        action="store_true",
-        default=False,
-        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
-    )
-    parser.add_argument(
-        "--lr_warmup_steps",
-        type=int,
-        default=10,
-        help="Number of steps for the warmup in the lr scheduler.",
-    )
-    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
-    parser.add_argument(
-        "--gradient_checkpointing",
-        action="store_true",
-        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
-    )
-    parser.add_argument("--selective_checkpointing", type=float, default=1.0)
-    parser.add_argument(
-        "--allow_tf32",
-        action="store_true",
-        help=("Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
-              " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"),
-    )
-    parser.add_argument(
-        "--mixed_precision",
-        type=str,
-        default=None,
-        choices=["no", "fp16", "bf16"],
-        help=(
-            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
-            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
-            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."),
-    )
-    parser.add_argument(
-        "--use_cpu_offload",
-        action="store_true",
-        help="Whether to use CPU offload for param & gradient & optimizer states.",
-    )
-
-    parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
-    parser.add_argument(
-        "--train_sp_batch_size",
-        type=int,
-        default=1,
-        help="Batch size for sequence parallel training",
-    )
-
-    parser.add_argument(
-        "--use_lora",
-        action="store_true",
-        default=False,
-        help="Whether to use LoRA for finetuning.",
-    )
-    parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.")
-    parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ")
-    parser.add_argument("--fsdp_sharding_startegy", default="full")
-    parser.add_argument("--multi_phased_distill_schedule", type=str, default=None)
-    parser.add_argument(
-        "--gradient_accumulation_steps",
-        type=int,
-        default=1,
-        help="Number of updates steps to accumulate before performing a backward/update pass.",
-    )
-
-    # lr_scheduler
-    parser.add_argument(
-        "--lr_scheduler",
-        type=str,
-        default="constant",
-        help=('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
-              ' "constant", "constant_with_warmup"]'),
-    )
-    parser.add_argument("--num_euler_timesteps", type=int, default=100)
-    parser.add_argument(
-        "--lr_num_cycles",
-        type=int,
-        default=1,
-        help="Number of cycles in the learning rate scheduler.",
-    )
-    parser.add_argument(
-        "--lr_power",
-        type=float,
-        default=1.0,
-        help="Power factor of the polynomial scheduler.",
-    )
-    parser.add_argument(
-        "--not_apply_cfg_solver",
-        action="store_true",
-        help="Whether to apply the cfg_solver.",
-    )
-    parser.add_argument("--distill_cfg", type=float, default=3.0, help="Distillation coefficient.")
-    # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
-    parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.")
-    parser.add_argument(
-        "--adv_weight",
-        type=float,
-        default=0.1,
-        help="The weight of the adversarial loss.",
-    )
-    parser.add_argument(
-        "--discriminator_head_stride",
-        type=int,
-        default=2,
-        help="The stride of the discriminator head.",
-    )
-    parser.add_argument(
-        "--linear_range",
-        type=float,
-        default=0.5,
-        help="Range for linear quadratic scheduler.",
-    )
-    parser.add_argument("--weight_decay", type=float, default=0.001, help="Weight decay to apply.")
-    parser.add_argument(
-        "--linear_quadratic_threshold",
-        type=float,
-        default=0.025,
-        help="The threshold of the linear quadratic scheduler.",
-    )
-    parser.add_argument(
-        "--master_weight_type",
-        type=str,
-        default="fp32",
-        help="Weight type to use - fp32 or bf16.",
-    )
-    args = parser.parse_args()
-    main(args)

fastvideo/models/flash_attn_no_pad.py

@@ -1,28 +0,0 @@
-from einops import rearrange
-from flash_attn import flash_attn_varlen_qkvpacked_func
-from flash_attn.bert_padding import pad_input, unpad_input
-
-
-def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None):
-    # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27
-    batch_size = qkv.shape[0]
-    seqlen = qkv.shape[1]
-    nheads = qkv.shape[-2]
-    x = rearrange(qkv, "b s three h d -> b s (three h d)")
-    x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask)
-
-    x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads)
-    output_unpad = flash_attn_varlen_qkvpacked_func(
-        x_unpad,
-        cu_seqlens,
-        max_s,
-        dropout_p,
-        softmax_scale=softmax_scale,
-        causal=causal,
-    )
-    output = rearrange(
-        pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen),
-        "b s (h d) -> b s h d",
-        h=nheads,
-    )
-    return output

fastvideo/models/hunyuan/constants.py

@@ -1,89 +0,0 @@
-import os
-
-import torch
-
-__all__ = [
-    "C_SCALE",
-    "PROMPT_TEMPLATE",
-    "MODEL_BASE",
-    "PRECISIONS",
-    "NORMALIZATION_TYPE",
-    "ACTIVATION_TYPE",
-    "VAE_PATH",
-    "TEXT_ENCODER_PATH",
-    "TOKENIZER_PATH",
-    "TEXT_PROJECTION",
-    "DATA_TYPE",
-    "NEGATIVE_PROMPT",
-]
-
-PRECISION_TO_TYPE = {
-    "fp32": torch.float32,
-    "fp16": torch.float16,
-    "bf16": torch.bfloat16,
-}
-
-# =================== Constant Values =====================
-# Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid
-# overflow error when tensorboard logging values.
-C_SCALE = 1_000_000_000_000_000
-
-# When using decoder-only models, we must provide a prompt template to instruct the text encoder
-# on how to generate the text.
-# --------------------------------------------------------------------
-PROMPT_TEMPLATE_ENCODE = (
-    "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, "
-    "quantity, text, spatial relationships of the objects and background:<|eot_id|>"
-    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>")
-PROMPT_TEMPLATE_ENCODE_VIDEO = (
-    "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
-    "1. The main content and theme of the video."
-    "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
-    "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
-    "4. background environment, light, style and atmosphere."
-    "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
-    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>")
-
-NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion"
-
-PROMPT_TEMPLATE = {
-    "dit-llm-encode": {
-        "template": PROMPT_TEMPLATE_ENCODE,
-        "crop_start": 36,
-    },
-    "dit-llm-encode-video": {
-        "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
-        "crop_start": 95,
-    },
-}
-
-# ======================= Model ======================
-PRECISIONS = {"fp32", "fp16", "bf16"}
-NORMALIZATION_TYPE = {"layer", "rms"}
-ACTIVATION_TYPE = {"relu", "silu", "gelu", "gelu_tanh"}
-
-# =================== Model Path =====================
-MODEL_BASE = os.getenv("MODEL_BASE", "./data/hunyuan")
-
-# =================== Data =======================
-DATA_TYPE = {"image", "video", "image_video"}
-
-# 3D VAE
-VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"}
-
-# Text Encoder
-TEXT_ENCODER_PATH = {
-    "clipL": f"{MODEL_BASE}/text_encoder_2",
-    "llm": f"{MODEL_BASE}/text_encoder",
-}
-
-# Tokenizer
-TOKENIZER_PATH = {
-    "clipL": f"{MODEL_BASE}/text_encoder_2",
-    "llm": f"{MODEL_BASE}/text_encoder",
-}
-
-TEXT_PROJECTION = {
-    "linear",  # Default, an nn.Linear() layer
-    "single_refiner",  # Single TokenRefiner. Refer to LI-DiT
-}

fastvideo/models/hunyuan/diffusion/__init__.py

@@ -1,3 +0,0 @@
-# ruff: noqa: F401
-from .pipelines import HunyuanVideoPipeline
-from .schedulers import FlowMatchDiscreteScheduler

fastvideo/models/hunyuan/diffusion/pipelines/__init__.py

@@ -1,2 +0,0 @@
-# ruff: noqa: F401
-from .pipeline_hunyuan_video import HunyuanVideoPipeline

fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py

@@ -1,931 +0,0 @@
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-#
-# Modified from diffusers==0.29.2
-#
-# ==============================================================================
-import inspect
-from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-import torch.distributed as dist
-import torch.nn.functional as F
-from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
-from diffusers.configuration_utils import FrozenDict
-from diffusers.image_processor import VaeImageProcessor
-from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin
-from diffusers.models import AutoencoderKL
-from diffusers.models.lora import adjust_lora_scale_text_encoder
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, deprecate, logging, replace_example_docstring,
-                             scale_lora_layers)
-from diffusers.utils.torch_utils import randn_tensor
-from einops import rearrange
-
-from fastvideo.utils.communications import all_gather
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-from ...constants import PRECISION_TO_TYPE
-from ...modules import HYVideoDiffusionTransformer
-from ...text_encoder import TextEncoder
-from ...vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-EXAMPLE_DOC_STRING = """"""
-
-
-def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
-    """
-    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
-    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
-    """
-    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
-    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
-    # rescale the results from guidance (fixes overexposure)
-    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
-    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
-    noise_cfg = (guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg)
-    return noise_cfg
-
-
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    """
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-@dataclass
-class HunyuanVideoPipelineOutput(BaseOutput):
-    videos: Union[torch.Tensor, np.ndarray]
-
-
-class HunyuanVideoPipeline(DiffusionPipeline):
-    r"""
-    Pipeline for text-to-video generation using HunyuanVideo.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Args:
-        vae ([`AutoencoderKL`]):
-            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
-        text_encoder ([`TextEncoder`]):
-            Frozen text-encoder.
-        text_encoder_2 ([`TextEncoder`]):
-            Frozen text-encoder_2.
-        transformer ([`HYVideoDiffusionTransformer`]):
-            A `HYVideoDiffusionTransformer` to denoise the encoded video latents.
-        scheduler ([`SchedulerMixin`]):
-            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
-    """
-
-    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
-    _optional_components = ["text_encoder_2"]
-    _exclude_from_cpu_offload = ["transformer"]
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        text_encoder: TextEncoder,
-        transformer: HYVideoDiffusionTransformer,
-        scheduler: KarrasDiffusionSchedulers,
-        text_encoder_2: Optional[TextEncoder] = None,
-        progress_bar_config: Dict[str, Any] = None,
-        args=None,
-    ):
-        super().__init__()
-
-        # ==========================================================================================
-        if progress_bar_config is None:
-            progress_bar_config = {}
-        if not hasattr(self, "_progress_bar_config"):
-            self._progress_bar_config = {}
-        self._progress_bar_config.update(progress_bar_config)
-
-        self.args = args
-        # ==========================================================================================
-
-        if (hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1):
-            deprecation_message = (
-                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
-                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
-                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
-                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
-                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
-                " file")
-            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
-            new_config = dict(scheduler.config)
-            new_config["steps_offset"] = 1
-            scheduler._internal_dict = FrozenDict(new_config)
-
-        if (hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True):
-            deprecation_message = (
-                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
-                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
-                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
-                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
-                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file")
-            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
-            new_config = dict(scheduler.config)
-            new_config["clip_sample"] = False
-            scheduler._internal_dict = FrozenDict(new_config)
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            transformer=transformer,
-            scheduler=scheduler,
-            text_encoder_2=text_encoder_2,
-        )
-        self.vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1)
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
-
-    def encode_prompt(
-        self,
-        prompt,
-        device,
-        num_videos_per_prompt,
-        do_classifier_free_guidance,
-        negative_prompt=None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        negative_attention_mask: Optional[torch.Tensor] = None,
-        lora_scale: Optional[float] = None,
-        clip_skip: Optional[int] = None,
-        text_encoder: Optional[TextEncoder] = None,
-        data_type: Optional[str] = "image",
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            device: (`torch.device`):
-                torch device
-            num_videos_per_prompt (`int`):
-                number of videos that should be generated per prompt
-            do_classifier_free_guidance (`bool`):
-                whether to use classifier free guidance or not
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the video generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            attention_mask (`torch.Tensor`, *optional*):
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            negative_attention_mask (`torch.Tensor`, *optional*):
-            lora_scale (`float`, *optional*):
-                A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
-            clip_skip (`int`, *optional*):
-                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
-                the output of the pre-final layer will be used for computing the prompt embeddings.
-            text_encoder (TextEncoder, *optional*):
-            data_type (`str`, *optional*):
-        """
-        if text_encoder is None:
-            text_encoder = self.text_encoder
-
-        # set lora scale so that monkey patched LoRA
-        # function of text encoder can correctly access it
-        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
-            self._lora_scale = lora_scale
-
-            # dynamically adjust the LoRA scale
-            if not USE_PEFT_BACKEND:
-                adjust_lora_scale_text_encoder(text_encoder.model, lora_scale)
-            else:
-                scale_lora_layers(text_encoder.model, lora_scale)
-
-        if prompt_embeds is None:
-            # textual inversion: process multi-vector tokens if necessary
-            if isinstance(self, TextualInversionLoaderMixin):
-                prompt = self.maybe_convert_prompt(prompt, text_encoder.tokenizer)
-
-            text_inputs = text_encoder.text2tokens(prompt, data_type=data_type)
-            if clip_skip is None:
-                prompt_outputs = text_encoder.encode(text_inputs, data_type=data_type, device=device)
-                prompt_embeds = prompt_outputs.hidden_state
-            else:
-                prompt_outputs = text_encoder.encode(
-                    text_inputs,
-                    output_hidden_states=True,
-                    data_type=data_type,
-                    device=device,
-                )
-                # Access the `hidden_states` first, that contains a tuple of
-                # all the hidden states from the encoder layers. Then index into
-                # the tuple to access the hidden states from the desired layer.
-                prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)]
-                # We also need to apply the final LayerNorm here to not mess with the
-                # representations. The `last_hidden_states` that we typically use for
-                # obtaining the final prompt representations passes through the LayerNorm
-                # layer.
-                prompt_embeds = text_encoder.model.text_model.final_layer_norm(prompt_embeds)
-
-            attention_mask = prompt_outputs.attention_mask
-            if attention_mask is not None:
-                attention_mask = attention_mask.to(device)
-                bs_embed, seq_len = attention_mask.shape
-                attention_mask = attention_mask.repeat(1, num_videos_per_prompt)
-                attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len)
-
-        if text_encoder is not None:
-            prompt_embeds_dtype = text_encoder.dtype
-        elif self.transformer is not None:
-            prompt_embeds_dtype = self.transformer.dtype
-        else:
-            prompt_embeds_dtype = prompt_embeds.dtype
-
-        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
-        if prompt_embeds.ndim == 2:
-            bs_embed, _ = prompt_embeds.shape
-            # duplicate text embeddings for each generation per prompt, using mps friendly method
-            prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt)
-            prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, -1)
-        else:
-            bs_embed, seq_len, _ = prompt_embeds.shape
-            # duplicate text embeddings for each generation per prompt, using mps friendly method
-            prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-            prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1)
-
-        return (
-            prompt_embeds,
-            negative_prompt_embeds,
-            attention_mask,
-            negative_attention_mask,
-        )
-
-    def decode_latents(self, latents, enable_tiling=True):
-        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
-        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
-
-        latents = 1 / self.vae.config.scaling_factor * latents
-        if enable_tiling:
-            self.vae.enable_tiling()
-        image = self.vae.decode(latents, return_dict=False)[0]
-        image = (image / 2 + 0.5).clamp(0, 1)
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-        if image.ndim == 4:
-            image = image.cpu().permute(0, 2, 3, 1).float()
-        else:
-            image = image.cpu().float()
-        return image
-
-    def prepare_extra_func_kwargs(self, func, kwargs):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-        extra_step_kwargs = {}
-
-        for k, v in kwargs.items():
-            accepts = k in set(inspect.signature(func).parameters.keys())
-            if accepts:
-                extra_step_kwargs[k] = v
-        return extra_step_kwargs
-
-    def check_inputs(
-        self,
-        prompt,
-        height,
-        width,
-        video_length,
-        callback_steps,
-        negative_prompt=None,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-        callback_on_step_end_tensor_inputs=None,
-        vae_ver="88-4c-sd",
-    ):
-        if height % 8 != 0 or width % 8 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
-        if video_length is not None:
-            if "884" in vae_ver:
-                if video_length != 1 and (video_length - 1) % 4 != 0:
-                    raise ValueError(f"`video_length` has to be 1 or a multiple of 4 but is {video_length}.")
-            elif "888" in vae_ver:
-                if video_length != 1 and (video_length - 1) % 8 != 0:
-                    raise ValueError(f"`video_length` has to be 1 or a multiple of 8 but is {video_length}.")
-
-        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
-            raise ValueError(f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
-                             f" {type(callback_steps)}.")
-        if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs
-                                                                      for k in callback_on_step_end_tensor_inputs):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two.")
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.")
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-        if negative_prompt is not None and negative_prompt_embeds is not None:
-            raise ValueError(f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
-                             f" {negative_prompt_embeds}. Please make sure to only forward one of the two.")
-
-        if prompt_embeds is not None and negative_prompt_embeds is not None:
-            if prompt_embeds.shape != negative_prompt_embeds.shape:
-                raise ValueError(
-                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
-                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
-                    f" {negative_prompt_embeds.shape}.")
-
-    def prepare_latents(
-        self,
-        batch_size,
-        num_channels_latents,
-        height,
-        width,
-        video_length,
-        dtype,
-        device,
-        generator,
-        latents=None,
-    ):
-        shape = (
-            batch_size,
-            num_channels_latents,
-            video_length,
-            int(height) // self.vae_scale_factor,
-            int(width) // self.vae_scale_factor,
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        else:
-            latents = latents.to(device)
-
-        # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler
-        if hasattr(self.scheduler, "init_noise_sigma"):
-            # scale the initial noise by the standard deviation required by the scheduler
-            latents = latents * self.scheduler.init_noise_sigma
-        return latents
-
-    # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
-    def get_guidance_scale_embedding(
-        self,
-        w: torch.Tensor,
-        embedding_dim: int = 512,
-        dtype: torch.dtype = torch.float32,
-    ) -> torch.Tensor:
-        """
-        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
-
-        Args:
-            w (`torch.Tensor`):
-                Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
-            embedding_dim (`int`, *optional*, defaults to 512):
-                Dimension of the embeddings to generate.
-            dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
-                Data type of the generated embeddings.
-
-        Returns:
-            `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
-        """
-        assert len(w.shape) == 1
-        w = w * 1000.0
-
-        half_dim = embedding_dim // 2
-        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
-        emb = w.to(dtype)[:, None] * emb[None, :]
-        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-        if embedding_dim % 2 == 1:  # zero pad
-            emb = torch.nn.functional.pad(emb, (0, 1))
-        assert emb.shape == (w.shape[0], embedding_dim)
-        return emb
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def guidance_rescale(self):
-        return self._guidance_rescale
-
-    @property
-    def clip_skip(self):
-        return self._clip_skip
-
-    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-    # corresponds to doing no classifier free guidance.
-    @property
-    def do_classifier_free_guidance(self):
-        # return self._guidance_scale > 1 and self.transformer.config.time_cond_proj_dim is None
-        return self._guidance_scale > 1
-
-    @property
-    def cross_attention_kwargs(self):
-        return self._cross_attention_kwargs
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]],
-        height: int,
-        width: int,
-        video_length: int,
-        data_type: str = "video",
-        num_inference_steps: int = 50,
-        timesteps: List[int] = None,
-        sigmas: List[float] = None,
-        guidance_scale: float = 7.5,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_videos_per_prompt: Optional[int] = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        negative_attention_mask: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        guidance_rescale: float = 0.0,
-        clip_skip: Optional[int] = None,
-        callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback,
-                                             MultiPipelineCallbacks, ]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        vae_ver: str = "88-4c-sd",
-        enable_tiling: bool = False,
-        enable_vae_sp: bool = False,
-        n_tokens: Optional[int] = None,
-        embedded_guidance_scale: Optional[float] = None,
-        mask_strategy: Optional[Dict[str, list]] = None,
-        **kwargs,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`):
-                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
-            height (`int`):
-                The height in pixels of the generated image.
-            width (`int`):
-                The width in pixels of the generated image.
-            video_length (`int`):
-                The number of frames in the generated video.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            timesteps (`List[int]`, *optional*):
-                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
-                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
-                passed will be used. Must be in descending order.
-            sigmas (`List[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            guidance_scale (`float`, *optional*, defaults to 7.5):
-                A higher guidance scale value encourages the model to generate images closely linked to the text
-                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
-                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            eta (`float`, *optional*, defaults to 0.0):
-                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
-                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
-                provided, text embeddings are generated from the `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
-                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
-
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a
-                plain tuple.
-            cross_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
-                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            guidance_rescale (`float`, *optional*, defaults to 0.0):
-                Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
-                Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
-                using zero terminal SNR.
-            clip_skip (`int`, *optional*):
-                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
-                the output of the pre-final layer will be used for computing the prompt embeddings.
-            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
-                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
-                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
-                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
-                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-
-        Examples:
-
-        Returns:
-            [`~HunyuanVideoPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned,
-                otherwise a `tuple` is returned where the first element is a list with the generated images and the
-                second element is a list of `bool`s indicating whether the corresponding generated image contains
-                "not-safe-for-work" (nsfw) content.
-        """
-        callback = kwargs.pop("callback", None)
-        callback_steps = kwargs.pop("callback_steps", None)
-
-        if callback is not None:
-            deprecate(
-                "callback",
-                "1.0.0",
-                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
-            )
-        if callback_steps is not None:
-            deprecate(
-                "callback_steps",
-                "1.0.0",
-                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
-            )
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        # 0. Default height and width to unet
-        # height = height or self.transformer.config.sample_size * self.vae_scale_factor
-        # width = width or self.transformer.config.sample_size * self.vae_scale_factor
-        # to deal with lora scaling and other possible forward hooks
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            height,
-            width,
-            video_length,
-            callback_steps,
-            negative_prompt,
-            prompt_embeds,
-            negative_prompt_embeds,
-            callback_on_step_end_tensor_inputs,
-            vae_ver=vae_ver,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._guidance_rescale = guidance_rescale
-        self._clip_skip = clip_skip
-        self._cross_attention_kwargs = cross_attention_kwargs
-        self._interrupt = False
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = (torch.device(f"cuda:{dist.get_rank()}") if dist.is_initialized() else self._execution_device)
-
-        # 3. Encode input prompt
-        lora_scale = (self.cross_attention_kwargs.get("scale", None)
-                      if self.cross_attention_kwargs is not None else None)
-
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            prompt_mask,
-            negative_prompt_mask,
-        ) = self.encode_prompt(
-            prompt,
-            device,
-            num_videos_per_prompt,
-            self.do_classifier_free_guidance,
-            negative_prompt,
-            prompt_embeds=prompt_embeds,
-            attention_mask=attention_mask,
-            negative_prompt_embeds=negative_prompt_embeds,
-            negative_attention_mask=negative_attention_mask,
-            lora_scale=lora_scale,
-            clip_skip=self.clip_skip,
-            data_type=data_type,
-        )
-        if self.text_encoder_2 is not None:
-            (
-                prompt_embeds_2,
-                negative_prompt_embeds_2,
-                prompt_mask_2,
-                negative_prompt_mask_2,
-            ) = self.encode_prompt(
-                prompt,
-                device,
-                num_videos_per_prompt,
-                self.do_classifier_free_guidance,
-                negative_prompt,
-                prompt_embeds=None,
-                attention_mask=None,
-                negative_prompt_embeds=None,
-                negative_attention_mask=None,
-                lora_scale=lora_scale,
-                clip_skip=self.clip_skip,
-                text_encoder=self.text_encoder_2,
-                data_type=data_type,
-            )
-        else:
-            prompt_embeds_2 = None
-            negative_prompt_embeds_2 = None
-            prompt_mask_2 = None
-            negative_prompt_mask_2 = None
-
-        # For classifier free guidance, we need to do two forward passes.
-        # Here we concatenate the unconditional and text embeddings into a single batch
-        # to avoid doing two forward passes
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
-            if prompt_mask is not None:
-                prompt_mask = torch.cat([negative_prompt_mask, prompt_mask])
-            if prompt_embeds_2 is not None:
-                prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2])
-            if prompt_mask_2 is not None:
-                prompt_mask_2 = torch.cat([negative_prompt_mask_2, prompt_mask_2])
-
-        # 4. Prepare timesteps
-        extra_set_timesteps_kwargs = self.prepare_extra_func_kwargs(self.scheduler.set_timesteps,
-                                                                    {"n_tokens": n_tokens})
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            device,
-            timesteps,
-            sigmas,
-            **extra_set_timesteps_kwargs,
-        )
-        if "884" in vae_ver:
-            video_length = (video_length - 1) // 4 + 1
-        elif "888" in vae_ver:
-            video_length = (video_length - 1) // 8 + 1
-        else:
-            video_length = video_length
-
-        # 5. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            video_length,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
-        if get_sequence_parallel_state():
-            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
-            latents = latents[:, :, rank, :, :, :]
-
-        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_func_kwargs(
-            self.scheduler.step,
-            {
-                "generator": generator,
-                "eta": eta
-            },
-        )
-
-        target_dtype = PRECISION_TO_TYPE[self.args.precision]
-        autocast_enabled = (target_dtype != torch.float32) and not self.args.disable_autocast
-        vae_dtype = PRECISION_TO_TYPE[self.args.vae_precision]
-        vae_autocast_enabled = (vae_dtype != torch.float32) and not self.args.disable_autocast
-
-        # 7. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        self._num_timesteps = len(timesteps)
-
-        def dict_to_3d_list(mask_strategy, t_max=50, l_max=60, h_max=24):
-            result = [[[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max)]
-            if mask_strategy is None:
-                return result
-            for key, value in mask_strategy.items():
-                t, l, h = map(int, key.split('_'))
-                result[t][l][h] = value
-            return result
-
-        mask_strategy = dict_to_3d_list(mask_strategy)
-        # if is_progress_bar:
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = (torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents)
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
-                t_expand = t.repeat(latent_model_input.shape[0])
-                guidance_expand = (torch.tensor(
-                    [embedded_guidance_scale] * latent_model_input.shape[0],
-                    dtype=torch.float32,
-                    device=device,
-                ).to(target_dtype) * 1000.0 if embedded_guidance_scale is not None else None)
-                # predict the noise residual
-                with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
-                    # concat prompt_embeds_2 and prompt_embeds. Mismatch fill with zeros
-                    if prompt_embeds_2.shape[-1] != prompt_embeds.shape[-1]:
-                        prompt_embeds_2 = F.pad(
-                            prompt_embeds_2,
-                            (0, prompt_embeds.shape[2] - prompt_embeds_2.shape[1]),
-                            value=0,
-                        ).unsqueeze(1)
-                    encoder_hidden_states = torch.cat([prompt_embeds_2, prompt_embeds], dim=1)
-                    noise_pred = self.transformer(  # For an input image (129, 192, 336) (1, 256, 256)
-                        latent_model_input,
-                        encoder_hidden_states,
-                        t_expand,
-                        prompt_mask,
-                        mask_strategy=mask_strategy[i],
-                        guidance=guidance_expand,
-                        return_dict=False,
-                    )[0]
-
-                # perform guidance
-                if self.do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
-                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
-                    noise_pred = rescale_noise_cfg(
-                        noise_pred,
-                        noise_pred_text,
-                        guidance_rescale=self.guidance_rescale,
-                    )
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    if progress_bar is not None:
-                        progress_bar.update()
-                    if callback is not None and i % callback_steps == 0:
-                        step_idx = i // getattr(self.scheduler, "order", 1)
-                        callback(step_idx, t, latents)
-
-        if get_sequence_parallel_state():
-            latents = all_gather(latents, dim=2)
-
-        if not output_type == "latent":
-            expand_temporal_dim = False
-            if len(latents.shape) == 4:
-                if isinstance(self.vae, AutoencoderKLCausal3D):
-                    latents = latents.unsqueeze(2)
-                    expand_temporal_dim = True
-            elif len(latents.shape) == 5:
-                pass
-            else:
-                raise ValueError(
-                    f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}.")
-
-            if (hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor):
-                latents = (latents / self.vae.config.scaling_factor + self.vae.config.shift_factor)
-            else:
-                latents = latents / self.vae.config.scaling_factor
-
-            with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
-                if enable_tiling:
-                    self.vae.enable_tiling()
-                if enable_vae_sp:
-                    self.vae.enable_parallel()
-                image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
-
-            if expand_temporal_dim or image.shape[2] == 1:
-                image = image.squeeze(2)
-
-        else:
-            image = latents
-
-        image = (image / 2 + 0.5).clamp(0, 1)
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
-        image = image.cpu().float()
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return image
-
-        return HunyuanVideoPipelineOutput(videos=image)

fastvideo/models/hunyuan/diffusion/schedulers/__init__.py

@@ -1,2 +0,0 @@
-# ruff: noqa: F401
-from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler

fastvideo/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py

@@ -1,239 +0,0 @@
-# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-#
-# Modified from diffusers==0.29.2
-#
-# ==============================================================================
-
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.schedulers.scheduling_utils import SchedulerMixin
-from diffusers.utils import BaseOutput, logging
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-@dataclass
-class FlowMatchDiscreteSchedulerOutput(BaseOutput):
-    """
-    Output class for the scheduler's `step` function output.
-
-    Args:
-        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
-            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
-            denoising loop.
-    """
-
-    prev_sample: torch.FloatTensor
-
-
-class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
-    """
-    Euler scheduler.
-
-    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
-    methods the library implements for all schedulers such as loading and saving.
-
-    Args:
-        num_train_timesteps (`int`, defaults to 1000):
-            The number of diffusion steps to train the model.
-        timestep_spacing (`str`, defaults to `"linspace"`):
-            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
-            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
-        shift (`float`, defaults to 1.0):
-            The shift value for the timestep schedule.
-        reverse (`bool`, defaults to `True`):
-            Whether to reverse the timestep schedule.
-    """
-
-    _compatibles = []
-    order = 1
-
-    @register_to_config
-    def __init__(
-        self,
-        num_train_timesteps: int = 1000,
-        shift: float = 1.0,
-        reverse: bool = True,
-        solver: str = "euler",
-        n_tokens: Optional[int] = None,
-    ):
-        sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
-
-        if not reverse:
-            sigmas = sigmas.flip(0)
-
-        self.sigmas = sigmas
-        # the value fed to model
-        self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32)
-
-        self._step_index = None
-        self._begin_index = None
-
-        self.supported_solver = ["euler"]
-        if solver not in self.supported_solver:
-            raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}")
-
-    @property
-    def step_index(self):
-        """
-        The index counter for current timestep. It will increase 1 after each scheduler step.
-        """
-        return self._step_index
-
-    @property
-    def begin_index(self):
-        """
-        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
-        """
-        return self._begin_index
-
-    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
-    def set_begin_index(self, begin_index: int = 0):
-        """
-        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
-
-        Args:
-            begin_index (`int`):
-                The begin index for the scheduler.
-        """
-        self._begin_index = begin_index
-
-    def _sigma_to_t(self, sigma):
-        return sigma * self.config.num_train_timesteps
-
-    def set_timesteps(
-        self,
-        num_inference_steps: int,
-        device: Union[str, torch.device] = None,
-        n_tokens: int = None,
-    ):
-        """
-        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
-
-        Args:
-            num_inference_steps (`int`):
-                The number of diffusion steps used when generating samples with a pre-trained model.
-            device (`str` or `torch.device`, *optional*):
-                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-            n_tokens (`int`, *optional*):
-                Number of tokens in the input sequence.
-        """
-        self.num_inference_steps = num_inference_steps
-
-        sigmas = torch.linspace(1, 0, num_inference_steps + 1)
-        sigmas = self.sd3_time_shift(sigmas)
-
-        if not self.config.reverse:
-            sigmas = 1 - sigmas
-
-        self.sigmas = sigmas
-        self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device)
-
-        # Reset step index
-        self._step_index = None
-
-    def index_for_timestep(self, timestep, schedule_timesteps=None):
-        if schedule_timesteps is None:
-            schedule_timesteps = self.timesteps
-
-        indices = (schedule_timesteps == timestep).nonzero()
-
-        # The sigma index that is taken for the **very** first `step`
-        # is always the second index (or the last index if there is only 1)
-        # This way we can ensure we don't accidentally skip a sigma in
-        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
-        pos = 1 if len(indices) > 1 else 0
-
-        return indices[pos].item()
-
-    def _init_step_index(self, timestep):
-        if self.begin_index is None:
-            if isinstance(timestep, torch.Tensor):
-                timestep = timestep.to(self.timesteps.device)
-            self._step_index = self.index_for_timestep(timestep)
-        else:
-            self._step_index = self._begin_index
-
-    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
-        return sample
-
-    def sd3_time_shift(self, t: torch.Tensor):
-        return (self.config.shift * t) / (1 + (self.config.shift - 1) * t)
-
-    def step(
-        self,
-        model_output: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
-        sample: torch.FloatTensor,
-        return_dict: bool = True,
-    ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
-        """
-        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
-        process from the learned model outputs (most often the predicted noise).
-
-        Args:
-            model_output (`torch.FloatTensor`):
-                The direct output from learned diffusion model.
-            timestep (`float`):
-                The current discrete timestep in the diffusion chain.
-            sample (`torch.FloatTensor`):
-                A current instance of a sample created by the diffusion process.
-            generator (`torch.Generator`, *optional*):
-                A random number generator.
-            n_tokens (`int`, *optional*):
-                Number of tokens in the input sequence.
-            return_dict (`bool`):
-                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
-                tuple.
-
-        Returns:
-            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
-                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
-                returned, otherwise a tuple is returned where the first element is the sample tensor.
-        """
-
-        if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor)
-                or isinstance(timestep, torch.LongTensor)):
-            raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-                              " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-                              " one of the `scheduler.timesteps` as a timestep."), )
-
-        if self.step_index is None:
-            self._init_step_index(timestep)
-
-        # Upcast to avoid precision issues when computing prev_sample
-        sample = sample.to(torch.float32)
-
-        dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
-
-        if self.config.solver == "euler":
-            prev_sample = sample + model_output.to(torch.float32) * dt
-        else:
-            raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}")
-
-        # upon completion increase step index by one
-        self._step_index += 1
-
-        if not return_dict:
-            return (prev_sample, )
-
-        return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
-
-    def __len__(self):
-        return self.config.num_train_timesteps

fastvideo/models/hunyuan/idle_config.py

@@ -1,380 +0,0 @@
-# ruff: noqa: F405, F403
-import argparse
-import re
-
-from .constants import *
-from .modules.models import HUNYUAN_VIDEO_CONFIG
-
-
-def parse_args(namespace=None):
-    parser = argparse.ArgumentParser(description="HunyuanVideo inference script")
-
-    parser = add_network_args(parser)
-    parser = add_extra_models_args(parser)
-    parser = add_denoise_schedule_args(parser)
-    parser = add_inference_args(parser)
-    parser = add_parallel_args(parser)
-
-    args = parser.parse_args(namespace=namespace)
-    args = sanity_check_args(args)
-
-    return args
-
-
-def add_network_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="HunyuanVideo network args")
-
-    # Main model
-    group.add_argument(
-        "--model",
-        type=str,
-        choices=list(HUNYUAN_VIDEO_CONFIG.keys()),
-        default="HYVideo-T/2-cfgdistill",
-    )
-    group.add_argument(
-        "--latent-channels",
-        type=str,
-        default=16,
-        help="Number of latent channels of DiT. If None, it will be determined by `vae`. If provided, "
-        "it still needs to match the latent channels of the VAE model.",
-    )
-    group.add_argument(
-        "--precision",
-        type=str,
-        default="bf16",
-        choices=PRECISIONS,
-        help="Precision mode. Options: fp32, fp16, bf16. Applied to the backbone model and optimizer.",
-    )
-
-    # RoPE
-    group.add_argument("--rope-theta", type=int, default=256, help="Theta used in RoPE.")
-    return parser
-
-
-def add_extra_models_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Extra models args, including vae, text encoders and tokenizers)")
-
-    # - VAE
-    group.add_argument(
-        "--vae",
-        type=str,
-        default="884-16c-hy",
-        choices=list(VAE_PATH),
-        help="Name of the VAE model.",
-    )
-    group.add_argument(
-        "--vae-precision",
-        type=str,
-        default="fp16",
-        choices=PRECISIONS,
-        help="Precision mode for the VAE model.",
-    )
-    group.add_argument(
-        "--vae-tiling",
-        action="store_true",
-        help="Enable tiling for the VAE model to save GPU memory.",
-    )
-    group.set_defaults(vae_tiling=True)
-
-    group.add_argument(
-        "--text-encoder",
-        type=str,
-        default="llm",
-        choices=list(TEXT_ENCODER_PATH),
-        help="Name of the text encoder model.",
-    )
-    group.add_argument(
-        "--text-encoder-precision",
-        type=str,
-        default="fp16",
-        choices=PRECISIONS,
-        help="Precision mode for the text encoder model.",
-    )
-    group.add_argument(
-        "--text-states-dim",
-        type=int,
-        default=4096,
-        help="Dimension of the text encoder hidden states.",
-    )
-    group.add_argument("--text-len", type=int, default=256, help="Maximum length of the text input.")
-    group.add_argument(
-        "--tokenizer",
-        type=str,
-        default="llm",
-        choices=list(TOKENIZER_PATH),
-        help="Name of the tokenizer model.",
-    )
-    group.add_argument(
-        "--prompt-template",
-        type=str,
-        default="dit-llm-encode",
-        choices=PROMPT_TEMPLATE,
-        help="Image prompt template for the decoder-only text encoder model.",
-    )
-    group.add_argument(
-        "--prompt-template-video",
-        type=str,
-        default="dit-llm-encode-video",
-        choices=PROMPT_TEMPLATE,
-        help="Video prompt template for the decoder-only text encoder model.",
-    )
-    group.add_argument(
-        "--hidden-state-skip-layer",
-        type=int,
-        default=2,
-        help="Skip layer for hidden states.",
-    )
-    group.add_argument(
-        "--apply-final-norm",
-        action="store_true",
-        help="Apply final normalization to the used text encoder hidden states.",
-    )
-
-    # - CLIP
-    group.add_argument(
-        "--text-encoder-2",
-        type=str,
-        default="clipL",
-        choices=list(TEXT_ENCODER_PATH),
-        help="Name of the second text encoder model.",
-    )
-    group.add_argument(
-        "--text-encoder-precision-2",
-        type=str,
-        default="fp16",
-        choices=PRECISIONS,
-        help="Precision mode for the second text encoder model.",
-    )
-    group.add_argument(
-        "--text-states-dim-2",
-        type=int,
-        default=768,
-        help="Dimension of the second text encoder hidden states.",
-    )
-    group.add_argument(
-        "--tokenizer-2",
-        type=str,
-        default="clipL",
-        choices=list(TOKENIZER_PATH),
-        help="Name of the second tokenizer model.",
-    )
-    group.add_argument(
-        "--text-len-2",
-        type=int,
-        default=77,
-        help="Maximum length of the second text input.",
-    )
-
-    return parser
-
-
-def add_denoise_schedule_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Denoise schedule args")
-
-    group.add_argument(
-        "--denoise-type",
-        type=str,
-        default="flow",
-        help="Denoise type for noised inputs.",
-    )
-
-    # Flow Matching
-    group.add_argument(
-        "--flow-shift",
-        type=float,
-        default=7.0,
-        help="Shift factor for flow matching schedulers.",
-    )
-    group.add_argument(
-        "--flow-reverse",
-        action="store_true",
-        help="If reverse, learning/sampling from t=1 -> t=0.",
-    )
-    group.add_argument(
-        "--flow-solver",
-        type=str,
-        default="euler",
-        help="Solver for flow matching.",
-    )
-    group.add_argument(
-        "--use-linear-quadratic-schedule",
-        action="store_true",
-        help="Use linear quadratic schedule for flow matching."
-        "Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)",
-    )
-    group.add_argument(
-        "--linear-schedule-end",
-        type=int,
-        default=25,
-        help="End step for linear quadratic schedule for flow matching.",
-    )
-
-    return parser
-
-
-def add_inference_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Inference args")
-
-    # ======================== Model loads ========================
-    group.add_argument(
-        "--model-base",
-        type=str,
-        default="ckpts",
-        help="Root path of all the models, including t2v models and extra models.",
-    )
-    group.add_argument(
-        "--dit-weight",
-        type=str,
-        default="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
-        help="Path to the HunyuanVideo model. If None, search the model in the args.model_root."
-        "1. If it is a file, load the model directly."
-        "2. If it is a directory, search the model in the directory. Support two types of models: "
-        "1) named `pytorch_model_*.pt`"
-        "2) named `*_model_states.pt`, where * can be `mp_rank_00`.",
-    )
-    group.add_argument(
-        "--model-resolution",
-        type=str,
-        default="540p",
-        choices=["540p", "720p"],
-        help="The resolution of the model.",
-    )
-    group.add_argument(
-        "--load-key",
-        type=str,
-        default="module",
-        help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.",
-    )
-    group.add_argument(
-        "--use-cpu-offload",
-        action="store_true",
-        help="Use CPU offload for the model load.",
-    )
-
-    # ======================== Inference general setting ========================
-    group.add_argument(
-        "--batch-size",
-        type=int,
-        default=1,
-        help="Batch size for inference and evaluation.",
-    )
-    group.add_argument(
-        "--infer-steps",
-        type=int,
-        default=50,
-        help="Number of denoising steps for inference.",
-    )
-    group.add_argument(
-        "--disable-autocast",
-        action="store_true",
-        help="Disable autocast for denoising loop and vae decoding in pipeline sampling.",
-    )
-    group.add_argument(
-        "--save-path",
-        type=str,
-        default="./results",
-        help="Path to save the generated samples.",
-    )
-    group.add_argument(
-        "--save-path-suffix",
-        type=str,
-        default="",
-        help="Suffix for the directory of saved samples.",
-    )
-    group.add_argument(
-        "--name-suffix",
-        type=str,
-        default="",
-        help="Suffix for the names of saved samples.",
-    )
-    group.add_argument(
-        "--num-videos",
-        type=int,
-        default=1,
-        help="Number of videos to generate for each prompt.",
-    )
-    # ---sample size---
-    group.add_argument(
-        "--video-size",
-        type=int,
-        nargs="+",
-        default=(720, 1280),
-        help="Video size for training. If a single value is provided, it will be used for both height "
-        "and width. If two values are provided, they will be used for height and width "
-        "respectively.",
-    )
-    group.add_argument(
-        "--video-length",
-        type=int,
-        default=129,
-        help="How many frames to sample from a video. if using 3d vae, the number should be 4n+1",
-    )
-    # --- prompt ---
-    group.add_argument(
-        "--prompt",
-        type=str,
-        default=None,
-        help="Prompt for sampling during evaluation.",
-    )
-    group.add_argument(
-        "--seed-type",
-        type=str,
-        default="auto",
-        choices=["file", "random", "fixed", "auto"],
-        help="Seed type for evaluation. If file, use the seed from the CSV file. If random, generate a "
-        "random seed. If fixed, use the fixed seed given by `--seed`. If auto, `csv` will use the "
-        "seed column if available, otherwise use the fixed `seed` value. `prompt` will use the "
-        "fixed `seed` value.",
-    )
-    group.add_argument("--seed", type=int, default=None, help="Seed for evaluation.")
-
-    # Classifier-Free Guidance
-    group.add_argument("--neg-prompt", type=str, default=None, help="Negative prompt for sampling.")
-    group.add_argument("--cfg-scale", type=float, default=1.0, help="Classifier free guidance scale.")
-    group.add_argument(
-        "--embedded-cfg-scale",
-        type=float,
-        default=6.0,
-        help="Embedded classifier free guidance scale.",
-    )
-
-    group.add_argument(
-        "--reproduce",
-        action="store_true",
-        help="Enable reproducibility by setting random seeds and deterministic algorithms.",
-    )
-
-    return parser
-
-
-def add_parallel_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Parallel args")
-
-    # ======================== Model loads ========================
-    group.add_argument(
-        "--ulysses-degree",
-        type=int,
-        default=1,
-        help="Ulysses degree.",
-    )
-    group.add_argument(
-        "--ring-degree",
-        type=int,
-        default=1,
-        help="Ring degree.",
-    )
-
-    return parser
-
-
-def sanity_check_args(args):
-    # VAE channels
-    vae_pattern = r"\d{2,3}-\d{1,2}c-\w+"
-    if not re.match(vae_pattern, args.vae):
-        raise ValueError(f"Invalid VAE model: {args.vae}. Must be in the format of '{vae_pattern}'.")
-    vae_channels = int(args.vae.split("-")[1][:-1])
-    if args.latent_channels is None:
-        args.latent_channels = vae_channels
-    if vae_channels != args.latent_channels:
-        raise ValueError(f"Latent channels ({args.latent_channels}) must match the VAE channels ({vae_channels}).")
-    return args

fastvideo/models/hunyuan/inference.py

@@ -1,482 +0,0 @@
-import os
-import random
-import time
-from pathlib import Path
-
-import torch
-from loguru import logger
-from safetensors.torch import load_file as safetensors_load_file
-
-from fastvideo.models.hunyuan.constants import NEGATIVE_PROMPT, PRECISION_TO_TYPE, PROMPT_TEMPLATE
-from fastvideo.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline
-from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler
-from fastvideo.models.hunyuan.modules import load_model
-from fastvideo.models.hunyuan.text_encoder import TextEncoder
-from fastvideo.models.hunyuan.utils.data_utils import align_to
-from fastvideo.models.hunyuan.vae import load_vae
-from fastvideo.utils.parallel_states import nccl_info
-
-
-class Inference:
-
-    def __init__(
-        self,
-        args,
-        vae,
-        vae_kwargs,
-        text_encoder,
-        model,
-        text_encoder_2=None,
-        pipeline=None,
-        use_cpu_offload=False,
-        device=None,
-        logger=None,
-        parallel_args=None,
-    ):
-        self.vae = vae
-        self.vae_kwargs = vae_kwargs
-
-        self.text_encoder = text_encoder
-        self.text_encoder_2 = text_encoder_2
-
-        self.model = model
-        self.pipeline = pipeline
-        self.use_cpu_offload = use_cpu_offload
-
-        self.args = args
-        self.device = (device if device is not None else "cuda" if torch.cuda.is_available() else "cpu")
-        self.logger = logger
-        self.parallel_args = parallel_args
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_path, args, device=None, **kwargs):
-        """
-        Initialize the Inference pipeline.
-
-        Args:
-            pretrained_model_path (str or pathlib.Path): The model path, including t2v, text encoder and vae checkpoints.
-            args (argparse.Namespace): The arguments for the pipeline.
-            device (int): The device for inference. Default is 0.
-        """
-        # ========================================================================
-        logger.info(f"Got text-to-video model root path: {pretrained_model_path}")
-
-        # ==================== Initialize Distributed Environment ================
-        if nccl_info.sp_size > 1:
-            device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}")
-        if device is None:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-
-        parallel_args = None  # {"ulysses_degree": args.ulysses_degree, "ring_degree": args.ring_degree}
-
-        # ======================== Get the args path =============================
-
-        # Disable gradient
-        torch.set_grad_enabled(False)
-
-        # =========================== Build main model ===========================
-        logger.info("Building model...")
-        factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]}
-        in_channels = args.latent_channels
-        out_channels = args.latent_channels
-
-        model = load_model(
-            args,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            factor_kwargs=factor_kwargs,
-        )
-        model = model.to(device)
-        model = Inference.load_state_dict(args, model, pretrained_model_path)
-        if args.enable_torch_compile:
-            model = torch.compile(model)
-        model.eval()
-
-        # ============================= Build extra models ========================
-        # VAE
-        vae, _, s_ratio, t_ratio = load_vae(
-            args.vae,
-            args.vae_precision,
-            logger=logger,
-            device=device if not args.use_cpu_offload else "cpu",
-        )
-        vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio}
-
-        # Text encoder
-        if args.prompt_template_video is not None:
-            crop_start = PROMPT_TEMPLATE[args.prompt_template_video].get("crop_start", 0)
-        elif args.prompt_template is not None:
-            crop_start = PROMPT_TEMPLATE[args.prompt_template].get("crop_start", 0)
-        else:
-            crop_start = 0
-        max_length = args.text_len + crop_start
-
-        # prompt_template
-        prompt_template = (PROMPT_TEMPLATE[args.prompt_template] if args.prompt_template is not None else None)
-
-        # prompt_template_video
-        prompt_template_video = (PROMPT_TEMPLATE[args.prompt_template_video]
-                                 if args.prompt_template_video is not None else None)
-
-        text_encoder = TextEncoder(
-            text_encoder_type=args.text_encoder,
-            max_length=max_length,
-            text_encoder_precision=args.text_encoder_precision,
-            tokenizer_type=args.tokenizer,
-            prompt_template=prompt_template,
-            prompt_template_video=prompt_template_video,
-            hidden_state_skip_layer=args.hidden_state_skip_layer,
-            apply_final_norm=args.apply_final_norm,
-            reproduce=args.reproduce,
-            logger=logger,
-            device=device if not args.use_cpu_offload else "cpu",
-        )
-        text_encoder_2 = None
-        if args.text_encoder_2 is not None:
-            text_encoder_2 = TextEncoder(
-                text_encoder_type=args.text_encoder_2,
-                max_length=args.text_len_2,
-                text_encoder_precision=args.text_encoder_precision_2,
-                tokenizer_type=args.tokenizer_2,
-                reproduce=args.reproduce,
-                logger=logger,
-                device=device if not args.use_cpu_offload else "cpu",
-            )
-
-        return cls(
-            args=args,
-            vae=vae,
-            vae_kwargs=vae_kwargs,
-            text_encoder=text_encoder,
-            text_encoder_2=text_encoder_2,
-            model=model,
-            use_cpu_offload=args.use_cpu_offload,
-            device=device,
-            logger=logger,
-            parallel_args=parallel_args,
-        )
-
-    @staticmethod
-    def load_state_dict(args, model, pretrained_model_path):
-        load_key = args.load_key
-        dit_weight = Path(args.dit_weight)
-
-        if dit_weight is None:
-            model_dir = pretrained_model_path / f"t2v_{args.model_resolution}"
-            files = list(model_dir.glob("*.pt"))
-            if len(files) == 0:
-                raise ValueError(f"No model weights found in {model_dir}")
-            if str(files[0]).startswith("pytorch_model_"):
-                model_path = dit_weight / f"pytorch_model_{load_key}.pt"
-                bare_model = True
-            elif any(str(f).endswith("_model_states.pt") for f in files):
-                files = [f for f in files if str(f).endswith("_model_states.pt")]
-                model_path = files[0]
-                if len(files) > 1:
-                    logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}")
-                bare_model = False
-            else:
-                raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: "
-                                 f"{list(map(str, files))}. When given a directory as --dit-weight, only "
-                                 f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and "
-                                 f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
-                                 f"specific weight file, please provide the full path to the file.")
-        else:
-            if dit_weight.is_dir():
-                files = list(dit_weight.glob("*.pt"))
-                if len(files) == 0:
-                    raise ValueError(f"No model weights found in {dit_weight}")
-                if str(files[0]).startswith("pytorch_model_"):
-                    model_path = dit_weight / f"pytorch_model_{load_key}.pt"
-                    bare_model = True
-                elif any(str(f).endswith("_model_states.pt") for f in files):
-                    files = [f for f in files if str(f).endswith("_model_states.pt")]
-                    model_path = files[0]
-                    if len(files) > 1:
-                        logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}")
-                    bare_model = False
-                else:
-                    raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: "
-                                     f"{list(map(str, files))}. When given a directory as --dit-weight, only "
-                                     f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and "
-                                     f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
-                                     f"specific weight file, please provide the full path to the file.")
-            elif dit_weight.is_file():
-                model_path = dit_weight
-                bare_model = "unknown"
-            else:
-                raise ValueError(f"Invalid model path: {dit_weight}")
-
-        if not model_path.exists():
-            raise ValueError(f"model_path not exists: {model_path}")
-        logger.info(f"Loading torch model {model_path}...")
-        if model_path.suffix == ".safetensors":
-            # Use safetensors library for .safetensors files
-            state_dict = safetensors_load_file(model_path)
-        elif model_path.suffix == ".pt":
-            # Use torch for .pt files
-            state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
-        else:
-            raise ValueError(f"Unsupported file format: {model_path}")
-
-        if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict):
-            bare_model = False
-        if bare_model is False:
-            if load_key in state_dict:
-                state_dict = state_dict[load_key]
-            else:
-                raise KeyError(f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint "
-                               f"are: {list(state_dict.keys())}.")
-        model.load_state_dict(state_dict, strict=True)
-        return model
-
-    @staticmethod
-    def parse_size(size):
-        if isinstance(size, int):
-            size = [size]
-        if not isinstance(size, (list, tuple)):
-            raise ValueError(f"Size must be an integer or (height, width), got {size}.")
-        if len(size) == 1:
-            size = [size[0], size[0]]
-        if len(size) != 2:
-            raise ValueError(f"Size must be an integer or (height, width), got {size}.")
-        return size
-
-
-class HunyuanVideoSampler(Inference):
-
-    def __init__(
-        self,
-        args,
-        vae,
-        vae_kwargs,
-        text_encoder,
-        model,
-        text_encoder_2=None,
-        pipeline=None,
-        use_cpu_offload=False,
-        device=0,
-        logger=None,
-        parallel_args=None,
-    ):
-        super().__init__(
-            args,
-            vae,
-            vae_kwargs,
-            text_encoder,
-            model,
-            text_encoder_2=text_encoder_2,
-            pipeline=pipeline,
-            use_cpu_offload=use_cpu_offload,
-            device=device,
-            logger=logger,
-            parallel_args=parallel_args,
-        )
-
-        self.pipeline = self.load_diffusion_pipeline(
-            args=args,
-            vae=self.vae,
-            text_encoder=self.text_encoder,
-            text_encoder_2=self.text_encoder_2,
-            model=self.model,
-            device=self.device,
-        )
-
-        self.default_negative_prompt = NEGATIVE_PROMPT
-
-    def load_diffusion_pipeline(
-        self,
-        args,
-        vae,
-        text_encoder,
-        text_encoder_2,
-        model,
-        scheduler=None,
-        device=None,
-        progress_bar_config=None,
-        data_type="video",
-    ):
-        """Load the denoising scheduler for inference."""
-        if scheduler is None:
-            if args.denoise_type == "flow":
-                scheduler = FlowMatchDiscreteScheduler(
-                    shift=args.flow_shift,
-                    reverse=args.flow_reverse,
-                    solver=args.flow_solver,
-                )
-            else:
-                raise ValueError(f"Invalid denoise type {args.denoise_type}")
-
-        pipeline = HunyuanVideoPipeline(
-            vae=vae,
-            text_encoder=text_encoder,
-            text_encoder_2=text_encoder_2,
-            transformer=model,
-            scheduler=scheduler,
-            progress_bar_config=progress_bar_config,
-            args=args,
-        )
-        if self.use_cpu_offload:
-            pipeline.enable_sequential_cpu_offload()
-        else:
-            pipeline = pipeline.to(device)
-
-        return pipeline
-
-    @torch.no_grad()
-    def predict(
-        self,
-        prompt,
-        height=192,
-        width=336,
-        video_length=129,
-        seed=None,
-        negative_prompt=None,
-        infer_steps=50,
-        guidance_scale=6,
-        flow_shift=5.0,
-        embedded_guidance_scale=None,
-        batch_size=1,
-        num_videos_per_prompt=1,
-        mask_strategy=None,
-        **kwargs,
-    ):
-        """
-        Predict the image/video from the given text.
-
-        Args:
-            prompt (str or List[str]): The input text.
-            kwargs:
-                height (int): The height of the output video. Default is 192.
-                width (int): The width of the output video. Default is 336.
-                video_length (int): The frame number of the output video. Default is 129.
-                seed (int or List[str]): The random seed for the generation. Default is a random integer.
-                negative_prompt (str or List[str]): The negative text prompt. Default is an empty string.
-                guidance_scale (float): The guidance scale for the generation. Default is 6.0.
-                num_images_per_prompt (int): The number of images per prompt. Default is 1.
-                infer_steps (int): The number of inference steps. Default is 100.
-        """
-
-        out_dict = dict()
-
-        # ========================================================================
-        # Arguments: seed
-        # ========================================================================
-        if isinstance(seed, torch.Tensor):
-            seed = seed.tolist()
-        if seed is None:
-            seeds = [random.randint(0, 1_000_000) for _ in range(batch_size * num_videos_per_prompt)]
-        elif isinstance(seed, int):
-            seeds = [seed + i for _ in range(batch_size) for i in range(num_videos_per_prompt)]
-        elif isinstance(seed, (list, tuple)):
-            if len(seed) == batch_size:
-                seeds = [int(seed[i]) + j for i in range(batch_size) for j in range(num_videos_per_prompt)]
-            elif len(seed) == batch_size * num_videos_per_prompt:
-                seeds = [int(s) for s in seed]
-            else:
-                raise ValueError(
-                    f"Length of seed must be equal to number of prompt(batch_size) or "
-                    f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}.")
-        else:
-            raise ValueError(f"Seed must be an integer, a list of integers, or None, got {seed}.")
-        # Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
-        generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds]
-        out_dict["seeds"] = seeds
-
-        # ========================================================================
-        # Arguments: target_width, target_height, target_video_length
-        # ========================================================================
-        if width <= 0 or height <= 0 or video_length <= 0:
-            raise ValueError(
-                f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}"
-            )
-        if (video_length - 1) % 4 != 0:
-            raise ValueError(f"`video_length-1` must be a multiple of 4, got {video_length}")
-
-        logger.info(f"Input (height, width, video_length) = ({height}, {width}, {video_length})")
-
-        target_height = align_to(height, 16)
-        target_width = align_to(width, 16)
-        target_video_length = video_length
-
-        out_dict["size"] = (target_height, target_width, target_video_length)
-
-        # ========================================================================
-        # Arguments: prompt, new_prompt, negative_prompt
-        # ========================================================================
-        if not isinstance(prompt, str):
-            raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
-        prompt = [prompt.strip()]
-
-        # negative prompt
-        if negative_prompt is None or negative_prompt == "":
-            negative_prompt = self.default_negative_prompt
-        if not isinstance(negative_prompt, str):
-            raise TypeError(f"`negative_prompt` must be a string, but got {type(negative_prompt)}")
-        negative_prompt = [negative_prompt.strip()]
-
-        # ========================================================================
-        # Scheduler
-        # ========================================================================
-        scheduler = FlowMatchDiscreteScheduler(
-            shift=flow_shift,
-            reverse=self.args.flow_reverse,
-            solver=self.args.flow_solver,
-        )
-        self.pipeline.scheduler = scheduler
-
-        if "884" in self.args.vae:
-            latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8]
-        elif "888" in self.args.vae:
-            latents_size = [(video_length - 1) // 8 + 1, height // 8, width // 8]
-        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
-
-        # ========================================================================
-        # Print infer args
-        # ========================================================================
-        debug_str = f"""
-                        height: {target_height}
-                         width: {target_width}
-                  video_length: {target_video_length}
-                        prompt: {prompt}
-                    neg_prompt: {negative_prompt}
-                          seed: {seed}
-                   infer_steps: {infer_steps}
-         num_videos_per_prompt: {num_videos_per_prompt}
-                guidance_scale: {guidance_scale}
-                      n_tokens: {n_tokens}
-                    flow_shift: {flow_shift}
-       embedded_guidance_scale: {embedded_guidance_scale}"""
-        logger.debug(debug_str)
-
-        # ========================================================================
-        # Pipeline inference
-        # ========================================================================
-        start_time = time.perf_counter()
-        samples = self.pipeline(
-            prompt=prompt,
-            height=target_height,
-            width=target_width,
-            video_length=target_video_length,
-            num_inference_steps=infer_steps,
-            guidance_scale=guidance_scale,
-            negative_prompt=negative_prompt,
-            num_videos_per_prompt=num_videos_per_prompt,
-            generator=generator,
-            output_type="pil",
-            n_tokens=n_tokens,
-            embedded_guidance_scale=embedded_guidance_scale,
-            data_type="video" if target_video_length > 1 else "image",
-            is_progress_bar=True,
-            vae_ver=self.args.vae,
-            enable_tiling=self.args.vae_tiling,
-            enable_vae_sp=self.args.vae_sp,
-            mask_strategy=mask_strategy,
-        )[0]
-        out_dict["samples"] = samples
-        out_dict["prompts"] = prompt
-
-        gen_time = time.perf_counter() - start_time
-        logger.info(f"Success, time: {gen_time}")
-
-        return out_dict

fastvideo/models/hunyuan/modules/__init__.py

@@ -1,25 +0,0 @@
-from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer
-
-
-def load_model(args, in_channels, out_channels, factor_kwargs):
-    """load hunyuan video model
-
-    Args:
-        args (dict): model args
-        in_channels (int): input channels number
-        out_channels (int): output channels number
-        factor_kwargs (dict): factor kwargs
-
-    Returns:
-        model (nn.Module): The hunyuan video model
-    """
-    if args.model in HUNYUAN_VIDEO_CONFIG.keys():
-        model = HYVideoDiffusionTransformer(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            **HUNYUAN_VIDEO_CONFIG[args.model],
-            **factor_kwargs,
-        )
-        return model
-    else:
-        raise NotImplementedError()

fastvideo/models/hunyuan/modules/activation_layers.py

@@ -1,23 +0,0 @@
-import torch.nn as nn
-
-
-def get_activation_layer(act_type):
-    """get activation layer
-
-    Args:
-        act_type (str): the activation type
-
-    Returns:
-        torch.nn.functional: the activation layer
-    """
-    if act_type == "gelu":
-        return lambda: nn.GELU()
-    elif act_type == "gelu_tanh":
-        # Approximate `tanh` requires torch >= 1.13
-        return lambda: nn.GELU(approximate="tanh")
-    elif act_type == "relu":
-        return nn.ReLU
-    elif act_type == "silu":
-        return nn.SiLU
-    else:
-        raise ValueError(f"Unknown activation type: {act_type}")

fastvideo/models/hunyuan/modules/attenion.py

@@ -1,124 +0,0 @@
-import torch
-import torch.nn.functional as F
-from einops import rearrange
-
-try:
-    from st_attn import sliding_tile_attention
-except ImportError:
-    print("Could not load Sliding Tile Attention.")
-    sliding_tile_attention = None
-
-from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
-from fastvideo.utils.communications import all_gather, all_to_all_4D
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-
-def attention(
-    q,
-    k,
-    v,
-    drop_rate=0,
-    attn_mask=None,
-    causal=False,
-):
-
-    qkv = torch.stack([q, k, v], dim=2)
-
-    if attn_mask is not None and attn_mask.dtype != torch.bool:
-        attn_mask = attn_mask.bool()
-
-    x = flash_attn_no_pad(qkv, attn_mask, causal=causal, dropout_p=drop_rate, softmax_scale=None)
-
-    b, s, a, d = x.shape
-    out = x.reshape(b, s, -1)
-    return out
-
-
-def tile(x, sp_size):
-    x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80)
-    return rearrange(x,
-                     "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-                     n_t=5,
-                     n_h=6,
-                     n_w=10,
-                     ts_t=6,
-                     ts_h=8,
-                     ts_w=8)
-
-
-def untile(x, sp_size):
-    x = rearrange(x,
-                  "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-                  n_t=5,
-                  n_h=6,
-                  n_w=10,
-                  ts_t=6,
-                  ts_h=8,
-                  ts_w=8)
-    return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80)
-
-
-def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=None):
-    query, encoder_query = q
-    key, encoder_key = k
-    value, encoder_value = v
-    text_length = text_mask.sum()
-
-    if get_sequence_parallel_state():
-        # batch_size, seq_len, attn_heads, head_dim
-        query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
-        key = all_to_all_4D(key, scatter_dim=2, gather_dim=1)
-        value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
-
-        def shrink_head(encoder_state, dim):
-            local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-            return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
-
-        encoder_query = shrink_head(encoder_query, dim=2)
-        encoder_key = shrink_head(encoder_key, dim=2)
-        encoder_value = shrink_head(encoder_value, dim=2)
-        # [b, s, h, d]
-
-    sequence_length = query.size(1)
-    encoder_sequence_length = encoder_query.size(1)
-
-    if mask_strategy[0] is not None:
-        query = torch.cat([tile(query, nccl_info.sp_size), encoder_query], dim=1).transpose(1, 2)
-        key = torch.cat([tile(key, nccl_info.sp_size), encoder_key], dim=1).transpose(1, 2)
-        value = torch.cat([tile(value, nccl_info.sp_size), encoder_value], dim=1).transpose(1, 2)
-
-        head_num = query.size(1)
-        current_rank = nccl_info.rank_within_group
-        start_head = current_rank * head_num
-        windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)]
-
-        hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2)
-    else:
-        query = torch.cat([query, encoder_query], dim=1)
-        key = torch.cat([key, encoder_key], dim=1)
-        value = torch.cat([value, encoder_value], dim=1)
-        # B, S, 3, H, D
-        qkv = torch.stack([query, key, value], dim=2)
-
-        attn_mask = F.pad(text_mask, (sequence_length, 0), value=True)
-        hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None)
-
-    hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length),
-                                                                          dim=1)
-
-    if mask_strategy[0] is not None:
-        hidden_states = untile(hidden_states, nccl_info.sp_size)
-
-    if get_sequence_parallel_state():
-        hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)
-        encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous()
-
-    hidden_states = hidden_states.to(query.dtype)
-    encoder_hidden_states = encoder_hidden_states.to(query.dtype)
-
-    attn = torch.cat([hidden_states, encoder_hidden_states], dim=1)
-
-    b, s, a, d = attn.shape
-    attn = attn.reshape(b, s, -1)
-
-    return attn

fastvideo/models/hunyuan/modules/embed_layers.py

@@ -1,150 +0,0 @@
-import math
-
-import torch
-import torch.nn as nn
-
-from ..utils.helpers import to_2tuple
-
-
-class PatchEmbed(nn.Module):
-    """2D Image to Patch Embedding
-
-    Image to Patch Embedding using Conv2d
-
-    A convolution based approach to patchifying a 2D image w/ embedding projection.
-
-    Based on the impl in https://github.com/google-research/vision_transformer
-
-    Hacked together by / Copyright 2020 Ross Wightman
-
-    Remove the _assert function in forward function to be compatible with multi-resolution images.
-    """
-
-    def __init__(
-        self,
-        patch_size=16,
-        in_chans=3,
-        embed_dim=768,
-        norm_layer=None,
-        flatten=True,
-        bias=True,
-        dtype=None,
-        device=None,
-    ):
-        factory_kwargs = {"dtype": dtype, "device": device}
-        super().__init__()
-        patch_size = to_2tuple(patch_size)
-        self.patch_size = patch_size
-        self.flatten = flatten
-
-        self.proj = nn.Conv3d(
-            in_chans,
-            embed_dim,
-            kernel_size=patch_size,
-            stride=patch_size,
-            bias=bias,
-            **factory_kwargs,
-        )
-        nn.init.xavier_uniform_(self.proj.weight.view(self.proj.weight.size(0), -1))
-        if bias:
-            nn.init.zeros_(self.proj.bias)
-
-        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
-
-    def forward(self, x):
-        x = self.proj(x)
-        if self.flatten:
-            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
-        x = self.norm(x)
-        return x
-
-
-class TextProjection(nn.Module):
-    """
-    Projects text embeddings. Also handles dropout for classifier-free guidance.
-
-    Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
-    """
-
-    def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None):
-        factory_kwargs = {"dtype": dtype, "device": device}
-        super().__init__()
-        self.linear_1 = nn.Linear(
-            in_features=in_channels,
-            out_features=hidden_size,
-            bias=True,
-            **factory_kwargs,
-        )
-        self.act_1 = act_layer()
-        self.linear_2 = nn.Linear(
-            in_features=hidden_size,
-            out_features=hidden_size,
-            bias=True,
-            **factory_kwargs,
-        )
-
-    def forward(self, caption):
-        hidden_states = self.linear_1(caption)
-        hidden_states = self.act_1(hidden_states)
-        hidden_states = self.linear_2(hidden_states)
-        return hidden_states
-
-
-def timestep_embedding(t, dim, max_period=10000):
-    """
-    Create sinusoidal timestep embeddings.
-
-    Args:
-        t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.
-        dim (int): the dimension of the output.
-        max_period (int): controls the minimum frequency of the embeddings.
-
-    Returns:
-        embedding (torch.Tensor): An (N, D) Tensor of positional embeddings.
-
-    .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-    """
-    half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) /
-                      half).to(device=t.device)
-    args = t[:, None].float() * freqs[None]
-    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-    if dim % 2:
-        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    return embedding
-
-
-class TimestepEmbedder(nn.Module):
-    """
-    Embeds scalar timesteps into vector representations.
-    """
-
-    def __init__(
-        self,
-        hidden_size,
-        act_layer,
-        frequency_embedding_size=256,
-        max_period=10000,
-        out_size=None,
-        dtype=None,
-        device=None,
-    ):
-        factory_kwargs = {"dtype": dtype, "device": device}
-        super().__init__()
-        self.frequency_embedding_size = frequency_embedding_size
-        self.max_period = max_period
-        if out_size is None:
-            out_size = hidden_size
-
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs),
-            act_layer(),
-            nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs),
-        )
-        nn.init.normal_(self.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.mlp[2].weight, std=0.02)
-
-    def forward(self, t):
-        t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype)
-        t_emb = self.mlp(t_freq)
-        return t_emb

fastvideo/models/hunyuan/modules/mlp_layers.py

@@ -1,107 +0,0 @@
-# Modified from timm library:
-# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13
-
-from functools import partial
-
-import torch
-import torch.nn as nn
-
-from ..utils.helpers import to_2tuple
-from .modulate_layers import modulate
-
-
-class MLP(nn.Module):
-    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
-
-    def __init__(
-        self,
-        in_channels,
-        hidden_channels=None,
-        out_features=None,
-        act_layer=nn.GELU,
-        norm_layer=None,
-        bias=True,
-        drop=0.0,
-        use_conv=False,
-        device=None,
-        dtype=None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        out_features = out_features or in_channels
-        hidden_channels = hidden_channels or in_channels
-        bias = to_2tuple(bias)
-        drop_probs = to_2tuple(drop)
-        linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
-
-        self.fc1 = linear_layer(in_channels, hidden_channels, bias=bias[0], **factory_kwargs)
-        self.act = act_layer()
-        self.drop1 = nn.Dropout(drop_probs[0])
-        self.norm = (norm_layer(hidden_channels, **factory_kwargs) if norm_layer is not None else nn.Identity())
-        self.fc2 = linear_layer(hidden_channels, out_features, bias=bias[1], **factory_kwargs)
-        self.drop2 = nn.Dropout(drop_probs[1])
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop1(x)
-        x = self.norm(x)
-        x = self.fc2(x)
-        x = self.drop2(x)
-        return x
-
-
-#
-class MLPEmbedder(nn.Module):
-    """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py"""
-
-    def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs)
-        self.silu = nn.SiLU()
-        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.out_layer(self.silu(self.in_layer(x)))
-
-
-class FinalLayer(nn.Module):
-    """The final layer of DiT."""
-
-    def __init__(self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-
-        # Just use LayerNorm for the final layer
-        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-        if isinstance(patch_size, int):
-            self.linear = nn.Linear(
-                hidden_size,
-                patch_size * patch_size * out_channels,
-                bias=True,
-                **factory_kwargs,
-            )
-        else:
-            self.linear = nn.Linear(
-                hidden_size,
-                patch_size[0] * patch_size[1] * patch_size[2] * out_channels,
-                bias=True,
-            )
-        nn.init.zeros_(self.linear.weight)
-        nn.init.zeros_(self.linear.bias)
-
-        # Here we don't distinguish between the modulate types. Just use the simple one.
-        self.adaLN_modulation = nn.Sequential(
-            act_layer(),
-            nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
-        )
-        # Zero-initialize the modulation
-        nn.init.zeros_(self.adaLN_modulation[1].weight)
-        nn.init.zeros_(self.adaLN_modulation[1].bias)
-
-    def forward(self, x, c):
-        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
-        x = modulate(self.norm_final(x), shift=shift, scale=scale)
-        x = self.linear(x)
-        return x

fastvideo/models/hunyuan/modules/models.py

@@ -1,666 +0,0 @@
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models import ModelMixin
-from einops import rearrange
-
-from fastvideo.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_embed
-from fastvideo.utils.parallel_states import nccl_info
-
-from .activation_layers import get_activation_layer
-from .attenion import parallel_attention
-from .embed_layers import PatchEmbed, TextProjection, TimestepEmbedder
-from .mlp_layers import MLP, FinalLayer, MLPEmbedder
-from .modulate_layers import ModulateDiT, apply_gate, modulate
-from .norm_layers import get_norm_layer
-from .posemb_layers import apply_rotary_emb
-from .token_refiner import SingleTokenRefiner
-
-
-class MMDoubleStreamBlock(nn.Module):
-    """
-    A multimodal dit block with separate modulation for
-    text and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206
-                                     (Flux.1): https://github.com/black-forest-labs/flux
-    """
-
-    def __init__(
-        self,
-        hidden_size: int,
-        heads_num: int,
-        mlp_width_ratio: float,
-        mlp_act_type: str = "gelu_tanh",
-        qk_norm: bool = True,
-        qk_norm_type: str = "rms",
-        qkv_bias: bool = False,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-
-        self.deterministic = False
-        self.heads_num = heads_num
-        head_dim = hidden_size // heads_num
-        mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
-
-        self.img_mod = ModulateDiT(
-            hidden_size,
-            factor=6,
-            act_layer=get_activation_layer("silu"),
-            **factory_kwargs,
-        )
-        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-
-        self.img_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs)
-        qk_norm_layer = get_norm_layer(qk_norm_type)
-        self.img_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-                                if qk_norm else nn.Identity())
-        self.img_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-                                if qk_norm else nn.Identity())
-        self.img_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs)
-
-        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-        self.img_mlp = MLP(
-            hidden_size,
-            mlp_hidden_dim,
-            act_layer=get_activation_layer(mlp_act_type),
-            bias=True,
-            **factory_kwargs,
-        )
-
-        self.txt_mod = ModulateDiT(
-            hidden_size,
-            factor=6,
-            act_layer=get_activation_layer("silu"),
-            **factory_kwargs,
-        )
-        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-
-        self.txt_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs)
-        self.txt_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-                                if qk_norm else nn.Identity())
-        self.txt_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-                                if qk_norm else nn.Identity())
-        self.txt_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs)
-
-        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-        self.txt_mlp = MLP(
-            hidden_size,
-            mlp_hidden_dim,
-            act_layer=get_activation_layer(mlp_act_type),
-            bias=True,
-            **factory_kwargs,
-        )
-        self.hybrid_seq_parallel_attn = None
-
-    def enable_deterministic(self):
-        self.deterministic = True
-
-    def disable_deterministic(self):
-        self.deterministic = False
-
-    def forward(
-        self,
-        img: torch.Tensor,
-        txt: torch.Tensor,
-        vec: torch.Tensor,
-        freqs_cis: tuple = None,
-        text_mask: torch.Tensor = None,
-        mask_strategy=None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        (
-            img_mod1_shift,
-            img_mod1_scale,
-            img_mod1_gate,
-            img_mod2_shift,
-            img_mod2_scale,
-            img_mod2_gate,
-        ) = self.img_mod(vec).chunk(6, dim=-1)
-        (
-            txt_mod1_shift,
-            txt_mod1_scale,
-            txt_mod1_gate,
-            txt_mod2_shift,
-            txt_mod2_scale,
-            txt_mod2_gate,
-        ) = self.txt_mod(vec).chunk(6, dim=-1)
-
-        # Prepare image for attention.
-        img_modulated = self.img_norm1(img)
-        img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale)
-        img_qkv = self.img_attn_qkv(img_modulated)
-        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
-        # Apply QK-Norm if needed
-        img_q = self.img_attn_q_norm(img_q).to(img_v)
-        img_k = self.img_attn_k_norm(img_k).to(img_v)
-
-        # Apply RoPE if needed.
-        if freqs_cis is not None:
-
-            def shrink_head(encoder_state, dim):
-                local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-                return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
-
-            freqs_cis = (
-                shrink_head(freqs_cis[0], dim=0),
-                shrink_head(freqs_cis[1], dim=0),
-            )
-
-            img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False)
-            assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape
-                    ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}"
-            img_q, img_k = img_qq, img_kk
-
-        # Prepare txt for attention.
-        txt_modulated = self.txt_norm1(txt)
-        txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale)
-        txt_qkv = self.txt_attn_qkv(txt_modulated)
-        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
-        # Apply QK-Norm if needed.
-        txt_q = self.txt_attn_q_norm(txt_q).to(txt_v)
-        txt_k = self.txt_attn_k_norm(txt_k).to(txt_v)
-
-        attn = parallel_attention(
-            (img_q, txt_q),
-            (img_k, txt_k),
-            (img_v, txt_v),
-            img_q_len=img_q.shape[1],
-            img_kv_len=img_k.shape[1],
-            text_mask=text_mask,
-            mask_strategy=mask_strategy,
-        )
-
-        # attention computation end
-
-        img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:]
-
-        # Calculate the img blocks.
-        img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate)
-        img = img + apply_gate(
-            self.img_mlp(modulate(self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale)),
-            gate=img_mod2_gate,
-        )
-
-        # Calculate the txt blocks.
-        txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate)
-        txt = txt + apply_gate(
-            self.txt_mlp(modulate(self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale)),
-            gate=txt_mod2_gate,
-        )
-        return img, txt
-
-
-class MMSingleStreamBlock(nn.Module):
-    """
-    A DiT block with parallel linear layers as described in
-    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
-    Also refer to (SD3): https://arxiv.org/abs/2403.03206
-                  (Flux.1): https://github.com/black-forest-labs/flux
-    """
-
-    def __init__(
-        self,
-        hidden_size: int,
-        heads_num: int,
-        mlp_width_ratio: float = 4.0,
-        mlp_act_type: str = "gelu_tanh",
-        qk_norm: bool = True,
-        qk_norm_type: str = "rms",
-        qk_scale: float = None,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-
-        self.deterministic = False
-        self.hidden_size = hidden_size
-        self.heads_num = heads_num
-        head_dim = hidden_size // heads_num
-        mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
-        self.mlp_hidden_dim = mlp_hidden_dim
-        self.scale = qk_scale or head_dim**-0.5
-
-        # qkv and mlp_in
-        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + mlp_hidden_dim, **factory_kwargs)
-        # proj and mlp_out
-        self.linear2 = nn.Linear(hidden_size + mlp_hidden_dim, hidden_size, **factory_kwargs)
-
-        qk_norm_layer = get_norm_layer(qk_norm_type)
-        self.q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-                       if qk_norm else nn.Identity())
-        self.k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-                       if qk_norm else nn.Identity())
-
-        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-
-        self.mlp_act = get_activation_layer(mlp_act_type)()
-        self.modulation = ModulateDiT(
-            hidden_size,
-            factor=3,
-            act_layer=get_activation_layer("silu"),
-            **factory_kwargs,
-        )
-        self.hybrid_seq_parallel_attn = None
-
-    def enable_deterministic(self):
-        self.deterministic = True
-
-    def disable_deterministic(self):
-        self.deterministic = False
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        vec: torch.Tensor,
-        txt_len: int,
-        freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
-        text_mask: torch.Tensor = None,
-        mask_strategy=None,
-    ) -> torch.Tensor:
-        mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
-        x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale)
-        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
-
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
-
-        # Apply QK-Norm if needed.
-        q = self.q_norm(q).to(v)
-        k = self.k_norm(k).to(v)
-
-        def shrink_head(encoder_state, dim):
-            local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-            return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
-
-        freqs_cis = (
-            shrink_head(freqs_cis[0], dim=0),
-            shrink_head(freqs_cis[1], dim=0),
-        )
-
-        img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :]
-        img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :]
-        img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :]
-        img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False)
-        assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape
-                ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}"
-        img_q, img_k = img_qq, img_kk
-
-        attn = parallel_attention(
-            (img_q, txt_q),
-            (img_k, txt_k),
-            (img_v, txt_v),
-            img_q_len=img_q.shape[1],
-            img_kv_len=img_k.shape[1],
-            text_mask=text_mask,
-            mask_strategy=mask_strategy,
-        )
-
-        # attention computation end
-
-        # Compute activation in mlp stream, cat again and run second linear layer.
-        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
-        return x + apply_gate(output, gate=mod_gate)
-
-
-class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
-    """
-    HunyuanVideo Transformer backbone
-
-    Inherited from ModelMixin and ConfigMixin for compatibility with diffusers' sampler StableDiffusionPipeline.
-
-    Reference:
-    [1] Flux.1: https://github.com/black-forest-labs/flux
-    [2] MMDiT: http://arxiv.org/abs/2403.03206
-
-    Parameters
-    ----------
-    args: argparse.Namespace
-        The arguments parsed by argparse.
-    patch_size: list
-        The size of the patch.
-    in_channels: int
-        The number of input channels.
-    out_channels: int
-        The number of output channels.
-    hidden_size: int
-        The hidden size of the transformer backbone.
-    heads_num: int
-        The number of attention heads.
-    mlp_width_ratio: float
-        The ratio of the hidden size of the MLP in the transformer block.
-    mlp_act_type: str
-        The activation function of the MLP in the transformer block.
-    depth_double_blocks: int
-        The number of transformer blocks in the double blocks.
-    depth_single_blocks: int
-        The number of transformer blocks in the single blocks.
-    rope_dim_list: list
-        The dimension of the rotary embedding for t, h, w.
-    qkv_bias: bool
-        Whether to use bias in the qkv linear layer.
-    qk_norm: bool
-        Whether to use qk norm.
-    qk_norm_type: str
-        The type of qk norm.
-    guidance_embed: bool
-        Whether to use guidance embedding for distillation.
-    text_projection: str
-        The type of the text projection, default is single_refiner.
-    use_attention_mask: bool
-        Whether to use attention mask for text encoder.
-    dtype: torch.dtype
-        The dtype of the model.
-    device: torch.device
-        The device of the model.
-    """
-
-    @register_to_config
-    def __init__(
-        self,
-        patch_size: list = [1, 2, 2],
-        in_channels: int = 4,  # Should be VAE.config.latent_channels.
-        out_channels: int = None,
-        hidden_size: int = 3072,
-        heads_num: int = 24,
-        mlp_width_ratio: float = 4.0,
-        mlp_act_type: str = "gelu_tanh",
-        mm_double_blocks_depth: int = 20,
-        mm_single_blocks_depth: int = 40,
-        rope_dim_list: List[int] = [16, 56, 56],
-        qkv_bias: bool = True,
-        qk_norm: bool = True,
-        qk_norm_type: str = "rms",
-        guidance_embed: bool = False,  # For modulation.
-        text_projection: str = "single_refiner",
-        use_attention_mask: bool = True,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        text_states_dim: int = 4096,
-        text_states_dim_2: int = 768,
-        rope_theta: int = 256,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-
-        self.patch_size = patch_size
-        self.in_channels = in_channels
-        self.out_channels = in_channels if out_channels is None else out_channels
-        self.unpatchify_channels = self.out_channels
-        self.guidance_embed = guidance_embed
-        self.rope_dim_list = rope_dim_list
-        self.rope_theta = rope_theta
-        # Text projection. Default to linear projection.
-        # Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831
-        self.use_attention_mask = use_attention_mask
-        self.text_projection = text_projection
-
-        if hidden_size % heads_num != 0:
-            raise ValueError(f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}")
-        pe_dim = hidden_size // heads_num
-        if sum(rope_dim_list) != pe_dim:
-            raise ValueError(f"Got {rope_dim_list} but expected positional dim {pe_dim}")
-        self.hidden_size = hidden_size
-        self.heads_num = heads_num
-
-        # image projection
-        self.img_in = PatchEmbed(self.patch_size, self.in_channels, self.hidden_size, **factory_kwargs)
-
-        # text projection
-        if self.text_projection == "linear":
-            self.txt_in = TextProjection(
-                self.config.text_states_dim,
-                self.hidden_size,
-                get_activation_layer("silu"),
-                **factory_kwargs,
-            )
-        elif self.text_projection == "single_refiner":
-            self.txt_in = SingleTokenRefiner(
-                self.config.text_states_dim,
-                hidden_size,
-                heads_num,
-                depth=2,
-                **factory_kwargs,
-            )
-        else:
-            raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}")
-
-        # time modulation
-        self.time_in = TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs)
-
-        # text modulation
-        self.vector_in = MLPEmbedder(self.config.text_states_dim_2, self.hidden_size, **factory_kwargs)
-
-        # guidance modulation
-        self.guidance_in = (TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs)
-                            if guidance_embed else None)
-
-        # double blocks
-        self.double_blocks = nn.ModuleList([
-            MMDoubleStreamBlock(
-                self.hidden_size,
-                self.heads_num,
-                mlp_width_ratio=mlp_width_ratio,
-                mlp_act_type=mlp_act_type,
-                qk_norm=qk_norm,
-                qk_norm_type=qk_norm_type,
-                qkv_bias=qkv_bias,
-                **factory_kwargs,
-            ) for _ in range(mm_double_blocks_depth)
-        ])
-
-        # single blocks
-        self.single_blocks = nn.ModuleList([
-            MMSingleStreamBlock(
-                self.hidden_size,
-                self.heads_num,
-                mlp_width_ratio=mlp_width_ratio,
-                mlp_act_type=mlp_act_type,
-                qk_norm=qk_norm,
-                qk_norm_type=qk_norm_type,
-                **factory_kwargs,
-            ) for _ in range(mm_single_blocks_depth)
-        ])
-
-        self.final_layer = FinalLayer(
-            self.hidden_size,
-            self.patch_size,
-            self.out_channels,
-            get_activation_layer("silu"),
-            **factory_kwargs,
-        )
-
-    def enable_deterministic(self):
-        for block in self.double_blocks:
-            block.enable_deterministic()
-        for block in self.single_blocks:
-            block.enable_deterministic()
-
-    def disable_deterministic(self):
-        for block in self.double_blocks:
-            block.disable_deterministic()
-        for block in self.single_blocks:
-            block.disable_deterministic()
-
-    def get_rotary_pos_embed(self, rope_sizes):
-        target_ndim = 3
-
-        head_dim = self.hidden_size // self.heads_num
-        rope_dim_list = self.rope_dim_list
-        if rope_dim_list is None:
-            rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
-        assert (sum(rope_dim_list) == head_dim), "sum(rope_dim_list) should equal to head_dim of attention layer"
-        freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
-            rope_dim_list,
-            rope_sizes,
-            theta=self.rope_theta,
-            use_real=True,
-            theta_rescale_factor=1,
-        )
-        return freqs_cos, freqs_sin
-        # x: torch.Tensor,
-        # t: torch.Tensor,  # Should be in range(0, 1000).
-        # text_states: torch.Tensor = None,
-        # text_mask: torch.Tensor = None,  # Now we don't use it.
-        # text_states_2: Optional[torch.Tensor] = None,  # Text embedding for modulation.
-        # guidance: torch.Tensor = None,  # Guidance for modulation, should be cfg_scale x 1000.
-        # return_dict: bool = True,
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        encoder_attention_mask: torch.Tensor,
-        mask_strategy=None,
-        output_features=False,
-        output_features_stride=8,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = False,
-        guidance=None,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
-        if guidance is None:
-            guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16)
-        if mask_strategy is None:
-            mask_strategy = [[None] * self.heads_num for _ in range(len(self.double_blocks) + len(self.single_blocks))]
-        img = x = hidden_states
-        text_mask = encoder_attention_mask
-        t = timestep
-        txt = encoder_hidden_states[:, 1:]
-        text_states_2 = encoder_hidden_states[:, 0, :self.config.text_states_dim_2]
-        _, _, ot, oh, ow = x.shape  # codespell:ignore
-        tt, th, tw = (
-            ot // self.patch_size[0],  # codespell:ignore
-            oh // self.patch_size[1],  # codespell:ignore
-            ow // self.patch_size[2],  # codespell:ignore
-        )
-        original_tt = nccl_info.sp_size * tt
-        freqs_cos, freqs_sin = self.get_rotary_pos_embed((original_tt, th, tw))
-        # Prepare modulation vectors.
-        vec = self.time_in(t)
-
-        # text modulation
-        vec = vec + self.vector_in(text_states_2)
-
-        # guidance modulation
-        if self.guidance_embed:
-            if guidance is None:
-                raise ValueError("Didn't get guidance strength for guidance distilled model.")
-
-            # our timestep_embedding is merged into guidance_in(TimestepEmbedder)
-            vec = vec + self.guidance_in(guidance)
-
-        # Embed image and text.
-        img = self.img_in(img)
-        if self.text_projection == "linear":
-            txt = self.txt_in(txt)
-        elif self.text_projection == "single_refiner":
-            txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None)
-        else:
-            raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}")
-
-        txt_seq_len = txt.shape[1]
-        img_seq_len = img.shape[1]
-
-        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
-        # --------------------- Pass through DiT blocks ------------------------
-
-        for index, block in enumerate(self.double_blocks):
-            double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]]
-            img, txt = block(*double_block_args)
-        # Merge txt and img to pass through single stream blocks.
-        x = torch.cat((img, txt), 1)
-        if output_features:
-            features_list = []
-        if len(self.single_blocks) > 0:
-            for index, block in enumerate(self.single_blocks):
-                single_block_args = [
-                    x,
-                    vec,
-                    txt_seq_len,
-                    (freqs_cos, freqs_sin),
-                    text_mask,
-                    mask_strategy[index + len(self.double_blocks)],
-                ]
-                x = block(*single_block_args)
-                if output_features and _ % output_features_stride == 0:
-                    features_list.append(x[:, :img_seq_len, ...])
-
-        img = x[:, :img_seq_len, ...]
-
-        # ---------------------------- Final layer ------------------------------
-        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
-
-        img = self.unpatchify(img, tt, th, tw)
-        assert not return_dict, "return_dict is not supported."
-        if output_features:
-            features_list = torch.stack(features_list, dim=0)
-        else:
-            features_list = None
-        return (img, features_list)
-
-    def unpatchify(self, x, t, h, w):
-        """
-        x: (N, T, patch_size**2 * C)
-        imgs: (N, H, W, C)
-        """
-        c = self.unpatchify_channels
-        pt, ph, pw = self.patch_size
-        assert t * h * w == x.shape[1]
-
-        x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
-        x = torch.einsum("nthwcopq->nctohpwq", x)
-        imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
-
-        return imgs
-
-    def params_count(self):
-        counts = {
-            "double":
-            sum([
-                sum(p.numel()
-                    for p in block.img_attn_qkv.parameters()) + sum(p.numel()
-                                                                    for p in block.img_attn_proj.parameters()) +
-                sum(p.numel() for p in block.img_mlp.parameters()) + sum(p.numel()
-                                                                         for p in block.txt_attn_qkv.parameters()) +
-                sum(p.numel() for p in block.txt_attn_proj.parameters()) + sum(p.numel()
-                                                                               for p in block.txt_mlp.parameters())
-                for block in self.double_blocks
-            ]),
-            "single":
-            sum([
-                sum(p.numel() for p in block.linear1.parameters()) + sum(p.numel() for p in block.linear2.parameters())
-                for block in self.single_blocks
-            ]),
-            "total":
-            sum(p.numel() for p in self.parameters()),
-        }
-        counts["attn+mlp"] = counts["double"] + counts["single"]
-        return counts
-
-
-#################################################################################
-#                             HunyuanVideo Configs                              #
-#################################################################################
-
-HUNYUAN_VIDEO_CONFIG = {
-    "HYVideo-T/2": {
-        "mm_double_blocks_depth": 20,
-        "mm_single_blocks_depth": 40,
-        "rope_dim_list": [16, 56, 56],
-        "hidden_size": 3072,
-        "heads_num": 24,
-        "mlp_width_ratio": 4,
-    },
-    "HYVideo-T/2-cfgdistill": {
-        "mm_double_blocks_depth": 20,
-        "mm_single_blocks_depth": 40,
-        "rope_dim_list": [16, 56, 56],
-        "hidden_size": 3072,
-        "heads_num": 24,
-        "mlp_width_ratio": 4,
-        "guidance_embed": True,
-    },
-}

fastvideo/models/hunyuan/modules/modulate_layers.py

@@ -1,152 +0,0 @@
-from typing import Callable
-
-import torch
-import torch.nn as nn
-
-
-class ModulateDiT(nn.Module):
-    """Modulation layer for DiT."""
-
-    def __init__(
-        self,
-        hidden_size: int,
-        factor: int,
-        act_layer: Callable,
-        dtype=None,
-        device=None,
-    ):
-        factory_kwargs = {"dtype": dtype, "device": device}
-        super().__init__()
-        self.act = act_layer()
-        self.linear = nn.Linear(hidden_size, factor * hidden_size, bias=True, **factory_kwargs)
-        # Zero-initialize the modulation
-        nn.init.zeros_(self.linear.weight)
-        nn.init.zeros_(self.linear.bias)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.linear(self.act(x))
-
-
-def modulate(x, shift=None, scale=None):
-    """modulate by shift and scale
-
-    Args:
-        x (torch.Tensor): input tensor.
-        shift (torch.Tensor, optional): shift tensor. Defaults to None.
-        scale (torch.Tensor, optional): scale tensor. Defaults to None.
-
-    Returns:
-        torch.Tensor: the output tensor after modulate.
-    """
-    if scale is None and shift is None:
-        return x
-    elif shift is None:
-        return x * (1 + scale.unsqueeze(1))
-    elif scale is None:
-        return x + shift.unsqueeze(1)
-    else:
-        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-def apply_gate(x, gate=None, tanh=False):
-    """AI is creating summary for apply_gate
-
-    Args:
-        x (torch.Tensor): input tensor.
-        gate (torch.Tensor, optional): gate tensor. Defaults to None.
-        tanh (bool, optional): whether to use tanh function. Defaults to False.
-
-    Returns:
-        torch.Tensor: the output tensor after apply gate.
-    """
-    if gate is None:
-        return x
-    if tanh:
-        return x * gate.unsqueeze(1).tanh()
-    else:
-        return x * gate.unsqueeze(1)
-
-
-def ckpt_wrapper(module):
-
-    def ckpt_forward(*inputs):
-        outputs = module(*inputs)
-        return outputs
-
-    return ckpt_forward
-
-
-class RMSNorm(nn.Module):
-
-    def __init__(
-        self,
-        dim: int,
-        elementwise_affine=True,
-        eps: float = 1e-6,
-        device=None,
-        dtype=None,
-    ):
-        """
-        Initialize the RMSNorm normalization layer.
-
-        Args:
-            dim (int): The dimension of the input tensor.
-            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
-
-        Attributes:
-            eps (float): A small value added to the denominator for numerical stability.
-            weight (nn.Parameter): Learnable scaling parameter.
-
-        """
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.eps = eps
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
-
-    def _norm(self, x):
-        """
-        Apply the RMSNorm normalization to the input tensor.
-
-        Args:
-            x (torch.Tensor): The input tensor.
-
-        Returns:
-            torch.Tensor: The normalized tensor.
-
-        """
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
-    def forward(self, x):
-        """
-        Forward pass through the RMSNorm layer.
-
-        Args:
-            x (torch.Tensor): The input tensor.
-
-        Returns:
-            torch.Tensor: The output tensor after applying RMSNorm.
-
-        """
-        output = self._norm(x.float()).type_as(x)
-        if hasattr(self, "weight"):
-            output = output * self.weight
-        return output
-
-
-def get_norm_layer(norm_layer):
-    """
-    Get the normalization layer.
-
-    Args:
-        norm_layer (str): The type of normalization layer.
-
-    Returns:
-        norm_layer (nn.Module): The normalization layer.
-    """
-    if norm_layer == "layer":
-        return nn.LayerNorm
-    elif norm_layer == "rms":
-        return RMSNorm
-    else:
-        raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")

fastvideo/models/hunyuan/modules/norm_layers.py

@@ -1,78 +0,0 @@
-import torch
-import torch.nn as nn
-
-
-class RMSNorm(nn.Module):
-
-    def __init__(
-        self,
-        dim: int,
-        elementwise_affine=True,
-        eps: float = 1e-6,
-        device=None,
-        dtype=None,
-    ):
-        """
-        Initialize the RMSNorm normalization layer.
-
-        Args:
-            dim (int): The dimension of the input tensor.
-            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
-
-        Attributes:
-            eps (float): A small value added to the denominator for numerical stability.
-            weight (nn.Parameter): Learnable scaling parameter.
-
-        """
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.eps = eps
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
-
-    def _norm(self, x):
-        """
-        Apply the RMSNorm normalization to the input tensor.
-
-        Args:
-            x (torch.Tensor): The input tensor.
-
-        Returns:
-            torch.Tensor: The normalized tensor.
-
-        """
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
-    def forward(self, x):
-        """
-        Forward pass through the RMSNorm layer.
-
-        Args:
-            x (torch.Tensor): The input tensor.
-
-        Returns:
-            torch.Tensor: The output tensor after applying RMSNorm.
-
-        """
-        output = self._norm(x.float()).type_as(x)
-        if hasattr(self, "weight"):
-            output = output * self.weight
-        return output
-
-
-def get_norm_layer(norm_layer):
-    """
-    Get the normalization layer.
-
-    Args:
-        norm_layer (str): The type of normalization layer.
-
-    Returns:
-        norm_layer (nn.Module): The normalization layer.
-    """
-    if norm_layer == "layer":
-        return nn.LayerNorm
-    elif norm_layer == "rms":
-        return RMSNorm
-    else:
-        raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")

fastvideo/models/hunyuan/modules/posemb_layers.py

@@ -1,289 +0,0 @@
-from typing import List, Tuple, Union
-
-import torch
-
-
-def _to_tuple(x, dim=2):
-    if isinstance(x, int):
-        return (x, ) * dim
-    elif len(x) == dim:
-        return x
-    else:
-        raise ValueError(f"Expected length {dim} or int, but got {x}")
-
-
-def get_meshgrid_nd(start, *args, dim=2):
-    """
-    Get n-D meshgrid with start, stop and num.
-
-    Args:
-        start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop,
-            step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num
-            should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in
-            n-tuples.
-        *args: See above.
-        dim (int): Dimension of the meshgrid. Defaults to 2.
-
-    Returns:
-        grid (np.ndarray): [dim, ...]
-    """
-    if len(args) == 0:
-        # start is grid_size
-        num = _to_tuple(start, dim=dim)
-        start = (0, ) * dim
-        stop = num
-    elif len(args) == 1:
-        # start is start, args[0] is stop, step is 1
-        start = _to_tuple(start, dim=dim)
-        stop = _to_tuple(args[0], dim=dim)
-        num = [stop[i] - start[i] for i in range(dim)]
-    elif len(args) == 2:
-        # start is start, args[0] is stop, args[1] is num
-        start = _to_tuple(start, dim=dim)  # Left-Top       eg: 12,0
-        stop = _to_tuple(args[0], dim=dim)  # Right-Bottom   eg: 20,32
-        num = _to_tuple(args[1], dim=dim)  # Target Size    eg: 32,124
-    else:
-        raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}")
-
-    # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False)
-    axis_grid = []
-    for i in range(dim):
-        a, b, n = start[i], stop[i], num[i]
-        g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n]
-        axis_grid.append(g)
-    grid = torch.meshgrid(*axis_grid, indexing="ij")  # dim x [W, H, D]
-    grid = torch.stack(grid, dim=0)  # [dim, W, H, D]
-
-    return grid
-
-
-#################################################################################
-#                   Rotary Positional Embedding Functions                       #
-#################################################################################
-# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80
-
-
-def reshape_for_broadcast(
-    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
-    x: torch.Tensor,
-    head_first=False,
-):
-    """
-    Reshape frequency tensor for broadcasting it with another tensor.
-
-    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
-    for the purpose of broadcasting the frequency tensor during element-wise operations.
-
-    Notes:
-        When using FlashMHAModified, head_first should be False.
-        When using Attention, head_first should be True.
-
-    Args:
-        freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped.
-        x (torch.Tensor): Target tensor for broadcasting compatibility.
-        head_first (bool): head dimension first (except batch dim) or not.
-
-    Returns:
-        torch.Tensor: Reshaped frequency tensor.
-
-    Raises:
-        AssertionError: If the frequency tensor doesn't match the expected shape.
-        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
-    """
-    ndim = x.ndim
-    assert 0 <= 1 < ndim
-
-    if isinstance(freqs_cis, tuple):
-        # freqs_cis: (cos, sin) in real space
-        if head_first:
-            assert freqs_cis[0].shape == (
-                x.shape[-2],
-                x.shape[-1],
-            ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
-            shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
-        else:
-            assert freqs_cis[0].shape == (
-                x.shape[1],
-                x.shape[-1],
-            ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
-            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
-        return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
-    else:
-        # freqs_cis: values in complex space
-        if head_first:
-            assert freqs_cis.shape == (
-                x.shape[-2],
-                x.shape[-1],
-            ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
-            shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
-        else:
-            assert freqs_cis.shape == (
-                x.shape[1],
-                x.shape[-1],
-            ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
-            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
-        return freqs_cis.view(*shape)
-
-
-def rotate_half(x):
-    x_real, x_imag = (x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1))  # [B, S, H, D//2]
-    return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
-
-
-def apply_rotary_emb(
-    xq: torch.Tensor,
-    xk: torch.Tensor,
-    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
-    head_first: bool = False,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Apply rotary embeddings to input tensors using the given frequency tensor.
-
-    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
-    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
-    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
-    returned as real tensors.
-
-    Args:
-        xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D]
-        xk (torch.Tensor): Key tensor to apply rotary embeddings.   [B, S, H, D]
-        freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential.
-        head_first (bool): head dimension first (except batch dim) or not.
-
-    Returns:
-        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
-
-    """
-    xk_out = None
-    if isinstance(freqs_cis, tuple):
-        cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first)  # [S, D]
-        cos, sin = cos.to(xq.device), sin.to(xq.device)
-        # real * cos - imag * sin
-        # imag * cos + real * sin
-        xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
-        xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk)
-    else:
-        # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex)
-        xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))  # [B, S, H, D//2]
-        freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device)  # [S, D//2] --> [1, S, 1, D//2]
-        # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin)
-        # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real)
-        xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq)
-        xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))  # [B, S, H, D//2]
-        xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk)
-
-    return xq_out, xk_out
-
-
-def get_nd_rotary_pos_embed(
-    rope_dim_list,
-    start,
-    *args,
-    theta=10000.0,
-    use_real=False,
-    theta_rescale_factor: Union[float, List[float]] = 1.0,
-    interpolation_factor: Union[float, List[float]] = 1.0,
-):
-    """
-    This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
-
-    Args:
-        rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n.
-            sum(rope_dim_list) should equal to head_dim of attention layer.
-        start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start,
-            args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num.
-        *args: See above.
-        theta (float): Scaling factor for frequency computation. Defaults to 10000.0.
-        use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers.
-            Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real
-            part and an imaginary part separately.
-        theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0.
-
-    Returns:
-        pos_embed (torch.Tensor): [HW, D/2]
-    """
-
-    grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list))  # [3, W, H, D] / [2, W, H]
-
-    if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float):
-        theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
-    elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
-        theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list)
-    assert len(theta_rescale_factor) == len(
-        rope_dim_list), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
-
-    if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float):
-        interpolation_factor = [interpolation_factor] * len(rope_dim_list)
-    elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
-        interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list)
-    assert len(interpolation_factor) == len(
-        rope_dim_list), "len(interpolation_factor) should equal to len(rope_dim_list)"
-
-    # use 1/ndim of dimensions to encode grid_axis
-    embs = []
-    for i in range(len(rope_dim_list)):
-        emb = get_1d_rotary_pos_embed(
-            rope_dim_list[i],
-            grid[i].reshape(-1),
-            theta,
-            use_real=use_real,
-            theta_rescale_factor=theta_rescale_factor[i],
-            interpolation_factor=interpolation_factor[i],
-        )  # 2 x [WHD, rope_dim_list[i]]
-        embs.append(emb)
-
-    if use_real:
-        cos = torch.cat([emb[0] for emb in embs], dim=1)  # (WHD, D/2)
-        sin = torch.cat([emb[1] for emb in embs], dim=1)  # (WHD, D/2)
-        return cos, sin
-    else:
-        emb = torch.cat(embs, dim=1)  # (WHD, D/2)
-        return emb
-
-
-def get_1d_rotary_pos_embed(
-    dim: int,
-    pos: Union[torch.FloatTensor, int],
-    theta: float = 10000.0,
-    use_real: bool = False,
-    theta_rescale_factor: float = 1.0,
-    interpolation_factor: float = 1.0,
-) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-    """
-    Precompute the frequency tensor for complex exponential (cis) with given dimensions.
-    (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
-
-    This function calculates a frequency tensor with complex exponential using the given dimension 'dim'
-    and the end index 'end'. The 'theta' parameter scales the frequencies.
-    The returned tensor contains complex values in complex64 data type.
-
-    Args:
-        dim (int): Dimension of the frequency tensor.
-        pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar
-        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
-        use_real (bool, optional): If True, return real part and imaginary part separately.
-                                   Otherwise, return complex numbers.
-        theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
-
-    Returns:
-        freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2]
-        freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
-    """
-    if isinstance(pos, int):
-        pos = torch.arange(pos).float()
-
-    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
-    # has some connection to NTK literature
-    if theta_rescale_factor != 1.0:
-        theta *= theta_rescale_factor**(dim / (dim - 2))
-
-    freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))  # [D/2]
-    # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}"
-    freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
-    if use_real:
-        freqs_cos = freqs.cos().repeat_interleave(2, dim=1)  # [S, D]
-        freqs_sin = freqs.sin().repeat_interleave(2, dim=1)  # [S, D]
-        return freqs_cos, freqs_sin
-    else:
-        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64     # [S, D/2]
-        return freqs_cis

fastvideo/models/hunyuan/modules/token_refiner.py

@@ -1,202 +0,0 @@
-from typing import Optional
-
-import torch
-import torch.nn as nn
-from einops import rearrange
-
-from .activation_layers import get_activation_layer
-from .attenion import attention
-from .embed_layers import TextProjection, TimestepEmbedder
-from .mlp_layers import MLP
-from .modulate_layers import apply_gate
-from .norm_layers import get_norm_layer
-
-
-class IndividualTokenRefinerBlock(nn.Module):
-
-    def __init__(
-        self,
-        hidden_size,
-        heads_num,
-        mlp_width_ratio: str = 4.0,
-        mlp_drop_rate: float = 0.0,
-        act_type: str = "silu",
-        qk_norm: bool = False,
-        qk_norm_type: str = "layer",
-        qkv_bias: bool = True,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.heads_num = heads_num
-        head_dim = hidden_size // heads_num
-        mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
-
-        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-        self.self_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs)
-        qk_norm_layer = get_norm_layer(qk_norm_type)
-        self.self_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-                                 if qk_norm else nn.Identity())
-        self.self_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-                                 if qk_norm else nn.Identity())
-        self.self_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs)
-
-        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs)
-        act_layer = get_activation_layer(act_type)
-        self.mlp = MLP(
-            in_channels=hidden_size,
-            hidden_channels=mlp_hidden_dim,
-            act_layer=act_layer,
-            drop=mlp_drop_rate,
-            **factory_kwargs,
-        )
-
-        self.adaLN_modulation = nn.Sequential(
-            act_layer(),
-            nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
-        )
-        # Zero-initialize the modulation
-        nn.init.zeros_(self.adaLN_modulation[1].weight)
-        nn.init.zeros_(self.adaLN_modulation[1].bias)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        c: torch.Tensor,  # timestep_aware_representations + context_aware_representations
-        attn_mask: torch.Tensor = None,
-    ):
-        gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1)
-
-        norm_x = self.norm1(x)
-        qkv = self.self_attn_qkv(norm_x)
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
-        # Apply QK-Norm if needed
-        q = self.self_attn_q_norm(q).to(v)
-        k = self.self_attn_k_norm(k).to(v)
-
-        # Self-Attention
-        attn = attention(q, k, v, attn_mask=attn_mask)
-
-        x = x + apply_gate(self.self_attn_proj(attn), gate_msa)
-
-        # FFN Layer
-        x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp)
-
-        return x
-
-
-class IndividualTokenRefiner(nn.Module):
-
-    def __init__(
-        self,
-        hidden_size,
-        heads_num,
-        depth,
-        mlp_width_ratio: float = 4.0,
-        mlp_drop_rate: float = 0.0,
-        act_type: str = "silu",
-        qk_norm: bool = False,
-        qk_norm_type: str = "layer",
-        qkv_bias: bool = True,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.blocks = nn.ModuleList([
-            IndividualTokenRefinerBlock(
-                hidden_size=hidden_size,
-                heads_num=heads_num,
-                mlp_width_ratio=mlp_width_ratio,
-                mlp_drop_rate=mlp_drop_rate,
-                act_type=act_type,
-                qk_norm=qk_norm,
-                qk_norm_type=qk_norm_type,
-                qkv_bias=qkv_bias,
-                **factory_kwargs,
-            ) for _ in range(depth)
-        ])
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        c: torch.LongTensor,
-        mask: Optional[torch.Tensor] = None,
-    ):
-        mask = mask.clone().bool()
-        # avoid attention weight become NaN
-        mask[:, 0] = True
-        for block in self.blocks:
-            x = block(x, c, mask)
-        return x
-
-
-class SingleTokenRefiner(nn.Module):
-    """
-    A single token refiner block for llm text embedding refine.
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        hidden_size,
-        heads_num,
-        depth,
-        mlp_width_ratio: float = 4.0,
-        mlp_drop_rate: float = 0.0,
-        act_type: str = "silu",
-        qk_norm: bool = False,
-        qk_norm_type: str = "layer",
-        qkv_bias: bool = True,
-        attn_mode: str = "torch",
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.attn_mode = attn_mode
-        assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner."
-
-        self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True, **factory_kwargs)
-
-        act_layer = get_activation_layer(act_type)
-        # Build timestep embedding layer
-        self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs)
-        # Build context embedding layer
-        self.c_embedder = TextProjection(in_channels, hidden_size, act_layer, **factory_kwargs)
-
-        self.individual_token_refiner = IndividualTokenRefiner(
-            hidden_size=hidden_size,
-            heads_num=heads_num,
-            depth=depth,
-            mlp_width_ratio=mlp_width_ratio,
-            mlp_drop_rate=mlp_drop_rate,
-            act_type=act_type,
-            qk_norm=qk_norm,
-            qk_norm_type=qk_norm_type,
-            qkv_bias=qkv_bias,
-            **factory_kwargs,
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        t: torch.LongTensor,
-        mask: Optional[torch.LongTensor] = None,
-    ):
-        timestep_aware_representations = self.t_embedder(t)
-
-        if mask is None:
-            context_aware_representations = x.mean(dim=1)
-        else:
-            mask_float = mask.float().unsqueeze(-1)  # [b, s1, 1]
-            context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
-        context_aware_representations = self.c_embedder(context_aware_representations)
-        c = timestep_aware_representations + context_aware_representations
-
-        x = self.input_embedder(x)
-
-        x = self.individual_token_refiner(x, c, mask)
-
-        return x

fastvideo/models/hunyuan/prompt_rewrite.py

@@ -1,52 +0,0 @@
-normal_mode_prompt = """Normal mode - Video Recaption Task:
-
-You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description.
-
-0. Preserve ALL information, including style words and technical terms.
-
-1. If the input is in Chinese, translate the entire description to English. 
-
-2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences.
-
-3. If the input does not include style, lighting, atmosphere, you can make reasonable associations.
-
-4. Output ALL must be in English.
-
-Given Input:
-input: "{input}"
-"""
-
-master_mode_prompt = """Master mode - Video Recaption Task:
-
-You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description.
-
-0. Preserve ALL information, including style words and technical terms.
-
-1. If the input is in Chinese, translate the entire description to English. 
-
-2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences.
-
-3. If the input does not include style, lighting, atmosphere, you can make reasonable associations.
-
-4. Output ALL must be in English.
-
-Given Input:
-input: "{input}"
-"""
-
-
-def get_rewrite_prompt(ori_prompt, mode="Normal"):
-    if mode == "Normal":
-        prompt = normal_mode_prompt.format(input=ori_prompt)
-    elif mode == "Master":
-        prompt = master_mode_prompt.format(input=ori_prompt)
-    else:
-        raise Exception("Only supports Normal and Master mode, but got {}".format(mode))
-    return prompt
-
-
-ori_prompt = "一只小狗在草地上奔跑。"
-normal_prompt = get_rewrite_prompt(ori_prompt, mode="Normal")
-master_prompt = get_rewrite_prompt(ori_prompt, mode="Master")
-
-# Then you can use the normal_prompt or master_prompt to access the hunyuan-large rewrite model to get the final prompt.

fastvideo/models/hunyuan/text_encoder/__init__.py

@@ -1,323 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-from transformers import AutoModel, AutoTokenizer, CLIPTextModel, CLIPTokenizer
-from transformers.utils import ModelOutput
-
-from ..constants import PRECISION_TO_TYPE, TEXT_ENCODER_PATH, TOKENIZER_PATH
-
-
-def use_default(value, default):
-    return value if value is not None else default
-
-
-def load_text_encoder(
-    text_encoder_type,
-    text_encoder_precision=None,
-    text_encoder_path=None,
-    logger=None,
-    device=None,
-):
-    if text_encoder_path is None:
-        text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type]
-    if logger is not None:
-        logger.info(f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}")
-
-    if text_encoder_type == "clipL":
-        text_encoder = CLIPTextModel.from_pretrained(text_encoder_path)
-        text_encoder.final_layer_norm = text_encoder.text_model.final_layer_norm
-    elif text_encoder_type == "llm":
-        text_encoder = AutoModel.from_pretrained(text_encoder_path, low_cpu_mem_usage=True)
-        text_encoder.final_layer_norm = text_encoder.norm
-    else:
-        raise ValueError(f"Unsupported text encoder type: {text_encoder_type}")
-    # from_pretrained will ensure that the model is in eval mode.
-
-    if text_encoder_precision is not None:
-        text_encoder = text_encoder.to(dtype=PRECISION_TO_TYPE[text_encoder_precision])
-
-    text_encoder.requires_grad_(False)
-
-    if logger is not None:
-        logger.info(f"Text encoder to dtype: {text_encoder.dtype}")
-
-    if device is not None:
-        text_encoder = text_encoder.to(device)
-
-    return text_encoder, text_encoder_path
-
-
-def load_tokenizer(tokenizer_type, tokenizer_path=None, padding_side="right", logger=None):
-    if tokenizer_path is None:
-        tokenizer_path = TOKENIZER_PATH[tokenizer_type]
-    if logger is not None:
-        logger.info(f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}")
-
-    if tokenizer_type == "clipL":
-        tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path, max_length=77)
-    elif tokenizer_type == "llm":
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, padding_side=padding_side)
-    else:
-        raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}")
-
-    return tokenizer, tokenizer_path
-
-
-@dataclass
-class TextEncoderModelOutput(ModelOutput):
-    """
-    Base class for model's outputs that also contains a pooling of the last hidden states.
-
-    Args:
-        hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
-            Sequence of hidden-states at the output of the last layer of the model.
-        attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
-        hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed):
-            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
-            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
-        text_outputs (`list`, *optional*, returned when `return_texts=True` is passed):
-            List of decoded texts.
-    """
-
-    hidden_state: torch.FloatTensor = None
-    attention_mask: Optional[torch.LongTensor] = None
-    hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None
-    text_outputs: Optional[list] = None
-
-
-class TextEncoder(nn.Module):
-
-    def __init__(
-        self,
-        text_encoder_type: str,
-        max_length: int,
-        text_encoder_precision: Optional[str] = None,
-        text_encoder_path: Optional[str] = None,
-        tokenizer_type: Optional[str] = None,
-        tokenizer_path: Optional[str] = None,
-        output_key: Optional[str] = None,
-        use_attention_mask: bool = True,
-        input_max_length: Optional[int] = None,
-        prompt_template: Optional[dict] = None,
-        prompt_template_video: Optional[dict] = None,
-        hidden_state_skip_layer: Optional[int] = None,
-        apply_final_norm: bool = False,
-        reproduce: bool = False,
-        logger=None,
-        device=None,
-    ):
-        super().__init__()
-        self.text_encoder_type = text_encoder_type
-        self.max_length = max_length
-        self.precision = text_encoder_precision
-        self.model_path = text_encoder_path
-        self.tokenizer_type = (tokenizer_type if tokenizer_type is not None else text_encoder_type)
-        self.tokenizer_path = (tokenizer_path if tokenizer_path is not None else text_encoder_path)
-        self.use_attention_mask = use_attention_mask
-        if prompt_template_video is not None:
-            assert (use_attention_mask is True), "Attention mask is True required when training videos."
-        self.input_max_length = (input_max_length if input_max_length is not None else max_length)
-        self.prompt_template = prompt_template
-        self.prompt_template_video = prompt_template_video
-        self.hidden_state_skip_layer = hidden_state_skip_layer
-        self.apply_final_norm = apply_final_norm
-        self.reproduce = reproduce
-        self.logger = logger
-
-        self.use_template = self.prompt_template is not None
-        if self.use_template:
-            assert (isinstance(self.prompt_template, dict) and "template" in self.prompt_template
-                    ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}"
-            assert "{}" in str(self.prompt_template["template"]), (
-                "`prompt_template['template']` must contain a placeholder `{}` for the input text, "
-                f"got {self.prompt_template['template']}")
-
-        self.use_video_template = self.prompt_template_video is not None
-        if self.use_video_template:
-            if self.prompt_template_video is not None:
-                assert (
-                    isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video
-                ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}"
-            assert "{}" in str(self.prompt_template_video["template"]), (
-                "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, "
-                f"got {self.prompt_template_video['template']}")
-
-        if "t5" in text_encoder_type:
-            self.output_key = output_key or "last_hidden_state"
-        elif "clip" in text_encoder_type:
-            self.output_key = output_key or "pooler_output"
-        elif "llm" in text_encoder_type or "glm" in text_encoder_type:
-            self.output_key = output_key or "last_hidden_state"
-        else:
-            raise ValueError(f"Unsupported text encoder type: {text_encoder_type}")
-
-        self.model, self.model_path = load_text_encoder(
-            text_encoder_type=self.text_encoder_type,
-            text_encoder_precision=self.precision,
-            text_encoder_path=self.model_path,
-            logger=self.logger,
-            device=device,
-        )
-        self.dtype = self.model.dtype
-        self.device = self.model.device
-
-        self.tokenizer, self.tokenizer_path = load_tokenizer(
-            tokenizer_type=self.tokenizer_type,
-            tokenizer_path=self.tokenizer_path,
-            padding_side="right",
-            logger=self.logger,
-        )
-
-    def __repr__(self):
-        return f"{self.text_encoder_type} ({self.precision} - {self.model_path})"
-
-    @staticmethod
-    def apply_text_to_template(text, template, prevent_empty_text=True):
-        """
-        Apply text to template.
-
-        Args:
-            text (str): Input text.
-            template (str or list): Template string or list of chat conversation.
-            prevent_empty_text (bool): If True, we will prevent the user text from being empty
-                by adding a space. Defaults to True.
-        """
-        if isinstance(template, str):
-            # Will send string to tokenizer. Used for llm
-            return template.format(text)
-        else:
-            raise TypeError(f"Unsupported template type: {type(template)}")
-
-    def text2tokens(self, text, data_type="image"):
-        """
-        Tokenize the input text.
-
-        Args:
-            text (str or list): Input text.
-        """
-        tokenize_input_type = "str"
-        if self.use_template:
-            if data_type == "image":
-                prompt_template = self.prompt_template["template"]
-            elif data_type == "video":
-                prompt_template = self.prompt_template_video["template"]
-            else:
-                raise ValueError(f"Unsupported data type: {data_type}")
-            if isinstance(text, (list, tuple)):
-                text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text]
-                if isinstance(text[0], list):
-                    tokenize_input_type = "list"
-            elif isinstance(text, str):
-                text = self.apply_text_to_template(text, prompt_template)
-                if isinstance(text, list):
-                    tokenize_input_type = "list"
-            else:
-                raise TypeError(f"Unsupported text type: {type(text)}")
-
-        kwargs = dict(
-            truncation=True,
-            max_length=self.max_length,
-            padding="max_length",
-            return_tensors="pt",
-        )
-        if tokenize_input_type == "str":
-            return self.tokenizer(
-                text,
-                return_length=False,
-                return_overflowing_tokens=False,
-                return_attention_mask=True,
-                **kwargs,
-            )
-        elif tokenize_input_type == "list":
-            return self.tokenizer.apply_chat_template(
-                text,
-                add_generation_prompt=True,
-                tokenize=True,
-                return_dict=True,
-                **kwargs,
-            )
-        else:
-            raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}")
-
-    def encode(
-        self,
-        batch_encoding,
-        use_attention_mask=None,
-        output_hidden_states=False,
-        do_sample=None,
-        hidden_state_skip_layer=None,
-        return_texts=False,
-        data_type="image",
-        device=None,
-    ):
-        """
-        Args:
-            batch_encoding (dict): Batch encoding from tokenizer.
-            use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask.
-                Defaults to None.
-            output_hidden_states (bool): Whether to output hidden states. If False, return the value of
-                self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer,
-                output_hidden_states will be set True. Defaults to False.
-            do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None.
-                When self.produce is False, do_sample is set to True by default.
-            hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer.
-                If None, self.output_key will be used. Defaults to None.
-            return_texts (bool): Whether to return the decoded texts. Defaults to False.
-        """
-        device = self.model.device if device is None else device
-        use_attention_mask = use_default(use_attention_mask, self.use_attention_mask)
-        hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer)
-        do_sample = use_default(do_sample, not self.reproduce)
-        attention_mask = (batch_encoding["attention_mask"].to(device) if use_attention_mask else None)
-        outputs = self.model(
-            input_ids=batch_encoding["input_ids"].to(device),
-            attention_mask=attention_mask,
-            output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None,
-        )
-        if hidden_state_skip_layer is not None:
-            last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)]
-            # Real last hidden state already has layer norm applied. So here we only apply it
-            # for intermediate layers.
-            if hidden_state_skip_layer > 0 and self.apply_final_norm:
-                last_hidden_state = self.model.final_layer_norm(last_hidden_state)
-        else:
-            last_hidden_state = outputs[self.output_key]
-
-        # Remove hidden states of instruction tokens, only keep prompt tokens.
-        if self.use_template:
-            if data_type == "image":
-                crop_start = self.prompt_template.get("crop_start", -1)
-            elif data_type == "video":
-                crop_start = self.prompt_template_video.get("crop_start", -1)
-            else:
-                raise ValueError(f"Unsupported data type: {data_type}")
-            if crop_start > 0:
-                last_hidden_state = last_hidden_state[:, crop_start:]
-                attention_mask = (attention_mask[:, crop_start:] if use_attention_mask else None)
-
-        if output_hidden_states:
-            return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states)
-        return TextEncoderModelOutput(last_hidden_state, attention_mask)
-
-    def forward(
-        self,
-        text,
-        use_attention_mask=None,
-        output_hidden_states=False,
-        do_sample=False,
-        hidden_state_skip_layer=None,
-        return_texts=False,
-    ):
-        batch_encoding = self.text2tokens(text)
-        return self.encode(
-            batch_encoding,
-            use_attention_mask=use_attention_mask,
-            output_hidden_states=output_hidden_states,
-            do_sample=do_sample,
-            hidden_state_skip_layer=hidden_state_skip_layer,
-            return_texts=return_texts,
-        )

fastvideo/models/hunyuan/utils/data_utils.py

@@ -1,14 +0,0 @@
-import math
-
-
-def align_to(value, alignment):
-    """align height, width according to alignment
-
-    Args:
-        value (int): height or width
-        alignment (int): target alignment factor
-
-    Returns:
-        int: the aligned value
-    """
-    return int(math.ceil(value / alignment) * alignment)

fastvideo/models/hunyuan/utils/file_utils.py

@@ -1,71 +0,0 @@
-import os
-from pathlib import Path
-
-import imageio
-import numpy as np
-import torch
-import torchvision
-from einops import rearrange
-
-CODE_SUFFIXES = {
-    ".py",  # Python codes
-    ".sh",  # Shell scripts
-    ".yaml",
-    ".yml",  # Configuration files
-}
-
-
-def safe_dir(path):
-    """
-    Create a directory (or the parent directory of a file) if it does not exist.
-
-    Args:
-        path (str or Path): Path to the directory.
-
-    Returns:
-        path (Path): Path object of the directory.
-    """
-    path = Path(path)
-    path.mkdir(exist_ok=True, parents=True)
-    return path
-
-
-def safe_file(path):
-    """
-    Create the parent directory of a file if it does not exist.
-
-    Args:
-        path (str or Path): Path to the file.
-
-    Returns:
-        path (Path): Path object of the file.
-    """
-    path = Path(path)
-    path.parent.mkdir(exist_ok=True, parents=True)
-    return path
-
-
-def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=1, fps=24):
-    """save videos by video tensor
-       copy from https://github.com/guoyww/AnimateDiff/blob/e92bd5671ba62c0d774a32951453e328018b7c5b/animatediff/utils/util.py#L61
-
-    Args:
-        videos (torch.Tensor): video tensor predicted by the model
-        path (str): path to save video
-        rescale (bool, optional): rescale the video tensor from [-1, 1] to  . Defaults to False.
-        n_rows (int, optional): Defaults to 1.
-        fps (int, optional): video save fps. Defaults to 8.
-    """
-    videos = rearrange(videos, "b c t h w -> t b c h w")
-    outputs = []
-    for x in videos:
-        x = torchvision.utils.make_grid(x, nrow=n_rows)
-        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-        if rescale:
-            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
-        x = torch.clamp(x, 0, 1)
-        x = (x * 255).numpy().astype(np.uint8)
-        outputs.append(x)
-
-    os.makedirs(os.path.dirname(path), exist_ok=True)
-    imageio.mimsave(path, outputs, fps=fps)

fastvideo/models/hunyuan/utils/helpers.py

@@ -1,41 +0,0 @@
-import collections.abc
-from itertools import repeat
-
-
-def _ntuple(n):
-
-    def parse(x):
-        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
-            x = tuple(x)
-            if len(x) == 1:
-                x = tuple(repeat(x[0], n))
-            return x
-        return tuple(repeat(x, n))
-
-    return parse
-
-
-to_1tuple = _ntuple(1)
-to_2tuple = _ntuple(2)
-to_3tuple = _ntuple(3)
-to_4tuple = _ntuple(4)
-
-
-def as_tuple(x):
-    if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
-        return tuple(x)
-    if x is None or isinstance(x, (int, float, str)):
-        return (x, )
-    else:
-        raise ValueError(f"Unknown type {type(x)}")
-
-
-def as_list_of_2tuple(x):
-    x = as_tuple(x)
-    if len(x) == 1:
-        x = (x[0], x[0])
-    assert len(x) % 2 == 0, f"Expect even length, got {len(x)}."
-    lst = []
-    for i in range(0, len(x), 2):
-        lst.append((x[i], x[i + 1]))
-    return lst

fastvideo/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py

@@ -1,41 +0,0 @@
-import argparse
-
-import torch
-from transformers import AutoProcessor, LlavaForConditionalGeneration
-
-
-def preprocess_text_encoder_tokenizer(args):
-
-    processor = AutoProcessor.from_pretrained(args.input_dir)
-    model = LlavaForConditionalGeneration.from_pretrained(
-        args.input_dir,
-        torch_dtype=torch.float16,
-        low_cpu_mem_usage=True,
-    ).to(0)
-
-    model.language_model.save_pretrained(f"{args.output_dir}")
-    processor.tokenizer.save_pretrained(f"{args.output_dir}")
-
-
-if __name__ == "__main__":
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--input_dir",
-        type=str,
-        required=True,
-        help="The path to the llava-llama-3-8b-v1_1-transformers.",
-    )
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default="",
-        help="The output path of the llava-llama-3-8b-text-encoder-tokenizer."
-        "if '', the parent dir of output will be the same as input dir.",
-    )
-    args = parser.parse_args()
-
-    if len(args.output_dir) == 0:
-        args.output_dir = "/".join(args.input_dir.split("/")[:-1])
-
-    preprocess_text_encoder_tokenizer(args)

fastvideo/models/hunyuan/vae/__init__.py

@@ -1,64 +0,0 @@
-from pathlib import Path
-
-import torch
-
-from ..constants import PRECISION_TO_TYPE, VAE_PATH
-from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D
-
-
-def load_vae(
-    vae_type: str = "884-16c-hy",
-    vae_precision: str = None,
-    sample_size: tuple = None,
-    vae_path: str = None,
-    logger=None,
-    device=None,
-):
-    """the function to load the 3D VAE model
-
-    Args:
-        vae_type (str): the type of the 3D VAE model. Defaults to "884-16c-hy".
-        vae_precision (str, optional): the precision to load vae. Defaults to None.
-        sample_size (tuple, optional): the tiling size. Defaults to None.
-        vae_path (str, optional): the path to vae. Defaults to None.
-        logger (_type_, optional): logger. Defaults to None.
-        device (_type_, optional): device to load vae. Defaults to None.
-    """
-    if vae_path is None:
-        vae_path = VAE_PATH[vae_type]
-
-    if logger is not None:
-        logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}")
-    config = AutoencoderKLCausal3D.load_config(vae_path)
-    if sample_size:
-        vae = AutoencoderKLCausal3D.from_config(config, sample_size=sample_size)
-    else:
-        vae = AutoencoderKLCausal3D.from_config(config)
-
-    vae_ckpt = Path(vae_path) / "pytorch_model.pt"
-    assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}"
-
-    ckpt = torch.load(vae_ckpt, map_location=vae.device)
-    if "state_dict" in ckpt:
-        ckpt = ckpt["state_dict"]
-    if any(k.startswith("vae.") for k in ckpt.keys()):
-        ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")}
-    vae.load_state_dict(ckpt)
-
-    spatial_compression_ratio = vae.config.spatial_compression_ratio
-    time_compression_ratio = vae.config.time_compression_ratio
-
-    if vae_precision is not None:
-        vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision])
-
-    vae.requires_grad_(False)
-
-    if logger is not None:
-        logger.info(f"VAE to dtype: {vae.dtype}")
-
-    if device is not None:
-        vae = vae.to(device)
-
-    vae.eval()
-
-    return vae, vae_path, spatial_compression_ratio, time_compression_ratio

fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py

@@ -1,764 +0,0 @@
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-#
-# Modified from diffusers==0.29.2
-#
-# ==============================================================================
-from dataclasses import dataclass
-from math import prod
-from typing import Dict, Optional, Tuple, Union
-
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-
-from fastvideo.utils.parallel_states import nccl_info
-
-try:
-    # This diffusers is modified and packed in the mirror.
-    from diffusers.loaders import FromOriginalVAEMixin
-except ImportError:
-    # Use this to be compatible with the original diffusers.
-    from diffusers.loaders.single_file_model import (
-        FromOriginalModelMixin as FromOriginalVAEMixin, )
-
-from diffusers.models.attention_processor import (ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, Attention,
-                                                  AttentionProcessor, AttnAddedKVProcessor, AttnProcessor)
-from diffusers.models.modeling_outputs import AutoencoderKLOutput
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.utils.accelerate_utils import apply_forward_hook
-
-from .vae import BaseOutput, DecoderCausal3D, DecoderOutput, DiagonalGaussianDistribution, EncoderCausal3D
-
-
-@dataclass
-class DecoderOutput2(BaseOutput):
-    sample: torch.FloatTensor
-    posterior: Optional[DiagonalGaussianDistribution] = None
-
-
-class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin):
-    r"""
-    A VAE model with KL loss for encoding images/videos into latents and decoding latent representations into images/videos.
-
-    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
-    for all models (such as downloading or saving).
-    """
-
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-        self,
-        in_channels: int = 3,
-        out_channels: int = 3,
-        down_block_types: Tuple[str] = ("DownEncoderBlockCausal3D", ),
-        up_block_types: Tuple[str] = ("UpDecoderBlockCausal3D", ),
-        block_out_channels: Tuple[int] = (64, ),
-        layers_per_block: int = 1,
-        act_fn: str = "silu",
-        latent_channels: int = 4,
-        norm_num_groups: int = 32,
-        sample_size: int = 32,
-        sample_tsize: int = 64,
-        scaling_factor: float = 0.18215,
-        force_upcast: float = True,
-        spatial_compression_ratio: int = 8,
-        time_compression_ratio: int = 4,
-        mid_block_add_attention: bool = True,
-    ):
-        super().__init__()
-
-        self.time_compression_ratio = time_compression_ratio
-
-        self.encoder = EncoderCausal3D(
-            in_channels=in_channels,
-            out_channels=latent_channels,
-            down_block_types=down_block_types,
-            block_out_channels=block_out_channels,
-            layers_per_block=layers_per_block,
-            act_fn=act_fn,
-            norm_num_groups=norm_num_groups,
-            double_z=True,
-            time_compression_ratio=time_compression_ratio,
-            spatial_compression_ratio=spatial_compression_ratio,
-            mid_block_add_attention=mid_block_add_attention,
-        )
-
-        self.decoder = DecoderCausal3D(
-            in_channels=latent_channels,
-            out_channels=out_channels,
-            up_block_types=up_block_types,
-            block_out_channels=block_out_channels,
-            layers_per_block=layers_per_block,
-            norm_num_groups=norm_num_groups,
-            act_fn=act_fn,
-            time_compression_ratio=time_compression_ratio,
-            spatial_compression_ratio=spatial_compression_ratio,
-            mid_block_add_attention=mid_block_add_attention,
-        )
-
-        self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1)
-        self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1)
-
-        self.use_slicing = False
-        self.use_spatial_tiling = False
-        self.use_temporal_tiling = False
-        self.use_parallel = False
-
-        # only relevant if vae tiling is enabled
-        self.tile_sample_min_tsize = sample_tsize
-        self.tile_latent_min_tsize = sample_tsize // time_compression_ratio
-
-        self.tile_sample_min_size = self.config.sample_size
-        sample_size = (self.config.sample_size[0] if isinstance(self.config.sample_size,
-                                                                (list, tuple)) else self.config.sample_size)
-        self.tile_latent_min_size = int(sample_size / (2**(len(self.config.block_out_channels) - 1)))
-        self.tile_overlap_factor = 0.25
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, (EncoderCausal3D, DecoderCausal3D)):
-            module.gradient_checkpointing = value
-
-    def enable_temporal_tiling(self, use_tiling: bool = True):
-        self.use_temporal_tiling = use_tiling
-
-    def disable_temporal_tiling(self):
-        self.enable_temporal_tiling(False)
-
-    def enable_spatial_tiling(self, use_tiling: bool = True):
-        self.use_spatial_tiling = use_tiling
-
-    def disable_spatial_tiling(self):
-        self.enable_spatial_tiling(False)
-
-    def enable_tiling(self, use_tiling: bool = True):
-        r"""
-        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
-        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
-        processing larger videos.
-        """
-        self.enable_spatial_tiling(use_tiling)
-        self.enable_temporal_tiling(use_tiling)
-
-    def disable_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
-        decoding in one step.
-        """
-        self.disable_spatial_tiling()
-        self.disable_temporal_tiling()
-
-    def enable_parallel(self):
-        r"""
-        Enable sequence parallelism for the model. This will allow the vae to decode (with tiling) in parallel.
-        """
-        self.use_parallel = True
-
-    def enable_slicing(self):
-        r"""
-        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
-        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        self.use_slicing = True
-
-    def disable_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
-        decoding in one step.
-        """
-        self.use_slicing = False
-
-    @property
-    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
-    def attn_processors(self) -> Dict[str, AttentionProcessor]:
-        r"""
-        Returns:
-            `dict` of attention processors: A dictionary containing all attention processors used in the model with
-            indexed by its weight name.
-        """
-        # set recursively
-        processors = {}
-
-        def fn_recursive_add_processors(
-            name: str,
-            module: torch.nn.Module,
-            processors: Dict[str, AttentionProcessor],
-        ):
-            if hasattr(module, "get_processor"):
-                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
-
-            for sub_name, child in module.named_children():
-                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
-
-            return processors
-
-        for name, module in self.named_children():
-            fn_recursive_add_processors(name, module, processors)
-
-        return processors
-
-    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
-    def set_attn_processor(
-        self,
-        processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]],
-        _remove_lora=False,
-    ):
-        r"""
-        Sets the attention processor to use to compute attention.
-
-        Parameters:
-            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
-                The instantiated processor class or a dictionary of processor classes that will be set as the processor
-                for **all** `Attention` layers.
-
-                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
-                processor. This is strongly recommended when setting trainable attention processors.
-
-        """
-        count = len(self.attn_processors.keys())
-
-        if isinstance(processor, dict) and len(processor) != count:
-            raise ValueError(
-                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
-                f" number of attention layers: {count}. Please make sure to pass {count} processor classes.")
-
-        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
-            if hasattr(module, "set_processor"):
-                if not isinstance(processor, dict):
-                    module.set_processor(processor, _remove_lora=_remove_lora)
-                else:
-                    module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora)
-
-            for sub_name, child in module.named_children():
-                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
-
-        for name, module in self.named_children():
-            fn_recursive_attn_processor(name, module, processor)
-
-    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
-    def set_default_attn_processor(self):
-        """
-        Disables custom attention processors and sets the default attention implementation.
-        """
-        if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
-            processor = AttnAddedKVProcessor()
-        elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
-            processor = AttnProcessor()
-        else:
-            raise ValueError(
-                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
-            )
-
-        self.set_attn_processor(processor, _remove_lora=True)
-
-    @apply_forward_hook
-    def encode(self,
-               x: torch.FloatTensor,
-               return_dict: bool = True) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
-        """
-        Encode a batch of images/videos into latents.
-
-        Args:
-            x (`torch.FloatTensor`): Input batch of images/videos.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
-
-        Returns:
-                The latent representations of the encoded images/videos. If `return_dict` is True, a
-                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
-        """
-        assert len(x.shape) == 5, "The input tensor should have 5 dimensions."
-
-        if self.use_temporal_tiling and x.shape[2] > self.tile_sample_min_tsize:
-            return self.temporal_tiled_encode(x, return_dict=return_dict)
-
-        if self.use_spatial_tiling and (x.shape[-1] > self.tile_sample_min_size
-                                        or x.shape[-2] > self.tile_sample_min_size):
-            return self.spatial_tiled_encode(x, return_dict=return_dict)
-
-        if self.use_slicing and x.shape[0] > 1:
-            encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
-            h = torch.cat(encoded_slices)
-        else:
-            h = self.encoder(x)
-
-        moments = self.quant_conv(h)
-        posterior = DiagonalGaussianDistribution(moments)
-
-        if not return_dict:
-            return (posterior, )
-
-        return AutoencoderKLOutput(latent_dist=posterior)
-
-    def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
-        assert len(z.shape) == 5, "The input tensor should have 5 dimensions."
-
-        if self.use_parallel:
-            return self.parallel_tiled_decode(z, return_dict=return_dict)
-
-        if self.use_temporal_tiling and z.shape[2] > self.tile_latent_min_tsize:
-            return self.temporal_tiled_decode(z, return_dict=return_dict)
-
-        if self.use_spatial_tiling and (z.shape[-1] > self.tile_latent_min_size
-                                        or z.shape[-2] > self.tile_latent_min_size):
-            return self.spatial_tiled_decode(z, return_dict=return_dict)
-
-        z = self.post_quant_conv(z)
-        dec = self.decoder(z)
-
-        if not return_dict:
-            return (dec, )
-
-        return DecoderOutput(sample=dec)
-
-    @apply_forward_hook
-    def decode(self,
-               z: torch.FloatTensor,
-               return_dict: bool = True,
-               generator=None) -> Union[DecoderOutput, torch.FloatTensor]:
-        """
-        Decode a batch of images/videos.
-
-        Args:
-            z (`torch.FloatTensor`): Input batch of latent vectors.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
-
-        Returns:
-            [`~models.vae.DecoderOutput`] or `tuple`:
-                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
-                returned.
-
-        """
-        if self.use_slicing and z.shape[0] > 1:
-            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
-            decoded = torch.cat(decoded_slices)
-        else:
-            decoded = self._decode(z).sample
-
-        if not return_dict:
-            return (decoded, )
-
-        return DecoderOutput(sample=decoded)
-
-    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
-        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
-        for y in range(blend_extent):
-            b[:, :, :,
-              y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent)
-        return b
-
-    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
-        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
-        for x in range(blend_extent):
-            b[:, :, :, :,
-              x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent)
-        return b
-
-    def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
-        blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
-        for x in range(blend_extent):
-            b[:, :,
-              x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :,
-                                                                                       x, :, :] * (x / blend_extent)
-        return b
-
-    def spatial_tiled_encode(
-        self,
-        x: torch.FloatTensor,
-        return_dict: bool = True,
-        return_moments: bool = False,
-    ) -> AutoencoderKLOutput:
-        r"""Encode a batch of images/videos using a tiled encoder.
-
-        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
-        steps. This is useful to keep memory use constant regardless of image/videos size. The end result of tiled encoding is
-        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
-        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
-        output, but they should be much less noticeable.
-
-        Args:
-            x (`torch.FloatTensor`): Input batch of images/videos.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
-
-        Returns:
-            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
-                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
-                `tuple` is returned.
-        """
-        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
-        row_limit = self.tile_latent_min_size - blend_extent
-
-        # Split video into tiles and encode them separately.
-        rows = []
-        for i in range(0, x.shape[-2], overlap_size):
-            row = []
-            for j in range(0, x.shape[-1], overlap_size):
-                tile = x[:, :, :, i:i + self.tile_sample_min_size, j:j + self.tile_sample_min_size, ]
-                tile = self.encoder(tile)
-                tile = self.quant_conv(tile)
-                row.append(tile)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                # blend the above tile and the left tile
-                # to the current tile and add the current tile to the result row
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=-1))
-
-        moments = torch.cat(result_rows, dim=-2)
-        if return_moments:
-            return moments
-
-        posterior = DiagonalGaussianDistribution(moments)
-        if not return_dict:
-            return (posterior, )
-
-        return AutoencoderKLOutput(latent_dist=posterior)
-
-    def spatial_tiled_decode(self,
-                             z: torch.FloatTensor,
-                             return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
-        r"""
-        Decode a batch of images/videos using a tiled decoder.
-
-        Args:
-            z (`torch.FloatTensor`): Input batch of latent vectors.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
-
-        Returns:
-            [`~models.vae.DecoderOutput`] or `tuple`:
-                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
-                returned.
-        """
-        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
-        row_limit = self.tile_sample_min_size - blend_extent
-
-        # Split z into overlapping tiles and decode them separately.
-        # The tiles have an overlap to avoid seams between tiles.
-        rows = []
-        for i in range(0, z.shape[-2], overlap_size):
-            row = []
-            for j in range(0, z.shape[-1], overlap_size):
-                tile = z[:, :, :, i:i + self.tile_latent_min_size, j:j + self.tile_latent_min_size, ]
-                tile = self.post_quant_conv(tile)
-                decoded = self.decoder(tile)
-                row.append(decoded)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                # blend the above tile and the left tile
-                # to the current tile and add the current tile to the result row
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=-1))
-
-        dec = torch.cat(result_rows, dim=-2)
-        if not return_dict:
-            return (dec, )
-
-        return DecoderOutput(sample=dec)
-
-    def temporal_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput:
-
-        B, C, T, H, W = x.shape
-        overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor)
-        t_limit = self.tile_latent_min_tsize - blend_extent
-
-        # Split the video into tiles and encode them separately.
-        row = []
-        for i in range(0, T, overlap_size):
-            tile = x[:, :, i:i + self.tile_sample_min_tsize + 1, :, :]
-            if self.use_spatial_tiling and (tile.shape[-1] > self.tile_sample_min_size
-                                            or tile.shape[-2] > self.tile_sample_min_size):
-                tile = self.spatial_tiled_encode(tile, return_moments=True)
-            else:
-                tile = self.encoder(tile)
-                tile = self.quant_conv(tile)
-            if i > 0:
-                tile = tile[:, :, 1:, :, :]
-            row.append(tile)
-        result_row = []
-        for i, tile in enumerate(row):
-            if i > 0:
-                tile = self.blend_t(row[i - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :t_limit, :, :])
-            else:
-                result_row.append(tile[:, :, :t_limit + 1, :, :])
-
-        moments = torch.cat(result_row, dim=2)
-        posterior = DiagonalGaussianDistribution(moments)
-
-        if not return_dict:
-            return (posterior, )
-
-        return AutoencoderKLOutput(latent_dist=posterior)
-
-    def temporal_tiled_decode(self,
-                              z: torch.FloatTensor,
-                              return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
-        # Split z into overlapping tiles and decode them separately.
-
-        B, C, T, H, W = z.shape
-        overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor)
-        t_limit = self.tile_sample_min_tsize - blend_extent
-
-        row = []
-        for i in range(0, T, overlap_size):
-            tile = z[:, :, i:i + self.tile_latent_min_tsize + 1, :, :]
-            if self.use_spatial_tiling and (tile.shape[-1] > self.tile_latent_min_size
-                                            or tile.shape[-2] > self.tile_latent_min_size):
-                decoded = self.spatial_tiled_decode(tile, return_dict=True).sample
-            else:
-                tile = self.post_quant_conv(tile)
-                decoded = self.decoder(tile)
-            if i > 0:
-                decoded = decoded[:, :, 1:, :, :]
-            row.append(decoded)
-        result_row = []
-        for i, tile in enumerate(row):
-            if i > 0:
-                tile = self.blend_t(row[i - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :t_limit, :, :])
-            else:
-                result_row.append(tile[:, :, :t_limit + 1, :, :])
-
-        dec = torch.cat(result_row, dim=2)
-        if not return_dict:
-            return (dec, )
-
-        return DecoderOutput(sample=dec)
-
-    def _parallel_data_generator(self, gathered_results, gathered_dim_metadata):
-        global_idx = 0
-        for i, per_rank_metadata in enumerate(gathered_dim_metadata):
-            _start_shape = 0
-            for shape in per_rank_metadata:
-                mul_shape = prod(shape)
-                yield (gathered_results[i, _start_shape:_start_shape + mul_shape].reshape(shape), global_idx)
-                _start_shape += mul_shape
-                global_idx += 1
-
-    def parallel_tiled_decode(self,
-                              z: torch.FloatTensor,
-                              return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
-        """
-        Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs
-        """
-        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
-        B, C, T, H, W = z.shape
-
-        # Calculate parameters
-        t_overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor))
-        t_blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor)
-        t_limit = self.tile_sample_min_tsize - t_blend_extent
-
-        s_overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
-        s_blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
-        s_row_limit = self.tile_sample_min_size - s_blend_extent
-
-        # Calculate tile dimensions
-        num_t_tiles = (T + t_overlap_size - 1) // t_overlap_size
-        num_h_tiles = (H + s_overlap_size - 1) // s_overlap_size
-        num_w_tiles = (W + s_overlap_size - 1) // s_overlap_size
-        total_spatial_tiles = num_h_tiles * num_w_tiles
-        total_tiles = num_t_tiles * total_spatial_tiles
-
-        # Calculate tiles per rank and padding
-        tiles_per_rank = (total_tiles + world_size - 1) // world_size
-        start_tile_idx = rank * tiles_per_rank
-        end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles)
-
-        local_results = []
-        local_dim_metadata = []
-        # Process assigned tiles
-        for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)):
-            # Convert flat index to 3D indices
-            t_idx = global_idx // total_spatial_tiles
-            spatial_idx = global_idx % total_spatial_tiles
-            h_idx = spatial_idx // num_w_tiles
-            w_idx = spatial_idx % num_w_tiles
-
-            # Calculate positions
-            t_start = t_idx * t_overlap_size
-            h_start = h_idx * s_overlap_size
-            w_start = w_idx * s_overlap_size
-
-            # Extract and process tile
-            tile = z[:, :, t_start:t_start + self.tile_latent_min_tsize + 1,
-                     h_start:h_start + self.tile_latent_min_size, w_start:w_start + self.tile_latent_min_size]
-
-            # Process tile
-            tile = self.post_quant_conv(tile)
-            decoded = self.decoder(tile)
-
-            if t_start > 0:
-                decoded = decoded[:, :, 1:, :, :]
-
-            # Store metadata
-            shape = decoded.shape
-            # Store decoded data (flattened)
-            decoded_flat = decoded.reshape(-1)
-            local_results.append(decoded_flat)
-            local_dim_metadata.append(shape)
-
-        results = torch.cat(local_results, dim=0).contiguous()
-        del local_results
-        torch.cuda.empty_cache()
-        # first gather size to pad the results
-        local_size = torch.tensor([results.size(0)], device=results.device, dtype=torch.int64)
-        all_sizes = [torch.zeros(1, device=results.device, dtype=torch.int64) for _ in range(world_size)]
-        dist.all_gather(all_sizes, local_size)
-        max_size = max(size.item() for size in all_sizes)
-        padded_results = torch.zeros(max_size, device=results.device)
-        padded_results[:results.size(0)] = results
-        del results
-        torch.cuda.empty_cache()
-        # Gather all results
-        gathered_dim_metadata = [None] * world_size
-        gathered_results = torch.zeros_like(padded_results).repeat(
-            world_size, *[1] * len(padded_results.shape)).contiguous(
-            )  # use contiguous to make sure it won't copy data in the following operations
-        dist.all_gather_into_tensor(gathered_results, padded_results)
-        dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
-        # Process gathered results
-        data = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] for _ in range(num_t_tiles)]
-        for current_data, global_idx in self._parallel_data_generator(gathered_results, gathered_dim_metadata):
-            t_idx = global_idx // total_spatial_tiles
-            spatial_idx = global_idx % total_spatial_tiles
-            h_idx = spatial_idx // num_w_tiles
-            w_idx = spatial_idx % num_w_tiles
-            data[t_idx][h_idx][w_idx] = current_data
-        # Merge results
-        result_slices = []
-        last_slice_data = None
-        for i, tem_data in enumerate(data):
-            slice_data = self._merge_spatial_tiles(tem_data, s_blend_extent, s_row_limit)
-            if i > 0:
-                slice_data = self.blend_t(last_slice_data, slice_data, t_blend_extent)
-                result_slices.append(slice_data[:, :, :t_limit, :, :])
-            else:
-                result_slices.append(slice_data[:, :, :t_limit + 1, :, :])
-            last_slice_data = slice_data
-        dec = torch.cat(result_slices, dim=2)
-
-        if not return_dict:
-            return (dec, )
-        return DecoderOutput(sample=dec)
-
-    def _merge_spatial_tiles(self, spatial_rows, blend_extent, row_limit):
-        """Helper function to merge spatial tiles with blending"""
-        result_rows = []
-        for i, row in enumerate(spatial_rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                if i > 0:
-                    tile = self.blend_v(spatial_rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=-1))
-        return torch.cat(result_rows, dim=-2)
-
-    def forward(
-        self,
-        sample: torch.FloatTensor,
-        sample_posterior: bool = False,
-        return_dict: bool = True,
-        return_posterior: bool = False,
-        generator: Optional[torch.Generator] = None,
-    ) -> Union[DecoderOutput2, torch.FloatTensor]:
-        r"""
-        Args:
-            sample (`torch.FloatTensor`): Input sample.
-            sample_posterior (`bool`, *optional*, defaults to `False`):
-                Whether to sample from the posterior.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
-        """
-        x = sample
-        posterior = self.encode(x).latent_dist
-        if sample_posterior:
-            z = posterior.sample(generator=generator)
-        else:
-            z = posterior.mode()
-        dec = self.decode(z).sample
-
-        if not return_dict:
-            if return_posterior:
-                return (dec, posterior)
-            else:
-                return (dec, )
-        if return_posterior:
-            return DecoderOutput2(sample=dec, posterior=posterior)
-        else:
-            return DecoderOutput2(sample=dec)
-
-    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
-    def fuse_qkv_projections(self):
-        """
-        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query,
-        key, value) are fused. For cross-attention modules, key and value projection matrices are fused.
-
-        <Tip warning={true}>
-
-        This API is 🧪 experimental.
-
-        </Tip>
-        """
-        self.original_attn_processors = None
-
-        for _, attn_processor in self.attn_processors.items():
-            if "Added" in str(attn_processor.__class__.__name__):
-                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
-
-        self.original_attn_processors = self.attn_processors
-
-        for module in self.modules():
-            if isinstance(module, Attention):
-                module.fuse_projections(fuse=True)
-
-    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
-    def unfuse_qkv_projections(self):
-        """Disables the fused QKV projection if enabled.
-
-        <Tip warning={true}>
-
-        This API is 🧪 experimental.
-
-        </Tip>
-
-        """
-        if self.original_attn_processors is not None:
-            self.set_attn_processor(self.original_attn_processors)

fastvideo/models/hunyuan/vae/unet_causal_3d_blocks.py

@@ -1,760 +0,0 @@
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-#
-# Modified from diffusers==0.29.2
-#
-# ==============================================================================
-
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-from diffusers.models.activations import get_activation
-from diffusers.models.attention_processor import Attention, SpatialNorm
-from diffusers.models.normalization import AdaGroupNorm, RMSNorm
-from diffusers.utils import logging
-from einops import rearrange
-from torch import nn
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None):
-    seq_len = n_frame * n_hw
-    mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
-    for i in range(seq_len):
-        i_frame = i // n_hw
-        mask[i, :(i_frame + 1) * n_hw] = 0
-    if batch_size is not None:
-        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
-    return mask
-
-
-class CausalConv3d(nn.Module):
-    """
-    Implements a causal 3D convolution layer where each position only depends on previous timesteps and current spatial locations.
-    This maintains temporal causality in video generation tasks.
-    """
-
-    def __init__(
-        self,
-        chan_in,
-        chan_out,
-        kernel_size: Union[int, Tuple[int, int, int]],
-        stride: Union[int, Tuple[int, int, int]] = 1,
-        dilation: Union[int, Tuple[int, int, int]] = 1,
-        pad_mode="replicate",
-        **kwargs,
-    ):
-        super().__init__()
-
-        self.pad_mode = pad_mode
-        padding = (
-            kernel_size // 2,
-            kernel_size // 2,
-            kernel_size // 2,
-            kernel_size // 2,
-            kernel_size - 1,
-            0,
-        )  # W, H, T
-        self.time_causal_padding = padding
-
-        self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
-
-    def forward(self, x):
-        x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
-        return self.conv(x)
-
-
-class UpsampleCausal3D(nn.Module):
-    """
-    A 3D upsampling layer with an optional convolution.
-    """
-
-    def __init__(
-            self,
-            channels: int,
-            use_conv: bool = False,
-            use_conv_transpose: bool = False,
-            out_channels: Optional[int] = None,
-            name: str = "conv",
-            kernel_size: Optional[int] = None,
-            padding=1,
-            norm_type=None,
-            eps=None,
-            elementwise_affine=None,
-            bias=True,
-            interpolate=True,
-            upsample_factor=(2, 2, 2),
-    ):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.use_conv_transpose = use_conv_transpose
-        self.name = name
-        self.interpolate = interpolate
-        self.upsample_factor = upsample_factor
-
-        if norm_type == "ln_norm":
-            self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
-        elif norm_type == "rms_norm":
-            self.norm = RMSNorm(channels, eps, elementwise_affine)
-        elif norm_type is None:
-            self.norm = None
-        else:
-            raise ValueError(f"unknown norm_type: {norm_type}")
-
-        conv = None
-        if use_conv_transpose:
-            raise NotImplementedError
-        elif use_conv:
-            if kernel_size is None:
-                kernel_size = 3
-            conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, bias=bias)
-
-        if name == "conv":
-            self.conv = conv
-        else:
-            self.Conv2d_0 = conv
-
-    def forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        output_size: Optional[int] = None,
-        scale: float = 1.0,
-    ) -> torch.FloatTensor:
-        assert hidden_states.shape[1] == self.channels
-
-        if self.norm is not None:
-            raise NotImplementedError
-
-        if self.use_conv_transpose:
-            return self.conv(hidden_states)
-
-        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
-        dtype = hidden_states.dtype
-        if dtype == torch.bfloat16:
-            hidden_states = hidden_states.to(torch.float32)
-
-        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
-        if hidden_states.shape[0] >= 64:
-            hidden_states = hidden_states.contiguous()
-
-        # if `output_size` is passed we force the interpolation output
-        # size and do not make use of `scale_factor=2`
-        if self.interpolate:
-            B, C, T, H, W = hidden_states.shape
-            first_h, other_h = hidden_states.split((1, T - 1), dim=2)
-            if output_size is None:
-                if T > 1:
-                    other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest")
-
-                first_h = first_h.squeeze(2)
-                first_h = F.interpolate(first_h, scale_factor=self.upsample_factor[1:], mode="nearest")
-                first_h = first_h.unsqueeze(2)
-            else:
-                raise NotImplementedError
-
-            if T > 1:
-                hidden_states = torch.cat((first_h, other_h), dim=2)
-            else:
-                hidden_states = first_h
-
-        # If the input is bfloat16, we cast back to bfloat16
-        if dtype == torch.bfloat16:
-            hidden_states = hidden_states.to(dtype)
-
-        if self.use_conv:
-            if self.name == "conv":
-                hidden_states = self.conv(hidden_states)
-            else:
-                hidden_states = self.Conv2d_0(hidden_states)
-
-        return hidden_states
-
-
-class DownsampleCausal3D(nn.Module):
-    """
-    A 3D downsampling layer with an optional convolution.
-    """
-
-    def __init__(
-        self,
-        channels: int,
-        use_conv: bool = False,
-        out_channels: Optional[int] = None,
-        padding: int = 1,
-        name: str = "conv",
-        kernel_size=3,
-        norm_type=None,
-        eps=None,
-        elementwise_affine=None,
-        bias=True,
-        stride=2,
-    ):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.padding = padding
-        stride = stride
-        self.name = name
-
-        if norm_type == "ln_norm":
-            self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
-        elif norm_type == "rms_norm":
-            self.norm = RMSNorm(channels, eps, elementwise_affine)
-        elif norm_type is None:
-            self.norm = None
-        else:
-            raise ValueError(f"unknown norm_type: {norm_type}")
-
-        if use_conv:
-            conv = CausalConv3d(
-                self.channels,
-                self.out_channels,
-                kernel_size=kernel_size,
-                stride=stride,
-                bias=bias,
-            )
-        else:
-            raise NotImplementedError
-
-        if name == "conv":
-            self.Conv2d_0 = conv
-            self.conv = conv
-        elif name == "Conv2d_0":
-            self.conv = conv
-        else:
-            self.conv = conv
-
-    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
-        assert hidden_states.shape[1] == self.channels
-
-        if self.norm is not None:
-            hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-
-        assert hidden_states.shape[1] == self.channels
-
-        hidden_states = self.conv(hidden_states)
-
-        return hidden_states
-
-
-class ResnetBlockCausal3D(nn.Module):
-    r"""
-    A Resnet block.
-    """
-
-    def __init__(
-        self,
-        *,
-        in_channels: int,
-        out_channels: Optional[int] = None,
-        conv_shortcut: bool = False,
-        dropout: float = 0.0,
-        temb_channels: int = 512,
-        groups: int = 32,
-        groups_out: Optional[int] = None,
-        pre_norm: bool = True,
-        eps: float = 1e-6,
-        non_linearity: str = "swish",
-        skip_time_act: bool = False,
-        # default, scale_shift, ada_group, spatial
-        time_embedding_norm: str = "default",
-        kernel: Optional[torch.FloatTensor] = None,
-        output_scale_factor: float = 1.0,
-        use_in_shortcut: Optional[bool] = None,
-        up: bool = False,
-        down: bool = False,
-        conv_shortcut_bias: bool = True,
-        conv_3d_out_channels: Optional[int] = None,
-    ):
-        super().__init__()
-        self.pre_norm = pre_norm
-        self.pre_norm = True
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-        self.up = up
-        self.down = down
-        self.output_scale_factor = output_scale_factor
-        self.time_embedding_norm = time_embedding_norm
-        self.skip_time_act = skip_time_act
-
-        linear_cls = nn.Linear
-
-        if groups_out is None:
-            groups_out = groups
-
-        if self.time_embedding_norm == "ada_group":
-            self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps)
-        elif self.time_embedding_norm == "spatial":
-            self.norm1 = SpatialNorm(in_channels, temb_channels)
-        else:
-            self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
-
-        self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, stride=1)
-
-        if temb_channels is not None:
-            if self.time_embedding_norm == "default":
-                self.time_emb_proj = linear_cls(temb_channels, out_channels)
-            elif self.time_embedding_norm == "scale_shift":
-                self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels)
-            elif (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"):
-                self.time_emb_proj = None
-            else:
-                raise ValueError(f"Unknown time_embedding_norm : {self.time_embedding_norm} ")
-        else:
-            self.time_emb_proj = None
-
-        if self.time_embedding_norm == "ada_group":
-            self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps)
-        elif self.time_embedding_norm == "spatial":
-            self.norm2 = SpatialNorm(out_channels, temb_channels)
-        else:
-            self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
-
-        self.dropout = torch.nn.Dropout(dropout)
-        conv_3d_out_channels = conv_3d_out_channels or out_channels
-        self.conv2 = CausalConv3d(out_channels, conv_3d_out_channels, kernel_size=3, stride=1)
-
-        self.nonlinearity = get_activation(non_linearity)
-
-        self.upsample = self.downsample = None
-        if self.up:
-            self.upsample = UpsampleCausal3D(in_channels, use_conv=False)
-        elif self.down:
-            self.downsample = DownsampleCausal3D(in_channels, use_conv=False, name="op")
-
-        self.use_in_shortcut = (self.in_channels != conv_3d_out_channels
-                                if use_in_shortcut is None else use_in_shortcut)
-
-        self.conv_shortcut = None
-        if self.use_in_shortcut:
-            self.conv_shortcut = CausalConv3d(
-                in_channels,
-                conv_3d_out_channels,
-                kernel_size=1,
-                stride=1,
-                bias=conv_shortcut_bias,
-            )
-
-    def forward(
-        self,
-        input_tensor: torch.FloatTensor,
-        temb: torch.FloatTensor,
-        scale: float = 1.0,
-    ) -> torch.FloatTensor:
-        hidden_states = input_tensor
-
-        if (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"):
-            hidden_states = self.norm1(hidden_states, temb)
-        else:
-            hidden_states = self.norm1(hidden_states)
-
-        hidden_states = self.nonlinearity(hidden_states)
-
-        if self.upsample is not None:
-            # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
-            if hidden_states.shape[0] >= 64:
-                input_tensor = input_tensor.contiguous()
-                hidden_states = hidden_states.contiguous()
-            input_tensor = self.upsample(input_tensor, scale=scale)
-            hidden_states = self.upsample(hidden_states, scale=scale)
-        elif self.downsample is not None:
-            input_tensor = self.downsample(input_tensor, scale=scale)
-            hidden_states = self.downsample(hidden_states, scale=scale)
-
-        hidden_states = self.conv1(hidden_states)
-
-        if self.time_emb_proj is not None:
-            if not self.skip_time_act:
-                temb = self.nonlinearity(temb)
-            temb = self.time_emb_proj(temb, scale)[:, :, None, None]
-
-        if temb is not None and self.time_embedding_norm == "default":
-            hidden_states = hidden_states + temb
-
-        if (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"):
-            hidden_states = self.norm2(hidden_states, temb)
-        else:
-            hidden_states = self.norm2(hidden_states)
-
-        if temb is not None and self.time_embedding_norm == "scale_shift":
-            scale, shift = torch.chunk(temb, 2, dim=1)
-            hidden_states = hidden_states * (1 + scale) + shift
-
-        hidden_states = self.nonlinearity(hidden_states)
-
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = self.conv2(hidden_states)
-
-        if self.conv_shortcut is not None:
-            input_tensor = self.conv_shortcut(input_tensor)
-
-        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
-
-        return output_tensor
-
-
-def get_down_block3d(
-    down_block_type: str,
-    num_layers: int,
-    in_channels: int,
-    out_channels: int,
-    temb_channels: int,
-    add_downsample: bool,
-    downsample_stride: int,
-    resnet_eps: float,
-    resnet_act_fn: str,
-    transformer_layers_per_block: int = 1,
-    num_attention_heads: Optional[int] = None,
-    resnet_groups: Optional[int] = None,
-    cross_attention_dim: Optional[int] = None,
-    downsample_padding: Optional[int] = None,
-    dual_cross_attention: bool = False,
-    use_linear_projection: bool = False,
-    only_cross_attention: bool = False,
-    upcast_attention: bool = False,
-    resnet_time_scale_shift: str = "default",
-    attention_type: str = "default",
-    resnet_skip_time_act: bool = False,
-    resnet_out_scale_factor: float = 1.0,
-    cross_attention_norm: Optional[str] = None,
-    attention_head_dim: Optional[int] = None,
-    downsample_type: Optional[str] = None,
-    dropout: float = 0.0,
-):
-    # If attn head dim is not defined, we default it to the number of heads
-    if attention_head_dim is None:
-        logger.warn(
-            f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
-        )
-        attention_head_dim = num_attention_heads
-
-    down_block_type = (down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type)
-    if down_block_type == "DownEncoderBlockCausal3D":
-        return DownEncoderBlockCausal3D(
-            num_layers=num_layers,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            dropout=dropout,
-            add_downsample=add_downsample,
-            downsample_stride=downsample_stride,
-            resnet_eps=resnet_eps,
-            resnet_act_fn=resnet_act_fn,
-            resnet_groups=resnet_groups,
-            downsample_padding=downsample_padding,
-            resnet_time_scale_shift=resnet_time_scale_shift,
-        )
-    raise ValueError(f"{down_block_type} does not exist.")
-
-
-def get_up_block3d(
-    up_block_type: str,
-    num_layers: int,
-    in_channels: int,
-    out_channels: int,
-    prev_output_channel: int,
-    temb_channels: int,
-    add_upsample: bool,
-    upsample_scale_factor: Tuple,
-    resnet_eps: float,
-    resnet_act_fn: str,
-    resolution_idx: Optional[int] = None,
-    transformer_layers_per_block: int = 1,
-    num_attention_heads: Optional[int] = None,
-    resnet_groups: Optional[int] = None,
-    cross_attention_dim: Optional[int] = None,
-    dual_cross_attention: bool = False,
-    use_linear_projection: bool = False,
-    only_cross_attention: bool = False,
-    upcast_attention: bool = False,
-    resnet_time_scale_shift: str = "default",
-    attention_type: str = "default",
-    resnet_skip_time_act: bool = False,
-    resnet_out_scale_factor: float = 1.0,
-    cross_attention_norm: Optional[str] = None,
-    attention_head_dim: Optional[int] = None,
-    upsample_type: Optional[str] = None,
-    dropout: float = 0.0,
-) -> nn.Module:
-    # If attn head dim is not defined, we default it to the number of heads
-    if attention_head_dim is None:
-        logger.warn(
-            f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
-        )
-        attention_head_dim = num_attention_heads
-
-    up_block_type = (up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type)
-    if up_block_type == "UpDecoderBlockCausal3D":
-        return UpDecoderBlockCausal3D(
-            num_layers=num_layers,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            resolution_idx=resolution_idx,
-            dropout=dropout,
-            add_upsample=add_upsample,
-            upsample_scale_factor=upsample_scale_factor,
-            resnet_eps=resnet_eps,
-            resnet_act_fn=resnet_act_fn,
-            resnet_groups=resnet_groups,
-            resnet_time_scale_shift=resnet_time_scale_shift,
-            temb_channels=temb_channels,
-        )
-    raise ValueError(f"{up_block_type} does not exist.")
-
-
-class UNetMidBlockCausal3D(nn.Module):
-    """
-    A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks.
-    """
-
-    def __init__(
-        self,
-        in_channels: int,
-        temb_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        resnet_eps: float = 1e-6,
-        resnet_time_scale_shift: str = "default",  # default, spatial
-        resnet_act_fn: str = "swish",
-        resnet_groups: int = 32,
-        attn_groups: Optional[int] = None,
-        resnet_pre_norm: bool = True,
-        add_attention: bool = True,
-        attention_head_dim: int = 1,
-        output_scale_factor: float = 1.0,
-    ):
-        super().__init__()
-        resnet_groups = (resnet_groups if resnet_groups is not None else min(in_channels // 4, 32))
-        self.add_attention = add_attention
-
-        if attn_groups is None:
-            attn_groups = (resnet_groups if resnet_time_scale_shift == "default" else None)
-
-        # there is always at least one resnet
-        resnets = [
-            ResnetBlockCausal3D(
-                in_channels=in_channels,
-                out_channels=in_channels,
-                temb_channels=temb_channels,
-                eps=resnet_eps,
-                groups=resnet_groups,
-                dropout=dropout,
-                time_embedding_norm=resnet_time_scale_shift,
-                non_linearity=resnet_act_fn,
-                output_scale_factor=output_scale_factor,
-                pre_norm=resnet_pre_norm,
-            )
-        ]
-        attentions = []
-
-        if attention_head_dim is None:
-            logger.warn(
-                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
-            )
-            attention_head_dim = in_channels
-
-        for _ in range(num_layers):
-            if self.add_attention:
-                attentions.append(
-                    Attention(
-                        in_channels,
-                        heads=in_channels // attention_head_dim,
-                        dim_head=attention_head_dim,
-                        rescale_output_factor=output_scale_factor,
-                        eps=resnet_eps,
-                        norm_num_groups=attn_groups,
-                        spatial_norm_dim=(temb_channels if resnet_time_scale_shift == "spatial" else None),
-                        residual_connection=True,
-                        bias=True,
-                        upcast_softmax=True,
-                        _from_deprecated_attn_block=True,
-                    ))
-            else:
-                attentions.append(None)
-
-            resnets.append(
-                ResnetBlockCausal3D(
-                    in_channels=in_channels,
-                    out_channels=in_channels,
-                    temb_channels=temb_channels,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                ))
-
-        self.attentions = nn.ModuleList(attentions)
-        self.resnets = nn.ModuleList(resnets)
-
-    def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
-        hidden_states = self.resnets[0](hidden_states, temb)
-        for attn, resnet in zip(self.attentions, self.resnets[1:]):
-            if attn is not None:
-                B, C, T, H, W = hidden_states.shape
-                hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c")
-                attention_mask = prepare_causal_attention_mask(T,
-                                                               H * W,
-                                                               hidden_states.dtype,
-                                                               hidden_states.device,
-                                                               batch_size=B)
-                hidden_states = attn(hidden_states, temb=temb, attention_mask=attention_mask)
-                hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W)
-            hidden_states = resnet(hidden_states, temb)
-
-        return hidden_states
-
-
-class DownEncoderBlockCausal3D(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        resnet_eps: float = 1e-6,
-        resnet_time_scale_shift: str = "default",
-        resnet_act_fn: str = "swish",
-        resnet_groups: int = 32,
-        resnet_pre_norm: bool = True,
-        output_scale_factor: float = 1.0,
-        add_downsample: bool = True,
-        downsample_stride: int = 2,
-        downsample_padding: int = 1,
-    ):
-        super().__init__()
-        resnets = []
-
-        for i in range(num_layers):
-            in_channels = in_channels if i == 0 else out_channels
-            resnets.append(
-                ResnetBlockCausal3D(
-                    in_channels=in_channels,
-                    out_channels=out_channels,
-                    temb_channels=None,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                ))
-
-        self.resnets = nn.ModuleList(resnets)
-
-        if add_downsample:
-            self.downsamplers = nn.ModuleList([
-                DownsampleCausal3D(
-                    out_channels,
-                    use_conv=True,
-                    out_channels=out_channels,
-                    padding=downsample_padding,
-                    name="op",
-                    stride=downsample_stride,
-                )
-            ])
-        else:
-            self.downsamplers = None
-
-    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
-        for resnet in self.resnets:
-            hidden_states = resnet(hidden_states, temb=None, scale=scale)
-
-        if self.downsamplers is not None:
-            for downsampler in self.downsamplers:
-                hidden_states = downsampler(hidden_states, scale)
-
-        return hidden_states
-
-
-class UpDecoderBlockCausal3D(nn.Module):
-
-    def __init__(
-            self,
-            in_channels: int,
-            out_channels: int,
-            resolution_idx: Optional[int] = None,
-            dropout: float = 0.0,
-            num_layers: int = 1,
-            resnet_eps: float = 1e-6,
-            resnet_time_scale_shift: str = "default",  # default, spatial
-            resnet_act_fn: str = "swish",
-            resnet_groups: int = 32,
-            resnet_pre_norm: bool = True,
-            output_scale_factor: float = 1.0,
-            add_upsample: bool = True,
-            upsample_scale_factor=(2, 2, 2),
-            temb_channels: Optional[int] = None,
-    ):
-        super().__init__()
-        resnets = []
-
-        for i in range(num_layers):
-            input_channels = in_channels if i == 0 else out_channels
-
-            resnets.append(
-                ResnetBlockCausal3D(
-                    in_channels=input_channels,
-                    out_channels=out_channels,
-                    temb_channels=temb_channels,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                ))
-
-        self.resnets = nn.ModuleList(resnets)
-
-        if add_upsample:
-            self.upsamplers = nn.ModuleList([
-                UpsampleCausal3D(
-                    out_channels,
-                    use_conv=True,
-                    out_channels=out_channels,
-                    upsample_factor=upsample_scale_factor,
-                )
-            ])
-        else:
-            self.upsamplers = None
-
-        self.resolution_idx = resolution_idx
-
-    def forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        temb: Optional[torch.FloatTensor] = None,
-        scale: float = 1.0,
-    ) -> torch.FloatTensor:
-        for resnet in self.resnets:
-            hidden_states = resnet(hidden_states, temb=temb, scale=scale)
-
-        if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states)
-
-        return hidden_states

fastvideo/models/hunyuan/vae/vae.py

@@ -1,342 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import numpy as np
-import torch
-import torch.nn as nn
-from diffusers.models.attention_processor import SpatialNorm
-from diffusers.utils import BaseOutput, is_torch_version
-from diffusers.utils.torch_utils import randn_tensor
-
-from .unet_causal_3d_blocks import CausalConv3d, UNetMidBlockCausal3D, get_down_block3d, get_up_block3d
-
-
-@dataclass
-class DecoderOutput(BaseOutput):
-    r"""
-    Output of decoding method.
-
-    Args:
-        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
-            The decoded output sample from the last layer of the model.
-    """
-
-    sample: torch.FloatTensor
-
-
-class EncoderCausal3D(nn.Module):
-    r"""
-    The `EncoderCausal3D` layer of a variational autoencoder that encodes its input into a latent representation.
-    """
-
-    def __init__(
-        self,
-        in_channels: int = 3,
-        out_channels: int = 3,
-        down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D", ),
-        block_out_channels: Tuple[int, ...] = (64, ),
-        layers_per_block: int = 2,
-        norm_num_groups: int = 32,
-        act_fn: str = "silu",
-        double_z: bool = True,
-        mid_block_add_attention=True,
-        time_compression_ratio: int = 4,
-        spatial_compression_ratio: int = 8,
-    ):
-        super().__init__()
-        self.layers_per_block = layers_per_block
-
-        self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1)
-        self.mid_block = None
-        self.down_blocks = nn.ModuleList([])
-
-        # down
-        output_channel = block_out_channels[0]
-        for i, down_block_type in enumerate(down_block_types):
-            input_channel = output_channel
-            output_channel = block_out_channels[i]
-            is_final_block = i == len(block_out_channels) - 1
-            num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio))
-            num_time_downsample_layers = int(np.log2(time_compression_ratio))
-
-            if time_compression_ratio == 4:
-                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
-                add_time_downsample = bool(i >= (len(block_out_channels) - 1 - num_time_downsample_layers)
-                                           and not is_final_block)
-            else:
-                raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.")
-
-            downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
-            downsample_stride_T = (2, ) if add_time_downsample else (1, )
-            downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
-            down_block = get_down_block3d(
-                down_block_type,
-                num_layers=self.layers_per_block,
-                in_channels=input_channel,
-                out_channels=output_channel,
-                add_downsample=bool(add_spatial_downsample or add_time_downsample),
-                downsample_stride=downsample_stride,
-                resnet_eps=1e-6,
-                downsample_padding=0,
-                resnet_act_fn=act_fn,
-                resnet_groups=norm_num_groups,
-                attention_head_dim=output_channel,
-                temb_channels=None,
-            )
-            self.down_blocks.append(down_block)
-
-        # mid
-        self.mid_block = UNetMidBlockCausal3D(
-            in_channels=block_out_channels[-1],
-            resnet_eps=1e-6,
-            resnet_act_fn=act_fn,
-            output_scale_factor=1,
-            resnet_time_scale_shift="default",
-            attention_head_dim=block_out_channels[-1],
-            resnet_groups=norm_num_groups,
-            temb_channels=None,
-            add_attention=mid_block_add_attention,
-        )
-
-        # out
-        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
-        self.conv_act = nn.SiLU()
-
-        conv_out_channels = 2 * out_channels if double_z else out_channels
-        self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3)
-
-    def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor:
-        r"""The forward method of the `EncoderCausal3D` class."""
-        assert len(sample.shape) == 5, "The input tensor should have 5 dimensions"
-
-        sample = self.conv_in(sample)
-
-        # down
-        for down_block in self.down_blocks:
-            sample = down_block(sample)
-
-        # middle
-        sample = self.mid_block(sample)
-
-        # post-process
-        sample = self.conv_norm_out(sample)
-        sample = self.conv_act(sample)
-        sample = self.conv_out(sample)
-
-        return sample
-
-
-class DecoderCausal3D(nn.Module):
-    r"""
-    The `DecoderCausal3D` layer of a variational autoencoder that decodes its latent representation into an output sample.
-    """
-
-    def __init__(
-        self,
-        in_channels: int = 3,
-        out_channels: int = 3,
-        up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D", ),
-        block_out_channels: Tuple[int, ...] = (64, ),
-        layers_per_block: int = 2,
-        norm_num_groups: int = 32,
-        act_fn: str = "silu",
-        norm_type: str = "group",  # group, spatial
-        mid_block_add_attention=True,
-        time_compression_ratio: int = 4,
-        spatial_compression_ratio: int = 8,
-    ):
-        super().__init__()
-        self.layers_per_block = layers_per_block
-
-        self.conv_in = CausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1)
-        self.mid_block = None
-        self.up_blocks = nn.ModuleList([])
-
-        temb_channels = in_channels if norm_type == "spatial" else None
-
-        # mid
-        self.mid_block = UNetMidBlockCausal3D(
-            in_channels=block_out_channels[-1],
-            resnet_eps=1e-6,
-            resnet_act_fn=act_fn,
-            output_scale_factor=1,
-            resnet_time_scale_shift="default" if norm_type == "group" else norm_type,
-            attention_head_dim=block_out_channels[-1],
-            resnet_groups=norm_num_groups,
-            temb_channels=temb_channels,
-            add_attention=mid_block_add_attention,
-        )
-
-        # up
-        reversed_block_out_channels = list(reversed(block_out_channels))
-        output_channel = reversed_block_out_channels[0]
-        for i, up_block_type in enumerate(up_block_types):
-            prev_output_channel = output_channel
-            output_channel = reversed_block_out_channels[i]
-            is_final_block = i == len(block_out_channels) - 1
-            num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio))
-            num_time_upsample_layers = int(np.log2(time_compression_ratio))
-
-            if time_compression_ratio == 4:
-                add_spatial_upsample = bool(i < num_spatial_upsample_layers)
-                add_time_upsample = bool(i >= len(block_out_channels) - 1 - num_time_upsample_layers
-                                         and not is_final_block)
-            else:
-                raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.")
-
-            upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1)
-            upsample_scale_factor_T = (2, ) if add_time_upsample else (1, )
-            upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW)
-            up_block = get_up_block3d(
-                up_block_type,
-                num_layers=self.layers_per_block + 1,
-                in_channels=prev_output_channel,
-                out_channels=output_channel,
-                prev_output_channel=None,
-                add_upsample=bool(add_spatial_upsample or add_time_upsample),
-                upsample_scale_factor=upsample_scale_factor,
-                resnet_eps=1e-6,
-                resnet_act_fn=act_fn,
-                resnet_groups=norm_num_groups,
-                attention_head_dim=output_channel,
-                temb_channels=temb_channels,
-                resnet_time_scale_shift=norm_type,
-            )
-            self.up_blocks.append(up_block)
-            prev_output_channel = output_channel
-
-        # out
-        if norm_type == "spatial":
-            self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels)
-        else:
-            self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
-        self.conv_act = nn.SiLU()
-        self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3)
-
-        self.gradient_checkpointing = False
-
-    def forward(
-        self,
-        sample: torch.FloatTensor,
-        latent_embeds: Optional[torch.FloatTensor] = None,
-    ) -> torch.FloatTensor:
-        r"""The forward method of the `DecoderCausal3D` class."""
-        assert len(sample.shape) == 5, "The input tensor should have 5 dimensions."
-
-        sample = self.conv_in(sample)
-
-        upscale_dtype = next(iter(self.up_blocks.parameters())).dtype
-        if self.training and self.gradient_checkpointing:
-
-            def create_custom_forward(module):
-
-                def custom_forward(*inputs):
-                    return module(*inputs)
-
-                return custom_forward
-
-            if is_torch_version(">=", "1.11.0"):
-                # middle
-                sample = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(self.mid_block),
-                    sample,
-                    latent_embeds,
-                    use_reentrant=False,
-                )
-                sample = sample.to(upscale_dtype)
-
-                # up
-                for up_block in self.up_blocks:
-                    sample = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(up_block),
-                        sample,
-                        latent_embeds,
-                        use_reentrant=False,
-                    )
-            else:
-                # middle
-                sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample, latent_embeds)
-                sample = sample.to(upscale_dtype)
-
-                # up
-                for up_block in self.up_blocks:
-                    sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds)
-        else:
-            # middle
-            sample = self.mid_block(sample, latent_embeds)
-            sample = sample.to(upscale_dtype)
-
-            # up
-            for up_block in self.up_blocks:
-                sample = up_block(sample, latent_embeds)
-
-        # post-process
-        if latent_embeds is None:
-            sample = self.conv_norm_out(sample)
-        else:
-            sample = self.conv_norm_out(sample, latent_embeds)
-        sample = self.conv_act(sample)
-        sample = self.conv_out(sample)
-
-        return sample
-
-
-class DiagonalGaussianDistribution(object):
-
-    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
-        if parameters.ndim == 3:
-            dim = 2  # (B, L, C)
-        elif parameters.ndim == 5 or parameters.ndim == 4:
-            dim = 1  # (B, C, T, H ,W) / (B, C, H, W)
-        else:
-            raise NotImplementedError
-        self.parameters = parameters
-        self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim)
-        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
-        self.deterministic = deterministic
-        self.std = torch.exp(0.5 * self.logvar)
-        self.var = torch.exp(self.logvar)
-        if self.deterministic:
-            self.var = self.std = torch.zeros_like(self.mean,
-                                                   device=self.parameters.device,
-                                                   dtype=self.parameters.dtype)
-
-    def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor:
-        # make sure sample is on the same device as the parameters and has same dtype
-        sample = randn_tensor(
-            self.mean.shape,
-            generator=generator,
-            device=self.parameters.device,
-            dtype=self.parameters.dtype,
-        )
-        x = self.mean + self.std * sample
-        return x
-
-    def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
-        if self.deterministic:
-            return torch.Tensor([0.0])
-        else:
-            reduce_dim = list(range(1, self.mean.ndim))
-            if other is None:
-                return 0.5 * torch.sum(
-                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
-                    dim=reduce_dim,
-                )
-            else:
-                return 0.5 * torch.sum(
-                    torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar +
-                    other.logvar,
-                    dim=reduce_dim,
-                )
-
-    def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor:
-        if self.deterministic:
-            return torch.Tensor([0.0])
-        logtwopi = np.log(2.0 * np.pi)
-        return 0.5 * torch.sum(
-            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
-            dim=dims,
-        )
-
-    def mode(self) -> torch.Tensor:
-        return self.mean

fastvideo/models/hunyuan_hf/modeling_hunyuan.py

@@ -1,836 +0,0 @@
-# Copyright 2024 The Hunyuan Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
-from diffusers.models.attention import FeedForward
-from diffusers.models.attention_processor import Attention, AttentionProcessor
-from diffusers.models.embeddings import (CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings,
-                                         get_1d_rotary_pos_embed)
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
-from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
-
-from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
-from fastvideo.utils.communications import all_gather, all_to_all_4D
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-def shrink_head(encoder_state, dim):
-    local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-    return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
-
-
-class HunyuanVideoAttnProcessor2_0:
-
-    def __init__(self):
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError(
-                "HunyuanVideoAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        image_rotary_emb: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-
-        sequence_length = hidden_states.size(1)
-        encoder_sequence_length = encoder_hidden_states.size(1)
-        if attn.add_q_proj is None and encoder_hidden_states is not None:
-            hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
-
-        # 1. QKV projections
-        query = attn.to_q(hidden_states)
-        key = attn.to_k(hidden_states)
-        value = attn.to_v(hidden_states)
-
-        query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-        key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-        value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-
-        # 2. QK normalization
-        if attn.norm_q is not None:
-            query = attn.norm_q(query).to(value)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key).to(value)
-
-        image_rotary_emb = (
-            shrink_head(image_rotary_emb[0], dim=0),
-            shrink_head(image_rotary_emb[1], dim=0),
-        )
-
-        # 3. Rotational positional embeddings applied to latent stream
-        if image_rotary_emb is not None:
-            from diffusers.models.embeddings import apply_rotary_emb
-
-            if attn.add_q_proj is None and encoder_hidden_states is not None:
-                query = torch.cat(
-                    [
-                        apply_rotary_emb(query[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb),
-                        query[:, :, -encoder_hidden_states.shape[1]:],
-                    ],
-                    dim=2,
-                )
-                key = torch.cat(
-                    [
-                        apply_rotary_emb(key[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb),
-                        key[:, :, -encoder_hidden_states.shape[1]:],
-                    ],
-                    dim=2,
-                )
-            else:
-                query = apply_rotary_emb(query, image_rotary_emb)
-                key = apply_rotary_emb(key, image_rotary_emb)
-
-        # 4. Encoder condition QKV projection and normalization
-        if attn.add_q_proj is not None and encoder_hidden_states is not None:
-            encoder_query = attn.add_q_proj(encoder_hidden_states)
-            encoder_key = attn.add_k_proj(encoder_hidden_states)
-            encoder_value = attn.add_v_proj(encoder_hidden_states)
-
-            encoder_query = encoder_query.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-            encoder_key = encoder_key.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-            encoder_value = encoder_value.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-
-            if attn.norm_added_q is not None:
-                encoder_query = attn.norm_added_q(encoder_query).to(encoder_value)
-            if attn.norm_added_k is not None:
-                encoder_key = attn.norm_added_k(encoder_key).to(encoder_value)
-
-            query = torch.cat([query, encoder_query], dim=2)
-            key = torch.cat([key, encoder_key], dim=2)
-            value = torch.cat([value, encoder_value], dim=2)
-
-        if get_sequence_parallel_state():
-            query_img, query_txt = query[:, :, :sequence_length, :], query[:, :, sequence_length:, :]
-            key_img, key_txt = key[:, :, :sequence_length, :], key[:, :, sequence_length:, :]
-            value_img, value_txt = value[:, :, :sequence_length, :], value[:, :, sequence_length:, :]
-            query_img = all_to_all_4D(query_img, scatter_dim=1, gather_dim=2)  #
-            key_img = all_to_all_4D(key_img, scatter_dim=1, gather_dim=2)
-            value_img = all_to_all_4D(value_img, scatter_dim=1, gather_dim=2)
-
-            query_txt = shrink_head(query_txt, dim=1)
-            key_txt = shrink_head(key_txt, dim=1)
-            value_txt = shrink_head(value_txt, dim=1)
-            query = torch.cat([query_img, query_txt], dim=2)
-            key = torch.cat([key_img, key_txt], dim=2)
-            value = torch.cat([value_img, value_txt], dim=2)
-
-        query = query.unsqueeze(2)
-        key = key.unsqueeze(2)
-        value = value.unsqueeze(2)
-        qkv = torch.cat([query, key, value], dim=2)
-        qkv = qkv.transpose(1, 3)
-
-        # 5. Attention
-        attention_mask = attention_mask[:, 0, :]
-        seq_len = qkv.shape[1]
-        attn_len = attention_mask.shape[1]
-        attention_mask = F.pad(attention_mask, (seq_len - attn_len, 0), value=True)
-
-        hidden_states = flash_attn_no_pad(qkv, attention_mask, causal=False, dropout_p=0.0, softmax_scale=None)
-
-        if get_sequence_parallel_state():
-            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
-                (sequence_length * nccl_info.sp_size, encoder_sequence_length), dim=1)
-            hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)
-            encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous()
-            hidden_states = hidden_states.flatten(2, 3)
-            hidden_states = hidden_states.to(query.dtype)
-            encoder_hidden_states = encoder_hidden_states.flatten(2, 3)
-            encoder_hidden_states = encoder_hidden_states.to(query.dtype)
-        else:
-            hidden_states = hidden_states.flatten(2, 3)
-            hidden_states = hidden_states.to(query.dtype)
-
-            # 6. Output projection
-            if encoder_hidden_states is not None:
-                hidden_states, encoder_hidden_states = (
-                    hidden_states[:, :-encoder_hidden_states.shape[1]],
-                    hidden_states[:, -encoder_hidden_states.shape[1]:],
-                )
-
-        if encoder_hidden_states is not None:
-            if getattr(attn, "to_out", None) is not None:
-                hidden_states = attn.to_out[0](hidden_states)
-                hidden_states = attn.to_out[1](hidden_states)
-
-            if getattr(attn, "to_add_out", None) is not None:
-                encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
-
-        return hidden_states, encoder_hidden_states
-
-
-class HunyuanVideoPatchEmbed(nn.Module):
-
-    def __init__(
-        self,
-        patch_size: Union[int, Tuple[int, int, int]] = 16,
-        in_chans: int = 3,
-        embed_dim: int = 768,
-    ) -> None:
-        super().__init__()
-
-        patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size
-        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.proj(hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)  # BCFHW -> BNC
-        return hidden_states
-
-
-class HunyuanVideoAdaNorm(nn.Module):
-
-    def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
-        super().__init__()
-
-        out_features = out_features or 2 * in_features
-        self.linear = nn.Linear(in_features, out_features)
-        self.nonlinearity = nn.SiLU()
-
-    def forward(self,
-                temb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        temb = self.linear(self.nonlinearity(temb))
-        gate_msa, gate_mlp = temb.chunk(2, dim=1)
-        gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
-        return gate_msa, gate_mlp
-
-
-class HunyuanVideoIndividualTokenRefinerBlock(nn.Module):
-
-    def __init__(
-        self,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        mlp_width_ratio: str = 4.0,
-        mlp_drop_rate: float = 0.0,
-        attention_bias: bool = True,
-    ) -> None:
-        super().__init__()
-
-        hidden_size = num_attention_heads * attention_head_dim
-
-        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
-        self.attn = Attention(
-            query_dim=hidden_size,
-            cross_attention_dim=None,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            bias=attention_bias,
-        )
-
-        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
-        self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate)
-
-        self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        norm_hidden_states = self.norm1(hidden_states)
-
-        attn_output = self.attn(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=None,
-            attention_mask=attention_mask,
-        )
-
-        gate_msa, gate_mlp = self.norm_out(temb)
-        hidden_states = hidden_states + attn_output * gate_msa
-
-        ff_output = self.ff(self.norm2(hidden_states))
-        hidden_states = hidden_states + ff_output * gate_mlp
-
-        return hidden_states
-
-
-class HunyuanVideoIndividualTokenRefiner(nn.Module):
-
-    def __init__(
-        self,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        num_layers: int,
-        mlp_width_ratio: float = 4.0,
-        mlp_drop_rate: float = 0.0,
-        attention_bias: bool = True,
-    ) -> None:
-        super().__init__()
-
-        self.refiner_blocks = nn.ModuleList([
-            HunyuanVideoIndividualTokenRefinerBlock(
-                num_attention_heads=num_attention_heads,
-                attention_head_dim=attention_head_dim,
-                mlp_width_ratio=mlp_width_ratio,
-                mlp_drop_rate=mlp_drop_rate,
-                attention_bias=attention_bias,
-            ) for _ in range(num_layers)
-        ])
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> None:
-        self_attn_mask = None
-        if attention_mask is not None:
-            batch_size = attention_mask.shape[0]
-            seq_len = attention_mask.shape[1]
-            attention_mask = attention_mask.to(hidden_states.device).bool()
-            self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
-            self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
-            self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
-            self_attn_mask[:, :, :, 0] = True
-
-        for block in self.refiner_blocks:
-            hidden_states = block(hidden_states, temb, self_attn_mask)
-
-        return hidden_states
-
-
-class HunyuanVideoTokenRefiner(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        num_layers: int,
-        mlp_ratio: float = 4.0,
-        mlp_drop_rate: float = 0.0,
-        attention_bias: bool = True,
-    ) -> None:
-        super().__init__()
-
-        hidden_size = num_attention_heads * attention_head_dim
-
-        self.time_text_embed = CombinedTimestepTextProjEmbeddings(embedding_dim=hidden_size,
-                                                                  pooled_projection_dim=in_channels)
-        self.proj_in = nn.Linear(in_channels, hidden_size, bias=True)
-        self.token_refiner = HunyuanVideoIndividualTokenRefiner(
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            num_layers=num_layers,
-            mlp_width_ratio=mlp_ratio,
-            mlp_drop_rate=mlp_drop_rate,
-            attention_bias=attention_bias,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        attention_mask: Optional[torch.LongTensor] = None,
-    ) -> torch.Tensor:
-        if attention_mask is None:
-            pooled_projections = hidden_states.mean(dim=1)
-        else:
-            original_dtype = hidden_states.dtype
-            mask_float = attention_mask.float().unsqueeze(-1)
-            pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
-            pooled_projections = pooled_projections.to(original_dtype)
-
-        temb = self.time_text_embed(timestep, pooled_projections)
-        hidden_states = self.proj_in(hidden_states)
-        hidden_states = self.token_refiner(hidden_states, temb, attention_mask)
-
-        return hidden_states
-
-
-class HunyuanVideoRotaryPosEmbed(nn.Module):
-
-    def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None:
-        super().__init__()
-
-        self.patch_size = patch_size
-        self.patch_size_t = patch_size_t
-        self.rope_dim = rope_dim
-        self.theta = theta
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
-        rope_sizes = [
-            num_frames * nccl_info.sp_size // self.patch_size_t, height // self.patch_size, width // self.patch_size
-        ]
-
-        axes_grids = []
-        for i in range(3):
-            # Note: The following line diverges from original behaviour. We create the grid on the device, whereas
-            # original implementation creates it on CPU and then moves it to device. This results in numerical
-            # differences in layerwise debugging outputs, but visually it is the same.
-            grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32)
-            axes_grids.append(grid)
-        grid = torch.meshgrid(*axes_grids, indexing="ij")  # [W, H, T]
-        grid = torch.stack(grid, dim=0)  # [3, W, H, T]
-
-        freqs = []
-        for i in range(3):
-            freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True)
-            freqs.append(freq)
-
-        freqs_cos = torch.cat([f[0] for f in freqs], dim=1)  # (W * H * T, D / 2)
-        freqs_sin = torch.cat([f[1] for f in freqs], dim=1)  # (W * H * T, D / 2)
-        return freqs_cos, freqs_sin
-
-
-class HunyuanVideoSingleTransformerBlock(nn.Module):
-
-    def __init__(
-        self,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        mlp_ratio: float = 4.0,
-        qk_norm: str = "rms_norm",
-    ) -> None:
-        super().__init__()
-
-        hidden_size = num_attention_heads * attention_head_dim
-        mlp_dim = int(hidden_size * mlp_ratio)
-
-        self.attn = Attention(
-            query_dim=hidden_size,
-            cross_attention_dim=None,
-            dim_head=attention_head_dim,
-            heads=num_attention_heads,
-            out_dim=hidden_size,
-            bias=True,
-            processor=HunyuanVideoAttnProcessor2_0(),
-            qk_norm=qk_norm,
-            eps=1e-6,
-            pre_only=True,
-        )
-
-        self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm")
-        self.proj_mlp = nn.Linear(hidden_size, mlp_dim)
-        self.act_mlp = nn.GELU(approximate="tanh")
-        self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-    ) -> torch.Tensor:
-        text_seq_length = encoder_hidden_states.shape[1]
-        hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
-
-        residual = hidden_states
-
-        # 1. Input normalization
-        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
-        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
-
-        norm_hidden_states, norm_encoder_hidden_states = (
-            norm_hidden_states[:, :-text_seq_length, :],
-            norm_hidden_states[:, -text_seq_length:, :],
-        )
-
-        # 2. Attention
-        attn_output, context_attn_output = self.attn(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            attention_mask=attention_mask,
-            image_rotary_emb=image_rotary_emb,
-        )
-        attn_output = torch.cat([attn_output, context_attn_output], dim=1)
-
-        # 3. Modulation and residual connection
-        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
-        hidden_states = gate.unsqueeze(1) * self.proj_out(hidden_states)
-        hidden_states = hidden_states + residual
-
-        hidden_states, encoder_hidden_states = (
-            hidden_states[:, :-text_seq_length, :],
-            hidden_states[:, -text_seq_length:, :],
-        )
-        return hidden_states, encoder_hidden_states
-
-
-class HunyuanVideoTransformerBlock(nn.Module):
-
-    def __init__(
-        self,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        mlp_ratio: float,
-        qk_norm: str = "rms_norm",
-    ) -> None:
-        super().__init__()
-
-        hidden_size = num_attention_heads * attention_head_dim
-
-        self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
-        self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
-
-        self.attn = Attention(
-            query_dim=hidden_size,
-            cross_attention_dim=None,
-            added_kv_proj_dim=hidden_size,
-            dim_head=attention_head_dim,
-            heads=num_attention_heads,
-            out_dim=hidden_size,
-            context_pre_only=False,
-            bias=True,
-            processor=HunyuanVideoAttnProcessor2_0(),
-            qk_norm=qk_norm,
-            eps=1e-6,
-        )
-
-        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
-
-        self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # 1. Input normalization
-        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
-        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
-            encoder_hidden_states, emb=temb)
-
-        # 2. Joint attention
-        attn_output, context_attn_output = self.attn(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            attention_mask=attention_mask,
-            image_rotary_emb=freqs_cis,
-        )
-
-        # 3. Modulation and residual connection
-        hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1)
-        encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1)
-
-        norm_hidden_states = self.norm2(hidden_states)
-        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
-
-        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
-
-        # 4. Feed-forward
-        ff_output = self.ff(norm_hidden_states)
-        context_ff_output = self.ff_context(norm_encoder_hidden_states)
-
-        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
-        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
-
-        return hidden_states, encoder_hidden_states
-
-
-class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
-    r"""
-    A Transformer model for video-like data used in [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo).
-
-    Args:
-        in_channels (`int`, defaults to `16`):
-            The number of channels in the input.
-        out_channels (`int`, defaults to `16`):
-            The number of channels in the output.
-        num_attention_heads (`int`, defaults to `24`):
-            The number of heads to use for multi-head attention.
-        attention_head_dim (`int`, defaults to `128`):
-            The number of channels in each head.
-        num_layers (`int`, defaults to `20`):
-            The number of layers of dual-stream blocks to use.
-        num_single_layers (`int`, defaults to `40`):
-            The number of layers of single-stream blocks to use.
-        num_refiner_layers (`int`, defaults to `2`):
-            The number of layers of refiner blocks to use.
-        mlp_ratio (`float`, defaults to `4.0`):
-            The ratio of the hidden layer size to the input size in the feedforward network.
-        patch_size (`int`, defaults to `2`):
-            The size of the spatial patches to use in the patch embedding layer.
-        patch_size_t (`int`, defaults to `1`):
-            The size of the tmeporal patches to use in the patch embedding layer.
-        qk_norm (`str`, defaults to `rms_norm`):
-            The normalization to use for the query and key projections in the attention layers.
-        guidance_embeds (`bool`, defaults to `True`):
-            Whether to use guidance embeddings in the model.
-        text_embed_dim (`int`, defaults to `4096`):
-            Input dimension of text embeddings from the text encoder.
-        pooled_projection_dim (`int`, defaults to `768`):
-            The dimension of the pooled projection of the text embeddings.
-        rope_theta (`float`, defaults to `256.0`):
-            The value of theta to use in the RoPE layer.
-        rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`):
-            The dimensions of the axes to use in the RoPE layer.
-    """
-
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-            self,
-            in_channels: int = 16,
-            out_channels: int = 16,
-            num_attention_heads: int = 24,
-            attention_head_dim: int = 128,
-            num_layers: int = 20,
-            num_single_layers: int = 40,
-            num_refiner_layers: int = 2,
-            mlp_ratio: float = 4.0,
-            patch_size: int = 2,
-            patch_size_t: int = 1,
-            qk_norm: str = "rms_norm",
-            guidance_embeds: bool = True,
-            text_embed_dim: int = 4096,
-            pooled_projection_dim: int = 768,
-            rope_theta: float = 256.0,
-            rope_axes_dim: Tuple[int] = (16, 56, 56),
-    ) -> None:
-        super().__init__()
-
-        inner_dim = num_attention_heads * attention_head_dim
-        out_channels = out_channels or in_channels
-
-        # 1. Latent and condition embedders
-        self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
-        self.context_embedder = HunyuanVideoTokenRefiner(text_embed_dim,
-                                                         num_attention_heads,
-                                                         attention_head_dim,
-                                                         num_layers=num_refiner_layers)
-        self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim)
-
-        # 2. RoPE
-        self.rope = HunyuanVideoRotaryPosEmbed(patch_size, patch_size_t, rope_axes_dim, rope_theta)
-
-        # 3. Dual stream transformer blocks
-        self.transformer_blocks = nn.ModuleList([
-            HunyuanVideoTransformerBlock(num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm)
-            for _ in range(num_layers)
-        ])
-
-        # 4. Single stream transformer blocks
-        self.single_transformer_blocks = nn.ModuleList([
-            HunyuanVideoSingleTransformerBlock(num_attention_heads,
-                                               attention_head_dim,
-                                               mlp_ratio=mlp_ratio,
-                                               qk_norm=qk_norm) for _ in range(num_single_layers)
-        ])
-
-        # 5. Output projection
-        self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
-        self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels)
-
-        self.gradient_checkpointing = False
-
-    @property
-    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
-    def attn_processors(self) -> Dict[str, AttentionProcessor]:
-        r"""
-        Returns:
-            `dict` of attention processors: A dictionary containing all attention processors used in the model with
-            indexed by its weight name.
-        """
-        # set recursively
-        processors = {}
-
-        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
-            if hasattr(module, "get_processor"):
-                processors[f"{name}.processor"] = module.get_processor()
-
-            for sub_name, child in module.named_children():
-                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
-
-            return processors
-
-        for name, module in self.named_children():
-            fn_recursive_add_processors(name, module, processors)
-
-        return processors
-
-    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
-    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
-        r"""
-        Sets the attention processor to use to compute attention.
-
-        Parameters:
-            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
-                The instantiated processor class or a dictionary of processor classes that will be set as the processor
-                for **all** `Attention` layers.
-
-                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
-                processor. This is strongly recommended when setting trainable attention processors.
-
-        """
-        count = len(self.attn_processors.keys())
-
-        if isinstance(processor, dict) and len(processor) != count:
-            raise ValueError(
-                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
-                f" number of attention layers: {count}. Please make sure to pass {count} processor classes.")
-
-        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
-            if hasattr(module, "set_processor"):
-                if not isinstance(processor, dict):
-                    module.set_processor(processor)
-                else:
-                    module.set_processor(processor.pop(f"{name}.processor"))
-
-            for sub_name, child in module.named_children():
-                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
-
-        for name, module in self.named_children():
-            fn_recursive_attn_processor(name, module, processor)
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if hasattr(module, "gradient_checkpointing"):
-            module.gradient_checkpointing = value
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        encoder_attention_mask: torch.Tensor,
-        guidance: torch.Tensor = None,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = True,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
-        if guidance is None:
-            guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16)
-
-        if attention_kwargs is not None:
-            attention_kwargs = attention_kwargs.copy()
-            lora_scale = attention_kwargs.pop("scale", 1.0)
-        else:
-            lora_scale = 1.0
-
-        if USE_PEFT_BACKEND:
-            # weight the lora layers by setting `lora_scale` for each PEFT layer
-            scale_lora_layers(self, lora_scale)
-        else:
-            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
-                logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.")
-
-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
-        p, p_t = self.config.patch_size, self.config.patch_size_t
-        post_patch_num_frames = num_frames // p_t
-        post_patch_height = height // p
-        post_patch_width = width // p
-
-        pooled_projections = encoder_hidden_states[:, 0, :self.config.pooled_projection_dim]
-        encoder_hidden_states = encoder_hidden_states[:, 1:]
-
-        # 1. RoPE
-        image_rotary_emb = self.rope(hidden_states)
-
-        # 2. Conditional embeddings
-        temb = self.time_text_embed(timestep, guidance, pooled_projections)
-        hidden_states = self.x_embedder(hidden_states)
-        encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask)
-
-        # 3. Attention mask preparation
-        latent_sequence_length = hidden_states.shape[1]
-        condition_sequence_length = encoder_hidden_states.shape[1]
-        sequence_length = latent_sequence_length + condition_sequence_length
-        attention_mask = torch.zeros(batch_size,
-                                     sequence_length,
-                                     sequence_length,
-                                     device=hidden_states.device,
-                                     dtype=torch.bool)  # [B, N, N]
-
-        effective_condition_sequence_length = encoder_attention_mask.sum(dim=1, dtype=torch.int)
-        effective_sequence_length = latent_sequence_length + effective_condition_sequence_length
-
-        for i in range(batch_size):
-            attention_mask[i, :effective_sequence_length[i], :effective_sequence_length[i]] = True
-
-        # 4. Transformer blocks
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-
-            def create_custom_forward(module, return_dict=None):
-
-                def custom_forward(*inputs):
-                    if return_dict is not None:
-                        return module(*inputs, return_dict=return_dict)
-                    else:
-                        return module(*inputs)
-
-                return custom_forward
-
-            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-
-            for block in self.transformer_blocks:
-                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    hidden_states,
-                    encoder_hidden_states,
-                    temb,
-                    attention_mask,
-                    image_rotary_emb,
-                    **ckpt_kwargs,
-                )
-
-            for block in self.single_transformer_blocks:
-                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    hidden_states,
-                    encoder_hidden_states,
-                    temb,
-                    attention_mask,
-                    image_rotary_emb,
-                    **ckpt_kwargs,
-                )
-
-        else:
-            for block in self.transformer_blocks:
-                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask,
-                                                             image_rotary_emb)
-
-            for block in self.single_transformer_blocks:
-                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask,
-                                                             image_rotary_emb)
-
-        # 5. Output projection
-        hidden_states = self.norm_out(hidden_states, temb)
-        hidden_states = self.proj_out(hidden_states)
-
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
-                                              -1, p_t, p, p)
-        hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
-        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
-
-        if USE_PEFT_BACKEND:
-            # remove `lora_scale` from each PEFT layer
-            unscale_lora_layers(self, lora_scale)
-
-        if not return_dict:
-            return (hidden_states, )
-
-        return Transformer2DModelOutput(sample=hidden_states)

fastvideo/models/hunyuan_hf/pipeline_hunyuan.py

@@ -1,691 +0,0 @@
-# Copyright 2024 The HunyuanVideo Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
-from diffusers.loaders import HunyuanVideoLoraLoaderMixin
-from diffusers.models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel
-from diffusers.pipelines.hunyuan_video.pipeline_output import HunyuanVideoPipelineOutput
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import logging, replace_example_docstring
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.video_processor import VideoProcessor
-from einops import rearrange
-from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast
-
-from fastvideo.utils.communications import all_gather
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```python
-        >>> import torch
-        >>> from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
-        >>> from diffusers.utils import export_to_video
-
-        >>> model_id = "tencent/HunyuanVideo"
-        >>> transformer = HunyuanVideoTransformer3DModel.from_pretrained(
-        ...     model_id, subfolder="transformer", torch_dtype=torch.bfloat16
-        ... )
-        >>> pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.float16)
-        >>> pipe.vae.enable_tiling()
-        >>> pipe.to("cuda")
-
-        >>> output = pipe(
-        ...     prompt="A cat walks on the grass, realistic",
-        ...     height=320,
-        ...     width=512,
-        ...     num_frames=61,
-        ...     num_inference_steps=30,
-        ... ).frames[0]
-        >>> export_to_video(output, "output.mp4", fps=15)
-        ```
-"""
-
-DEFAULT_PROMPT_TEMPLATE = {
-    "template": ("<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
-                 "1. The main content and theme of the video."
-                 "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
-                 "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
-                 "4. background environment, light, style and atmosphere."
-                 "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
-                 "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"),
-    "crop_start":
-    95,
-}
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class HunyuanVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin):
-    r"""
-    Pipeline for text-to-video generation using HunyuanVideo.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Args:
-        text_encoder ([`LlamaModel`]):
-            [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers).
-        tokenizer_2 (`LlamaTokenizer`):
-            Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers).
-        transformer ([`HunyuanVideoTransformer3DModel`]):
-            Conditional Transformer to denoise the encoded image latents.
-        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae ([`AutoencoderKLHunyuanVideo`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
-        text_encoder_2 ([`CLIPTextModel`]):
-            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
-            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
-        tokenizer_2 (`CLIPTokenizer`):
-            Tokenizer of class
-            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
-    """
-
-    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
-    _callback_tensor_inputs = ["latents", "prompt_embeds"]
-
-    def __init__(
-        self,
-        text_encoder: LlamaModel,
-        tokenizer: LlamaTokenizerFast,
-        transformer: HunyuanVideoTransformer3DModel,
-        vae: AutoencoderKLHunyuanVideo,
-        scheduler: FlowMatchEulerDiscreteScheduler,
-        text_encoder_2: CLIPTextModel,
-        tokenizer_2: CLIPTokenizer,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler,
-            text_encoder_2=text_encoder_2,
-            tokenizer_2=tokenizer_2,
-        )
-
-        self.vae_scale_factor_temporal = (self.vae.temporal_compression_ratio
-                                          if hasattr(self, "vae") and self.vae is not None else 4)
-        self.vae_scale_factor_spatial = (self.vae.spatial_compression_ratio
-                                         if hasattr(self, "vae") and self.vae is not None else 8)
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
-
-    def _get_llama_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]],
-        prompt_template: Dict[str, Any],
-        num_videos_per_prompt: int = 1,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-        max_sequence_length: int = 256,
-        num_hidden_layers_to_skip: int = 2,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
-
-        prompt = [prompt_template["template"].format(p) for p in prompt]
-
-        crop_start = prompt_template.get("crop_start", None)
-        if crop_start is None:
-            prompt_template_input = self.tokenizer(
-                prompt_template["template"],
-                padding="max_length",
-                return_tensors="pt",
-                return_length=False,
-                return_overflowing_tokens=False,
-                return_attention_mask=False,
-            )
-            crop_start = prompt_template_input["input_ids"].shape[-1]
-            # Remove <|eot_id|> token and placeholder {}
-            crop_start -= 2
-
-        max_sequence_length += crop_start
-        text_inputs = self.tokenizer(
-            prompt,
-            max_length=max_sequence_length,
-            padding="max_length",
-            truncation=True,
-            return_tensors="pt",
-            return_length=False,
-            return_overflowing_tokens=False,
-            return_attention_mask=True,
-        )
-        text_input_ids = text_inputs.input_ids.to(device=device)
-        prompt_attention_mask = text_inputs.attention_mask.to(device=device)
-
-        prompt_embeds = self.text_encoder(
-            input_ids=text_input_ids,
-            attention_mask=prompt_attention_mask,
-            output_hidden_states=True,
-        ).hidden_states[-(num_hidden_layers_to_skip + 1)]
-        prompt_embeds = prompt_embeds.to(dtype=dtype)
-
-        if crop_start is not None and crop_start > 0:
-            prompt_embeds = prompt_embeds[:, crop_start:]
-            prompt_attention_mask = prompt_attention_mask[:, crop_start:]
-
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat(1, num_videos_per_prompt)
-        prompt_attention_mask = prompt_attention_mask.view(batch_size * num_videos_per_prompt, seq_len)
-
-        return prompt_embeds, prompt_attention_mask
-
-    def _get_clip_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]],
-        num_videos_per_prompt: int = 1,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-        max_sequence_length: int = 77,
-    ) -> torch.Tensor:
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder_2.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
-
-        text_inputs = self.tokenizer_2(
-            prompt,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            return_tensors="pt",
-        )
-
-        text_input_ids = text_inputs.input_ids
-        untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1])
-            logger.warning("The following part of your input was truncated because CLIP can only handle sequences up to"
-                           f" {max_sequence_length} tokens: {removed_text}")
-
-        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output
-
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, -1)
-
-        return prompt_embeds
-
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        prompt_2: Union[str, List[str]] = None,
-        prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
-        num_videos_per_prompt: int = 1,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        pooled_prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-        max_sequence_length: int = 256,
-    ):
-
-        if prompt_embeds is None:
-            prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds(
-                prompt,
-                prompt_template,
-                num_videos_per_prompt,
-                device=device,
-                dtype=dtype,
-                max_sequence_length=max_sequence_length,
-            )
-
-        if pooled_prompt_embeds is None:
-            if prompt_2 is None and pooled_prompt_embeds is None:
-                prompt_2 = prompt
-            pooled_prompt_embeds = self._get_clip_prompt_embeds(
-                prompt,
-                num_videos_per_prompt,
-                device=device,
-                dtype=dtype,
-                max_sequence_length=77,
-            )
-
-        return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask
-
-    def check_inputs(
-        self,
-        prompt,
-        prompt_2,
-        height,
-        width,
-        prompt_embeds=None,
-        callback_on_step_end_tensor_inputs=None,
-        prompt_template=None,
-    ):
-        if height % 16 != 0 or width % 16 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
-
-        if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs
-                                                                      for k in callback_on_step_end_tensor_inputs):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two.")
-        elif prompt_2 is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two.")
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.")
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
-            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
-
-        if prompt_template is not None:
-            if not isinstance(prompt_template, dict):
-                raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}")
-            if "template" not in prompt_template:
-                raise ValueError(
-                    f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}")
-
-    def prepare_latents(
-        self,
-        batch_size: int,
-        num_channels_latents: 32,
-        height: int = 720,
-        width: int = 1280,
-        num_frames: int = 129,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-
-        shape = (
-            batch_size,
-            num_channels_latents,
-            num_frames,
-            int(height) // self.vae_scale_factor_spatial,
-            int(width) // self.vae_scale_factor_spatial,
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
-
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        return latents
-
-    def enable_vae_slicing(self):
-        r"""
-        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
-        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        self.vae.enable_slicing()
-
-    def disable_vae_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_slicing()
-
-    def enable_vae_tiling(self):
-        r"""
-        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
-        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
-        processing larger images.
-        """
-        self.vae.enable_tiling()
-
-    def disable_vae_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_tiling()
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def attention_kwargs(self):
-        return self._attention_kwargs
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        prompt_2: Union[str, List[str]] = None,
-        height: int = 720,
-        width: int = 1280,
-        num_frames: int = 129,
-        num_inference_steps: int = 50,
-        sigmas: List[float] = None,
-        guidance_scale: float = 6.0,
-        num_videos_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        pooled_prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback,
-                                             MultiPipelineCallbacks]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
-        max_sequence_length: int = 256,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            prompt_2 (`str` or `List[str]`, *optional*):
-                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
-                will be used instead.
-            height (`int`, defaults to `720`):
-                The height in pixels of the generated image.
-            width (`int`, defaults to `1280`):
-                The width in pixels of the generated image.
-            num_frames (`int`, defaults to `129`):
-                The number of frames in the generated video.
-            num_inference_steps (`int`, defaults to `50`):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`List[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            guidance_scale (`float`, defaults to `6.0`):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality. Note that the only available HunyuanVideo model is
-                CFG-distilled, which means that traditional guidance between unconditional and conditional latent is
-                not applied.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
-                provided, text embeddings are generated from the `prompt` input argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            clip_skip (`int`, *optional*):
-                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
-                the output of the pre-final layer will be used for computing the prompt embeddings.
-            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
-                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
-                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
-                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
-                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-
-        Examples:
-
-        Returns:
-            [`~HunyuanVideoPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned, otherwise a `tuple` is returned
-                where the first element is a list with the generated images and the second element is a list of `bool`s
-                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
-        """
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            prompt_2,
-            height,
-            width,
-            prompt_embeds,
-            callback_on_step_end_tensor_inputs,
-            prompt_template,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._interrupt = False
-
-        device = self._execution_device
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # 3. Encode input prompt
-        prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt(
-            prompt=prompt,
-            prompt_2=prompt,
-            prompt_template=prompt_template,
-            num_videos_per_prompt=num_videos_per_prompt,
-            prompt_embeds=prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            device=device,
-            max_sequence_length=max_sequence_length,
-        )
-
-        transformer_dtype = self.transformer.dtype
-        prompt_embeds = prompt_embeds.to(transformer_dtype)
-        prompt_attention_mask = prompt_attention_mask.to(transformer_dtype)
-        if pooled_prompt_embeds is not None:
-            pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype)
-
-        # 4. Prepare timesteps
-        sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            device,
-            sigmas=sigmas,
-        )
-
-        # 5. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels
-        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_latent_frames,
-            torch.float32,
-            device,
-            generator,
-            latents,
-        )
-        # check sequence_parallel
-        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
-        if get_sequence_parallel_state():
-            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
-            latents = latents[:, :, rank, :, :, :]
-
-        # 6. Prepare guidance condition
-        guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
-
-        # 7. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        self._num_timesteps = len(timesteps)
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                latent_model_input = latents.to(transformer_dtype)
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latents.shape[0]).to(latents.dtype)
-                if pooled_prompt_embeds.shape[-1] != prompt_embeds.shape[-1]:
-                    pooled_prompt_embeds_padding = F.pad(
-                        pooled_prompt_embeds,
-                        (0, prompt_embeds.shape[2] - pooled_prompt_embeds.shape[1]),
-                        value=0,
-                    ).unsqueeze(1)
-                encoder_hidden_states = torch.cat([pooled_prompt_embeds_padding, prompt_embeds], dim=1)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    encoder_hidden_states=encoder_hidden_states,  # [1, 257, 4096]
-                    timestep=timestep,
-                    encoder_attention_mask=prompt_attention_mask,
-                    guidance=guidance,
-                    attention_kwargs=attention_kwargs,
-                    return_dict=False,
-                )[0]
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-        if get_sequence_parallel_state():
-            latents = all_gather(latents, dim=2)
-
-        if not output_type == "latent":
-            latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor
-            video = self.vae.decode(latents, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
-        else:
-            video = latents
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (video, )
-
-        return HunyuanVideoPipelineOutput(frames=video)

fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py

@@ -1,361 +0,0 @@
-import argparse
-import os
-
-import torch
-from safetensors.torch import save_file
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--diffusers_path", required=True, type=str)
-parser.add_argument("--transformer_path", type=str, default=None, help="Path to save transformer model")
-parser.add_argument("--vae_encoder_path", type=str, default=None, help="Path to save VAE encoder model")
-parser.add_argument("--vae_decoder_path", type=str, default=None, help="Path to save VAE decoder model")
-
-args = parser.parse_args()
-
-
-def reverse_scale_shift(weight, dim):
-    scale, shift = weight.chunk(2, dim=0)
-    new_weight = torch.cat([shift, scale], dim=0)
-    return new_weight
-
-
-def reverse_proj_gate(weight):
-    gate, proj = weight.chunk(2, dim=0)
-    new_weight = torch.cat([proj, gate], dim=0)
-    return new_weight
-
-
-def convert_diffusers_transformer_to_mochi(state_dict):
-    original_state_dict = state_dict.copy()
-    new_state_dict = {}
-
-    # Convert patch_embed
-    new_state_dict["x_embedder.proj.weight"] = original_state_dict.pop("patch_embed.proj.weight")
-    new_state_dict["x_embedder.proj.bias"] = original_state_dict.pop("patch_embed.proj.bias")
-
-    # Convert time_embed
-    new_state_dict["t_embedder.mlp.0.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.weight")
-    new_state_dict["t_embedder.mlp.0.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.bias")
-    new_state_dict["t_embedder.mlp.2.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.weight")
-    new_state_dict["t_embedder.mlp.2.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.bias")
-    new_state_dict["t5_y_embedder.to_kv.weight"] = original_state_dict.pop("time_embed.pooler.to_kv.weight")
-    new_state_dict["t5_y_embedder.to_kv.bias"] = original_state_dict.pop("time_embed.pooler.to_kv.bias")
-    new_state_dict["t5_y_embedder.to_q.weight"] = original_state_dict.pop("time_embed.pooler.to_q.weight")
-    new_state_dict["t5_y_embedder.to_q.bias"] = original_state_dict.pop("time_embed.pooler.to_q.bias")
-    new_state_dict["t5_y_embedder.to_out.weight"] = original_state_dict.pop("time_embed.pooler.to_out.weight")
-    new_state_dict["t5_y_embedder.to_out.bias"] = original_state_dict.pop("time_embed.pooler.to_out.bias")
-    new_state_dict["t5_yproj.weight"] = original_state_dict.pop("time_embed.caption_proj.weight")
-    new_state_dict["t5_yproj.bias"] = original_state_dict.pop("time_embed.caption_proj.bias")
-
-    # Convert transformer blocks
-    num_layers = 48
-    for i in range(num_layers):
-        block_prefix = f"transformer_blocks.{i}."
-        new_prefix = f"blocks.{i}."
-
-        # norm1
-        new_state_dict[new_prefix + "mod_x.weight"] = original_state_dict.pop(block_prefix + "norm1.linear.weight")
-        new_state_dict[new_prefix + "mod_x.bias"] = original_state_dict.pop(block_prefix + "norm1.linear.bias")
-
-        if i < num_layers - 1:
-            new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                  "norm1_context.linear.weight")
-            new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(block_prefix +
-                                                                                "norm1_context.linear.bias")
-        else:
-            new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                  "norm1_context.linear_1.weight")
-            new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(block_prefix +
-                                                                                "norm1_context.linear_1.bias")
-
-        # Visual attention
-        q = original_state_dict.pop(block_prefix + "attn1.to_q.weight")
-        k = original_state_dict.pop(block_prefix + "attn1.to_k.weight")
-        v = original_state_dict.pop(block_prefix + "attn1.to_v.weight")
-        qkv_weight = torch.cat([q, k, v], dim=0)
-        new_state_dict[new_prefix + "attn.qkv_x.weight"] = qkv_weight
-
-        new_state_dict[new_prefix + "attn.q_norm_x.weight"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.norm_q.weight")
-        new_state_dict[new_prefix + "attn.k_norm_x.weight"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.norm_k.weight")
-        new_state_dict[new_prefix + "attn.proj_x.weight"] = original_state_dict.pop(block_prefix +
-                                                                                    "attn1.to_out.0.weight")
-        new_state_dict[new_prefix + "attn.proj_x.bias"] = original_state_dict.pop(block_prefix + "attn1.to_out.0.bias")
-
-        # Context attention
-        q = original_state_dict.pop(block_prefix + "attn1.add_q_proj.weight")
-        k = original_state_dict.pop(block_prefix + "attn1.add_k_proj.weight")
-        v = original_state_dict.pop(block_prefix + "attn1.add_v_proj.weight")
-        qkv_weight = torch.cat([q, k, v], dim=0)
-        new_state_dict[new_prefix + "attn.qkv_y.weight"] = qkv_weight
-
-        new_state_dict[new_prefix + "attn.q_norm_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.norm_added_q.weight")
-        new_state_dict[new_prefix + "attn.k_norm_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.norm_added_k.weight")
-        if i < num_layers - 1:
-            new_state_dict[new_prefix + "attn.proj_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                        "attn1.to_add_out.weight")
-            new_state_dict[new_prefix + "attn.proj_y.bias"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.to_add_out.bias")
-
-        # MLP
-        new_state_dict[new_prefix + "mlp_x.w1.weight"] = reverse_proj_gate(
-            original_state_dict.pop(block_prefix + "ff.net.0.proj.weight"))
-        new_state_dict[new_prefix + "mlp_x.w2.weight"] = original_state_dict.pop(block_prefix + "ff.net.2.weight")
-        if i < num_layers - 1:
-            new_state_dict[new_prefix + "mlp_y.w1.weight"] = reverse_proj_gate(
-                original_state_dict.pop(block_prefix + "ff_context.net.0.proj.weight"))
-            new_state_dict[new_prefix + "mlp_y.w2.weight"] = original_state_dict.pop(block_prefix +
-                                                                                     "ff_context.net.2.weight")
-
-    # Output layers
-    new_state_dict["final_layer.mod.weight"] = reverse_scale_shift(original_state_dict.pop("norm_out.linear.weight"),
-                                                                   dim=0)
-    new_state_dict["final_layer.mod.bias"] = reverse_scale_shift(original_state_dict.pop("norm_out.linear.bias"), dim=0)
-    new_state_dict["final_layer.linear.weight"] = original_state_dict.pop("proj_out.weight")
-    new_state_dict["final_layer.linear.bias"] = original_state_dict.pop("proj_out.bias")
-
-    new_state_dict["pos_frequencies"] = original_state_dict.pop("pos_frequencies")
-
-    print("Remaining Keys:", original_state_dict.keys())
-
-    return new_state_dict
-
-
-def convert_diffusers_vae_to_mochi(state_dict):
-    original_state_dict = state_dict.copy()
-    encoder_state_dict = {}
-    decoder_state_dict = {}
-
-    # Convert encoder
-    prefix = "encoder."
-
-    encoder_state_dict["layers.0.weight"] = original_state_dict.pop(f"{prefix}proj_in.weight")
-    encoder_state_dict["layers.0.bias"] = original_state_dict.pop(f"{prefix}proj_in.bias")
-
-    # Convert block_in
-    for i in range(3):
-        encoder_state_dict[f"layers.{i+1}.stack.0.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+1}.stack.0.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias")
-        encoder_state_dict[f"layers.{i+1}.stack.2.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv1.conv.weight")
-        encoder_state_dict[f"layers.{i+1}.stack.2.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv1.conv.bias")
-        encoder_state_dict[f"layers.{i+1}.stack.3.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+1}.stack.3.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias")
-        encoder_state_dict[f"layers.{i+1}.stack.5.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv2.conv.weight")
-        encoder_state_dict[f"layers.{i+1}.stack.5.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv2.conv.bias")
-
-    # Convert down_blocks
-    down_block_layers = [3, 4, 6]
-    for block in range(3):
-        encoder_state_dict[f"layers.{block+4}.layers.0.weight"] = original_state_dict.pop(
-            f"{prefix}down_blocks.{block}.conv_in.conv.weight")
-        encoder_state_dict[f"layers.{block+4}.layers.0.bias"] = original_state_dict.pop(
-            f"{prefix}down_blocks.{block}.conv_in.conv.bias")
-
-        for i in range(down_block_layers[block]):
-            # Convert resnets
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias")
-
-            # Convert attentions
-            q = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_q.weight")
-            k = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_k.weight")
-            v = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_v.weight")
-            qkv_weight = torch.cat([q, k, v], dim=0)
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.qkv.weight"] = qkv_weight
-
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.bias")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.bias")
-
-    # Convert block_out
-    for i in range(3):
-        encoder_state_dict[f"layers.{i+7}.stack.0.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+7}.stack.0.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias")
-        encoder_state_dict[f"layers.{i+7}.stack.2.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv1.conv.weight")
-        encoder_state_dict[f"layers.{i+7}.stack.2.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv1.conv.bias")
-        encoder_state_dict[f"layers.{i+7}.stack.3.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+7}.stack.3.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias")
-        encoder_state_dict[f"layers.{i+7}.stack.5.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv2.conv.weight")
-        encoder_state_dict[f"layers.{i+7}.stack.5.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv2.conv.bias")
-
-        q = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_q.weight")
-        k = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_k.weight")
-        v = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_v.weight")
-        qkv_weight = torch.cat([q, k, v], dim=0)
-        encoder_state_dict[f"layers.{i+7}.attn_block.attn.qkv.weight"] = qkv_weight
-
-        encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.attentions.{i}.to_out.0.weight")
-        encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.attentions.{i}.to_out.0.bias")
-        encoder_state_dict[f"layers.{i+7}.attn_block.norm.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.norms.{i}.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+7}.attn_block.norm.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.norms.{i}.norm_layer.bias")
-
-    # Convert output layers
-    encoder_state_dict["output_norm.weight"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.weight")
-    encoder_state_dict["output_norm.bias"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.bias")
-    encoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight")
-
-    # Convert decoder
-    prefix = "decoder."
-
-    decoder_state_dict["blocks.0.0.weight"] = original_state_dict.pop(f"{prefix}conv_in.weight")
-    decoder_state_dict["blocks.0.0.bias"] = original_state_dict.pop(f"{prefix}conv_in.bias")
-
-    # Convert block_in
-    for i in range(3):
-        decoder_state_dict[f"blocks.0.{i+1}.stack.0.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.0.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.2.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv1.conv.weight")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.2.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv1.conv.bias")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.3.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.3.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.5.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv2.conv.weight")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.5.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv2.conv.bias")
-
-    # Convert up_blocks
-    up_block_layers = [6, 4, 3]
-    for block in range(3):
-        for i in range(up_block_layers[block]):
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.weight"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.bias"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.weight"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.bias"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.weight"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.bias"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.weight"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.bias"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias")
-        decoder_state_dict[f"blocks.{block+1}.proj.weight"] = original_state_dict.pop(
-            f"{prefix}up_blocks.{block}.proj.weight")
-        decoder_state_dict[f"blocks.{block+1}.proj.bias"] = original_state_dict.pop(
-            f"{prefix}up_blocks.{block}.proj.bias")
-
-    # Convert block_out
-    for i in range(3):
-        decoder_state_dict[f"blocks.4.{i}.stack.0.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight")
-        decoder_state_dict[f"blocks.4.{i}.stack.0.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias")
-        decoder_state_dict[f"blocks.4.{i}.stack.2.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv1.conv.weight")
-        decoder_state_dict[f"blocks.4.{i}.stack.2.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv1.conv.bias")
-        decoder_state_dict[f"blocks.4.{i}.stack.3.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight")
-        decoder_state_dict[f"blocks.4.{i}.stack.3.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias")
-        decoder_state_dict[f"blocks.4.{i}.stack.5.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv2.conv.weight")
-        decoder_state_dict[f"blocks.4.{i}.stack.5.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv2.conv.bias")
-
-    # Convert output layers
-    decoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight")
-    decoder_state_dict["output_proj.bias"] = original_state_dict.pop(f"{prefix}proj_out.bias")
-
-    return encoder_state_dict, decoder_state_dict
-
-
-def ensure_safetensors_extension(path):
-    if not path.endswith(".safetensors"):
-        path = path + ".safetensors"
-    return path
-
-
-def ensure_directory_exists(path):
-    directory = os.path.dirname(path)
-    if directory:
-        os.makedirs(directory, exist_ok=True)
-
-
-def main(args):
-    from diffusers import MochiPipeline
-
-    pipe = MochiPipeline.from_pretrained(args.diffusers_path)
-
-    if args.transformer_path:
-        transformer_path = ensure_safetensors_extension(args.transformer_path)
-        ensure_directory_exists(transformer_path)
-
-        print("Converting transformer model...")
-        transformer_state_dict = convert_diffusers_transformer_to_mochi(pipe.transformer.state_dict())
-        save_file(transformer_state_dict, transformer_path)
-        print(f"Saved transformer to {transformer_path}")
-
-    if args.vae_encoder_path and args.vae_decoder_path:
-        encoder_path = ensure_safetensors_extension(args.vae_encoder_path)
-        decoder_path = ensure_safetensors_extension(args.vae_decoder_path)
-
-        ensure_directory_exists(encoder_path)
-        ensure_directory_exists(decoder_path)
-
-        print("Converting VAE models...")
-        encoder_state_dict, decoder_state_dict = convert_diffusers_vae_to_mochi(pipe.vae.state_dict())
-
-        save_file(encoder_state_dict, encoder_path)
-        print(f"Saved VAE encoder to {encoder_path}")
-
-        save_file(decoder_state_dict, decoder_path)
-        print(f"Saved VAE decoder to {decoder_path}")
-    elif args.vae_encoder_path or args.vae_decoder_path:
-        print("Warning: Both VAE encoder and decoder paths must be specified to convert VAE models.")
-
-
-if __name__ == "__main__":
-    main(args)

fastvideo/models/mochi_hf/mochi_latents_utils.py

@@ -1,45 +0,0 @@
-import torch
-
-mochi_latents_mean = torch.tensor([
-    -0.06730895953510081,
-    -0.038011381506090416,
-    -0.07477820912866141,
-    -0.05565264470995561,
-    0.012767231469026969,
-    -0.04703542746246419,
-    0.043896967884726704,
-    -0.09346305707025976,
-    -0.09918314763016893,
-    -0.008729793427399178,
-    -0.011931556316503654,
-    -0.0321993391887285,
-]).view(1, 12, 1, 1, 1)
-mochi_latents_std = torch.tensor([
-    0.9263795028493863,
-    0.9248894543193766,
-    0.9393059390890617,
-    0.959253732819592,
-    0.8244560132752793,
-    0.917259975397747,
-    0.9294154431013696,
-    1.3720942357788521,
-    0.881393668867029,
-    0.9168315692124348,
-    0.9185249279345552,
-    0.9274757570805041,
-]).view(1, 12, 1, 1, 1)
-mochi_scaling_factor = 1.0
-
-
-def normalize_dit_input(model_type, latents):
-    if model_type == "mochi":
-        latents_mean = mochi_latents_mean.to(latents.device, latents.dtype)
-        latents_std = mochi_latents_std.to(latents.device, latents.dtype)
-        latents = (latents - latents_mean) / latents_std
-        return latents
-    elif model_type == "hunyuan_hf":
-        return latents * 0.476986
-    elif model_type == "hunyuan":
-        return latents * 0.476986
-    else:
-        raise NotImplementedError(f"model_type {model_type} not supported")

fastvideo/models/mochi_hf/modeling_mochi.py

@@ -1,663 +0,0 @@
-# Copyright 2024 The Genmo team and The HuggingFace Team.
-# All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import Any, Dict, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.loaders import PeftAdapterMixin
-from diffusers.models.attention import FeedForward as HF_FeedForward
-from diffusers.models.attention_processor import Attention
-from diffusers.models.embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.normalization import AdaLayerNormContinuous
-from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
-from diffusers.utils.torch_utils import maybe_allow_in_graph
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
-
-from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
-from fastvideo.models.mochi_hf.norm import (MochiLayerNormContinuous, MochiModulatedRMSNorm, MochiRMSNorm,
-                                            MochiRMSNormZero)
-from fastvideo.utils.communications import all_gather, all_to_all_4D
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-class FeedForward(HF_FeedForward):
-
-    def __init__(
-        self,
-        dim: int,
-        dim_out: Optional[int] = None,
-        mult: int = 4,
-        dropout: float = 0.0,
-        activation_fn: str = "geglu",
-        final_dropout: bool = False,
-        inner_dim=None,
-        bias: bool = True,
-    ):
-        super().__init__(dim, dim_out, mult, dropout, activation_fn, final_dropout, inner_dim, bias)
-        assert activation_fn == "swiglu"
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.net[0].proj(hidden_states)
-        hidden_states, gate = hidden_states.chunk(2, dim=-1)
-
-        return self.net[2](LigerSiLUMulFunction.apply(gate, hidden_states))
-
-
-class MochiAttention(nn.Module):
-
-    def __init__(
-        self,
-        query_dim: int,
-        processor: "MochiAttnProcessor2_0",
-        heads: int = 8,
-        dim_head: int = 64,
-        dropout: float = 0.0,
-        bias: bool = False,
-        added_kv_proj_dim: Optional[int] = None,
-        added_proj_bias: Optional[bool] = True,
-        out_dim: int = None,
-        out_context_dim: int = None,
-        out_bias: bool = True,
-        context_pre_only: bool = False,
-        eps: float = 1e-5,
-    ):
-        super().__init__()
-        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
-        self.out_dim = out_dim if out_dim is not None else query_dim
-        self.out_context_dim = out_context_dim if out_context_dim else query_dim
-        self.context_pre_only = context_pre_only
-
-        self.heads = out_dim // dim_head if out_dim is not None else heads
-
-        self.norm_q = MochiRMSNorm(dim_head, eps)
-        self.norm_k = MochiRMSNorm(dim_head, eps)
-        self.norm_added_q = MochiRMSNorm(dim_head, eps)
-        self.norm_added_k = MochiRMSNorm(dim_head, eps)
-
-        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
-
-        self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
-        self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
-        if self.context_pre_only is not None:
-            self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
-
-        self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
-        self.to_out.append(nn.Dropout(dropout))
-
-        if not self.context_pre_only:
-            self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
-
-        self.processor = processor
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        **kwargs,
-    ):
-        return self.processor(
-            self,
-            hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            attention_mask=attention_mask,
-            **kwargs,
-        )
-
-
-class MochiAttnProcessor2_0:
-    """Attention processor used in Mochi."""
-
-    def __init__(self):
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        image_rotary_emb: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        # [b, s, h * d]
-        query = attn.to_q(hidden_states)
-        key = attn.to_k(hidden_states)
-        value = attn.to_v(hidden_states)
-
-        # [b, s, h=24, d=128]
-        query = query.unflatten(2, (attn.heads, -1))
-        key = key.unflatten(2, (attn.heads, -1))
-        value = value.unflatten(2, (attn.heads, -1))
-
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-        # [b, 256, h * d]
-        encoder_query = attn.add_q_proj(encoder_hidden_states)
-        encoder_key = attn.add_k_proj(encoder_hidden_states)
-        encoder_value = attn.add_v_proj(encoder_hidden_states)
-
-        # [b, 256, h=24, d=128]
-        encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
-        encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
-        encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
-
-        if attn.norm_added_q is not None:
-            encoder_query = attn.norm_added_q(encoder_query)
-        if attn.norm_added_k is not None:
-            encoder_key = attn.norm_added_k(encoder_key)
-
-        if image_rotary_emb is not None:
-            freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1]
-        # shard the head dimension
-        if get_sequence_parallel_state():
-            # B, S, H, D to (S, B,) H, D
-            # batch_size, seq_len, attn_heads, head_dim
-            query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
-            key = all_to_all_4D(key, scatter_dim=2, gather_dim=1)
-            value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
-
-            def shrink_head(encoder_state, dim):
-                local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-                return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
-
-            encoder_query = shrink_head(encoder_query, dim=2)
-            encoder_key = shrink_head(encoder_key, dim=2)
-            encoder_value = shrink_head(encoder_value, dim=2)
-            if image_rotary_emb is not None:
-                freqs_cos = shrink_head(freqs_cos, dim=1)
-                freqs_sin = shrink_head(freqs_sin, dim=1)
-
-        if image_rotary_emb is not None:
-
-            def apply_rotary_emb(x, freqs_cos, freqs_sin):
-                x_even = x[..., 0::2].float()
-                x_odd = x[..., 1::2].float()
-                cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
-                sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
-
-                return torch.stack([cos, sin], dim=-1).flatten(-2)
-
-            query = apply_rotary_emb(query, freqs_cos, freqs_sin)
-            key = apply_rotary_emb(key, freqs_cos, freqs_sin)
-
-        # query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
-        # encoder_query, encoder_key, encoder_value = (
-        #     encoder_query.transpose(1, 2),
-        #     encoder_key.transpose(1, 2),
-        #     encoder_value.transpose(1, 2),
-        # )
-        # [b, s, h, d]
-        sequence_length = query.size(1)
-        encoder_sequence_length = encoder_query.size(1)
-
-        # H
-        query = torch.cat([query, encoder_query], dim=1).unsqueeze(2)
-        key = torch.cat([key, encoder_key], dim=1).unsqueeze(2)
-        value = torch.cat([value, encoder_value], dim=1).unsqueeze(2)
-        # B, S, 3, H, D
-        qkv = torch.cat([query, key, value], dim=2)
-
-        attn_mask = encoder_attention_mask[:, :].bool()
-        attn_mask = F.pad(attn_mask, (sequence_length, 0), value=True)
-        hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None)
-
-        # hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask = None, dropout_p=0.0, is_causal=False)
-
-        # valid_lengths = encoder_attention_mask.sum(dim=1) + sequence_length
-        # def no_padding_mask(score, b, h, q_idx, kv_idx):
-        #     return torch.where(kv_idx < valid_lengths[b],score,  -float("inf"))
-
-        # hidden_states = flex_attention(query, key, value, score_mod=no_padding_mask)
-        if get_sequence_parallel_state():
-            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
-                (sequence_length, encoder_sequence_length), dim=1)
-            # B, S, H, D
-            hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)
-            encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous()
-            hidden_states = hidden_states.flatten(2, 3)
-            hidden_states = hidden_states.to(query.dtype)
-            encoder_hidden_states = encoder_hidden_states.flatten(2, 3)
-            encoder_hidden_states = encoder_hidden_states.to(query.dtype)
-        else:
-            hidden_states = hidden_states.flatten(2, 3)
-            hidden_states = hidden_states.to(query.dtype)
-
-            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
-                (sequence_length, encoder_sequence_length), dim=1)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        if hasattr(attn, "to_add_out"):
-            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
-
-        return hidden_states, encoder_hidden_states
-
-
-@maybe_allow_in_graph
-class MochiTransformerBlock(nn.Module):
-    r"""
-    Transformer block used in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
-
-    Args:
-        dim (`int`):
-            The number of channels in the input and output.
-        num_attention_heads (`int`):
-            The number of heads to use for multi-head attention.
-        attention_head_dim (`int`):
-            The number of channels in each head.
-        qk_norm (`str`, defaults to `"rms_norm"`):
-            The normalization layer to use.
-        activation_fn (`str`, defaults to `"swiglu"`):
-            Activation function to use in feed-forward.
-        context_pre_only (`bool`, defaults to `False`):
-            Whether or not to process context-related conditions with additional layers.
-        eps (`float`, defaults to `1e-6`):
-            Epsilon value for normalization layers.
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        pooled_projection_dim: int,
-        qk_norm: str = "rms_norm",
-        activation_fn: str = "swiglu",
-        context_pre_only: bool = False,
-        eps: float = 1e-6,
-    ) -> None:
-        super().__init__()
-
-        self.context_pre_only = context_pre_only
-        self.ff_inner_dim = (4 * dim * 2) // 3
-        self.ff_context_inner_dim = (4 * pooled_projection_dim * 2) // 3
-
-        self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False)
-
-        if not context_pre_only:
-            self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False)
-        else:
-            self.norm1_context = MochiLayerNormContinuous(
-                embedding_dim=pooled_projection_dim,
-                conditioning_embedding_dim=dim,
-                eps=eps,
-            )
-
-        self.attn1 = MochiAttention(
-            query_dim=dim,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            bias=False,
-            added_kv_proj_dim=pooled_projection_dim,
-            added_proj_bias=False,
-            out_dim=dim,
-            out_context_dim=pooled_projection_dim,
-            context_pre_only=context_pre_only,
-            processor=MochiAttnProcessor2_0(),
-            eps=1e-5,
-        )
-
-        # TODO(aryan): norm_context layers are not needed when `context_pre_only` is True
-        self.norm2 = MochiModulatedRMSNorm(eps=eps)
-        self.norm2_context = (MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None)
-
-        self.norm3 = MochiModulatedRMSNorm(eps)
-        self.norm3_context = (MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None)
-
-        self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False)
-        self.ff_context = None
-        if not context_pre_only:
-            self.ff_context = FeedForward(
-                pooled_projection_dim,
-                inner_dim=self.ff_context_inner_dim,
-                activation_fn=activation_fn,
-                bias=False,
-            )
-
-        self.norm4 = MochiModulatedRMSNorm(eps=eps)
-        self.norm4_context = MochiModulatedRMSNorm(eps=eps)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        temb: torch.Tensor,
-        image_rotary_emb: Optional[torch.Tensor] = None,
-        output_attn=False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
-
-        if not self.context_pre_only:
-            (
-                norm_encoder_hidden_states,
-                enc_gate_msa,
-                enc_scale_mlp,
-                enc_gate_mlp,
-            ) = self.norm1_context(encoder_hidden_states, temb)
-        else:
-            norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)
-
-        attn_hidden_states, context_attn_hidden_states = self.attn1(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            image_rotary_emb=image_rotary_emb,
-            encoder_attention_mask=encoder_attention_mask,
-        )
-
-        hidden_states = hidden_states + self.norm2(attn_hidden_states, torch.tanh(gate_msa).unsqueeze(1))
-        norm_hidden_states = self.norm3(hidden_states, (1 + scale_mlp.unsqueeze(1).to(torch.float32)))
-        ff_output = self.ff(norm_hidden_states)
-        hidden_states = hidden_states + self.norm4(ff_output, torch.tanh(gate_mlp).unsqueeze(1))
-
-        if not self.context_pre_only:
-            encoder_hidden_states = encoder_hidden_states + self.norm2_context(context_attn_hidden_states,
-                                                                               torch.tanh(enc_gate_msa).unsqueeze(1))
-            norm_encoder_hidden_states = self.norm3_context(
-                encoder_hidden_states,
-                (1 + enc_scale_mlp.unsqueeze(1).to(torch.float32)),
-            )
-            context_ff_output = self.ff_context(norm_encoder_hidden_states)
-            encoder_hidden_states = encoder_hidden_states + self.norm4_context(context_ff_output,
-                                                                               torch.tanh(enc_gate_mlp).unsqueeze(1))
-
-        if not output_attn:
-            attn_hidden_states = None
-        return hidden_states, encoder_hidden_states, attn_hidden_states
-
-
-class MochiRoPE(nn.Module):
-    r"""
-    RoPE implementation used in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
-
-    Args:
-        base_height (`int`, defaults to `192`):
-            Base height used to compute interpolation scale for rotary positional embeddings.
-        base_width (`int`, defaults to `192`):
-            Base width used to compute interpolation scale for rotary positional embeddings.
-    """
-
-    def __init__(self, base_height: int = 192, base_width: int = 192) -> None:
-        super().__init__()
-
-        self.target_area = base_height * base_width
-
-    def _centers(self, start, stop, num, device, dtype) -> torch.Tensor:
-        edges = torch.linspace(start, stop, num + 1, device=device, dtype=dtype)
-        return (edges[:-1] + edges[1:]) / 2
-
-    def _get_positions(
-        self,
-        num_frames: int,
-        height: int,
-        width: int,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ) -> torch.Tensor:
-        scale = (self.target_area / (height * width))**0.5
-        t = torch.arange(num_frames * nccl_info.sp_size, device=device, dtype=dtype)
-        h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
-        w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
-
-        grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
-
-        positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
-        return positions
-
-    def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
-        with torch.autocast(freqs.device.type, enabled=False):
-            # Always run ROPE freqs computation in FP32
-            freqs = torch.einsum(
-                "nd,dhf->nhf",  # codespell:ignore
-                pos.to(torch.float32),  # codespell:ignore
-                freqs.to(torch.float32))
-        freqs_cos = torch.cos(freqs)
-        freqs_sin = torch.sin(freqs)
-        return freqs_cos, freqs_sin
-
-    def forward(
-        self,
-        pos_frequencies: torch.Tensor,
-        num_frames: int,
-        height: int,
-        width: int,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        pos = self._get_positions(num_frames, height, width, device, dtype)
-        rope_cos, rope_sin = self._create_rope(pos_frequencies, pos)
-        return rope_cos, rope_sin
-
-
-@maybe_allow_in_graph
-class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
-    r"""
-    A Transformer model for video-like data introduced in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
-
-    Args:
-        patch_size (`int`, defaults to `2`):
-            The size of the patches to use in the patch embedding layer.
-        num_attention_heads (`int`, defaults to `24`):
-            The number of heads to use for multi-head attention.
-        attention_head_dim (`int`, defaults to `128`):
-            The number of channels in each head.
-        num_layers (`int`, defaults to `48`):
-            The number of layers of Transformer blocks to use.
-        in_channels (`int`, defaults to `12`):
-            The number of channels in the input.
-        out_channels (`int`, *optional*, defaults to `None`):
-            The number of channels in the output.
-        qk_norm (`str`, defaults to `"rms_norm"`):
-            The normalization layer to use.
-        text_embed_dim (`int`, defaults to `4096`):
-            Input dimension of text embeddings from the text encoder.
-        time_embed_dim (`int`, defaults to `256`):
-            Output dimension of timestep embeddings.
-        activation_fn (`str`, defaults to `"swiglu"`):
-            Activation function to use in feed-forward.
-        max_sequence_length (`int`, defaults to `256`):
-            The maximum sequence length of text embeddings supported.
-    """
-
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-        self,
-        patch_size: int = 2,
-        num_attention_heads: int = 24,
-        attention_head_dim: int = 128,
-        num_layers: int = 48,
-        pooled_projection_dim: int = 1536,
-        in_channels: int = 12,
-        out_channels: Optional[int] = None,
-        qk_norm: str = "rms_norm",
-        text_embed_dim: int = 4096,
-        time_embed_dim: int = 256,
-        activation_fn: str = "swiglu",
-        max_sequence_length: int = 256,
-    ) -> None:
-        super().__init__()
-
-        inner_dim = num_attention_heads * attention_head_dim
-        out_channels = out_channels or in_channels
-
-        self.patch_embed = PatchEmbed(
-            patch_size=patch_size,
-            in_channels=in_channels,
-            embed_dim=inner_dim,
-            pos_embed_type=None,
-        )
-
-        self.time_embed = MochiCombinedTimestepCaptionEmbedding(
-            embedding_dim=inner_dim,
-            pooled_projection_dim=pooled_projection_dim,
-            text_embed_dim=text_embed_dim,
-            time_embed_dim=time_embed_dim,
-            num_attention_heads=8,
-        )
-
-        self.pos_frequencies = nn.Parameter(torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0))
-        self.rope = MochiRoPE()
-
-        self.transformer_blocks = nn.ModuleList([
-            MochiTransformerBlock(
-                dim=inner_dim,
-                num_attention_heads=num_attention_heads,
-                attention_head_dim=attention_head_dim,
-                pooled_projection_dim=pooled_projection_dim,
-                qk_norm=qk_norm,
-                activation_fn=activation_fn,
-                context_pre_only=i == num_layers - 1,
-            ) for i in range(num_layers)
-        ])
-
-        self.norm_out = AdaLayerNormContinuous(
-            inner_dim,
-            inner_dim,
-            elementwise_affine=False,
-            eps=1e-6,
-            norm_type="layer_norm",
-        )
-        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
-
-        self.gradient_checkpointing = False
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if hasattr(module, "gradient_checkpointing"):
-            module.gradient_checkpointing = value
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        encoder_attention_mask: torch.Tensor,
-        output_features=False,
-        output_features_stride=8,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = False,
-    ) -> torch.Tensor:
-        assert (return_dict is False), "return_dict is not supported in MochiTransformer3DModel"
-
-        if attention_kwargs is not None:
-            attention_kwargs = attention_kwargs.copy()
-            lora_scale = attention_kwargs.pop("scale", 1.0)
-        else:
-            lora_scale = 1.0
-
-        if USE_PEFT_BACKEND:
-            # weight the lora layers by setting `lora_scale` for each PEFT layer
-            scale_lora_layers(self, lora_scale)
-        else:
-            if (attention_kwargs is not None and attention_kwargs.get("scale", None) is not None):
-                logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.")
-
-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
-        p = self.config.patch_size
-
-        post_patch_height = height // p
-        post_patch_width = width // p
-        # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux
-        timestep = 1000 - timestep
-        temb, encoder_hidden_states = self.time_embed(
-            timestep,
-            encoder_hidden_states,
-            encoder_attention_mask,
-            hidden_dtype=hidden_states.dtype,
-        )
-
-        hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
-        hidden_states = self.patch_embed(hidden_states)
-        hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2)
-
-        image_rotary_emb = self.rope(
-            self.pos_frequencies,
-            num_frames,
-            post_patch_height,
-            post_patch_width,
-            device=hidden_states.device,
-            dtype=torch.float32,
-        )
-        attn_outputs_list = []
-        for i, block in enumerate(self.transformer_blocks):
-            if self.gradient_checkpointing:
-
-                def create_custom_forward(module):
-
-                    def custom_forward(*inputs):
-                        return module(*inputs)
-
-                    return custom_forward
-
-                ckpt_kwargs: Dict[str, Any] = ({"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {})
-                (
-                    hidden_states,
-                    encoder_hidden_states,
-                    attn_outputs,
-                ) = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    hidden_states,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                    temb,
-                    image_rotary_emb,
-                    output_features,
-                    **ckpt_kwargs,
-                )
-            else:
-                hidden_states, encoder_hidden_states, attn_outputs = block(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=encoder_attention_mask,
-                    temb=temb,
-                    image_rotary_emb=image_rotary_emb,
-                    output_attn=output_features,
-                )
-            if i % output_features_stride == 0:
-                attn_outputs_list.append(attn_outputs)
-
-        hidden_states = self.norm_out(hidden_states, temb)
-        hidden_states = self.proj_out(hidden_states)
-
-        hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1)
-        hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5)
-        output = hidden_states.reshape(batch_size, -1, num_frames, height, width)
-
-        if USE_PEFT_BACKEND:
-            # remove `lora_scale` from each PEFT layer
-            unscale_lora_layers(self, lora_scale)
-
-        if not output_features:
-            attn_outputs_list = None
-        else:
-            attn_outputs_list = torch.stack(attn_outputs_list, dim=0)
-        # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux
-        return (-output, attn_outputs_list)

fastvideo/models/mochi_hf/norm.py

@@ -1,128 +0,0 @@
-# Copyright 2024 The Genmo team and The HuggingFace Team.
-# All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Tuple
-
-import torch
-import torch.nn as nn
-
-
-class MochiModulatedRMSNorm(nn.Module):
-
-    def __init__(self, eps: float):
-        super().__init__()
-
-        self.eps = eps
-
-    def forward(self, hidden_states, scale=None):
-        hidden_states_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
-        if scale is not None:
-            hidden_states = hidden_states * scale
-
-        hidden_states = hidden_states.to(hidden_states_dtype)
-
-        return hidden_states
-
-
-class MochiRMSNorm(nn.Module):
-
-    def __init__(self, dim, eps: float, elementwise_affine=True):
-        super().__init__()
-
-        self.eps = eps
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.ones(dim))
-        else:
-            self.weight = None
-
-    def forward(self, hidden_states):
-        hidden_states_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
-        if self.weight is not None:
-            # convert into half-precision if necessary
-            if self.weight.dtype in [torch.float16, torch.bfloat16]:
-                hidden_states = hidden_states.to(self.weight.dtype)
-            hidden_states = hidden_states * self.weight
-        hidden_states = hidden_states.to(hidden_states_dtype)
-
-        return hidden_states
-
-
-class MochiLayerNormContinuous(nn.Module):
-
-    def __init__(
-        self,
-        embedding_dim: int,
-        conditioning_embedding_dim: int,
-        eps=1e-5,
-        bias=True,
-    ):
-        super().__init__()
-
-        # AdaLN
-        self.silu = nn.SiLU()
-        self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias)
-        self.norm = MochiModulatedRMSNorm(eps=eps)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        conditioning_embedding: torch.Tensor,
-    ) -> torch.Tensor:
-        input_dtype = x.dtype
-
-        # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
-        scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype))
-        x = self.norm(x, (1 + scale.unsqueeze(1).to(torch.float32)))
-
-        return x.to(input_dtype)
-
-
-class MochiRMSNormZero(nn.Module):
-    r"""
-    Adaptive RMS Norm used in Mochi.
-    Parameters:
-        embedding_dim (`int`): The size of each embedding vector.
-    """
-
-    def __init__(
-        self,
-        embedding_dim: int,
-        hidden_dim: int,
-        eps: float = 1e-5,
-        elementwise_affine: bool = False,
-    ) -> None:
-        super().__init__()
-
-        self.silu = nn.SiLU()
-        self.linear = nn.Linear(embedding_dim, hidden_dim)
-        self.norm = MochiModulatedRMSNorm(eps=eps)
-
-    def forward(self, hidden_states: torch.Tensor,
-                emb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        hidden_states_dtype = hidden_states.dtype
-
-        emb = self.linear(self.silu(emb))
-        scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
-
-        hidden_states = self.norm(hidden_states, (1 + scale_msa[:, None].to(torch.float32)))
-        hidden_states = hidden_states.to(hidden_states_dtype)
-
-        return hidden_states, gate_msa, scale_mlp, gate_mlp

fastvideo/models/mochi_hf/pipeline_mochi.py

@@ -1,757 +0,0 @@
-# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import copy
-import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
-from diffusers.loaders import Mochi1LoraLoaderMixin
-from diffusers.models.autoencoders import AutoencoderKL
-from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.video_processor import VideoProcessor
-from einops import rearrange
-from transformers import T5EncoderModel, T5TokenizerFast
-
-from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
-from fastvideo.utils.communications import all_gather
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-if is_torch_xla_available():
-    import torch_xla.core.xla_model as xm
-
-    XLA_AVAILABLE = True
-else:
-    XLA_AVAILABLE = False
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> import torch
-        >>> from diffusers import MochiPipeline
-        >>> from diffusers.utils import export_to_video
-
-        >>> pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", torch_dtype=torch.bfloat16)
-        >>> pipe.to("cuda")
-        >>> prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k."
-        >>> frames = pipe(prompt, num_inference_steps=28, guidance_scale=3.5).frames[0]
-        >>> export_to_video(frames, "mochi.mp4")
-        ```
-"""
-
-
-def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.16,
-):
-    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-    b = base_shift - m * base_seq_len
-    mu = image_seq_len * m + b
-    return mu
-
-
-# from: https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
-def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
-    if linear_steps is None:
-        linear_steps = num_steps // 2
-    linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
-    threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
-    quadratic_steps = num_steps - linear_steps
-    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
-    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
-    const = quadratic_coef * (linear_steps**2)
-    quadratic_sigma_schedule = [
-        quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps)
-    ]
-    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule
-    sigma_schedule = [1.0 - x for x in sigma_schedule]
-    return sigma_schedule
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
-    r"""
-    The mochi pipeline for text-to-video generation.
-
-    Reference: https://github.com/genmoai/models
-
-    Args:
-        transformer ([`MochiTransformer3DModel`]):
-            Conditional Transformer architecture to denoise the encoded video latents.
-        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae ([`AutoencoderKL`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
-        text_encoder ([`T5EncoderModel`]):
-            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
-            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
-        tokenizer (`CLIPTokenizer`):
-            Tokenizer of class
-            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
-        tokenizer (`T5TokenizerFast`):
-            Second Tokenizer of class
-            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
-    """
-
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-    _optional_components = []
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-
-    def __init__(
-        self,
-        scheduler: FlowMatchEulerDiscreteScheduler,
-        vae: AutoencoderKL,
-        text_encoder: T5EncoderModel,
-        tokenizer: T5TokenizerFast,
-        transformer: MochiTransformer3DModel,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler,
-        )
-        self.vae_spatial_scale_factor = 8
-        self.vae_temporal_scale_factor = 6
-        self.patch_size = 2
-
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_scale_factor)
-        self.tokenizer_max_length = (self.tokenizer.model_max_length
-                                     if hasattr(self, "tokenizer") and self.tokenizer is not None else 77)
-        self.default_height = 480
-        self.default_width = 848
-
-    # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds
-    def _get_t5_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]] = None,
-        num_videos_per_prompt: int = 1,
-        max_sequence_length: int = 256,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
-
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids
-        prompt_attention_mask = text_inputs.attention_mask
-        prompt_attention_mask = prompt_attention_mask.bool().to(device)
-
-        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1])
-            logger.warning("The following part of your input was truncated because `max_sequence_length` is set to "
-                           f" {max_sequence_length} tokens: {removed_text}")
-
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0]
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
-
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-
-        prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1)
-
-        return prompt_embeds, prompt_attention_mask
-
-    # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        do_classifier_free_guidance: bool = True,
-        num_videos_per_prompt: int = 1,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
-        max_sequence_length: int = 256,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
-                Whether to use classifier free guidance or not.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            device: (`torch.device`, *optional*):
-                torch device
-            dtype: (`torch.dtype`, *optional*):
-                torch dtype
-        """
-        device = device or self._execution_device
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        if prompt is not None:
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        if prompt_embeds is None:
-            prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds(
-                prompt=prompt,
-                num_videos_per_prompt=num_videos_per_prompt,
-                max_sequence_length=max_sequence_length,
-                device=device,
-                dtype=dtype,
-            )
-
-        if do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt = (batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt)
-
-            if prompt is not None and type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}.")
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`.")
-
-            (
-                negative_prompt_embeds,
-                negative_prompt_attention_mask,
-            ) = self._get_t5_prompt_embeds(
-                prompt=negative_prompt,
-                num_videos_per_prompt=num_videos_per_prompt,
-                max_sequence_length=max_sequence_length,
-                device=device,
-                dtype=dtype,
-            )
-
-        return (
-            prompt_embeds,
-            prompt_attention_mask,
-            negative_prompt_embeds,
-            negative_prompt_attention_mask,
-        )
-
-    def check_inputs(
-        self,
-        prompt,
-        height,
-        width,
-        callback_on_step_end_tensor_inputs=None,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-        prompt_attention_mask=None,
-        negative_prompt_attention_mask=None,
-    ):
-        if height % 8 != 0 or width % 8 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
-        if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs
-                                                                      for k in callback_on_step_end_tensor_inputs):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two.")
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.")
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-        if prompt_embeds is not None and prompt_attention_mask is None:
-            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
-
-        if (negative_prompt_embeds is not None and negative_prompt_attention_mask is None):
-            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
-
-        if prompt_embeds is not None and negative_prompt_embeds is not None:
-            if prompt_embeds.shape != negative_prompt_embeds.shape:
-                raise ValueError(
-                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
-                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
-                    f" {negative_prompt_embeds.shape}.")
-            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
-                raise ValueError(
-                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
-                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
-                    f" {negative_prompt_attention_mask.shape}.")
-
-    def enable_vae_slicing(self):
-        r"""
-        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
-        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        self.vae.enable_slicing()
-
-    def disable_vae_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_slicing()
-
-    def enable_vae_tiling(self):
-        r"""
-        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
-        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
-        processing larger images.
-        """
-        self.vae.enable_tiling()
-
-    def disable_vae_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_tiling()
-
-    def prepare_latents(
-        self,
-        batch_size,
-        num_channels_latents,
-        height,
-        width,
-        num_frames,
-        dtype,
-        device,
-        generator,
-        latents=None,
-    ):
-        height = height // self.vae_spatial_scale_factor
-        width = width // self.vae_spatial_scale_factor
-        num_frames = (num_frames - 1) // self.vae_temporal_scale_factor + 1
-
-        shape = (batch_size, num_channels_latents, num_frames, height, width)
-
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
-
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
-        latents = latents.to(dtype)
-        return latents
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1.0
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def attention_kwargs(self):
-        return self._attention_kwargs
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_frames: int = 19,
-        num_inference_steps: int = 64,
-        timesteps: List[int] = None,
-        guidance_scale: float = 4.5,
-        num_videos_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        max_sequence_length: int = 256,
-        return_all_states=False,
-    ):
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_frames (`int`, defaults to 16):
-                The number of video frames to generate
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            timesteps (`List[int]`, *optional*):
-                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
-                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
-                passed will be used. Must be in descending order.
-            guidance_scale (`float`, defaults to `4.5`):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of videos to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will ge generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            prompt_attention_mask (`torch.Tensor`, *optional*):
-                Pre-generated attention mask for text embeddings.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
-                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
-            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
-                Pre-generated attention mask for negative text embeddings.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.mochi.MochiPipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to `256`):
-                Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.mochi.MochiPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`~pipelines.mochi.MochiPipelineOutput`] is returned, otherwise a `tuple`
-                is returned where the first element is a list with the generated images.
-        """
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        height = height or self.default_height
-        width = width or self.default_width
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt=prompt,
-            height=height,
-            width=width,
-            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._interrupt = False
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-
-        # 3. Prepare text embeddings
-        (
-            prompt_embeds,
-            prompt_attention_mask,
-            negative_prompt_embeds,
-            negative_prompt_attention_mask,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            do_classifier_free_guidance=self.do_classifier_free_guidance,
-            num_videos_per_prompt=num_videos_per_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-            max_sequence_length=max_sequence_length,
-            device=device,
-        )
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
-
-        # 4. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_frames,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
-        if get_sequence_parallel_state():
-            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
-            latents = latents[:, :, rank, :, :, :]
-
-        original_noise = copy.deepcopy(latents)
-        # 5. Prepare timestep
-        # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
-        threshold_noise = 0.025
-        sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise)
-        sigmas = np.array(sigmas)
-        # check if of type FlowMatchEulerDiscreteScheduler
-        if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
-            timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler,
-                num_inference_steps,
-                device,
-                timesteps,
-                sigmas,
-            )
-        else:
-            timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler,
-                num_inference_steps,
-                device,
-            )
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-        self._num_timesteps = len(timesteps)
-
-        # 6. Denoising loop
-        self._progress_bar_config = {"disable": nccl_info.rank_within_group != 0}
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                latent_model_input = (torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents)
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    encoder_hidden_states=prompt_embeds,
-                    timestep=timestep,
-                    encoder_attention_mask=prompt_attention_mask,
-                    attention_kwargs=attention_kwargs,
-                    return_dict=False,
-                )[0]
-
-                # Mochi CFG + Sampling runs in FP32
-                noise_pred = noise_pred.to(torch.float32)
-                if self.do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents_dtype = latents.dtype
-                latents = self.scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0]
-                latents = latents.to(latents_dtype)
-
-                if latents.dtype != latents_dtype:
-                    if torch.backends.mps.is_available():
-                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
-                        latents = latents.to(latents_dtype)
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        if get_sequence_parallel_state():
-            latents = all_gather(latents, dim=2)
-            # latents_shape = list(latents.shape)
-            # full_shape = [latents_shape[0] * world_size] + latents_shape[1:]
-            # all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device)
-            # torch.distributed.all_gather_into_tensor(all_latents, latents)
-            # latents_list = list(all_latents.chunk(world_size, dim=0))
-            # latents = torch.cat(latents_list, dim=2)
-
-        if output_type == "latent":
-            video = latents
-        else:
-            # unscale/denormalize the latents
-            # denormalize with the mean and std if available and not None
-            has_latents_mean = (hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None)
-            has_latents_std = (hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None)
-            if has_latents_mean and has_latents_std:
-                latents_mean = (torch.tensor(self.vae.config.latents_mean).view(1, 12, 1, 1,
-                                                                                1).to(latents.device, latents.dtype))
-                latents_std = (torch.tensor(self.vae.config.latents_std).view(1, 12, 1, 1,
-                                                                              1).to(latents.device, latents.dtype))
-                latents = (latents * latents_std / self.vae.config.scaling_factor + latents_mean)
-            else:
-                latents = latents / self.vae.config.scaling_factor
-
-            video = self.vae.decode(latents, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-        if return_all_states:
-            # Pay extra attention here:
-            # prompt_embeds with shape torch.Size([2, 256]), where prompt_embeds[1] is the prompt_embeds for the actual prompt
-            # prompt_embeds[0] is for negative prompt
-            return original_noise, video, latents, prompt_embeds, prompt_attention_mask
-
-        if not return_dict:
-            return (video, )
-
-        return MochiPipelineOutput(frames=video)

fastvideo/models/stepvideo/__init__.py

@@ -1,7 +0,0 @@
-import os
-
-os.environ["NCCL_DEBUG"] = "ERROR"
-
-from .diffusion.scheduler import *
-from .diffusion.video_pipeline import *
-from .modules.model import *

fastvideo/models/stepvideo/__version__.py

@@ -1 +0,0 @@
-__version__ = "0.1.0"

fastvideo/models/stepvideo/config.py

@@ -1,174 +0,0 @@
-import argparse
-
-
-def parse_args(namespace=None):
-    parser = argparse.ArgumentParser(description="StepVideo inference script")
-
-    parser = add_extra_models_args(parser)
-    parser = add_denoise_schedule_args(parser)
-    parser = add_inference_args(parser)
-    parser = add_parallel_args(parser)
-
-    args = parser.parse_args(namespace=namespace)
-
-    return args
-
-
-def add_extra_models_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Extra models args, including vae, text encoders and tokenizers)")
-
-    group.add_argument(
-        "--vae_url",
-        type=str,
-        default='127.0.0.1',
-        help="vae url.",
-    )
-    group.add_argument(
-        "--caption_url",
-        type=str,
-        default='127.0.0.1',
-        help="caption url.",
-    )
-
-    return parser
-
-
-def add_denoise_schedule_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Denoise schedule args")
-
-    # Flow Matching
-    group.add_argument(
-        "--time_shift",
-        type=float,
-        default=7.0,
-        help="Shift factor for flow matching schedulers.",
-    )
-    group.add_argument(
-        "--flow_reverse",
-        action="store_true",
-        help="If reverse, learning/sampling from t=1 -> t=0.",
-    )
-    group.add_argument(
-        "--flow_solver",
-        type=str,
-        default="euler",
-        help="Solver for flow matching.",
-    )
-
-    return parser
-
-
-def add_inference_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Inference args")
-
-    # ======================== Model loads ========================
-    group.add_argument(
-        "--model_dir",
-        type=str,
-        default="./ckpts",
-        help="Root path of all the models, including t2v models and extra models.",
-    )
-    group.add_argument(
-        "--model_resolution",
-        type=str,
-        default="540p",
-        choices=["540p"],
-        help="Root path of all the models, including t2v models and extra models.",
-    )
-    group.add_argument(
-        "--use-cpu-offload",
-        action="store_true",
-        help="Use CPU offload for the model load.",
-    )
-
-    # ======================== Inference general setting ========================
-    group.add_argument(
-        "--batch_size",
-        type=int,
-        default=1,
-        help="Batch size for inference and evaluation.",
-    )
-    group.add_argument(
-        "--infer_steps",
-        type=int,
-        default=50,
-        help="Number of denoising steps for inference.",
-    )
-    group.add_argument(
-        "--save_path",
-        type=str,
-        default="./results",
-        help="Path to save the generated samples.",
-    )
-    group.add_argument(
-        "--name_suffix",
-        type=str,
-        default="",
-        help="Suffix for the names of saved samples.",
-    )
-    group.add_argument(
-        "--num_videos",
-        type=int,
-        default=1,
-        help="Number of videos to generate for each prompt.",
-    )
-    # ---sample size---
-    group.add_argument(
-        "--num_frames",
-        type=int,
-        default=204,
-        help="How many frames to sample from a video. ",
-    )
-    group.add_argument(
-        "--height",
-        type=int,
-        default=544,
-        help="The height of video sample",
-    )
-    group.add_argument(
-        "--width",
-        type=int,
-        default=992,
-        help="The width of video sample",
-    )
-    # --- prompt ---
-    group.add_argument(
-        "--prompt",
-        type=str,
-        default=None,
-        help="Prompt for sampling during evaluation.",
-    )
-    group.add_argument("--seed", type=int, default=1234, help="Seed for evaluation.")
-
-    # Classifier-Free Guidance
-    group.add_argument("--pos_magic",
-                       type=str,
-                       default="超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。",
-                       help="Positive magic prompt for sampling.")
-    group.add_argument("--neg_magic",
-                       type=str,
-                       default="画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。",
-                       help="Negative magic prompt for sampling.")
-    group.add_argument("--cfg_scale", type=float, default=9.0, help="Classifier free guidance scale.")
-
-    return parser
-
-
-def add_parallel_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Parallel args")
-
-    # ======================== Model loads ========================
-    group.add_argument(
-        "--ulysses_degree",
-        type=int,
-        default=8,
-        help="Ulysses degree.",
-    )
-    group.add_argument(
-        "--ring_degree",
-        type=int,
-        default=1,
-        help="Ulysses degree.",
-    )
-
-    return parser

fastvideo/models/stepvideo/diffusion/scheduler.py

@@ -1,220 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.schedulers.scheduling_utils import SchedulerMixin
-from diffusers.utils import BaseOutput, logging
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-@dataclass
-class FlowMatchDiscreteSchedulerOutput(BaseOutput):
-    """
-    Output class for the scheduler's `step` function output.
-
-    Args:
-        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
-            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
-            denoising loop.
-    """
-
-    prev_sample: torch.FloatTensor
-
-
-class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
-    """
-    Euler scheduler.
-
-    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
-    methods the library implements for all schedulers such as loading and saving.
-
-    Args:
-        num_train_timesteps (`int`, defaults to 1000):
-            The number of diffusion steps to train the model.
-        timestep_spacing (`str`, defaults to `"linspace"`):
-            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
-            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
-        reverse (`bool`, defaults to `True`):
-            Whether to reverse the timestep schedule.
-    """
-
-    _compatibles = []
-    order = 1
-
-    @register_to_config
-    def __init__(
-        self,
-        num_train_timesteps: int = 1000,
-        reverse: bool = False,
-        solver: str = "euler",
-        device: Union[str, torch.device] = None,
-    ):
-        sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
-
-        if not reverse:
-            sigmas = sigmas.flip(0)
-
-        self.sigmas = sigmas
-        # the value fed to model
-        self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32)
-
-        self._step_index = None
-        self._begin_index = None
-
-        self.device = device
-
-        self.supported_solver = ["euler"]
-        if solver not in self.supported_solver:
-            raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}")
-
-    @property
-    def step_index(self):
-        """
-        The index counter for current timestep. It will increase 1 after each scheduler step.
-        """
-        return self._step_index
-
-    @property
-    def begin_index(self):
-        """
-        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
-        """
-        return self._begin_index
-
-    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
-    def set_begin_index(self, begin_index: int = 0):
-        """
-        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
-
-        Args:
-            begin_index (`int`):
-                The begin index for the scheduler.
-        """
-        self._begin_index = begin_index
-
-    def _sigma_to_t(self, sigma):
-        return sigma * self.config.num_train_timesteps
-
-    def set_timesteps(
-        self,
-        num_inference_steps: int,
-        time_shift: float = 13.0,
-        device: Union[str, torch.device] = None,
-    ):
-        """
-        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
-
-        Args:
-            num_inference_steps (`int`):
-                The number of diffusion steps used when generating samples with a pre-trained model.
-            device (`str` or `torch.device`, *optional*):
-                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-            n_tokens (`int`, *optional*):
-                Number of tokens in the input sequence.
-        """
-        device = device or self.device
-        self.num_inference_steps = num_inference_steps
-
-        sigmas = torch.linspace(1, 0, num_inference_steps + 1, device=device)
-        sigmas = self.sd3_time_shift(sigmas, time_shift)
-
-        if not self.config.reverse:
-            sigmas = 1 - sigmas
-
-        self.sigmas = sigmas
-        self.timesteps = sigmas[:-1]
-
-        # Reset step index
-        self._step_index = None
-
-    def index_for_timestep(self, timestep, schedule_timesteps=None):
-        if schedule_timesteps is None:
-            schedule_timesteps = self.timesteps
-
-        indices = (schedule_timesteps == timestep).nonzero()
-
-        # The sigma index that is taken for the **very** first `step`
-        # is always the second index (or the last index if there is only 1)
-        # This way we can ensure we don't accidentally skip a sigma in
-        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
-        pos = 1 if len(indices) > 1 else 0
-
-        return indices[pos].item()
-
-    def _init_step_index(self, timestep):
-        if self.begin_index is None:
-            if isinstance(timestep, torch.Tensor):
-                timestep = timestep.to(self.timesteps.device)
-            self._step_index = self.index_for_timestep(timestep)
-        else:
-            self._step_index = self._begin_index
-
-    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
-        return sample
-
-    def sd3_time_shift(self, t: torch.Tensor, time_shift: float = 13.0):
-        return (time_shift * t) / (1 + (time_shift - 1) * t)
-
-    def step(
-        self,
-        model_output: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
-        sample: torch.FloatTensor,
-        return_dict: bool = False,
-    ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
-        """
-        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
-        process from the learned model outputs (most often the predicted noise).
-
-        Args:
-            model_output (`torch.FloatTensor`):
-                The direct output from learned diffusion model.
-            timestep (`float`):
-                The current discrete timestep in the diffusion chain.
-            sample (`torch.FloatTensor`):
-                A current instance of a sample created by the diffusion process.
-            generator (`torch.Generator`, *optional*):
-                A random number generator.
-            n_tokens (`int`, *optional*):
-                Number of tokens in the input sequence.
-            return_dict (`bool`):
-                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
-                tuple.
-
-        Returns:
-            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
-                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
-                returned, otherwise a tuple is returned where the first element is the sample tensor.
-        """
-
-        if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor)
-                or isinstance(timestep, torch.LongTensor)):
-            raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-                              " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-                              " one of the `scheduler.timesteps` as a timestep."), )
-
-        if self.step_index is None:
-            self._init_step_index(timestep)
-
-        # Upcast to avoid precision issues when computing prev_sample
-        sample = sample.to(torch.float32)
-
-        dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
-
-        if self.config.solver == "euler":
-            prev_sample = sample + model_output.to(torch.float32) * dt
-        else:
-            raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}")
-
-        # upon completion increase step index by one
-        self._step_index += 1
-
-        if not return_dict:
-            return prev_sample
-
-        return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
-
-    def __len__(self):
-        return self.config.num_train_timesteps

fastvideo/models/stepvideo/diffusion/video_pipeline.py

@@ -1,325 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-
-import asyncio
-import pickle
-from dataclasses import dataclass
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-import torch
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.utils import BaseOutput
-
-from fastvideo.models.stepvideo.diffusion.scheduler import FlowMatchDiscreteScheduler
-from fastvideo.models.stepvideo.modules.model import StepVideoModel
-from fastvideo.models.stepvideo.utils import VideoProcessor
-
-
-def call_api_gen(url, api, port=8080):
-    url = f"http://{url}:{port}/{api}-api"
-    import aiohttp
-
-    async def _fn(samples, *args, **kwargs):
-        if api == 'vae':
-            data = {
-                "samples": samples,
-            }
-        elif api == 'caption':
-            data = {
-                "prompts": samples,
-            }
-        else:
-            raise Exception(f"Not supported api: {api}...")
-
-        async with aiohttp.ClientSession() as sess:
-            data_bytes = pickle.dumps(data)
-            async with sess.get(url, data=data_bytes, timeout=12000) as response:
-                result = bytearray()
-                while not response.content.at_eof():
-                    chunk = await response.content.read(1024)
-                    result += chunk
-                response_data = pickle.loads(result)
-        return response_data
-
-    return _fn
-
-
-@dataclass
-class StepVideoPipelineOutput(BaseOutput):
-    video: Union[torch.Tensor, np.ndarray]
-
-
-class StepVideoPipeline(DiffusionPipeline):
-    r"""
-    Pipeline for text-to-video generation using StepVideo.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Args:
-        transformer ([`StepVideoModel`]):
-            Conditional Transformer to denoise the encoded image latents.
-        scheduler ([`FlowMatchDiscreteScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae_url:
-            remote vae server's url.
-        caption_url:
-            remote caption (stepllm and clip) server's url.
-    """
-
-    def __init__(
-        self,
-        transformer: StepVideoModel,
-        scheduler: FlowMatchDiscreteScheduler,
-        vae_url: str = '127.0.0.1',
-        caption_url: str = '127.0.0.1',
-        save_path: str = './results',
-        name_suffix: str = '',
-    ):
-        super().__init__()
-
-        self.register_modules(
-            transformer=transformer,
-            scheduler=scheduler,
-        )
-
-        self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 8
-        self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 16
-        self.video_processor = VideoProcessor(save_path, name_suffix)
-
-        self.vae_url = vae_url
-        self.caption_url = caption_url
-        self.setup_api(self.vae_url, self.caption_url)
-
-    def setup_api(self, vae_url, caption_url):
-        self.vae_url = vae_url
-        self.caption_url = caption_url
-        self.caption = call_api_gen(caption_url, 'caption')
-        self.vae = call_api_gen(vae_url, 'vae')
-        return self
-
-    def encode_prompt(
-        self,
-        prompt: str,
-        neg_magic: str = '',
-        pos_magic: str = '',
-    ):
-        device = self._execution_device
-        prompts = [prompt + pos_magic]
-        bs = len(prompts)
-        prompts += [neg_magic] * bs
-
-        data = asyncio.run(self.caption(prompts))
-        prompt_embeds, prompt_attention_mask, clip_embedding = data['y'].to(device), data['y_mask'].to(
-            device), data['clip_embedding'].to(device)
-
-        return prompt_embeds, clip_embedding, prompt_attention_mask
-
-    def decode_vae(self, samples):
-        samples = asyncio.run(self.vae(samples.cpu()))
-        return samples
-
-    def check_inputs(self, num_frames, width, height):
-        num_frames = max(num_frames // 17 * 17, 1)
-        width = max(width // 16 * 16, 16)
-        height = max(height // 16 * 16, 16)
-        return num_frames, width, height
-
-    def prepare_latents(
-        self,
-        batch_size: int,
-        num_channels_latents: 64,
-        height: int = 544,
-        width: int = 992,
-        num_frames: int = 204,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-
-        num_frames, width, height = self.check_inputs(num_frames, width, height)
-        shape = (
-            batch_size,
-            max(num_frames // 17 * 3, 1),
-            num_channels_latents,
-            int(height) // self.vae_scale_factor_spatial,
-            int(width) // self.vae_scale_factor_spatial,
-        )  # b,f,c,h,w
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
-
-        if generator is None:
-            generator = torch.Generator(device=self._execution_device)
-
-        latents = torch.randn(shape, generator=generator, device=device, dtype=dtype)
-        return latents
-
-    @torch.inference_mode()
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        height: int = 544,
-        width: int = 992,
-        num_frames: int = 204,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 9.0,
-        time_shift: float = 13.0,
-        neg_magic: str = "",
-        pos_magic: str = "",
-        num_videos_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "mp4",
-        output_file_name: Optional[str] = "",
-        return_dict: bool = True,
-        mask_strategy: Optional[Dict[str, list]] = None,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            height (`int`, defaults to `544`):
-                The height in pixels of the generated image.
-            width (`int`, defaults to `992`):
-                The width in pixels of the generated image.
-            num_frames (`int`, defaults to `204`):
-                The number of frames in the generated video.
-            num_inference_steps (`int`, defaults to `50`):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            guidance_scale (`float`, defaults to `9.0`):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality. 
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            output_file_name(`str`, *optional*`):
-                The output mp4 file name.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`StepVideoPipelineOutput`] instead of a plain tuple.
-
-        Examples:
-
-        Returns:
-            [`~StepVideoPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`StepVideoPipelineOutput`] is returned, otherwise a `tuple` is returned
-                where the first element is a list with the generated images and the second element is a list of `bool`s
-                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
-        """
-
-        # 1. Check inputs. Raise error if not correct
-        device = self._execution_device
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # 3. Encode input prompt
-        prompt_embeds, prompt_embeds_2, prompt_attention_mask = self.encode_prompt(
-            prompt=prompt,
-            neg_magic=neg_magic,
-            pos_magic=pos_magic,
-        )
-
-        transformer_dtype = self.transformer.dtype
-        prompt_embeds = prompt_embeds.to(transformer_dtype)
-        prompt_attention_mask = prompt_attention_mask.to(transformer_dtype)
-        prompt_embeds_2 = prompt_embeds_2.to(transformer_dtype)
-
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps=num_inference_steps, time_shift=time_shift, device=device)
-
-        # 5. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_frames,
-            torch.bfloat16,
-            device,
-            generator,
-            latents,
-        )
-
-        def dict_to_3d_list(best_masks, t_max=50, l_max=48, h_max=48):
-            result = [[[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max)]
-            if best_masks is None:
-                return result
-            for key, value in best_masks.items():
-                timestep, layer, head = map(int, key.split('_'))
-                result[timestep][layer][head] = value
-            return result
-
-        mask_strategy = dict_to_3d_list(mask_strategy)
-
-        #best_mask_selections = None
-        # 7. Denoising loop
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(self.scheduler.timesteps):
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-                latent_model_input = latent_model_input.to(transformer_dtype)
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    timestep=timestep,
-                    encoder_hidden_states=prompt_embeds,
-                    encoder_attention_mask=prompt_attention_mask,
-                    encoder_hidden_states_2=prompt_embeds_2,
-                    return_dict=False,
-                    mask_strategy=mask_strategy[i],
-                )
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(model_output=noise_pred, timestep=t, sample=latents)
-
-                progress_bar.update()
-
-        if not torch.distributed.is_initialized() or int(torch.distributed.get_rank()) == 0:
-            if not output_type == "latent":
-                video = self.decode_vae(latents)
-                video = self.video_processor.postprocess_video(video,
-                                                               output_file_name=output_file_name,
-                                                               output_type=output_type)
-            else:
-                video = latents
-
-            # Offload all models
-            self.maybe_free_model_hooks()
-
-            if not return_dict:
-                return (video, )
-
-            return StepVideoPipelineOutput(video=video)

fastvideo/models/stepvideo/modules/attentions.py

@@ -1,96 +0,0 @@
-import torch
-import torch.nn as nn
-from einops import rearrange
-from flash_attn import flash_attn_func
-
-try:
-    from st_attn import sliding_tile_attention
-except ImportError:
-    print("Could not load Sliding Tile Attention.")
-    sliding_tile_attention = None
-
-from fastvideo.utils.communications import all_to_all_4D
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-
-class Attention(nn.Module):
-
-    def __init__(self):
-        super().__init__()
-
-    def attn_processor(self, attn_type):
-        if attn_type == 'torch':
-            return self.torch_attn_func
-        elif attn_type == 'parallel':
-            return self.parallel_attn_func
-        else:
-            raise Exception('Not supported attention type...')
-
-    def tile(self, x, sp_size):
-        x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=36 // sp_size, h=48, w=48)
-        return rearrange(x,
-                         "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-                         n_t=6,
-                         n_h=6,
-                         n_w=6,
-                         ts_t=6,
-                         ts_h=8,
-                         ts_w=8)
-
-    def untile(self, x, sp_size):
-        x = rearrange(x,
-                      "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-                      n_t=6,
-                      n_h=6,
-                      n_w=6,
-                      ts_t=6,
-                      ts_h=8,
-                      ts_w=8)
-        return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=36 // sp_size, h=48, w=48)
-
-    def torch_attn_func(self, q, k, v, attn_mask=None, causal=False, drop_rate=0.0, **kwargs):
-
-        if attn_mask is not None and attn_mask.dtype != torch.bool:
-            attn_mask = attn_mask.to(q.dtype)
-
-        if attn_mask is not None and attn_mask.ndim == 3:  ## no head
-            n_heads = q.shape[2]
-            attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)
-
-        q, k, v = map(lambda x: rearrange(x, 'b s h d -> b h s d'), (q, k, v))
-        x = torch.nn.functional.scaled_dot_product_attention(q,
-                                                             k,
-                                                             v,
-                                                             attn_mask=attn_mask,
-                                                             dropout_p=drop_rate,
-                                                             is_causal=causal)
-        x = rearrange(x, 'b h s d -> b s h d')
-        return x
-
-    def parallel_attn_func(self, q, k, v, causal=False, mask_strategy=None, **kwargs):
-        if get_sequence_parallel_state():
-            q = all_to_all_4D(q, scatter_dim=2, gather_dim=1)
-            k = all_to_all_4D(k, scatter_dim=2, gather_dim=1)
-            v = all_to_all_4D(v, scatter_dim=2, gather_dim=1)
-
-        if mask_strategy[0] is not None:
-            q = self.tile(q, nccl_info.sp_size).transpose(1, 2).contiguous()
-            k = self.tile(k, nccl_info.sp_size).transpose(1, 2).contiguous()
-            v = self.tile(v, nccl_info.sp_size).transpose(1, 2).contiguous()
-
-            head_num = q.size(1)  # 48 // sp_size
-            current_rank = nccl_info.rank_within_group
-
-            start_head = current_rank * head_num
-            windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)]
-
-            x = sliding_tile_attention(q, k, v, windows, 0, False).transpose(1, 2).contiguous()
-            x = self.untile(x, nccl_info.sp_size)
-        else:
-            x = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)
-
-        if get_sequence_parallel_state():
-            x = all_to_all_4D(x, scatter_dim=1, gather_dim=2)
-
-        x = x.to(q.dtype)
-        return x

fastvideo/models/stepvideo/modules/blocks.py

@@ -1,296 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-# ==============================================================================
-from typing import Optional
-
-import torch
-import torch.nn as nn
-from einops import rearrange
-
-from fastvideo.models.stepvideo.modules.attentions import Attention
-from fastvideo.models.stepvideo.modules.normalization import RMSNorm
-from fastvideo.models.stepvideo.modules.rope import RoPE3D
-
-
-class SelfAttention(Attention):
-
-    def __init__(self, hidden_dim, head_dim, bias=False, with_rope=True, with_qk_norm=True, attn_type='torch'):
-        super().__init__()
-        self.head_dim = head_dim
-        self.n_heads = hidden_dim // head_dim
-
-        self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=bias)
-        self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias)
-
-        self.with_rope = with_rope
-        self.with_qk_norm = with_qk_norm
-        if self.with_qk_norm:
-            self.q_norm = RMSNorm(head_dim, elementwise_affine=True)
-            self.k_norm = RMSNorm(head_dim, elementwise_affine=True)
-
-        if self.with_rope:
-            self.rope_3d = RoPE3D(freq=1e4, F0=1.0, scaling_factor=1.0)
-            self.rope_ch_split = [64, 32, 32]
-
-        self.core_attention = self.attn_processor(attn_type=attn_type)
-        self.parallel = attn_type == 'parallel'
-
-    def apply_rope3d(self, x, fhw_positions, rope_ch_split, parallel=True):
-        x = self.rope_3d(x, fhw_positions, rope_ch_split, parallel)
-        return x
-
-    def forward(self, x, cu_seqlens=None, max_seqlen=None, rope_positions=None, attn_mask=None, mask_strategy=None):
-        xqkv = self.wqkv(x)
-        xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim)
-
-        xq, xk, xv = torch.split(xqkv, [self.head_dim] * 3, dim=-1)  ## seq_len, n, dim
-
-        if self.with_qk_norm:
-            xq = self.q_norm(xq)
-            xk = self.k_norm(xk)
-
-        if self.with_rope:
-            xq = self.apply_rope3d(xq, rope_positions, self.rope_ch_split, parallel=self.parallel)
-            xk = self.apply_rope3d(xk, rope_positions, self.rope_ch_split, parallel=self.parallel)
-
-        output = self.core_attention(xq,
-                                     xk,
-                                     xv,
-                                     cu_seqlens=cu_seqlens,
-                                     max_seqlen=max_seqlen,
-                                     attn_mask=attn_mask,
-                                     mask_strategy=mask_strategy)
-        output = rearrange(output, 'b s h d -> b s (h d)')
-        output = self.wo(output)
-
-        return output
-
-
-class CrossAttention(Attention):
-
-    def __init__(self, hidden_dim, head_dim, bias=False, with_qk_norm=True, attn_type='torch'):
-        super().__init__()
-        self.head_dim = head_dim
-        self.n_heads = hidden_dim // head_dim
-
-        self.wq = nn.Linear(hidden_dim, hidden_dim, bias=bias)
-        self.wkv = nn.Linear(hidden_dim, hidden_dim * 2, bias=bias)
-        self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias)
-
-        self.with_qk_norm = with_qk_norm
-        if self.with_qk_norm:
-            self.q_norm = RMSNorm(head_dim, elementwise_affine=True)
-            self.k_norm = RMSNorm(head_dim, elementwise_affine=True)
-
-        self.core_attention = self.attn_processor(attn_type=attn_type)
-
-    def forward(self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, attn_mask=None):
-        xq = self.wq(x)
-        xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim)
-
-        xkv = self.wkv(encoder_hidden_states)
-        xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim)
-
-        xk, xv = torch.split(xkv, [self.head_dim] * 2, dim=-1)  ## seq_len, n, dim
-
-        if self.with_qk_norm:
-            xq = self.q_norm(xq)
-            xk = self.k_norm(xk)
-
-        output = self.core_attention(xq, xk, xv, attn_mask=attn_mask)
-
-        output = rearrange(output, 'b s h d -> b s (h d)')
-        output = self.wo(output)
-
-        return output
-
-
-class GELU(nn.Module):
-    r"""
-    GELU activation function with tanh approximation support with `approximate="tanh"`.
-
-    Parameters:
-        dim_in (`int`): The number of channels in the input.
-        dim_out (`int`): The number of channels in the output.
-        approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
-        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
-    """
-
-    def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True):
-        super().__init__()
-        self.proj = nn.Linear(dim_in, dim_out, bias=bias)
-        self.approximate = approximate
-
-    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
-        return torch.nn.functional.gelu(gate, approximate=self.approximate)
-
-    def forward(self, hidden_states):
-        hidden_states = self.proj(hidden_states)
-        hidden_states = self.gelu(hidden_states)
-        return hidden_states
-
-
-class FeedForward(nn.Module):
-
-    def __init__(
-        self,
-        dim: int,
-        inner_dim: Optional[int] = None,
-        dim_out: Optional[int] = None,
-        mult: int = 4,
-        bias: bool = False,
-    ):
-        super().__init__()
-        inner_dim = dim * mult if inner_dim is None else inner_dim
-        dim_out = dim if dim_out is None else dim_out
-        self.net = nn.ModuleList([
-            GELU(dim, inner_dim, approximate="tanh", bias=bias),
-            nn.Identity(),
-            nn.Linear(inner_dim, dim_out, bias=bias)
-        ])
-
-    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
-        for module in self.net:
-            hidden_states = module(hidden_states)
-        return hidden_states
-
-
-def modulate(x, scale, shift):
-    x = x * (1 + scale) + shift
-    return x
-
-
-def gate(x, gate):
-    x = gate * x
-    return x
-
-
-class StepVideoTransformerBlock(nn.Module):
-    r"""
-    A basic Transformer block.
-
-    Parameters:
-        dim (`int`): The number of channels in the input and output.
-        num_attention_heads (`int`): The number of heads to use for multi-head attention.
-        attention_head_dim (`int`): The number of channels in each head.
-        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
-        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
-        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
-        num_embeds_ada_norm (:
-            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
-        attention_bias (:
-            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
-        only_cross_attention (`bool`, *optional*):
-            Whether to use only cross-attention layers. In this case two cross attention layers are used.
-        double_self_attention (`bool`, *optional*):
-            Whether to use two self-attention layers. In this case no cross attention layers are used.
-        upcast_attention (`bool`, *optional*):
-            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
-        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
-            Whether to use learnable elementwise affine parameters for normalization.
-        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
-            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
-        final_dropout (`bool` *optional*, defaults to False):
-            Whether to apply a final dropout after the last feed-forward layer.
-        attention_type (`str`, *optional*, defaults to `"default"`):
-            The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
-        positional_embeddings (`str`, *optional*, defaults to `None`):
-            The type of positional embeddings to apply to.
-        num_positional_embeddings (`int`, *optional*, defaults to `None`):
-            The maximum number of positional embeddings to apply.
-    """
-
-    def __init__(self,
-                 dim: int,
-                 attention_head_dim: int,
-                 norm_eps: float = 1e-5,
-                 ff_inner_dim: Optional[int] = None,
-                 ff_bias: bool = False,
-                 attention_type: str = 'parallel'):
-        super().__init__()
-        self.dim = dim
-        self.norm1 = nn.LayerNorm(dim, eps=norm_eps)
-        self.attn1 = SelfAttention(dim,
-                                   attention_head_dim,
-                                   bias=False,
-                                   with_rope=True,
-                                   with_qk_norm=True,
-                                   attn_type=attention_type)
-
-        self.norm2 = nn.LayerNorm(dim, eps=norm_eps)
-        self.attn2 = CrossAttention(dim, attention_head_dim, bias=False, with_qk_norm=True, attn_type='torch')
-
-        self.ff = FeedForward(dim=dim, inner_dim=ff_inner_dim, dim_out=dim, bias=ff_bias)
-
-        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
-
-    @torch.no_grad()
-    def forward(self,
-                q: torch.Tensor,
-                kv: Optional[torch.Tensor] = None,
-                timestep: Optional[torch.LongTensor] = None,
-                attn_mask=None,
-                rope_positions: list = None,
-                mask_strategy=None) -> torch.Tensor:
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (torch.clone(chunk) for chunk in (
-            self.scale_shift_table[None] + timestep.reshape(-1, 6, self.dim)).chunk(6, dim=1))
-
-        scale_shift_q = modulate(self.norm1(q), scale_msa, shift_msa)
-
-        attn_q = self.attn1(scale_shift_q, rope_positions=rope_positions, mask_strategy=mask_strategy)
-
-        q = gate(attn_q, gate_msa) + q
-
-        attn_q = self.attn2(q, kv, attn_mask)
-
-        q = attn_q + q
-
-        scale_shift_q = modulate(self.norm2(q), scale_mlp, shift_mlp)
-
-        ff_output = self.ff(scale_shift_q)
-
-        q = gate(ff_output, gate_mlp) + q
-
-        return q
-
-
-class PatchEmbed(nn.Module):
-    """2D Image to Patch Embedding"""
-
-    def __init__(
-        self,
-        patch_size=64,
-        in_channels=3,
-        embed_dim=768,
-        layer_norm=False,
-        flatten=True,
-        bias=True,
-    ):
-        super().__init__()
-
-        self.flatten = flatten
-        self.layer_norm = layer_norm
-
-        self.proj = nn.Conv2d(in_channels,
-                              embed_dim,
-                              kernel_size=(patch_size, patch_size),
-                              stride=patch_size,
-                              bias=bias)
-
-    def forward(self, latent):
-        latent = self.proj(latent).to(latent.dtype)
-        if self.flatten:
-            latent = latent.flatten(2).transpose(1, 2)  # BCHW -> BNC
-        if self.layer_norm:
-            latent = self.norm(latent)
-
-        return latent