run lint

README.md

@@ -8,7 +8,7 @@ It features a clean, consistent API that works across popular video models, maki
 With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
 </p>
 
 <div align="center">

docs/source/contributing/developer_env/runpod.md

@@ -8,14 +8,12 @@ You can easily use the FastVideo Docker image as a custom container on [RunPod](
 
 Choose a GPU that supports CUDA 12.4
 
-Pick 1 or 2 L40S GPU(s)
-
 ![RunPod CUDA selection](../../_static/images/runpod_cuda.png)
 
 When creating your pod template, use this image:
 
 ```
-ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
+ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
 ```
 
 Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):

docs/source/getting_started/installation.md

@@ -117,4 +117,4 @@ If you're planning to contribute to FastVideo please see the following page:
 
 If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
 
-You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg) for additional support.

docs/source/inference/add_pipeline.md

@@ -12,7 +12,7 @@ This guide explains how to implement a custom diffusion pipeline in FastVideo, l
 4. **Register Your Pipeline** - Make it discoverable by the framework
 5. **Configure Your Pipeline** - (Coming soon)
 
-Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
 
 ## Step 1: Pipeline Modules
 

docs/source/inference/cli.md

@@ -27,7 +27,7 @@ fastvideo generate --help
 ### Hardware Configuration
 
 - `--num-gpus {NUM_GPUS}`: Number of GPUs to use
-- `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder, should not be larger than 1 if text encoder offload is enabled, as layerwise offload + prefetch is faster)
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (Typically should match the number of GPUs)
 - `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)
 
 #### Video Configuration
@@ -68,7 +68,7 @@ Example configuration file (config.json):
     "output_path": "outputs/",
     "num_gpus": 2,
     "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
     "num_frames": 45,
     "height": 720,
     "width": 1280,
@@ -102,7 +102,7 @@ prompt: "A beautiful woman in a red dress walking down a street"
 output_path: "outputs/"
 num_gpus: 2
 sp_size: 2
-tp_size: 1
+tp_size: 2
 num_frames: 45
 height: 720
 width: 1280

docs/source/inference/inference_quick_start.md

@@ -121,4 +121,4 @@ If the generated video doesn't match your prompt:
 - Learn about using [Optimizations](#inference-optimizations)
 - See [Examples](../examples/examples_inference_index.md) for more usage scenarios
 - Join our [Community Discord](https://discord.gg/JA7cksDz86).
-- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).

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

@@ -24,7 +24,6 @@ training_args=(
   --num_height 480
   --num_width 832
   --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )
 
 # Parallel arguments

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

@@ -1,11 +1,13 @@
 #!/bin/bash
 #SBATCH --job-name=i2v
 #SBATCH --partition=main
+#SBATCH --qos=hao
 #SBATCH --nodes=4
 #SBATCH --ntasks=4
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
+#SBATCH --nodelist=fs-mbz-gpu-[100-850]
 #SBATCH --mem=1440G
 #SBATCH --output=i2v_output/i2v_%j.out
 #SBATCH --error=i2v_output/i2v_%j.err
@@ -58,7 +60,6 @@ training_args=(
   --num_height 480
   --num_width 832
   --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )
 
 # Parallel arguments

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

@@ -24,14 +24,13 @@ training_args=(
   --num_height 480
   --num_width 832
   --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )
 
 # Parallel arguments
 parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size 8
-  --tp_size 1
+  --tp_size 8
   --hsdp_replicate_dim 1
   --hsdp_shard_dim 8
 )

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

@@ -1,11 +1,13 @@
 #!/bin/bash
 #SBATCH --job-name=i2v
 #SBATCH --partition=main
+#SBATCH --qos=hao
 #SBATCH --nodes=4
 #SBATCH --ntasks=4
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
+#SBATCH --nodelist=fs-mbz-gpu-[100-850]
 #SBATCH --mem=1440G
 #SBATCH --output=i2v_output/i2v_%j.out
 #SBATCH --error=i2v_output/i2v_%j.err
@@ -58,14 +60,13 @@ training_args=(
   --num_height 480
   --num_width 832
   --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )
 
 # Parallel arguments
 parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size $NUM_GPUS
-  --tp_size 1
+  --tp_size $NUM_GPUS
   --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
   --hsdp_shard_dim $NUM_GPUS
 )

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

@@ -24,14 +24,13 @@ training_args=(
   --num_height 480
   --num_width 832
   --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )
 
 # Parallel arguments
 parallel_args=(
   --num_gpus $NUM_GPUS 
   --sp_size $NUM_GPUS 
-  --tp_size 1
+  --tp_size $NUM_GPUS
   --hsdp_replicate_dim 1
   --hsdp_shard_dim $NUM_GPUS
 )

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

@@ -1,11 +1,13 @@
 #!/bin/bash
 #SBATCH --job-name=t2v
 #SBATCH --partition=main
+#SBATCH --qos=hao
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
+#SBATCH --nodelist=fs-mbz-gpu-[100-850]
 #SBATCH --mem=1440G
 #SBATCH --output=t2v_output/t2v_%j.out
 #SBATCH --error=t2v_output/t2v_%j.err
@@ -55,14 +57,13 @@ training_args=(
   --num_height 480
   --num_width 832
   --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )
 
 # Parallel arguments
 parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size 4
-  --tp_size 1
+  --tp_size 4
   --hsdp_replicate_dim 2
   --hsdp_shard_dim 4
 )

fastvideo/v1/configs/models/dits/__init__.py

@@ -1,5 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.v1.configs.models.dits.cosmos import CosmosConfig
 from fastvideo.v1.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.v1.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.v1.configs.models.dits.wanvideo import WanVideoConfig
 
-__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"]
+__all__ = [
+    "HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig", "CosmosConfig"
+]

fastvideo/v1/configs/models/dits/cosmos.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_transformer_blocks(n: str, m) -> bool:
+    return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class CosmosArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_transformer_blocks])
+
+    _param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embed\.(.*)$": r"patch_embed.\1",
+            r"^time_embed\.time_proj\.(.*)$": r"time_embed.time_proj.\1",
+            r"^time_embed\.t_embedder\.(.*)$": r"time_embed.t_embedder.\1",
+            r"^time_embed\.norm\.(.*)$": r"time_embed.norm.\1",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
+            r"transformer_blocks.\1.attn1.norm_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
+            r"transformer_blocks.\1.attn1.norm_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.norm_q\.(.*)$":
+            r"transformer_blocks.\1.attn2.norm_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.norm_k\.(.*)$":
+            r"transformer_blocks.\1.attn2.norm_k.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$":
+            r"transformer_blocks.\1.ff.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$":
+            r"transformer_blocks.\1.ff.fc_out.\2",
+            r"^norm_out\.(.*)$": r"norm_out.\1",
+            r"^proj_out\.(.*)$": r"proj_out.\1",
+        })
+
+    _lora_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.(.*)$":
+            r"transformer_blocks.\1.ff.\2",
+        })
+
+    # Cosmos-specific config parameters based on transformer_cosmos.py
+    in_channels: int = 16
+    out_channels: int = 16
+    num_attention_heads: int = 16
+    attention_head_dim: int = 128
+    num_layers: int = 28
+    mlp_ratio: float = 4.0
+    text_embed_dim: int = 1024
+    adaln_lora_dim: int = 256
+    max_size: tuple[int, int, int] = (128, 240, 240)
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    rope_scale: tuple[float, float, float] = (1.0, 4.0, 4.0)
+    concat_padding_mask: bool = True
+    extra_pos_embed_type: str | None = None
+    qk_norm: str = "rms_norm"
+    eps: float = 1e-6
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.in_channels
+
+
+@dataclass
+class CosmosConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=CosmosArchConfig)
+    prefix: str = "Cosmos"

fastvideo/v1/configs/sample/cosmos.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.sample.base import CacheParams
+
+
+@dataclass
+class CosmosTeaCacheParams(CacheParams):
+    cache_type: str = "teacache"
+    teacache_thresh: float = 0.0
+    use_ret_steps: bool = True
+    ret_steps_coeffs: list[float] = field(default_factory=list)
+    non_ret_steps_coeffs: list[float] = field(default_factory=list)
+
+    @property
+    def coefficients(self) -> list[float]:
+        if self.use_ret_steps:
+            return self.ret_steps_coeffs
+        else:
+            return self.non_ret_steps_coeffs
+
+    @property
+    def ret_steps(self) -> int:
+        if self.use_ret_steps:
+            return 5 * 2
+        else:
+            return 1 * 2
+
+    def get_cutoff_steps(self, num_inference_steps: int) -> int:
+        if self.use_ret_steps:
+            return num_inference_steps * 2
+        else:
+            return num_inference_steps * 2 - 2

fastvideo/v1/distributed/parallel_state.py

@@ -23,6 +23,7 @@ If you only need to use the distributed environment without model parallelism,
  you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
+import gc
 import os
 import pickle
 import weakref
@@ -1015,6 +1016,15 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     if shutdown_ray:
         import ray  # Lazy import Ray
         ray.shutdown()
+    gc.collect()
+    from fastvideo.v1.platforms import current_platform
+    if not current_platform.is_cpu():
+        torch.cuda.empty_cache()
+    try:
+        torch._C._host_emptyCache()
+    except AttributeError:
+        logger.warning(
+            "torch._C._host_emptyCache() only available in Pytorch >=2.5")
 
 
 def in_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup,

fastvideo/v1/entrypoints/video_generator.py

@@ -6,6 +6,7 @@ This module provides a consolidated interface for generating videos using
 diffusion models.
 """
 
+import gc
 import math
 import os
 import time
@@ -276,3 +277,5 @@ class VideoGenerator:
         """
         self.executor.shutdown()
         del self.executor
+        gc.collect()
+        torch.cuda.empty_cache()

fastvideo/v1/fastvideo_args.py

@@ -292,7 +292,7 @@ class FastVideoArgs:
             assert self.sp_size != -1, "sp_size must be set for training"
 
         if self.tp_size == -1:
-            self.tp_size = 1
+            self.tp_size = self.num_gpus
         if self.sp_size == -1:
             self.sp_size = self.num_gpus
         if self.hsdp_shard_dim == -1:
@@ -305,6 +305,11 @@ class FastVideoArgs:
         if self.num_gpus < max(self.tp_size, self.sp_size):
             self.num_gpus = max(self.tp_size, self.sp_size)
 
+        if self.tp_size != self.sp_size:
+            raise ValueError(
+                f"tp_size ({self.tp_size}) must be equal to sp_size ({self.sp_size})"
+            )
+
         if self.enable_torch_compile and self.num_gpus > 1:
             logger.warning(
                 "Currently torch compile does not work with multi-gpu. Setting enable_torch_compile to False"

fastvideo/v1/layers/rotary_embedding.py

@@ -44,6 +44,60 @@ def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
     return x.flatten(-2)
 
 
+def apply_rotary_emb(
+    x: torch.Tensor,
+    freqs_cis: torch.Tensor | tuple[torch.Tensor],
+    use_real: bool = True,
+    use_real_unbind_dim: int = -1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+    to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+    reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+    tensors contain rotary embeddings and are returned as real tensors.
+
+    Args:
+        x (`torch.Tensor`):
+            Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
+        freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+    """
+    if use_real:
+        cos, sin = freqs_cis  # [S, D]
+        cos = cos[None, None]
+        sin = sin[None, None]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+
+        if use_real_unbind_dim == -1:
+            # Used for flux, cogvideox, hunyuan-dit
+            x_real, x_imag = x.reshape(*x.shape[:-1], -1,
+                                       2).unbind(-1)  # [B, S, H, D//2]
+            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+        elif use_real_unbind_dim == -2:
+            # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+            x_real, x_imag = x.reshape(*x.shape[:-1], 2,
+                                       -1).unbind(-2)  # [B, S, H, D//2]
+            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+        else:
+            raise ValueError(
+                f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2."
+            )
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+        return out
+    else:
+        # used for lumina
+        x_rotated = torch.view_as_complex(x.float().reshape(
+            *x.shape[:-1], -1, 2))
+        freqs_cis = freqs_cis.unsqueeze(2)
+        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+        return x_out.type_as(x)
+
+
 def _apply_rotary_emb(
     x: torch.Tensor,
     cos: torch.Tensor,

fastvideo/v1/layers/visual_embedding.py

@@ -173,3 +173,80 @@ def unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor:
     imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
 
     return imgs
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+    Args
+        timesteps (torch.Tensor):
+            a 1-D Tensor of N indices, one per batch element. These may be fractional.
+        embedding_dim (int):
+            the dimension of the output.
+        flip_sin_to_cos (bool):
+            Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+        downscale_freq_shift (float):
+            Controls the delta between frequencies between dimensions
+        scale (float):
+            Scaling factor applied to the embeddings.
+        max_period (int):
+            Controls the maximum frequency of the embeddings
+    Returns
+        torch.Tensor: an [N x dim] Tensor of positional embeddings.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+class Timesteps(nn.Module):
+
+    def __init__(self,
+                 num_channels: int,
+                 flip_sin_to_cos: bool,
+                 downscale_freq_shift: float,
+                 scale: int = 1):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        t_emb = get_timestep_embedding(
+            timesteps,
+            self.num_channels,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+            scale=self.scale,
+        )
+        return t_emb

fastvideo/v1/models/dits/cosmos.py

@@ -0,0 +1,731 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from fastvideo.v1.attention import DistributedAttention, LocalAttention
+from fastvideo.v1.configs.models.dits.cosmos import CosmosConfig
+from fastvideo.v1.forward_context import get_forward_context
+from fastvideo.v1.layers.layernorm import RMSNorm
+from fastvideo.v1.layers.linear import ReplicatedLinear
+from fastvideo.v1.layers.mlp import MLP
+from fastvideo.v1.layers.rotary_embedding import apply_rotary_emb
+from fastvideo.v1.layers.visual_embedding import Timesteps
+from fastvideo.v1.models.dits.base import BaseDiT
+from fastvideo.v1.platforms import AttentionBackendEnum
+
+
+class CosmosPatchEmbed(nn.Module):
+
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 patch_size: tuple[int, int, int],
+                 bias: bool = True) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+
+        self.proj = nn.Linear(in_channels * patch_size[0] * patch_size[1] *
+                              patch_size[2],
+                              out_channels,
+                              bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        hidden_states = hidden_states.reshape(batch_size, num_channels,
+                                              num_frames // p_t, p_t,
+                                              height // p_h, p_h, width // p_w,
+                                              p_w)
+        hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5,
+                                              7).flatten(4, 7)
+        hidden_states = self.proj(hidden_states)
+        return hidden_states
+
+
+class CosmosTimestepEmbedding(nn.Module):
+
+    def __init__(self, in_features: int, out_features: int) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(in_features, out_features, bias=False)
+        self.activation = nn.SiLU()
+        self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False)
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        emb = self.linear_1(timesteps)
+        emb = self.activation(emb)
+        emb = self.linear_2(emb)
+        return emb
+
+
+class CosmosEmbedding(nn.Module):
+
+    def __init__(self, embedding_dim: int, condition_dim: int) -> None:
+        super().__init__()
+
+        self.time_proj = Timesteps(embedding_dim,
+                                   flip_sin_to_cos=True,
+                                   downscale_freq_shift=0.0)
+        self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim)
+        self.norm = RMSNorm(embedding_dim, eps=1e-6)
+
+    def forward(self, hidden_states: torch.Tensor,
+                timestep: torch.LongTensor) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep).type_as(hidden_states)
+        temb = self.t_embedder(timesteps_proj)
+        embedded_timestep = self.norm(timesteps_proj)
+        return temb, embedded_timestep
+
+
+class CosmosAdaLayerNorm(nn.Module):
+
+    def __init__(self, in_features: int, hidden_features: int) -> None:
+        super().__init__()
+        self.embedding_dim = in_features
+
+        self.activation = nn.SiLU()
+        self.norm = nn.LayerNorm(in_features,
+                                 elementwise_affine=False,
+                                 eps=1e-6)
+        self.linear_1 = nn.Linear(in_features, hidden_features, bias=False)
+        self.linear_2 = nn.Linear(hidden_features, 2 * in_features, bias=False)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                embedded_timestep: torch.Tensor,
+                temb: torch.Tensor | None = None) -> torch.Tensor:
+        embedded_timestep = self.activation(embedded_timestep)
+        embedded_timestep = self.linear_1(embedded_timestep)
+        embedded_timestep = self.linear_2(embedded_timestep)
+
+        if temb is not None:
+            embedded_timestep = embedded_timestep + temb[..., :2 *
+                                                         self.embedding_dim]
+
+        shift, scale = embedded_timestep.chunk(2, dim=-1)
+        hidden_states = self.norm(hidden_states)
+
+        if embedded_timestep.ndim == 2:
+            shift, scale = (x.unsqueeze(1) for x in (shift, scale))
+
+        hidden_states = hidden_states * (1 + scale) + shift
+        return hidden_states
+
+
+class CosmosAdaLayerNormZero(nn.Module):
+
+    def __init__(self,
+                 in_features: int,
+                 hidden_features: int | None = None) -> None:
+        super().__init__()
+
+        self.norm = nn.LayerNorm(in_features,
+                                 elementwise_affine=False,
+                                 eps=1e-6)
+        self.activation = nn.SiLU()
+
+        if hidden_features is None:
+            self.linear_1 = nn.Identity()
+        else:
+            self.linear_1 = nn.Linear(in_features, hidden_features, bias=False)
+
+        self.linear_2 = nn.Linear(
+            hidden_features if hidden_features is not None else in_features,
+            3 * in_features,
+            bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        embedded_timestep: torch.Tensor,
+        temb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        embedded_timestep = self.activation(embedded_timestep)
+        embedded_timestep = self.linear_1(embedded_timestep)
+        embedded_timestep = self.linear_2(embedded_timestep)
+
+        if temb is not None:
+            embedded_timestep = embedded_timestep + temb
+
+        shift, scale, gate = embedded_timestep.chunk(3, dim=-1)
+        hidden_states = self.norm(hidden_states)
+
+        if embedded_timestep.ndim == 2:
+            shift, scale, gate = (x.unsqueeze(1) for x in (shift, scale, gate))
+
+        hidden_states = hidden_states * (1 + scale) + shift
+        return hidden_states, gate
+
+
+class CosmosSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 qk_norm=True,
+                 eps=1e-6,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
+                 prefix: str = "") -> None:
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.to_q = ReplicatedLinear(dim, dim, bias=False)
+        self.to_k = ReplicatedLinear(dim, dim, bias=False)
+        self.to_v = ReplicatedLinear(dim, dim, bias=False)
+        self.to_out = ReplicatedLinear(dim, dim, bias=False)
+        self.norm_q = RMSNorm(self.head_dim,
+                              eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = RMSNorm(self.head_dim,
+                              eps=eps) if qk_norm else nn.Identity()
+
+        # Attention mechanism
+        self.attn = DistributedAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+            prefix=prefix)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor | None = None,
+                attention_mask: torch.Tensor | None = None,
+                image_rotary_emb: torch.Tensor | None = None) -> torch.Tensor:
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        # Get QKV
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(encoder_hidden_states)
+        value, _ = self.to_v(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        query = query.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
+        key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
+        value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
+
+        # Apply normalization
+        if self.norm_q is not None:
+            query = self.norm_q.forward_native(query)
+        if self.norm_k is not None:
+            key = self.norm_k.forward_native(key)
+
+        # Apply RoPE if provided
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query,
+                                     image_rotary_emb,
+                                     use_real=True,
+                                     use_real_unbind_dim=-2)
+            key = apply_rotary_emb(key,
+                                   image_rotary_emb,
+                                   use_real=True,
+                                   use_real_unbind_dim=-2)
+
+        # Attention computation
+        attn_output, _ = self.attn(query, key, value)
+        # attn_output = attn_output.flatten(2)
+        attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)
+
+        # Output projection
+        attn_output, _ = self.to_out(attn_output)
+        return attn_output
+
+
+class CosmosCrossAttention(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 cross_attention_dim: int,
+                 num_heads: int,
+                 qk_norm=True,
+                 eps=1e-6,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
+                 prefix: str = "") -> None:
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.cross_attention_dim = cross_attention_dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.to_q = ReplicatedLinear(dim, dim, bias=False)
+        self.to_k = ReplicatedLinear(cross_attention_dim, dim, bias=False)
+        self.to_v = ReplicatedLinear(cross_attention_dim, dim, bias=False)
+        self.to_out = ReplicatedLinear(dim, dim, bias=False)
+        self.norm_q = RMSNorm(self.head_dim,
+                              eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = RMSNorm(self.head_dim,
+                              eps=eps) if qk_norm else nn.Identity()
+
+        # Attention mechanism
+        self.attn = LocalAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=supported_attention_backends)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor,
+                attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+
+        # Get QKV
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(encoder_hidden_states)
+        value, _ = self.to_v(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        query = query.unflatten(2, (self.num_heads, -1))
+        key = key.unflatten(2, (self.num_heads, -1))
+        value = value.unflatten(2, (self.num_heads, -1))
+
+        # Apply normalization
+        if self.norm_q is not None:
+            query = self.norm_q.forward_native(query)
+        if self.norm_k is not None:
+            key = self.norm_k.forward_native(key)
+
+        # Attention computation
+        attn_output = self.attn(query, key, value)
+        attn_output = attn_output.flatten(2, 3).type_as(query)
+
+        # Output projection
+        attn_output, _ = self.to_out(attn_output)
+        return attn_output
+
+
+class CosmosTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        cross_attention_dim: int,
+        mlp_ratio: float = 4.0,
+        adaln_lora_dim: int = 256,
+        qk_norm: str = "rms_norm",
+        out_bias: bool = False,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...]
+        | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size,
+                                            hidden_features=adaln_lora_dim)
+        self.attn1 = CosmosSelfAttention(
+            dim=hidden_size,
+            num_heads=num_attention_heads,
+            qk_norm=(qk_norm == "rms_norm"),
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn1")
+
+        self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size,
+                                            hidden_features=adaln_lora_dim)
+        self.attn2 = CosmosCrossAttention(
+            dim=hidden_size,
+            cross_attention_dim=cross_attention_dim,
+            num_heads=num_attention_heads,
+            qk_norm=(qk_norm == "rms_norm"),
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn2")
+
+        self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size,
+                                            hidden_features=adaln_lora_dim)
+        self.ff = MLP(hidden_size,
+                      int(hidden_size * mlp_ratio),
+                      act_type="gelu",
+                      bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        embedded_timestep: torch.Tensor,
+        temb: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        extra_pos_emb: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if extra_pos_emb is not None:
+            hidden_states = hidden_states + extra_pos_emb
+
+        # 1. Self Attention
+        norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep,
+                                              temb)
+        attn_output = self.attn1(norm_hidden_states,
+                                 image_rotary_emb=image_rotary_emb)
+        hidden_states = hidden_states + gate * attn_output
+
+        # 2. Cross Attention
+        norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep,
+                                              temb)
+        attn_output = self.attn2(norm_hidden_states,
+                                 encoder_hidden_states=encoder_hidden_states,
+                                 attention_mask=attention_mask)
+        hidden_states = hidden_states + gate * attn_output
+
+        # 3. Feed Forward
+        norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep,
+                                              temb)
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + gate * ff_output
+
+        return hidden_states
+
+
+class CosmosRotaryPosEmbed(nn.Module):
+
+    def __init__(
+            self,
+            hidden_size: int,
+            max_size: tuple[int, int, int] = (128, 240, 240),
+            patch_size: tuple[int, int, int] = (1, 2, 2),
+            base_fps: int = 24,
+            rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0),
+    ) -> None:
+        super().__init__()
+
+        self.max_size = [
+            size // patch
+            for size, patch in zip(max_size, patch_size, strict=False)
+        ]
+        self.patch_size = patch_size
+        self.base_fps = base_fps
+
+        self.dim_h = hidden_size // 6 * 2
+        self.dim_w = hidden_size // 6 * 2
+        self.dim_t = hidden_size - self.dim_h - self.dim_w
+
+        self.h_ntk_factor = rope_scale[1]**(self.dim_h / (self.dim_h - 2))
+        self.w_ntk_factor = rope_scale[2]**(self.dim_w / (self.dim_w - 2))
+        self.t_ntk_factor = rope_scale[0]**(self.dim_t / (self.dim_t - 2))
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        pe_size = [
+            num_frames // self.patch_size[0], height // self.patch_size[1],
+            width // self.patch_size[2]
+        ]
+        device = hidden_states.device
+
+        h_theta = 10000.0 * self.h_ntk_factor
+        w_theta = 10000.0 * self.w_ntk_factor
+        t_theta = 10000.0 * self.t_ntk_factor
+
+        seq = torch.arange(max(self.max_size),
+                           device=device,
+                           dtype=torch.float32)
+        dim_h_range = (
+            torch.arange(0, self.dim_h, 2, device=device,
+                         dtype=torch.float32)[:(self.dim_h // 2)] / self.dim_h)
+        dim_w_range = (
+            torch.arange(0, self.dim_w, 2, device=device,
+                         dtype=torch.float32)[:(self.dim_w // 2)] / self.dim_w)
+        dim_t_range = (
+            torch.arange(0, self.dim_t, 2, device=device,
+                         dtype=torch.float32)[:(self.dim_t // 2)] / self.dim_t)
+        h_spatial_freqs = 1.0 / (h_theta**dim_h_range)
+        w_spatial_freqs = 1.0 / (w_theta**dim_w_range)
+        temporal_freqs = 1.0 / (t_theta**dim_t_range)
+
+        emb_h = torch.outer(seq[:pe_size[1]],
+                            h_spatial_freqs)[None, :, None, :].repeat(
+                                pe_size[0], 1, pe_size[2], 1)
+        emb_w = torch.outer(seq[:pe_size[2]],
+                            w_spatial_freqs)[None, None, :, :].repeat(
+                                pe_size[0], pe_size[1], 1, 1)
+
+        # Apply sequence scaling in temporal dimension
+        if fps is None:
+            # Images
+            emb_t = torch.outer(seq[:pe_size[0]], temporal_freqs)
+        else:
+            # Videos
+            emb_t = torch.outer(seq[:pe_size[0]] / fps * self.base_fps,
+                                temporal_freqs)
+
+        emb_t = emb_t[:, None, None, :].repeat(1, pe_size[1], pe_size[2], 1)
+        freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0,
+                                                                     2).float()
+        cos = torch.cos(freqs)
+        sin = torch.sin(freqs)
+        return cos, sin
+
+
+class CosmosLearnablePositionalEmbed(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        max_size: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+
+        self.max_size = [
+            size // patch
+            for size, patch in zip(max_size, patch_size, strict=False)
+        ]
+        self.patch_size = patch_size
+        self.eps = eps
+
+        self.pos_emb_t = nn.Parameter(torch.zeros(self.max_size[0],
+                                                  hidden_size))
+        self.pos_emb_h = nn.Parameter(torch.zeros(self.max_size[1],
+                                                  hidden_size))
+        self.pos_emb_w = nn.Parameter(torch.zeros(self.max_size[2],
+                                                  hidden_size))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        pe_size = [
+            num_frames // self.patch_size[0], height // self.patch_size[1],
+            width // self.patch_size[2]
+        ]
+
+        emb_t = self.pos_emb_t[:pe_size[0]][None, :, None, None, :].repeat(
+            batch_size, 1, pe_size[1], pe_size[2], 1)
+        emb_h = self.pos_emb_h[:pe_size[1]][None, None, :, None, :].repeat(
+            batch_size, pe_size[0], 1, pe_size[2], 1)
+        emb_w = self.pos_emb_w[:pe_size[2]][None, None, None, :, :].repeat(
+            batch_size, pe_size[0], pe_size[1], 1, 1)
+        emb = emb_t + emb_h + emb_w
+        emb = emb.flatten(1, 3)
+
+        norm = torch.linalg.vector_norm(emb,
+                                        dim=-1,
+                                        keepdim=True,
+                                        dtype=torch.float32)
+        norm = torch.add(self.eps,
+                         norm,
+                         alpha=np.sqrt(norm.numel() / emb.numel()))
+        return (emb / norm).type_as(hidden_states)
+
+
+class CosmosTransformer3DModel(BaseDiT):
+    _fsdp_shard_conditions = CosmosConfig()._fsdp_shard_conditions
+    _compile_conditions = CosmosConfig()._compile_conditions
+    _supported_attention_backends = CosmosConfig()._supported_attention_backends
+    _param_names_mapping = CosmosConfig()._param_names_mapping
+    _lora_param_names_mapping = CosmosConfig()._lora_param_names_mapping
+
+    def __init__(self, config: CosmosConfig, hf_config: dict[str, Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels
+        self.num_channels_latents = config.num_channels_latents
+        self.patch_size = config.patch_size
+        self.max_size = config.max_size
+        self.rope_scale = config.rope_scale
+        self.concat_padding_mask = config.concat_padding_mask
+        self.extra_pos_embed_type = config.extra_pos_embed_type
+
+        # 1. Patch Embedding
+        patch_embed_in_channels = config.in_channels + 1 if config.concat_padding_mask else config.in_channels
+        self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels,
+                                            inner_dim,
+                                            config.patch_size,
+                                            bias=False)
+
+        # 2. Positional Embedding
+        self.rope = CosmosRotaryPosEmbed(hidden_size=config.attention_head_dim,
+                                         max_size=config.max_size,
+                                         patch_size=config.patch_size,
+                                         rope_scale=config.rope_scale)
+
+        self.learnable_pos_embed = None
+        if config.extra_pos_embed_type == "learnable":
+            self.learnable_pos_embed = CosmosLearnablePositionalEmbed(
+                hidden_size=inner_dim,
+                max_size=config.max_size,
+                patch_size=config.patch_size,
+            )
+
+        # 3. Time Embedding
+        self.time_embed = CosmosEmbedding(inner_dim, inner_dim)
+
+        # 4. Transformer Blocks
+        self.transformer_blocks = nn.ModuleList([
+            CosmosTransformerBlock(
+                num_attention_heads=config.num_attention_heads,
+                attention_head_dim=config.attention_head_dim,
+                cross_attention_dim=config.text_embed_dim,
+                mlp_ratio=config.mlp_ratio,
+                adaln_lora_dim=config.adaln_lora_dim,
+                qk_norm=config.qk_norm,
+                out_bias=False,
+                supported_attention_backends=self._supported_attention_backends,
+                prefix=f"{config.prefix}.transformer_blocks.{i}",
+            ) for i in range(config.num_layers)
+        ])
+
+        # 5. Output norm & projection
+        self.norm_out = CosmosAdaLayerNorm(inner_dim, config.adaln_lora_dim)
+        self.proj_out = nn.Linear(inner_dim,
+                                  config.out_channels *
+                                  math.prod(config.patch_size),
+                                  bias=False)
+
+        self.gradient_checkpointing = False
+
+        # For TeaCache
+        self.previous_e0_even = None
+        self.previous_e0_odd = None
+        self.previous_residual_even = None
+        self.previous_residual_odd = None
+        self.is_even = True
+        self.should_calc_even = True
+        self.should_calc_odd = True
+        self.accumulated_rel_l1_distance_even = 0
+        self.accumulated_rel_l1_distance_odd = 0
+        self.cnt = 0
+        self.__post_init__()
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                timestep: torch.Tensor,
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+                attention_mask: torch.Tensor | None = None,
+                fps: int | None = None,
+                condition_mask: torch.Tensor | None = None,
+                padding_mask: torch.Tensor | None = None,
+                **kwargs) -> torch.Tensor:
+        forward_batch = get_forward_context().forward_batch
+        enable_teacache = forward_batch is not None and forward_batch.enable_teacache
+
+        orig_dtype = hidden_states.dtype
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+
+        # 1. Concatenate padding mask if needed & prepare attention mask
+        if condition_mask is not None:
+            hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
+
+        if self.concat_padding_mask and padding_mask is not None:
+            from torchvision import transforms
+            padding_mask = transforms.functional.resize(
+                padding_mask,
+                list(hidden_states.shape[-2:]),
+                interpolation=transforms.InterpolationMode.NEAREST)
+            hidden_states = torch.cat([
+                hidden_states,
+                padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1,
+                                                 1)
+            ],
+                                      dim=1)
+            # # Resize padding mask to match hidden states spatial dimensions
+            # padding_mask_resized = F.interpolate(
+            #     padding_mask.float().unsqueeze(1),
+            #     size=(height, width),
+            #     mode='nearest'
+            # ).squeeze(1)
+            # hidden_states = torch.cat(
+            #     [hidden_states, padding_mask_resized.unsqueeze(1).unsqueeze(2).repeat(1, 1, num_frames, 1, 1)], dim=1
+            # )
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(
+                1)  # [B, 1, 1, S]
+
+        # 2. Generate positional embeddings
+        image_rotary_emb = self.rope(hidden_states, fps=fps)
+        extra_pos_emb = self.learnable_pos_embed(
+            hidden_states) if self.extra_pos_embed_type == "learnable" else None
+
+        # 3. Patchify input
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+        hidden_states = self.patch_embed(hidden_states)
+        hidden_states = hidden_states.flatten(
+            1, 3)  # [B, T, H, W, C] -> [B, THW, C] codespell:ignore
+
+        # 4. Timestep embeddings
+        if timestep.ndim == 1:
+            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
+        elif timestep.ndim == 5:
+            assert timestep.shape == (batch_size, 1, num_frames, 1, 1), (
+                f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}"
+            )
+            timestep = timestep.flatten()
+            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
+            # We can do this because num_frames == post_patch_num_frames, as p_t is 1
+            temb, embedded_timestep = (
+                x.view(batch_size, post_patch_num_frames, 1, 1,
+                       -1).expand(-1, -1, post_patch_height, post_patch_width,
+                                  -1).flatten(1, 3)
+                for x in (temb, embedded_timestep)
+            )  # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C] codespell:ignore
+        else:
+            raise ValueError(f"Unsupported timestep shape: {timestep.shape}")
+
+        # 6. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for block in self.transformer_blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    embedded_timestep,
+                    temb,
+                    image_rotary_emb,
+                    extra_pos_emb,
+                    attention_mask,
+                )
+        else:
+            for block in self.transformer_blocks:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    embedded_timestep=embedded_timestep,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    extra_pos_emb=extra_pos_emb,
+                    attention_mask=attention_mask,
+                )
+
+        # 7. Output norm & projection & unpatchify
+        hidden_states = self.norm_out(hidden_states, embedded_timestep, temb)
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
+        hidden_states = hidden_states.unflatten(
+            1, (post_patch_num_frames, post_patch_height, post_patch_width))
+        # NOTE: The permutation order here is not the inverse operation of what happens when patching as usually expected.
+        # It might be a source of confusion to the reader, but this is correct
+        hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
+        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return hidden_states

fastvideo/v1/models/registry.py

@@ -23,7 +23,8 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "HunyuanVideoTransformer3DModel":
     ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
-    "StepVideoModel": ("dits", "stepvideo", "StepVideoModel")
+    "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
+    "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel")
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {

fastvideo/v1/models/vaes/common.py

@@ -239,6 +239,7 @@ class ParallelTiledVAE(ABC):
 
         results = torch.cat(local_results, dim=0).contiguous()
         del local_results
+        torch.cuda.empty_cache()
         # first gather size to pad the results
         local_size = torch.tensor([results.size(0)],
                                   device=results.device,
@@ -252,7 +253,7 @@ class ParallelTiledVAE(ABC):
         padded_results = torch.zeros(max_size, device=results.device)
         padded_results[:results.size(0)] = results
         del results
-
+        torch.cuda.empty_cache()
         # Gather all results
         gathered_dim_metadata = [None] * world_size
         gathered_results = torch.zeros_like(padded_results).repeat(

fastvideo/v1/pipelines/stages/encoding.py

@@ -136,6 +136,7 @@ class EncodingStage(PipelineStage):
             self.maybe_free_model_hooks()
 
         self.vae.to("cpu")
+        torch.cuda.empty_cache()
 
         return batch
 

fastvideo/v1/pipelines/stages/image_encoding.py

@@ -5,6 +5,8 @@ Image encoding stages for I2V diffusion pipelines.
 This module contains implementations of image encoding stages for diffusion pipelines.
 """
 
+import torch
+
 from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
@@ -66,6 +68,7 @@ class ImageEncodingStage(PipelineStage):
 
         if fastvideo_args.use_cpu_offload:
             self.image_encoder.to('cpu')
+            torch.cuda.empty_cache()
 
         return batch
 

fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py

@@ -105,7 +105,7 @@ def run_training():
         "--num_latent_t", "8",
         "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
         "--sp_size", NUM_GPUS_PER_NODE_TRAINING,
-        "--tp_size", 1,
+        "--tp_size", NUM_GPUS_PER_NODE_TRAINING,
         "--hsdp_replicate_dim", "1",
         "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
         "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,

fastvideo/v1/tests/ssim/README.md

@@ -24,7 +24,7 @@ FastHunyuan-diffusers: {
 "flow_shift": 17,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": true,
 "fps": 24
 }
@@ -41,7 +41,7 @@ Wan2.1-T2V-1.3B-Diffusers: {
 "flow_shift": 7.0,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": True,
 "fps": 24,
 "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
@@ -60,7 +60,7 @@ Wan2.1-I2V-14B-480P-Diffusers: {
 "flow_shift": 7.0,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": True,
 "fps": 24,
 "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",

fastvideo/v1/tests/ssim/test_inference_similarity.py

@@ -33,7 +33,7 @@ HUNYUAN_PARAMS = {
     "flow_shift": 17,
     "seed": 1024,
     "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
     "vae_sp": True,
     "fps": 24,
 }
@@ -50,7 +50,7 @@ WAN_T2V_PARAMS = {
     "flow_shift": 7.0,
     "seed": 1024,
     "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
     "vae_sp": True,
     "fps": 24,
     "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
@@ -69,7 +69,7 @@ WAN_I2V_PARAMS = {
     "flow_shift": 7.0,
     "seed": 1024,
     "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
     "vae_sp": True,
     "fps": 24,
     "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
@@ -238,7 +238,7 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
         logger.error("Failed to write SSIM results to file")
 
     min_acceptable_ssim = 0.97
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"
+    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"
 
 @pytest.mark.parametrize("prompt", TEST_PROMPTS)
 @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
@@ -337,5 +337,5 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     if not success:
         logger.error("Failed to write SSIM results to file")
 
-    min_acceptable_ssim = 0.93
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"
+    min_acceptable_ssim = 0.95
+    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"

fastvideo/v1/tests/training/Vanilla/test_training_loss.py

@@ -43,7 +43,7 @@ def run_worker():
         "--num_latent_t", "4",
         "--num_gpus", "4",
         "--sp_size", "4",
-        "--tp_size", "1",
+        "--tp_size", "4",
         "--hsdp_replicate_dim", "1",
         "--hsdp_shard_dim", "4",
         "--train_sp_batch_size", "1",

fastvideo/v1/tests/transformers/test_cosmos2.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import numpy as np
+import pytest
+import torch
+from diffusers import CosmosTransformer3DModel
+
+from fastvideo.v1.configs.pipelines import PipelineConfig
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import TransformerLoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.configs.models.dits import CosmosConfig
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29504"
+
+BASE_MODEL_PATH = "nvidia/Cosmos-Predict2-2B-Text2Image"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_cosmos2_transformer():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
+                         use_cpu_offload=False,
+                         pipeline_config=PipelineConfig(dit_config=CosmosConfig(), dit_precision=precision_str))
+    args.device = device
+
+    loader = TransformerLoader()
+    model2 = loader.load(TRANSFORMER_PATH, "", args).to(device, dtype=precision)
+
+    model1 = CosmosTransformer3DModel.from_pretrained(
+        TRANSFORMER_PATH, device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+
+    total_params = sum(p.numel() for p in model1.parameters())
+    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
+    weight_sum_model1 = sum(
+        p.to(torch.float64).sum().item() for p in model1.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model1 = weight_sum_model1 / total_params
+    logger.info("Model 1 weight sum: %s", weight_sum_model1)
+    logger.info("Model 1 weight mean: %s", weight_mean_model1)
+
+    # Calculate weight sum for model2 (converting to float64 to avoid overflow)
+    total_params_model2 = sum(p.numel() for p in model2.parameters())
+    weight_sum_model2 = sum(
+        p.to(torch.float64).sum().item() for p in model2.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model2 = weight_sum_model2 / total_params_model2
+    logger.info("Model 2 weight sum: %s", weight_sum_model2)
+    logger.info("Model 2 weight mean: %s", weight_mean_model2)
+
+    weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2)
+    logger.info("Weight sum difference: %s", weight_sum_diff)
+    weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2)
+    logger.info("Weight mean difference: %s", weight_mean_diff)
+
+    # Set both models to eval mode
+    model1 = model1.eval()
+    model2 = model2.eval()
+
+    # Create identical inputs for both models
+    batch_size = 1
+    seq_len = 30
+
+    # Video latents [B, C, T, H, W] - Cosmos2 specific dimensions
+    hidden_states = torch.randn(batch_size,
+                                16,
+                                1,  # Single frame for image generation
+                                32,  # Height patches
+                                32,  # Width patches
+                                device=device,
+                                dtype=precision)
+
+    # Text embeddings [B, L, D] - Cosmos2 uses T5 embeddings with 1024 dim
+    encoder_hidden_states = torch.randn(batch_size,
+                                        seq_len,
+                                        1024,  # T5 embedding dimension
+                                        device=device,
+                                        dtype=precision)
+
+    # Timestep
+    timestep = torch.tensor([500], device=device, dtype=precision)
+
+    # padding mask
+    padding_mask = hidden_states.new_zeros(1, 1, 32, 32, device=device, dtype=precision)
+    # print(padding_mask.shape)
+
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+
+    with torch.amp.autocast('cuda', dtype=precision):
+        output1 = model1(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+            padding_mask=padding_mask,
+            return_dict=False,
+        )[0]
+        with set_forward_context(
+                current_timestep=0,
+                attn_metadata=None,
+                forward_batch=forward_batch,
+        ):
+            output2 = model2(hidden_states=hidden_states,
+                             encoder_hidden_states=encoder_hidden_states,
+                             timestep=timestep,
+                             padding_mask=padding_mask)
+
+    # Check if outputs have the same shape
+    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
+
+    # Check if outputs are similar (allowing for small numerical differences)
+    max_diff = torch.max(torch.abs(output1 - output2))
+    mean_diff = torch.mean(torch.abs(output1 - output2))
+    logger.info("Max Diff: %s", max_diff.item())
+    logger.info("Mean Diff: %s", mean_diff.item())
+    assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}"
+    # mean diff
+    assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}"
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_cosmos2_transformer_video2world():
+    """Test Cosmos2 Video2World variant"""
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    
+    # Use Video2World model path
+    base_model_path = "nvidia/Cosmos-Predict2-2B-Video2World"
+    model_path = maybe_download_model(base_model_path,
+                                      local_dir=os.path.join(
+                                          'data', base_model_path))
+    transformer_path = os.path.join(model_path, "transformer")
+    
+    args = FastVideoArgs(model_path=transformer_path,
+                         use_cpu_offload=False,
+                         pipeline_config=PipelineConfig(dit_config=CosmosConfig(), dit_precision=precision_str))
+    args.device = device
+
+    loader = TransformerLoader()
+    model2 = loader.load(transformer_path, "", args).to(device, dtype=precision)
+
+    model1 = CosmosTransformer3DModel.from_pretrained(
+        transformer_path, device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+
+    # Set both models to eval mode
+    model1 = model1.eval()
+    model2 = model2.eval()
+
+    # Create identical inputs for both models
+    batch_size = 1
+    seq_len = 30
+
+    # Video latents [B, C, T, H, W] - Video2World has additional condition channel
+    hidden_states = torch.randn(batch_size,
+                                17,  # 16 + 1 for condition channel
+                                8,   # Multiple frames for video
+                                32,  # Height patches
+                                32,  # Width patches
+                                device=device,
+                                dtype=precision)
+
+    # Text embeddings [B, L, D]
+    encoder_hidden_states = torch.randn(batch_size,
+                                        seq_len,
+                                        1024,  # T5 embedding dimension
+                                        device=device,
+                                        dtype=precision)
+
+    # Timestep
+    timestep = torch.tensor([500], device=device, dtype=precision)
+
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+
+    with torch.amp.autocast('cuda', dtype=precision):
+        output1 = model1(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+            return_dict=False,
+        )[0]
+        with set_forward_context(
+                current_timestep=0,
+                attn_metadata=None,
+                forward_batch=forward_batch,
+        ):
+            output2 = model2(hidden_states=hidden_states,
+                             encoder_hidden_states=encoder_hidden_states,
+                             timestep=timestep)
+
+    # Check if outputs have the same shape
+    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
+
+    # Check if outputs are similar (allowing for small numerical differences)
+    max_diff = torch.max(torch.abs(output1 - output2))
+    mean_diff = torch.mean(torch.abs(output1 - output2))
+    logger.info("Max Diff: %s", max_diff.item())
+    logger.info("Mean Diff: %s", mean_diff.item())
+    assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}"
+    # mean diff
+    assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}"

fastvideo/v1/training/training_pipeline.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+import gc
 import math
 import os
 import time
@@ -705,3 +706,5 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         # Re-enable gradients for training
         training_args.inference_mode = False
         transformer.train()
+        gc.collect()
+        torch.cuda.empty_cache()

fastvideo/v1/worker/gpu_worker.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import contextlib
 import faulthandler
+import gc
 import multiprocessing as mp
 import os
 import signal
@@ -68,6 +69,8 @@ class Worker:
         torch.cuda.set_device(self.device)
 
         # _check_if_gpu_supports_dtype(self.model_config.dtype)
+        gc.collect()
+        torch.cuda.empty_cache()
         self.init_gpu_memory = torch.cuda.mem_get_info()[0]
 
         os.environ["MASTER_ADDR"] = "localhost"
@@ -99,6 +102,9 @@ class Worker:
         if hasattr(self, 'pipeline') and self.pipeline is not None:
             # Clean up pipeline resources if needed
             pass
+        # Release CUDA resources
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         # Destroy the distributed environment
         cleanup_dist_env_and_memory(shutdown_ray=False)
@@ -127,6 +133,8 @@ class Worker:
 
                 # Handle regular RPC calls
                 if method_name == 'execute_forward':
+                    gc.collect()
+                    torch.cuda.empty_cache()
                     forward_batch = recv_rpc['kwargs']['forward_batch']
                     fastvideo_args = recv_rpc['kwargs']['fastvideo_args']
                     output_batch = self.execute_forward(forward_batch,

scripts/finetune/finetune_v1.sh

@@ -20,7 +20,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --train_batch_size=4 \
     --num_latent_t 20 \
     --sp_size 4 \
-    --tp_size 1 \
+    --tp_size 4 \
     --hsdp_replicate_dim 1 \
     --hsdp_shard_dim 4 \
     --num_gpus $NUM_GPUS \

scripts/inference/v1_inference_fasthunyuan.sh

@@ -3,10 +3,13 @@
 num_gpus=4
 export MODEL_BASE=FastVideo/FastHunyuan-Diffusers
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
     --model-path $MODEL_BASE \
     --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 720 \
     --width 1280 \

scripts/inference/v1_inference_hunyuan.sh

@@ -4,10 +4,13 @@ num_gpus=4
 export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
     --model-path $MODEL_BASE \
     --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 720 \
     --width 1280 \

scripts/inference/v1_inference_hunyuan_STA.sh

@@ -5,10 +5,13 @@ export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_hunyuan.json
 export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
     --model-path $MODEL_BASE \
     --sp-size ${num_gpus} \
-    --tp-size 1 \
+    --tp-size ${num_gpus} \
     --height 768 \
     --width 1280 \
     --num-frames 117 \

scripts/inference/v1_inference_wan.sh

@@ -4,10 +4,13 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_BACKEND=
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
     --model-path $MODEL_BASE \
     --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 480 \
     --width 832 \

scripts/inference/v1_inference_wan_STA.sh

@@ -4,10 +4,13 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
 export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
     --model-path $MODEL_BASE \
     --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 768 \
     --width 1280 \

scripts/inference/v1_inference_wan_VSA.sh

@@ -4,11 +4,14 @@ num_gpus=1
 export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
 # change model path to local dir if you want to inference using your checkpoint
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo 
+# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
     --model-path $MODEL_BASE \
     --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 448 \
     --width 832 \

scripts/inference/v1_inference_wan_i2v.sh

@@ -4,6 +4,9 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_BACKEND=
 export MODEL_BASE=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
     --model-path $MODEL_BASE \
     --sp-size $num_gpus \