pd

.github/ISSUE_TEMPLATE/config.yml

@@ -1 +0,0 @@
-blank_issues_enabled: false

README.md

@@ -120,49 +120,51 @@ Then you can run the finetune with:
 ```
 bash scripts/finetune/finetune_mochi.sh # for mochi
 ```
-**Note that for finetuning, we did not tune the hyperparameters in the provided script.**
+**Note that for finetuning, we did not tune the hyperparameters in the provided script**
 ### ⚡ Lora Finetune 
-
-Hunyuan supports Lora fine-tuning of videos up to 720p. Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through:
+Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through:
 ```bash
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight --local_dir=data/Hunyuan-Black-Myth-Wukong-lora-weight --repo_type=model
 ```
-#### Minimum Hardware Requirement
-- 40 GB GPU memory each for 2 GPUs with lora.
-- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.  
-
-
 Currently, both Mochi and Hunyuan models support Lora finetuning through diffusers. To generate personalized videos from your own dataset, you'll need to follow three main steps: dataset preparation, finetuning, and inference.
 
 #### Dataset Preparation
 We provide scripts to better help you get started to train on your own characters!  
-You can run this to organize your dataset to get the videos2caption.json before preprocess. Specify your video folder and corresponding caption folder (caption files should be .txt files and have the same name with its video):
+You can run this to organize your dataset to get the videos2caption.json before preprocess. Specify your video folder and corresponding caption folder(Caption files should be .txt files and have the same name with its video):
 ```
 python scripts/dataset_preparation/prepare_json_file.py --video_dir data/input_videos/ --prompt_dir data/captions/ --output_path data/output_folder/videos2caption.json --verbose
 ```
 Also, we provide script to resize your videos:
 ```
-python scripts/data_preprocess/resize_videos.py 
+python scripts/data_preprocess/resize_videos.py \
+  --input_dir data/raw_videos/ \
+  --output_dir data/resized_videos/ \
+  --width 1280 \
+  --height 720 \
+  --fps 30
 ```
 #### Finetuning
 After basic dataset preparation and preprocess, you can start to finetune your model using Lora:
 ```
 bash scripts/finetune/finetune_hunyuan_hf_lora.sh
+bash scripts/finetune/finetune_mochi_lora.sh
 ```
 #### Inference
-For inference with Lora checkpoint, you can run the following scripts with additional parameter `--lora_checkpoint_dir`:
+For inference with Lora checkpoint, you can run the following scripts with Additional parameter --lora_checkpoint_dir:
 ```
 bash scripts/inference/inference_hunyuan_hf.sh 
+bash scripts/inference/inference_mochi_hf.sh 
 ```
-**We also provide scripts for Mochi in the same directory.**
-
+#### Minimum Hardware Requirement
+- 40 GB GPU memory each for 2 GPUs with lora
+- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.
 #### Finetune with Both Image and Video
 Our codebase support finetuning with both image and video. 
 ```bash
 bash scripts/finetune/finetune_hunyuan.sh
 bash scripts/finetune/finetune_mochi_lora_mix.sh
 ```
-For Image-Video Mixture Fine-tuning, make sure to enable the `--group_frame` option in your script.
+For Image-Video Mixture Fine-tuning, make sure to enable the --group_frame option in your script.
 
 ## 📑 Development Plan
 

fastvideo/models/hunyuan/modules/attenion.py

@@ -32,7 +32,7 @@ def attention(
     return out
 
 
-def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask):
+def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, sample_step, layer_id):
     # 1GPU torch.Size([1, 11264, 24, 128]) tensor([    0, 11275, 11520], device='cuda:0', dtype=torch.int32)
     # 2GPU torch.Size([1, 5632, 24, 128]) tensor([   0, 5643, 5888], device='cuda:0', dtype=torch.int32)
     query, encoder_query = q
@@ -57,6 +57,90 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask):
     sequence_length = query.size(1)
     encoder_sequence_length = encoder_query.size(1)
 
+    # attn_map plot
+    num_heads = query.size(2) 
+    map_q = query.transpose(1, 2)  # [B, H, S, D]
+    map_k = torch.cat([key, encoder_key], dim=1).transpose(1, 2)    # [B, H, S+T, D]
+    d_k = map_q.size(-1)
+
+    shape = (16, 45, 45)
+    q_coords = torch.tensor([get_block(i, shape=shape) for i in range(sequence_length)],
+                           device='cuda', dtype=torch.int16)
+    k_coords = torch.tensor([get_block(i, shape=shape) for i in range(img_kv_len)],
+                           device='cuda', dtype=torch.int16)
+
+    diffs = (q_coords.unsqueeze(1) - k_coords.unsqueeze(0)).abs()  # [seq_len, kv_len, 3]
+    mask_t_diff = 6
+    mask_x_diff = 12
+    mask_y_diff = 12
+    
+    mask_t = diffs[..., 0] <= mask_t_diff
+    mask_x = diffs[..., 1] <= mask_x_diff
+    mask_y = diffs[..., 2] <= mask_y_diff
+    valid_mask_2d = mask_t & mask_x & mask_y  # [seq_len, kv_len]
+    del mask_t, mask_x, mask_y, diffs
+    
+    # Pad mask for text part (all True)
+    full_valid_mask = F.pad(valid_mask_2d, (0, encoder_sequence_length), value=True)  # [seq_len, kv_len+text_len]
+    img_mask_density = valid_mask_2d.float().mean().item()
+    full_mask_density = full_valid_mask.float().mean().item()
+    del valid_mask_2d
+
+    chunk_size = 512
+    
+    save_attn = False #(sample_step == 1) and (layer_id == 59)
+    if save_attn:
+        attn_map_cumulated = torch.zeros((sequence_length, img_kv_len + encoder_sequence_length), 
+                                       dtype=torch.float32, device='cuda')
+    
+    # 逐head计算
+    for head_idx in range(num_heads):
+        current_q = map_q[:, head_idx:head_idx+1].to(dtype=torch.float32)  # [B, 1, S, D]
+        current_k = map_k[:, head_idx:head_idx+1].to(dtype=torch.float32)  # [B, 1, S+T, D]
+
+        valid_score_total = 0.0
+        all_score_total = 0.0
+        
+        if save_attn:
+            head_attn_map = torch.zeros((sequence_length, img_kv_len + encoder_sequence_length), 
+                                      dtype=torch.float32, device='cuda')
+        
+        # 分块计算
+        for i in range(0, current_q.size(2), chunk_size):
+            chunk_end = min(i + chunk_size, current_q.size(2))
+            q_chunk = current_q[:, :, i:chunk_end]  # [B, 1, chunk_size, D]
+            
+            scores_32 = torch.matmul(
+                q_chunk, 
+                current_k.transpose(-2, -1)
+            ) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32, device='cuda'))
+            
+            attn_weights = F.softmax(scores_32, dim=-1)
+            
+            chunk_valid_mask = full_valid_mask[i:chunk_end].unsqueeze(0).unsqueeze(0)
+            valid_score_sum = (attn_weights * chunk_valid_mask).sum()
+            all_score_sum = attn_weights.sum()
+            
+            valid_score_total += valid_score_sum.item()
+            all_score_total += all_score_sum.item()
+            
+            # For last layer, accumulate attention weights
+            if save_attn:
+                head_attn_map[i:chunk_end] = attn_weights.squeeze(0).squeeze(0)
+            
+            del scores_32, attn_weights
+            torch.cuda.empty_cache()
+        
+        recall = valid_score_total / (all_score_total + 1e-9)
+        print(f"step{sample_step}_layer{layer_id}_head{head_idx}_window_diff_{mask_t_diff}_{mask_x_diff}_{mask_y_diff}_shape_{shape[0]}_{shape[1]}_{shape[2]}_img{img_mask_density:.4f}_full{full_mask_density:.4f}'s recall:", recall)
+        
+        if save_attn:
+            attn_map_cumulated += head_attn_map
+    
+    if save_attn:
+        attn_map_avg = attn_map_cumulated / num_heads
+        torch.save(attn_map_avg, f'logs/attn_map_step{sample_step}_layer{layer_id}.pt')
+    
     # Hint: please check encoder_query.shape
     query = torch.cat([query, encoder_query], dim=1)
     key = torch.cat([key, encoder_key], dim=1)
@@ -70,7 +154,7 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_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 get_sequence_parallel_state():
@@ -88,3 +172,12 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask):
     attn = attn.reshape(b, s, -1)
 
     return attn
+
+def get_block(idx, shape=(16, 30, 30)):
+    t_size, x_size, y_size = shape
+    xy = x_size * y_size
+    t = idx // xy
+    r = idx % xy
+    x = r // y_size
+    y = r % y_size
+    return t, x, y

fastvideo/models/hunyuan/modules/models.py

@@ -38,6 +38,7 @@ class MMDoubleStreamBlock(nn.Module):
         qkv_bias: bool = False,
         dtype: Optional[torch.dtype] = None,
         device: Optional[torch.device] = None,
+        layer_id: int = 0,
     ):
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -124,6 +125,8 @@ class MMDoubleStreamBlock(nn.Module):
             **factory_kwargs,
         )
         self.hybrid_seq_parallel_attn = None
+        self.sample_step = 0
+        self.layer_id = layer_id
 
     def enable_deterministic(self):
         self.deterministic = True
@@ -207,6 +210,10 @@ class MMDoubleStreamBlock(nn.Module):
         txt_q = self.txt_attn_q_norm(txt_q).to(txt_v)
         txt_k = self.txt_attn_k_norm(txt_k).to(txt_v)
 
+        print("DOUBLE====")
+        print(img_q.shape, txt_q.shape)
+        
+        self.sample_step += 1
         attn = parallel_attention(
             (img_q, txt_q),
             (img_k, txt_k),
@@ -214,6 +221,8 @@ class MMDoubleStreamBlock(nn.Module):
             img_q_len=img_q.shape[1],
             img_kv_len=img_k.shape[1],
             text_mask=text_mask,
+            sample_step=self.sample_step,
+            layer_id=self.layer_id,
         )
 
         # attention computation end
@@ -264,6 +273,7 @@ class MMSingleStreamBlock(nn.Module):
         qk_scale: float = None,
         dtype: Optional[torch.dtype] = None,
         device: Optional[torch.device] = None,
+        layer_id: int = 0,
     ):
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -304,6 +314,8 @@ class MMSingleStreamBlock(nn.Module):
             **factory_kwargs,
         )
         self.hybrid_seq_parallel_attn = None
+        self.sample_step = 0
+        self.layer_id = layer_id
 
     def enable_deterministic(self):
         self.deterministic = True
@@ -354,6 +366,7 @@ class MMSingleStreamBlock(nn.Module):
         ), 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
 
+        self.sample_step += 1
         attn = parallel_attention(
             (img_q, txt_q),
             (img_k, txt_k),
@@ -361,6 +374,8 @@ class MMSingleStreamBlock(nn.Module):
             img_q_len=img_q.shape[1],
             img_kv_len=img_k.shape[1],
             text_mask=text_mask,
+            sample_step=self.sample_step,
+            layer_id=self.layer_id,
         )
 
         # attention computation end
@@ -521,6 +536,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
                 qk_norm=qk_norm,
                 qk_norm_type=qk_norm_type,
                 qkv_bias=qkv_bias,
+                layer_id=_,
                 **factory_kwargs,
             ) for _ in range(mm_double_blocks_depth)
         ])
@@ -534,6 +550,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
                 mlp_act_type=mlp_act_type,
                 qk_norm=qk_norm,
                 qk_norm_type=qk_norm_type,
+                layer_id=_+mm_double_blocks_depth,
                 **factory_kwargs,
             ) for _ in range(mm_single_blocks_depth)
         ])

fastvideo/ops/modulate/k_modulate.py

@@ -1,102 +0,0 @@
-import triton
-import triton.language as tl
-
-CONFIG_LIST = [
-    triton.Config({"BLOCK_M": 256, "BLOCK_N": 32}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 128, "BLOCK_N": 32}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 32}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 64}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 256}, num_stages=2, num_warps=4),
-]
-
-
-@triton.autotune(
-    configs=CONFIG_LIST,
-    key=["M", "N"],
-)
-@triton.jit
-def _modulate_fwd(
-        x_ptr,  # *Pointer* to first input vector.
-        output_ptr,  # *Pointer* to output vector.
-        scale_ptr,
-        shift_ptr,
-        m_stride,
-        s_stride,
-        M,
-        N,
-        seq_len,
-        BLOCK_M: tl.constexpr,  # Number of elements each program should process.
-        BLOCK_N: tl.constexpr,
-        # NOTE: `constexpr` so it can be used as a shape value.
-):
-    row_id = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
-    rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M)
-    s_rows = (row_id // seq_len) * BLOCK_M
-    col_id = tl.program_id(axis=1)
-    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
-
-    x_ptrs = x_ptr + rows[:, None] * m_stride + cols[None, :]
-    scale_ptrs = scale_ptr + s_rows * s_stride + cols[None, :]
-    shift_ptrs = shift_ptr + s_rows * s_stride + cols[None, :]
-
-    col_mask = cols[None, :] < N
-    block_mask = (rows[:, None] < M) & col_mask
-    s_block_mask = col_mask
-    x = tl.load(x_ptrs, mask=block_mask, other=0.0)
-    scale = tl.load(scale_ptrs, mask=s_block_mask, other=0.0)
-    shift = tl.load(shift_ptrs, mask=s_block_mask, other=0.0)
-
-    output = x * (1 + scale) + shift
-    # Write x + y back to DRAM.
-    tl.store(output_ptr + rows[:, None] * m_stride + cols[None, :], output, mask=block_mask)
-
-
-@triton.autotune(
-    configs=CONFIG_LIST,
-    key=["M", "N"],
-)
-@triton.jit
-def _modulate_bwd(
-        dx_ptr,  # *Pointer* to first input vector.
-        x_ptr,
-        dy_ptr,  # *Pointer* to output vector.
-        scale_ptr,
-        dscale_ptr,
-        m_stride,
-        s_stride,
-        M,
-        N,
-        seq_len,
-        BLOCK_M: tl.constexpr,  # Number of elements each program should process.
-        BLOCK_N: tl.constexpr,
-        # NOTE: `constexpr` so it can be used as a shape value.
-):
-    row_id = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
-    rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M)
-    s_rows = (row_id // seq_len) * BLOCK_M
-    col_id = tl.program_id(axis=1)
-    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
-
-    x_ptrs = x_ptr + rows[:, None] * m_stride + cols[None, :]
-    dy_ptrs = dy_ptr + rows[:, None] * m_stride + cols[None, :]
-    dx_ptrs = dx_ptr + rows[:, None] * m_stride + cols[None, :]
-    dscale_ptrs = dscale_ptr + rows[:, None] * m_stride + cols[None, :]
-
-    scale_ptrs = scale_ptr + s_rows * s_stride + cols[None, :]
-
-    col_mask = cols[None, :] < N
-    block_mask = (rows[:, None] < M) & col_mask
-    s_block_mask = col_mask
-    x = tl.load(x_ptrs, mask=block_mask, other=0.0)
-    dy = tl.load(dy_ptrs, mask=block_mask, other=0.0)
-    scale = tl.load(scale_ptrs, mask=s_block_mask, other=0.0)
-
-    dx = dy * (1 + scale)
-    dscale = dy * x
-    # Write x + y back to DRAM.
-    tl.store(dx_ptrs, dx, mask=block_mask)
-    tl.store(dscale_ptrs, dscale, mask=block_mask)

fastvideo/ops/modulate/modulate.py

@@ -1,63 +0,0 @@
-import torch
-import triton
-
-from fastvideo.ops.modulate.k_modulate import _modulate_fwd, _modulate_bwd
-
-
-class _FusedModulate(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x, scale, shift):
-        y = torch.empty_like(x)
-        batch, seq_len, dim = x.shape
-        M = batch * seq_len
-        N = dim
-        x = x.view(-1, dim).contiguous()
-        scale = scale.view(-1, dim).contiguous()
-        shift = shift.view(-1, dim).contiguous()
-
-        def grid(meta):
-            return (
-                triton.cdiv(batch * seq_len, meta["BLOCK_M"]),
-                triton.cdiv(dim, meta["BLOCK_N"]),
-            )
-
-        _modulate_fwd[grid](x, y, scale, shift, x.stride(0), scale.stride(0), M, N, seq_len)
-
-        ctx.save_for_backward(x, scale)
-        ctx.batch = batch
-        ctx.seq_len = seq_len
-        ctx.dim = dim
-        return y
-
-    @staticmethod
-    def backward(ctx, dy):  # pragma: no cover  # this is covered, but called directly from C++
-        x, scale = ctx.saved_tensors
-
-        batch, seq_len, dim = ctx.batch, ctx.seq_len, ctx.dim
-        M = batch * seq_len
-        N = dim
-
-        # allocate output
-        dy = dy.contiguous()
-        dx = torch.empty_like(dy)
-        dscale = torch.empty_like(dy)
-        dshift = torch.sum(dy, dim=1)
-
-        def grid(meta):
-            return (
-                triton.cdiv(batch * seq_len, meta["BLOCK_M"]),
-                triton.cdiv(dim, meta["BLOCK_N"]),
-            )
-
-        _modulate_bwd[grid](dx, x, dy, scale, dscale, x.stride(0), scale.stride(0), M, N, seq_len)
-
-        dscale = torch.sum(dscale, dim=1)
-        return dx, dscale, dshift
-
-
-def fused_modulate(
-        x: torch.Tensor,
-        scale: torch.Tensor,
-        shift: torch.Tensor,
-) -> torch.Tensor:
-    return _FusedModulate.apply(x, scale, shift)

fastvideo/sample/sample_t2v_cogvideo.py

@@ -1,306 +0,0 @@
-"""
-This script demonstrates how to generate a video using the CogVideoX model with the Hugging Face `diffusers` pipeline.
-The script supports different types of video generation, including text-to-video (t2v), image-to-video (i2v),
-and video-to-video (v2v), depending on the input data and different weight.
-
-- text-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
-- video-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
-- image-to-video: THUDM/CogVideoX-5b-I2V or THUDM/CogVideoX1.5-5b-I2V
-
-Running the Script:
-To run the script, use the following command with appropriate arguments:
-
-```bash
-$ python cli_demo.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX1.5-5b --generate_type "t2v"
-```
-
-Additional options are available to specify the model path, guidance scale, number of inference steps, video generation type, and output paths.
-"""
-
-import argparse
-import logging
-import os
-from typing import Literal, Optional
-
-import torch
-from diffusers import (
-    CogVideoXDPMScheduler,
-    CogVideoXImageToVideoPipeline,
-    CogVideoXPipeline,
-    CogVideoXVideoToVideoPipeline,
-)
-from diffusers.utils import export_to_video, load_image, load_video
-
-
-logging.basicConfig(level=logging.INFO)
-
-# Recommended resolution for each model (width, height)
-RESOLUTION_MAP = {
-    # cogvideox1.5-*
-    "cogvideox1.5-5b-i2v": (1360, 768),
-    "cogvideox1.5-5b": (1360, 768),
-    # cogvideox-*
-    "cogvideox-5b-i2v": (720, 480),
-    "cogvideox-5b": (720, 480),
-    "cogvideox-2b": (720, 480),
-}
-
-
-def generate_video(
-    prompt: str,
-    model_path: str,
-    lora_path: str = None,
-    lora_rank: int = 128,
-    num_frames: int = 81,
-    width: Optional[int] = None,
-    height: Optional[int] = None,
-    output_path: str = "./output.mp4",
-    image_or_video_path: str = "",
-    num_inference_steps: int = 50,
-    guidance_scale: float = 6.0,
-    num_videos_per_prompt: int = 1,
-    dtype: torch.dtype = torch.bfloat16,
-    generate_type: str = Literal[
-        "t2v", "i2v", "v2v"
-    ],  # i2v: image to video, v2v: video to video
-    seed: int = 42,
-    fps: int = 16,
-):
-    """
-    Generates a video based on the given prompt and saves it to the specified path.
-
-    Parameters:
-    - prompt (str): The description of the video to be generated.
-    - model_path (str): The path of the pre-trained model to be used.
-    - lora_path (str): The path of the LoRA weights to be used.
-    - lora_rank (int): The rank of the LoRA weights.
-    - output_path (str): The path where the generated video will be saved.
-    - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality.
-    - num_frames (int): Number of frames to generate. CogVideoX1.0 generates 49 frames for 6 seconds at 8 fps, while CogVideoX1.5 produces either 81 or 161 frames, corresponding to 5 seconds or 10 seconds at 16 fps.
-    - width (int): The width of the generated video, applicable only for CogVideoX1.5-5B-I2V
-    - height (int): The height of the generated video, applicable only for CogVideoX1.5-5B-I2V
-    - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt.
-    - num_videos_per_prompt (int): Number of videos to generate per prompt.
-    - dtype (torch.dtype): The data type for computation (default is torch.bfloat16).
-    - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').·
-    - seed (int): The seed for reproducibility.
-    - fps (int): The frames per second for the generated video.
-    """
-
-    # 1.  Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16).
-    # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload()
-    # function to use Multi GPUs.
-
-    image = None
-    video = None
-
-    model_name = model_path.split("/")[-1].lower()
-    desired_resolution = RESOLUTION_MAP[model_name]
-    if width is None or height is None:
-        width, height = desired_resolution
-        logging.info(
-            f"\033[1mUsing default resolution {desired_resolution} for {model_name}\033[0m"
-        )
-    elif (width, height) != desired_resolution:
-        if generate_type == "i2v":
-            # For i2v models, use user-defined width and height
-            logging.warning(
-                f"\033[1;31mThe width({width}) and height({height}) are not recommended for {model_name}. The best resolution is {desired_resolution}.\033[0m"
-            )
-        else:
-            # Otherwise, use the recommended width and height
-            logging.warning(
-                f"\033[1;31m{model_name} is not supported for custom resolution. Setting back to default resolution {desired_resolution}.\033[0m"
-            )
-            width, height = desired_resolution
-
-    if generate_type == "i2v":
-        pipe = CogVideoXImageToVideoPipeline.from_pretrained(
-            model_path, torch_dtype=dtype
-        )
-        image = load_image(image=image_or_video_path)
-    elif generate_type == "t2v":
-        pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype)
-    else:
-        pipe = CogVideoXVideoToVideoPipeline.from_pretrained(
-            model_path, torch_dtype=dtype
-        )
-        video = load_video(image_or_video_path)
-
-    # If you're using with lora, add this code
-    if lora_path:
-        pipe.load_lora_weights(
-            lora_path,
-            weight_name="pytorch_lora_weights.safetensors",
-            adapter_name="test_1",
-        )
-        pipe.fuse_lora(lora_scale=1 / lora_rank)
-
-    # 2. Set Scheduler.
-    # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`.
-    # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B.
-    # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V.
-
-    # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-    pipe.scheduler = CogVideoXDPMScheduler.from_config(
-        pipe.scheduler.config, timestep_spacing="trailing"
-    )
-
-    # 3. Enable CPU offload for the model.
-    # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference
-    # and enable to("cuda")
-
-    # pipe.to("cuda")
-    pipe.enable_sequential_cpu_offload()
-    pipe.vae.enable_slicing()
-    pipe.vae.enable_tiling()
-
-    # 4. Generate the video frames based on the prompt.
-    # `num_frames` is the Number of frames to generate.
-    if generate_type == "i2v":
-        video_generate = pipe(
-            height=height,
-            width=width,
-            prompt=prompt,
-            image=image,
-            # The path of the image, the resolution of video will be the same as the image for CogVideoX1.5-5B-I2V, otherwise it will be 720 * 480
-            num_videos_per_prompt=num_videos_per_prompt,  # Number of videos to generate per prompt
-            num_inference_steps=num_inference_steps,  # Number of inference steps
-            num_frames=num_frames,  # Number of frames to generate
-            use_dynamic_cfg=True,  # This id used for DPM scheduler, for DDIM scheduler, it should be False
-            guidance_scale=guidance_scale,
-            generator=torch.Generator().manual_seed(
-                seed
-            ),  # Set the seed for reproducibility
-        ).frames[0]
-    elif generate_type == "t2v":
-        video_generate = pipe(
-            height=height,
-            width=width,
-            prompt=prompt,
-            num_videos_per_prompt=num_videos_per_prompt,
-            num_inference_steps=num_inference_steps,
-            num_frames=num_frames,
-            use_dynamic_cfg=True,
-            guidance_scale=guidance_scale,
-            generator=torch.Generator().manual_seed(seed),
-        ).frames[0]
-    else:
-        video_generate = pipe(
-            height=height,
-            width=width,
-            prompt=prompt,
-            video=video,  # The path of the video to be used as the background of the video
-            num_videos_per_prompt=num_videos_per_prompt,
-            num_inference_steps=num_inference_steps,
-            num_frames=num_frames,
-            use_dynamic_cfg=True,
-            guidance_scale=guidance_scale,
-            generator=torch.Generator().manual_seed(
-                seed
-            ),  # Set the seed for reproducibility
-        ).frames[0]
-    export_to_video(video_generate, output_path, fps=fps)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Generate a video from a text prompt using CogVideoX"
-    )
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        required=True,
-        help="The description of the video to be generated",
-    )
-    parser.add_argument(
-        "--image_or_video_path",
-        type=str,
-        default=None,
-        help="The path of the image to be used as the background of the video",
-    )
-    parser.add_argument(
-        "--model_path",
-        type=str,
-        default="THUDM/CogVideoX1.5-5B",
-        help="Path of the pre-trained model use",
-    )
-    parser.add_argument(
-        "--lora_path",
-        type=str,
-        default=None,
-        help="The path of the LoRA weights to be used",
-    )
-    parser.add_argument(
-        "--lora_rank", type=int, default=128, help="The rank of the LoRA weights"
-    )
-    parser.add_argument(
-        "--output_path",
-        type=str,
-        default="./output.mp4",
-        help="The path save generated video",
-    )
-    parser.add_argument(
-        "--guidance_scale",
-        type=float,
-        default=6.0,
-        help="The scale for classifier-free guidance",
-    )
-    parser.add_argument(
-        "--num_inference_steps", type=int, default=50, help="Inference steps"
-    )
-    parser.add_argument(
-        "--num_frames",
-        type=int,
-        default=81,
-        help="Number of steps for the inference process",
-    )
-    parser.add_argument(
-        "--width", type=int, default=None, help="The width of the generated video"
-    )
-    parser.add_argument(
-        "--height", type=int, default=None, help="The height of the generated video"
-    )
-    parser.add_argument(
-        "--fps",
-        type=int,
-        default=16,
-        help="The frames per second for the generated video",
-    )
-    parser.add_argument(
-        "--num_videos_per_prompt",
-        type=int,
-        default=1,
-        help="Number of videos to generate per prompt",
-    )
-    parser.add_argument(
-        "--generate_type", type=str, default="t2v", help="The type of video generation"
-    )
-    parser.add_argument(
-        "--dtype", type=str, default="bfloat16", help="The data type for computation"
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42, help="The seed for reproducibility"
-    )
-
-    args = parser.parse_args()
-    dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16
-    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
-    generate_video(
-        prompt=args.prompt,
-        model_path=args.model_path,
-        lora_path=args.lora_path,
-        lora_rank=args.lora_rank,
-        output_path=args.output_path,
-        num_frames=args.num_frames,
-        width=args.width,
-        height=args.height,
-        image_or_video_path=args.image_or_video_path,
-        num_inference_steps=args.num_inference_steps,
-        guidance_scale=args.guidance_scale,
-        num_videos_per_prompt=args.num_videos_per_prompt,
-        dtype=dtype,
-        generate_type=args.generate_type,
-        seed=args.seed,
-        fps=args.fps,
-    )

fastvideo/sample/sample_t2v_hunyuan.py

@@ -50,6 +50,7 @@ def main(args):
         prompts = f.readlines()
 
     for prompt in prompts:
+        print("#####PROMPT#####", prompt)
         outputs = hunyuan_video_sampler.predict(
             prompt=prompt,
             height=args.height,

scripts/inference/inference_hunyuan_sw.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+num_gpus=1
+export MODEL_BASE=data/hunyuan
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+    fastvideo/sample/sample_t2v_hunyuan.py \
+    --height 720 \
+    --width 720 \
+    --num_frames 61 \
+    --num_inference_steps 6 \
+    --guidance_scale 1 \
+    --embedded_cfg_scale 6 \
+    --flow_shift 17 \
+    --flow-reverse \
+    --prompt ./assets/prompt.txt \
+    --seed 1024 \
+    --output_path outputs_video/hunyuan/sw/ \
+    --model_path $MODEL_BASE \
+    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+    --vae-sp

tests/test_kernels/benchmark.py

@@ -1,297 +0,0 @@
-# Copyright (c) 2023, Tri Dao.
-"""Useful functions for writing test code."""
-
-import torch
-import torch.utils.benchmark as benchmark
-
-
-def get_abs_err(x, y):
-    return (x - y).flatten().abs().max().item()
-
-
-def get_err_ratio(x, y):
-    err = (x - y).flatten().square().mean().sqrt().item()
-    base = x.flatten().square().mean().sqrt().item()
-    return err / base
-
-
-def assert_close(prefix, ref, tri, ratio):
-    msg = f"{prefix} diff: {get_abs_err(ref, tri):.6f} ratio: {get_err_ratio(ref, tri):.6f}"
-    print(msg)
-    assert get_err_ratio(ref, tri) < ratio, msg
-
-
-def benchmark_forward(
-    fn,
-    *inputs,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the forward pass of an arbitrary function."""
-    if verbose:
-        print(desc, "- Forward pass")
-
-    def amp_wrapper(*inputs, **kwinputs):
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            fn(*inputs, **kwinputs)
-
-    t = benchmark.Timer(
-        stmt="fn_amp(*inputs, **kwinputs)",
-        globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs},
-        num_threads=torch.get_num_threads(),
-    )
-    m = t.timeit(repeats)
-    if verbose:
-        print(m)
-    return t, m
-
-
-def benchmark_backward(
-    fn,
-    *inputs,
-    grad=None,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the backward pass of an arbitrary function."""
-    if verbose:
-        print(desc, "- Backward pass")
-    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-        y = fn(*inputs, **kwinputs)
-        if type(y) is tuple:
-            y = y[0]
-    if grad is None:
-        grad = torch.randn_like(y)
-    else:
-        if grad.shape != y.shape:
-            raise RuntimeError("Grad shape does not match output shape")
-
-    def f(*inputs, y, grad):
-        # Set .grad to None to avoid extra operation of gradient accumulation
-        for x in inputs:
-            if isinstance(x, torch.Tensor):
-                x.grad = None
-        y.backward(grad, retain_graph=True)
-
-    t = benchmark.Timer(
-        stmt="f(*inputs, y=y, grad=grad)",
-        globals={"f": f, "inputs": inputs, "y": y, "grad": grad},
-        num_threads=torch.get_num_threads(),
-    )
-    m = t.timeit(repeats)
-    if verbose:
-        print(m)
-    return t, m
-
-
-def benchmark_combined(
-    fn,
-    *inputs,
-    grad=None,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
-    if verbose:
-        print(desc, "- Forward + Backward pass")
-    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-        y = fn(*inputs, **kwinputs)
-        if type(y) is tuple:
-            y = y[0]
-    if grad is None:
-        grad = torch.randn_like(y)
-    else:
-        if grad.shape != y.shape:
-            raise RuntimeError("Grad shape does not match output shape")
-
-    def f(grad, *inputs, **kwinputs):
-        for x in inputs:
-            if isinstance(x, torch.Tensor):
-                x.grad = None
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            y = fn(*inputs, **kwinputs)
-            if type(y) is tuple:
-                y = y[0]
-        y.backward(grad, retain_graph=True)
-
-    t = benchmark.Timer(
-        stmt="f(grad, *inputs, **kwinputs)",
-        globals={
-            "f": f,
-            "fn": fn,
-            "inputs": inputs,
-            "grad": grad,
-            "kwinputs": kwinputs,
-        },
-        num_threads=torch.get_num_threads(),
-    )
-    m = t.timeit(repeats)
-    if verbose:
-        print(m)
-    return t, m
-
-
-def benchmark_fwd_bwd(
-    fn,
-    *inputs,
-    grad=None,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
-    return (
-        benchmark_forward(
-            fn,
-            *inputs,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-        benchmark_backward(
-            fn,
-            *inputs,
-            grad=grad,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-    )
-
-
-def benchmark_all(
-    fn,
-    *inputs,
-    grad=None,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
-    return (
-        benchmark_forward(
-            fn,
-            *inputs,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-        benchmark_backward(
-            fn,
-            *inputs,
-            grad=grad,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-        benchmark_combined(
-            fn,
-            *inputs,
-            grad=grad,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-    )
-
-
-def pytorch_profiler(
-    fn,
-    *inputs,
-    trace_filename=None,
-    backward=False,
-    amp=False,
-    amp_dtype=torch.float16,
-    cpu=False,
-    verbose=True,
-    **kwinputs,
-):
-    """Wrap benchmark functions in Pytorch profiler to see CUDA information."""
-    if backward:
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            out = fn(*inputs, **kwinputs)
-            if type(out) is tuple:
-                out = out[0]
-            g = torch.randn_like(out)
-    for _ in range(30):  # Warm up
-        if backward:
-            for x in inputs:
-                if isinstance(x, torch.Tensor):
-                    x.grad = None
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            out = fn(*inputs, **kwinputs)
-            if type(out) is tuple:
-                out = out[0]
-        # Backward should be done outside autocast
-        if backward:
-            out.backward(g, retain_graph=True)
-    activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [
-        torch.profiler.ProfilerActivity.CUDA
-    ]
-    with torch.profiler.profile(
-        activities=activities,
-        record_shapes=True,
-        # profile_memory=True,
-        with_stack=True,
-    ) as prof:
-        if backward:
-            for x in inputs:
-                if isinstance(x, torch.Tensor):
-                    x.grad = None
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            out = fn(*inputs, **kwinputs)
-            if type(out) is tuple:
-                out = out[0]
-        if backward:
-            out.backward(g, retain_graph=True)
-    if verbose:
-        # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50))
-        print(prof.key_averages().table(row_limit=50))
-    if trace_filename is not None:
-        prof.export_chrome_trace(trace_filename)
-
-
-def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs):
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    torch.cuda.synchronize()
-    fn(*inputs, **kwinputs)
-    torch.cuda.synchronize()
-    mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000)
-    if verbose:
-        print(f"{desc} max memory: {mem}GB")
-    torch.cuda.empty_cache()
-    return mem

tests/test_kernels/benchmark_modulate.py

@@ -1,83 +0,0 @@
-import torch
-from torch import nn
-
-from benchmark import (
-    benchmark_combined,
-    benchmark_forward,
-    benchmark_backward,
-    assert_close,
-)
-from fastvideo.ops.modulate.modulate import fused_modulate
-
-
-def torch_modulate(x, scale, shift):
-    return x * (scale.unsqueeze(1) + 1) + shift.unsqueeze(1)
-
-
-# from flash_attn import flash_attn_func
-
-
-def time_fwd(func, *args, **kwargs):
-    time_fb = benchmark_forward(func, *args, **kwargs)
-    return time_fb[1].mean
-
-
-def time_fwd_bwd(func, *args, **kwargs):
-    time_fb = benchmark_combined(func, *args, **kwargs)
-    return time_fb[1].mean
-
-
-def time_bwd(func, *args, **kwargs):
-    time_fb = benchmark_backward(func, *args, **kwargs)
-    return time_fb[1].mean
-
-
-device = "cuda"
-dtype = torch.bfloat16
-
-
-batch_sizes = [1, 8, 32]
-seq_lengths = [128, 512, 1024]
-hidden_dims = [768, 1024, 2048]
-
-# methods = (["torch", "triton", "thunderkitten"])
-methods = ["torch", "triton"]
-time_f = {}
-time_b = {}
-time_f_b = {}
-speed_f = {}
-speed_b = {}
-speed_f_b = {}
-for B in batch_sizes:
-    for T in seq_lengths:
-        for D in hidden_dims:
-            config = (B, T, D)
-
-            norm_func = nn.LayerNorm(D, elementwise_affine=False, eps=1e-6)
-            x = torch.randn(B, T, D, device="cuda", requires_grad=True, dtype=dtype)
-            x = norm_func(x.to(torch.float32)).to(dtype)
-            shift = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype)
-            scale = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype)
-            # test torch
-            o_ref = torch_modulate(x, scale, shift)
-            o_ref.sum().backward(retain_graph=True)
-            f_b = time_fwd_bwd(torch_modulate, x, scale, shift, verbose=False)
-            time_f_b[config, "torch"] = f_b
-            # test triton
-            o2 = fused_modulate(x, scale, shift)
-            o2.sum().backward(retain_graph=True)
-            f_b = time_fwd_bwd(fused_modulate, x, scale, shift, verbose=False)
-            time_f_b[config, "triton"] = f_b
-            # test if the results are close
-            assert_close("  o", o_ref, o2, 0.005)
-            # time_f_b[config, "thunderkitten"] = f_b
-
-            print(f"### batch size={B}, seq length={T}, B={B}, hidden dim={D} ###")
-            for method in methods:
-                # time_f_b[config, method] = time_f[config, method] + time_b[config, method]
-                print(
-                    f"{method:>50} fwd + bwd:\t {time_f_b[config, method]*1000:>6.4f} ms "
-                )
-
-# with open('flash2_attn_time.plk', 'wb') as fp:
-#     pickle.dump((speed_f, speed_b, speed_f_b), fp, protocol=pickle.HIGHEST_PROTOCOL)