syn

.github/ISSUE_TEMPLATE/config.yml

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

fastvideo/models/flash_attn_no_pad.py

@@ -1,36 +1,48 @@
-from einops import rearrange
+try: 
+    from flash_attn_interface import flash_attn_func as flash_attn_func_v3 
+    has_v3 = True
+except: 
+    pass
 from flash_attn import flash_attn_varlen_qkvpacked_func
 from flash_attn.bert_padding import pad_input, unpad_input
+from einops import rearrange
 
 
-def flash_attn_no_pad(qkv,
-                      key_padding_mask,
-                      causal=False,
-                      dropout_p=0.0,
-                      softmax_scale=None):
+def flash_attn_no_pad(
+    qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None, attn_impl="fa3",
+):
     # 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,
+        x, key_padding_mask
     )
+
+    x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads)
+    if attn_impl == "fa3":
+        assert has_v3
+        q, k, v = x_unpad[:, 0].unsqueeze(0), x_unpad[:, 1].unsqueeze(0), x_unpad[:, 2].unsqueeze(0)
+        assert dropout_p == 0.0
+        output_unpad = flash_attn_func_v3(
+            q,
+            k,
+            v,
+        )[0].squeeze(0)
+    else:
+        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),
+        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,
     )

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

@@ -606,6 +606,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
         enable_vae_sp: bool = False,
         n_tokens: Optional[int] = None,
         embedded_guidance_scale: Optional[float] = None,
+        only_save_latent: bool = False,
         **kwargs,
     ):
         r"""
@@ -978,8 +979,13 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                 latents = (latents / self.vae.config.scaling_factor +
                            self.vae.config.shift_factor)
             else:
-                latents = latents / self.vae.config.scaling_factor
-
+                latents = latents / self.vae.config.scaling_factor # 1, 16, 32, 90, 160
+                
+            if only_save_latent:
+                latents = latents.cpu().float()
+                self.maybe_free_model_hooks()
+                return latents
+            
             with torch.autocast(device_type="cuda",
                                 dtype=vae_dtype,
                                 enabled=vae_autocast_enabled):
@@ -993,7 +999,6 @@ class HunyuanVideoPipeline(DiffusionPipeline):
 
             if expand_temporal_dim or image.shape[2] == 1:
                 image = image.squeeze(2)
-
         else:
             image = latents
 

fastvideo/models/hunyuan/inference.py

@@ -369,6 +369,7 @@ class HunyuanVideoSampler(Inference):
         embedded_guidance_scale=None,
         batch_size=1,
         num_videos_per_prompt=1,
+        only_save_latent=False,
         **kwargs,
     ):
         """
@@ -524,7 +525,10 @@ class HunyuanVideoSampler(Inference):
             vae_ver=self.args.vae,
             enable_tiling=self.args.vae_tiling,
             enable_vae_sp=self.args.vae_sp,
+            only_save_latent=only_save_latent,
         )[0]
+        if only_save_latent:
+            return samples
         out_dict["samples"] = samples
         out_dict["prompts"] = prompt
 

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/generate_synthetic_hunyuan.py

@@ -0,0 +1,286 @@
+import argparse
+import os
+from pathlib import Path
+
+import imageio
+import numpy as np
+import torch
+import torch.distributed as dist
+import torchvision
+from einops import rearrange
+import json
+from fastvideo.models.hunyuan.inference import HunyuanVideoSampler
+from fastvideo.utils.parallel_states import (
+    initialize_sequence_parallel_state, nccl_info)
+
+
+def initialize_distributed():
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    print("world_size", world_size)
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend="nccl",
+                            init_method="env://",
+                            world_size=world_size,
+                            rank=local_rank)
+    initialize_sequence_parallel_state(world_size)
+
+
+def main(args):
+    print(nccl_info.sp_size)
+    print(args)
+    models_root_path = Path(args.model_path)
+    if not models_root_path.exists():
+        raise ValueError(f"`models_root` not exists: {models_root_path}")
+
+    save_path = args.output_path
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+    # Handle both JSON and TXT files
+    prompts = []
+    if args.prompt_json.endswith('.json'):
+        with open(args.prompt_json) as f:
+            prompt_json = json.load(f)
+            prompts = [(item["caption"], item["latent_path"]) for item in prompt_json]
+    else:
+        with open(args.prompt_json) as f:
+            for line in f:
+                prompt = line.strip()
+                if prompt:
+                    # Use first 50 chars of prompt as filename
+                    file_name = "".join(c for c in prompt[:50] if c.isalnum() or c.isspace())
+                    file_name = file_name.strip().replace(" ", "_")
+                    prompts.append((prompt, file_name))
+
+    # Filter for unprocessed prompts in latent-only mode
+    if args.only_save_latent:
+        filtered_prompts = []
+        for prompt, file_name in prompts:
+            latent_path = os.path.join(args.output_path, "latent", f"{file_name}.pt")
+            if not os.path.exists(latent_path):
+                filtered_prompts.append((prompt, file_name))
+            else:
+                print(f"Latent file exists for {file_name}, skipping...")
+        
+        prompts = filtered_prompts
+        if not prompts:
+            print("All prompts processed. Exiting...")
+            return
+
+    hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(
+        models_root_path, args=args)
+    args = hunyuan_video_sampler.args
+
+    prompt_start_end_idx = [int(i) for i in args.prompt_start_end_idx.split(",")]
+    prompts = prompts[prompt_start_end_idx[0]:prompt_start_end_idx[1]]
+    
+    for prompt, file_name in prompts:
+        outputs = hunyuan_video_sampler.predict(
+            prompt=prompt,
+            height=args.height,
+            width=args.width,
+            video_length=args.num_frames,
+            seed=args.seed,
+            negative_prompt=args.neg_prompt,
+            infer_steps=args.num_inference_steps,
+            guidance_scale=args.guidance_scale,
+            num_videos_per_prompt=args.num_videos,
+            flow_shift=args.flow_shift,
+            batch_size=args.batch_size,
+            embedded_guidance_scale=args.embedded_cfg_scale,
+            only_save_latent=args.only_save_latent,
+        )
+        if args.only_save_latent:
+            os.makedirs(os.path.join(args.output_path, "latent"), exist_ok=True)
+            latent_path = os.path.join(args.output_path, "latent", f"{file_name}.pt")
+            torch.save(outputs.to(torch.bfloat16), latent_path)
+        else:
+            videos = rearrange(outputs["samples"], "b c t h w -> t b c h w")
+            outputs = []
+            for x in videos:
+                x = torchvision.utils.make_grid(x, nrow=6)
+                x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                outputs.append((x * 255).numpy().astype(np.uint8))
+            os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
+            imageio.mimsave(os.path.join(args.output_path, f"{file_name}.mp4"),
+                            outputs,
+                            fps=args.fps)
+
+#                 
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # Basic parameters
+    parser.add_argument("--prompt_json", type=str, required=True, help="prompt file for inference")
+    parser.add_argument("--num_frames", type=int, default=16)
+    parser.add_argument("--height", type=int, default=256)
+    parser.add_argument("--width", type=int, default=256)
+    parser.add_argument("--num_inference_steps", type=int, default=50)
+    parser.add_argument("--model_path", type=str, default="data/hunyuan")
+    parser.add_argument("--output_path", type=str, default="./outputs/video")
+    parser.add_argument("--fps", type=int, default=24)
+
+    # Additional parameters
+    parser.add_argument(
+        "--prompt_start_end_idx",
+        type=str,
+        default="0,99",
+        help=" 0 stands for starting layer index, 9 for stride",
+    )
+    parser.add_argument(
+        "--only_save_latent",
+        type=bool,
+        default=False,
+        help="Whether save latent or not.",
+    )
+    parser.add_argument(
+        "--denoise-type",
+        type=str,
+        default="flow",
+        help="Denoise type for noised inputs.",
+    )
+    parser.add_argument("--seed",
+                        type=int,
+                        default=None,
+                        help="Seed for evaluation.")
+    parser.add_argument("--neg_prompt",
+                        type=str,
+                        default=None,
+                        help="Negative prompt for sampling.")
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=1.0,
+        help="Classifier free guidance scale.",
+    )
+    parser.add_argument(
+        "--embedded_cfg_scale",
+        type=float,
+        default=6.0,
+        help="Embedded classifier free guidance scale.",
+    )
+    parser.add_argument("--flow_shift",
+                        type=int,
+                        default=7,
+                        help="Flow shift parameter.")
+    parser.add_argument("--batch_size",
+                        type=int,
+                        default=1,
+                        help="Batch size for inference.")
+    parser.add_argument(
+        "--num_videos",
+        type=int,
+        default=1,
+        help="Number of videos to generate per prompt.",
+    )
+    parser.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.",
+    )
+    parser.add_argument(
+        "--use-cpu-offload",
+        action="store_true",
+        help="Use CPU offload for the model load.",
+    )
+    parser.add_argument(
+        "--dit-weight",
+        type=str,
+        default=
+        "data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
+    )
+    parser.add_argument(
+        "--reproduce",
+        action="store_true",
+        help=
+        "Enable reproducibility by setting random seeds and deterministic algorithms.",
+    )
+    parser.add_argument(
+        "--disable-autocast",
+        action="store_true",
+        help=
+        "Disable autocast for denoising loop and vae decoding in pipeline sampling.",
+    )
+
+    # Flow Matching
+    parser.add_argument(
+        "--flow-reverse",
+        action="store_true",
+        help="If reverse, learning/sampling from t=1 -> t=0.",
+    )
+    parser.add_argument("--flow-solver",
+                        type=str,
+                        default="euler",
+                        help="Solver for flow matching.")
+    parser.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)",
+    )
+    parser.add_argument(
+        "--linear-schedule-end",
+        type=int,
+        default=25,
+        help="End step for linear quadratic schedule for flow matching.",
+    )
+
+    # Model parameters
+    parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill")
+    parser.add_argument("--latent-channels", type=int, default=16)
+    parser.add_argument("--precision",
+                        type=str,
+                        default="bf16",
+                        choices=["fp32", "fp16", "bf16"])
+    parser.add_argument("--rope-theta",
+                        type=int,
+                        default=256,
+                        help="Theta used in RoPE.")
+
+    parser.add_argument("--vae", type=str, default="884-16c-hy")
+    parser.add_argument("--vae-precision",
+                        type=str,
+                        default="fp16",
+                        choices=["fp32", "fp16", "bf16"])
+    parser.add_argument("--vae-tiling", action="store_true", default=True)
+    parser.add_argument("--vae-sp", action="store_true", default=False)
+
+    parser.add_argument("--text-encoder", type=str, default="llm")
+    parser.add_argument(
+        "--text-encoder-precision",
+        type=str,
+        default="fp16",
+        choices=["fp32", "fp16", "bf16"],
+    )
+    parser.add_argument("--text-states-dim", type=int, default=4096)
+    parser.add_argument("--text-len", type=int, default=256)
+    parser.add_argument("--tokenizer", type=str, default="llm")
+    parser.add_argument("--prompt-template",
+                        type=str,
+                        default="dit-llm-encode")
+    parser.add_argument("--prompt-template-video",
+                        type=str,
+                        default="dit-llm-encode-video")
+    parser.add_argument("--hidden-state-skip-layer", type=int, default=2)
+    parser.add_argument("--apply-final-norm", action="store_true")
+
+    parser.add_argument("--text-encoder-2", type=str, default="clipL")
+    parser.add_argument(
+        "--text-encoder-precision-2",
+        type=str,
+        default="fp16",
+        choices=["fp32", "fp16", "bf16"],
+    )
+    parser.add_argument("--text-states-dim-2", type=int, default=768)
+    parser.add_argument("--tokenizer-2", type=str, default="clipL")
+    parser.add_argument("--text-len-2", type=int, default=77)
+
+    args = parser.parse_args()
+    # process for vae sequence parallel
+    if args.vae_sp and not args.vae_tiling:
+        raise ValueError(
+            "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
+        )
+    main(args)

fastvideo/sample/generate_synthetic_mochi.py

@@ -1,16 +1,16 @@
-import argparse
 import json
-import os
 
-import torch
 import torch.distributed as dist
-from diffusers.utils import export_to_video
-
+import torch
 from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
+import os
+from diffusers.utils import export_to_video
+import argparse
 
 
-def generate_video_and_latent(pipe, prompt, height, width, num_frames,
-                              num_inference_steps, guidance_scale):
+def generate_video_and_latent(
+    pipe, prompt, height, width, num_frames, num_inference_steps, guidance_scale
+):
     # Set the random seed for reproducibility
     generator = torch.Generator("cuda").manual_seed(12345)
     # Generate videos from the input prompt
@@ -25,8 +25,7 @@ def generate_video_and_latent(pipe, prompt, height, width, num_frames,
         output_type="latent_and_video",
     )
     # prompt_embed has negative prompt at index 0
-    return noise[0], video[0], latent[0], prompt_embed[
-        1], prompt_attention_mask[1]
+    return noise[0], video[0], latent[0], prompt_embed[1], prompt_attention_mask[1]
 
     # return dummy tensor to debug first
     # return torch.zeros(1, 3, 480, 848), torch.zeros(1, 256, 16, 16)
@@ -40,22 +39,19 @@ if __name__ == "__main__":
     parser.add_argument("--num_inference_steps", type=int, default=64)
     parser.add_argument("--guidance_scale", type=float, default=4.5)
     parser.add_argument("--model_path", type=str, default="data/mochi")
-    parser.add_argument("--prompt_path",
-                        type=str,
-                        default="data/dummyVid/videos2caption.json")
-    parser.add_argument("--dataset_output_dir",
-                        type=str,
-                        default="data/dummySynthetic")
+    parser.add_argument(
+        "--prompt_path", type=str, default="data/dummyVid/videos2caption.json"
+    )
+    parser.add_argument("--dataset_output_dir", type=str, default="data/dummySynthetic")
     args = parser.parse_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)
     torch.cuda.set_device(local_rank)
-    dist.init_process_group(backend="nccl",
-                            init_method="env://",
-                            world_size=world_size,
-                            rank=local_rank)
+    dist.init_process_group(
+        backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
+    )
 
     if not isinstance(args.prompt_path, list):
         args.prompt_path = [args.prompt_path]
@@ -63,8 +59,7 @@ if __name__ == "__main__":
         text_prompt = open(args.prompt_path[0], "r").readlines()
         text_prompt = [i.strip() for i in text_prompt]
 
-    pipe = MochiPipeline.from_pretrained(args.model_path,
-                                         torch_dtype=torch.bfloat16)
+    pipe = MochiPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
     pipe.enable_vae_tiling()
     pipe.enable_model_cpu_offload(gpu_id=local_rank)
     # make dir if not exist
@@ -73,10 +68,10 @@ if __name__ == "__main__":
     os.makedirs(os.path.join(args.dataset_output_dir, "noise"), exist_ok=True)
     os.makedirs(os.path.join(args.dataset_output_dir, "video"), exist_ok=True)
     os.makedirs(os.path.join(args.dataset_output_dir, "latent"), exist_ok=True)
-    os.makedirs(os.path.join(args.dataset_output_dir, "prompt_embed"),
-                exist_ok=True)
-    os.makedirs(os.path.join(args.dataset_output_dir, "prompt_attention_mask"),
-                exist_ok=True)
+    os.makedirs(os.path.join(args.dataset_output_dir, "prompt_embed"), exist_ok=True)
+    os.makedirs(
+        os.path.join(args.dataset_output_dir, "prompt_attention_mask"), exist_ok=True
+    )
     data = []
     for i, prompt in enumerate(text_prompt):
         if i % world_size != local_rank:
@@ -98,17 +93,17 @@ if __name__ == "__main__":
         )
         # save latent
         video_name = str(i)
-        noise_path = os.path.join(args.dataset_output_dir, "noise",
-                                  video_name + ".pt")
-        latent_path = os.path.join(args.dataset_output_dir, "latent",
-                                   video_name + ".pt")
-        prompt_embed_path = os.path.join(args.dataset_output_dir,
-                                         "prompt_embed", video_name + ".pt")
-        video_path = os.path.join(args.dataset_output_dir, "video",
-                                  video_name + ".mp4")
-        prompt_attention_mask_path = os.path.join(args.dataset_output_dir,
-                                                  "prompt_attention_mask",
-                                                  video_name + ".pt")
+        noise_path = os.path.join(args.dataset_output_dir, "noise", video_name + ".pt")
+        latent_path = os.path.join(
+            args.dataset_output_dir, "latent", video_name + ".pt"
+        )
+        prompt_embed_path = os.path.join(
+            args.dataset_output_dir, "prompt_embed", video_name + ".pt"
+        )
+        video_path = os.path.join(args.dataset_output_dir, "video", video_name + ".mp4")
+        prompt_attention_mask_path = os.path.join(
+            args.dataset_output_dir, "prompt_attention_mask", video_name + ".pt"
+        )
         # save latent
         torch.save(noise, noise_path)
         torch.save(latent, latent_path)
@@ -132,6 +127,7 @@ if __name__ == "__main__":
     # save json
     if local_rank == 0:
         all_data = [item for sublist in gathered_data for item in sublist]
-        with open(os.path.join(args.dataset_output_dir, "videos2caption.json"),
-                  "w") as f:
+        with open(
+            os.path.join(args.dataset_output_dir, "videos2caption.json"), "w"
+        ) as f:
             json.dump(all_data, f, indent=4)

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

local_scripts/fastvideo/Hunyuan_lora.slurm

@@ -0,0 +1,62 @@
+#!/bin/bash
+#SBATCH --job-name=fastvideo
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log/slurm_lora_1e4.out
+#SBATCH --error=slurm_log/slurm_lora_1e4.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+export MASTER_PORT=29803
+torchrun --nnodes 1 --nproc_per_node 8 --master_port $MASTER_PORT \
+    fastvideo/train.py \
+    --seed 1024 \
+    --pretrained_model_name_or_path ~/data/hunyuan_diffusers \
+    --model_type hunyuan_hf \
+    --cache_dir data/.cache \
+    --data_json_path ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+    --validation_prompt_dir ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/validation \
+    --gradient_checkpointing \
+    --train_batch_size 8 \
+    --num_latent_t 32 \
+    --sp_size 8 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps 1 \
+    --max_train_steps 8000 \
+    --learning_rate 1e-4 \
+    --mixed_precision bf16 \
+    --checkpointing_steps 100 \
+    --validation_steps 100 \
+    --validation_sampling_steps 50 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir data/outputs/SBA_lora_1e4_r32 \
+    --tracker_project_name SBA \
+    --num_frames 125 \
+    --num_width 1280 \
+    --num_height 768 \
+    --validation_guidance_scale "1.0" \
+    --shift 7 \
+    --use_lora \
+    --lora_rank 32 \
+    --lora_alpha 32

local_scripts/fastvideo/Hunyuan_ori_finetune.slurm

@@ -0,0 +1,100 @@
+#!/bin/bash
+#SBATCH --job-name=fv_syn_hunyuan
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log/slurm_ori_ft_syn.out
+#SBATCH --error=slurm_log/slurm_ori_ft_syn.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+export MASTER_PORT=29803
+torchrun --nnodes 1 --nproc_per_node 6 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/hunyuan \
+    --dit_model_name_or_path data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt\
+    --model_type "hunyuan" \
+    --cache_dir data/.cache \
+    --data_json_path ../FastVideo-Internal/data/synthetic_MixKit/videos2caption.json \
+    --validation_prompt_dir ../FastVideo-Internal/data/synthetic_MixKit/validation \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 30 \
+    --sp_size 6 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=2 \
+    --max_train_steps=2000 \
+    --learning_rate=1e-5 \
+    --mixed_precision=bf16 \
+    --checkpointing_steps=200 \
+    --validation_steps 100 \
+    --validation_sampling_steps 50 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir=data/outputs/SBA_hunyuan_ori_ft_1e5 \
+    --tracker_project_name SBA-Syn \
+    --num_frames 117 \
+    --num_height 768 \
+    --num_width 1280 \
+    --shift 7 \
+    --validation_guidance_scale "1.0" \
+    --training_guidance "6.0" \
+    --run_name "SBA_hunyuan_ori_ft_1e5_g6" \
+
+torchrun --nnodes 1 --nproc_per_node 6 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/hunyuan \
+    --dit_model_name_or_path data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt\
+    --model_type "hunyuan" \
+    --cache_dir data/.cache \
+    --data_json_path ../FastVideo-Internal/data/synthetic_MixKit/videos2caption.json \
+    --validation_prompt_dir ../FastVideo-Internal/data/synthetic_MixKit/validation \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 30 \
+    --sp_size 6 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=2 \
+    --max_train_steps=2000 \
+    --learning_rate=1e-5 \
+    --mixed_precision=bf16 \
+    --checkpointing_steps=200 \
+    --validation_steps 100 \
+    --validation_sampling_steps 50 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir=data/outputs/SBA_hunyuan_ori_ft_1e5 \
+    --tracker_project_name SBA-Syn \
+    --num_frames 117 \
+    --num_height 768 \
+    --num_width 1280 \
+    --shift 7 \
+    --validation_guidance_scale "1.0" \
+    --training_guidance "1.0" \
+    --run_name "SBA_hunyuan_ori_ft_1e5_g1" \

local_scripts/get_videos2caption.py

@@ -0,0 +1,12 @@
+import json
+src_json = "../HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json"
+des_json = "./videos2caption.json"
+syn_num = 8
+with open(src_json) as f:
+    data = json.load(f)
+data = data[:syn_num]
+# from IPython import embed
+# embed()
+
+with open(des_json, "w") as f2:
+    json.dump(data, f2, indent=4, ensure_ascii=False)

local_scripts/gpu_reservation.slurm

@@ -0,0 +1,65 @@
+#!/bin/bash
+#SBATCH --job-name=hyhf-lora-5e6-r32
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=512G
+#SBATCH --output=slurm_log/slurm.out
+#SBATCH --error=slurm_log/slurm.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+cd ~/yongqi/FastVideo
+export MASTER_PORT=29803
+torchrun --nnodes 1 --nproc_per_node 8 --master_port $MASTER_PORT \
+    fastvideo/train.py \
+    --seed 1024 \
+    --pretrained_model_name_or_path ~/data/hunyuan_diffusers \
+    --model_type hunyuan_hf \
+    --cache_dir data/.cache \
+    --data_json_path ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+    --validation_prompt_dir ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/validation \
+    --gradient_checkpointing \
+    --train_batch_size 4 \
+    --num_latent_t 32 \
+    --sp_size 8 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps 1 \
+    --max_train_steps 8000 \
+    --learning_rate 5e-6 \
+    --mixed_precision bf16 \
+    --checkpointing_steps 200 \
+    --validation_steps 100 \
+    --validation_sampling_steps 50 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir data/outputs/SBA_lora_5e6_r32 \
+    --tracker_project_name SBA \
+    --num_frames 125 \
+    --num_width 1280 \
+    --num_height 768 \
+    --validation_guidance_scale "1.0" \
+    --use_lora \
+    --lora_rank 32 \
+    --lora_alpha 64
+
+echo "Run completed at:- "
+date

local_scripts/synthetic_inference/exp-1.slurm

@@ -0,0 +1,60 @@
+#!/bin/bash
+#SBATCH --job-name=syn-1
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-1.out
+#SBATCH --error=slurm_log_last_dance/syn-1.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+
+# Calculate prompts per GPU
+prompt_start_idx_all=0
+num_promtps_all=200
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ./assets/prompt_vb.txt \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-10.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-10
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-10.out
+#SBATCH --error=slurm_log_last_dance/syn-10.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=2160
+num_promtps_all=240
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-11.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-11
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-11.out
+#SBATCH --error=slurm_log_last_dance/syn-11.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=2400
+num_promtps_all=240
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-2.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-2
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-2.out
+#SBATCH --error=slurm_log_last_dance/syn-2.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=198
+num_promtps_all=200
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ./assets/prompt_vb.txt \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-3.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-3
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-3.out
+#SBATCH --error=slurm_log_last_dance/syn-3.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=396
+num_promtps_all=200
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ./assets/prompt_vb.txt \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-4.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-4
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-4.out
+#SBATCH --error=slurm_log_last_dance/syn-4.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=594
+num_promtps_all=200
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ./assets/prompt_vb.txt \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-5.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-5
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-5.out
+#SBATCH --error=slurm_log_last_dance/syn-5.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=960
+num_promtps_all=240
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-6.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-6
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-6.out
+#SBATCH --error=slurm_log_last_dance/syn-6.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=1200
+num_promtps_all=240
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-7.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-7
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-7.out
+#SBATCH --error=slurm_log_last_dance/syn-7.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=1440
+num_promtps_all=240
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-8.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-8
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-8.out
+#SBATCH --error=slurm_log_last_dance/syn-8.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=1680
+num_promtps_all=240
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

local_scripts/synthetic_inference/exp-9.slurm

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=syn-9
+#SBATCH --partition=mbzuai
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=960G
+#SBATCH --output=slurm_log_last_dance/syn-9.out
+#SBATCH --error=slurm_log_last_dance/syn-9.err
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+conda init
+source ~/conda/miniconda/bin/activate
+PYTHON_VIRTUAL_ENVIRONMENT=fastvideo
+conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+cd /mbz/users/hao.zhang/yongqi/FastVideo
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Calculate prompts per GPU
+prompt_start_idx_all=1920
+num_promtps_all=240
+num_gpus=8
+prompts_per_gpu=$((num_promtps_all / num_gpus))
+export MODEL_BASE=data/hunyuan
+
+for gpu_id in $(seq 0 $((num_gpus-1))); do
+    start_idx=$((prompt_start_idx_all + gpu_id * prompts_per_gpu))
+    end_idx=$((start_idx + prompts_per_gpu))
+    prompt_start_end_idx_single_gpu="${start_idx},${end_idx}"
+    
+    CUDA_VISIBLE_DEVICES=$gpu_id python fastvideo/sample/generate_synthetic_hunyuan.py \
+        --height 768 \
+        --width 1280 \
+        --num_frames 117 \
+        --num_inference_steps 50 \
+        --guidance_scale 1 \
+        --embedded_cfg_scale 6 \
+        --flow_shift 7 \
+        --flow-reverse \
+        --prompt_json ~/data/HD-MixKit-Hunyuan-Distill-125x768x1280/videos2caption.json \
+        --seed 1024 \
+        --output_path outputs_latents/ \
+        --model_path $MODEL_BASE \
+        --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+        --only_save_latent True \
+        --prompt_start_end_idx ${prompt_start_end_idx_single_gpu} &
+done
+
+# Wait for all background jobs to complete
+wait

scripts/inference/inference_hunyuan.sh

@@ -18,3 +18,22 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     --model_path $MODEL_BASE \
     --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
     --vae-sp
+
+num_gpus=6
+export MODEL_BASE=data/hunyuan
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+    fastvideo/sample/sample_t2v_hunyuan.py \
+    --height 768 \
+    --width 1280 \
+    --num_frames 117 \
+    --num_inference_steps 50 \
+    --guidance_scale 1 \
+    --embedded_cfg_scale 8 \
+    --flow_shift 7 \
+    --flow-reverse \
+    --prompt ./assets/prompt_mixkit.txt \
+    --seed 1024 \
+    --output_path outputs_video/mixkit-8/ \
+    --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)