Update preprocess_pipeline_ode_trajectory.py

examples/inference/basic/basic_hy15.py

@@ -0,0 +1,42 @@
+from fastvideo import VideoGenerator
+import json
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_hy15"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        # image_encoder_cpu_offload=False,
+    )
+
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16)
+
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16)
+
+
+if __name__ == "__main__":
+    main()

fastvideo/attention/backends/flash_attn.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
+import torch.nn.functional as F
 from flash_attn import flash_attn_func as flash_attn_2_func
-
+from dataclasses import dataclass
 try:
     from flash_attn_interface import flash_attn_func as flash_attn_3_func
 
@@ -46,6 +47,29 @@ class FlashAttentionBackend(AttentionBackend):
         raise NotImplementedError
 
 
+@dataclass
+class FlashAttnMetadata(AttentionMetadata):
+    current_timestep: int
+    attn_mask: torch.Tensor | None = None
+
+
+class FlashAttnMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        current_timestep: int,
+        attn_mask: torch.Tensor,
+    ) -> FlashAttnMetadata:
+        return FlashAttnMetadata(current_timestep=current_timestep,
+                                 attn_mask=attn_mask)
+
+
 class FlashAttentionImpl(AttentionImpl):
 
     def __init__(
@@ -66,12 +90,27 @@ class FlashAttentionImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata,
+        attn_metadata: FlashAttnMetadata,
     ):
-        output = flash_attn_func(
-            query,  # type: ignore[no-untyped-call]
-            key,
-            value,
-            softmax_scale=self.softmax_scale,
-            causal=self.causal)
+        if attn_metadata is not None and hasattr(
+                attn_metadata,
+                "attn_mask") and attn_metadata.attn_mask is not None:
+            from fastvideo.attention.utils.flash_attn_no_pad import flash_attn_no_pad
+            attn_mask = attn_metadata.attn_mask
+            qkv = torch.stack([query, key, value], dim=2)
+
+            attn_mask = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0),
+                              value=True)
+            output = flash_attn_no_pad(qkv,
+                                       attn_mask,
+                                       causal=False,
+                                       dropout_p=0,
+                                       softmax_scale=None)
+        else:
+            output = flash_attn_func(
+                query,  # type: ignore[no-untyped-call]
+                key,
+                value,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal)
         return output

fastvideo/attention/backends/sdpa.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
-
+from dataclasses import dataclass
 from fastvideo.attention.backends.abstract import (  # FlashAttentionMetadata,
-    AttentionBackend, AttentionImpl, AttentionMetadata)
+    AttentionBackend, AttentionImpl, AttentionMetadata,
+    AttentionMetadataBuilder)
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
@@ -30,6 +31,29 @@ class SDPABackend(AttentionBackend):
     #     return FlashAttentionMetadata
 
 
+@dataclass
+class SDPAMetadata(AttentionMetadata):
+    current_timestep: int
+    attn_mask: torch.Tensor | None = None
+
+
+class SDPAMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        current_timestep: int,
+        attn_mask: torch.Tensor,
+    ) -> SDPAMetadata:
+        return SDPAMetadata(current_timestep=current_timestep,
+                            attn_mask=attn_mask)
+
+
 class SDPAImpl(AttentionImpl):
 
     def __init__(
@@ -51,14 +75,15 @@ class SDPAImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata,
+        attn_metadata: SDPAMetadata,
     ) -> torch.Tensor:
         # transpose to bs, heads, seq_len, head_dim
         query = query.transpose(1, 2)
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
+        attn_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         attn_kwargs = {
-            "attn_mask": None,
+            "attn_mask": attn_mask,
             "dropout_p": self.dropout,
             "is_causal": self.causal,
             "scale": self.softmax_scale

fastvideo/attention/layer.py

@@ -108,6 +108,7 @@ class DistributedAttention(nn.Module):
         # Since mask is [batch, full_seq_len], it's already in the correct format
 
         # LOAY TODO, instead of slicing repeatedly maintain an original qkv and rewrite into that
+        valid_seq_len = None
         if attention_mask is not None:
             valid_seq_len = (attention_mask[0] == 1).sum().item()
             qkv = qkv[:, :valid_seq_len, :, :]
@@ -140,8 +141,9 @@ class DistributedAttention(nn.Module):
         # Redistribute back if using sequence parallelism
         replicated_output = None
         if replicated_q is not None:
-            replicated_output = output[:, seq_len * world_size:]
-            output = output[:, :seq_len * world_size]
+            split_idx = seq_len * world_size if valid_seq_len is None else valid_seq_len
+            replicated_output = output[:, split_idx:]
+            output = output[:, :split_idx]
             # TODO: make this asynchronous
             replicated_output = sequence_model_parallel_all_gather(
                 replicated_output.contiguous(), dim=2)

fastvideo/attention/utils/flash_attn_no_pad.py

@@ -0,0 +1,99 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE
+#
+# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any
+# output and results there from are provided "AS IS" without any express or implied warranties of
+# any kind including any warranties of title, merchantability, noninfringement, course of dealing,
+# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the
+# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of
+# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a
+# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise
+# of rights and permissions under this agreement.
+# See the License for the specific language governing permissions and limitations under the License.
+
+from einops import rearrange
+
+
+def flash_attn_no_pad(qkv,
+                      key_padding_mask,
+                      causal=False,
+                      dropout_p=0.0,
+                      softmax_scale=None,
+                      deterministic=False):
+    from flash_attn import flash_attn_varlen_qkvpacked_func
+    from flash_attn.bert_padding import pad_input, unpad_input
+    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,
+        deterministic=deterministic,
+    )
+    output = rearrange(
+        pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices,
+                  batch_size, seqlen),
+        "b s (h d) -> b s h d",
+        h=nheads,
+    )
+    return output
+
+
+def flash_attn_no_pad_v3(qkv,
+                         key_padding_mask,
+                         causal=False,
+                         dropout_p=0.0,
+                         softmax_scale=None,
+                         deterministic=False):
+    from flash_attn.bert_padding import pad_input, unpad_input
+    from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3
+
+    if flash_attn_varlen_func_v3 is None:
+        raise ImportError("FlashAttention V3 backend not available")
+
+    batch_size, seqlen, _, nheads, head_dim = qkv.shape
+    query, key, value = qkv.unbind(dim=2)
+
+    query_unpad, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
+        rearrange(query, "b s h d -> b s (h d)"), key_padding_mask)
+    key_unpad, _, cu_seqlens_k, _, _ = unpad_input(
+        rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
+    value_unpad, _, _, _, _ = unpad_input(
+        rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
+
+    query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads)
+    key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads)
+    value_unpad = rearrange(value_unpad, "nnz (h d) -> nnz h d", h=nheads)
+
+    output_unpad = flash_attn_varlen_func_v3(query_unpad,
+                                             key_unpad,
+                                             value_unpad,
+                                             cu_seqlens_q,
+                                             cu_seqlens_k,
+                                             max_seqlen_q,
+                                             max_seqlen_q,
+                                             softmax_scale=softmax_scale,
+                                             causal=causal,
+                                             deterministic=deterministic)
+
+    output = rearrange(pad_input(
+        rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size,
+        seqlen),
+                       "b s (h d) -> b s h d",
+                       h=nheads)
+    return output

fastvideo/configs/models/dits/__init__.py

@@ -1,10 +1,11 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
+from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 
 __all__ = [
-    "HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig",
-    "CosmosVideoConfig", "Cosmos25VideoConfig"
+    "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
+    "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig"
 ]

fastvideo/configs/models/dits/base.py

@@ -23,6 +23,8 @@ class DiTArchConfig(ArchConfig):
     hidden_size: int = 0
     num_attention_heads: int = 0
     num_channels_latents: int = 0
+    in_channels: int = 0
+    out_channels: int = 0
     exclude_lora_layers: list[str] = field(default_factory=list)
     boundary_ratio: float | None = None
 

fastvideo/configs/models/dits/hunyuanvideo15.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_double_block(n: str, m) -> bool:
+    return "double_blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+def is_refiner_block(n: str, m) -> bool:
+    return "refiner" in n and str.isdigit(n.split(".")[-1])
+
+
+def is_txt_in(n: str, m) -> bool:
+    return n.split(".")[-1] == "txt_in"
+
+
+@dataclass
+class HunyuanVideo15ArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_double_block, is_refiner_block])
+
+    _compile_conditions: list = field(
+        default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
+            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
+            r"txt_in.t_embedder.mlp.fc_in.\1",
+            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
+            r"txt_in.t_embedder.mlp.fc_out.\1",
+            r"^context_embedder\.proj_in\.(.*)$":
+            r"txt_in.input_embedder.\1",
+            r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$":
+            r"txt_in.c_embedder.fc_in.\1",
+            r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$":
+            r"txt_in.c_embedder.fc_out.\1",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$":
+            r"txt_in.refiner_blocks.\1.norm1.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$":
+            r"txt_in.refiner_blocks.\1.norm2.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+            (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 0, 3),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+            (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 1, 3),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+            (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 2, 3),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
+            r"txt_in.refiner_blocks.\1.self_attn_proj.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
+            r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
+            r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$":
+            r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
+
+            # 2. txt_in_2 mapping:
+            r"^context_embedder_2\.(.*)$":
+            r"txt_in_2.\1",
+
+            # 3. x_embedder mapping:
+            r"^x_embedder\.proj\.(.*)$":
+            r"img_in.proj.\1",
+
+            # 4. Top-level time_text_embed mappings:
+            r"^time_embed\.timestep_embedder\.linear_1\.(.*)$":
+            r"time_in.timestep_embedder.mlp.fc_in.\1",
+            r"^time_embed\.timestep_embedder\.linear_2\.(.*)$":
+            r"time_in.timestep_embedder.mlp.fc_out.\1",
+            r"^time_embed\.timestep_embedder_r\.linear_1\.(.*)$":
+            r"time_in.timestep_embedder_r.mlp.fc_in.\1",
+            r"^time_embed\.timestep_embedder_r\.linear_2\.(.*)$":
+            r"time_in.timestep_embedder_r.mlp.fc_out.\1",
+
+            # 5. transformer_blocks mapping:
+            r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$":
+            r"double_blocks.\1.img_mod.linear.\2",
+            r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$":
+            r"double_blocks.\1.txt_mod.linear.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
+            r"double_blocks.\1.img_attn_q_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
+            r"double_blocks.\1.img_attn_k_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
+            r"double_blocks.\1.img_attn_proj.\2",
+            # Corrected: merge attn.to_add_out into the main projection.
+            r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$":
+            r"double_blocks.\1.txt_attn_proj.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$":
+            r"double_blocks.\1.txt_attn_q_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$":
+            r"double_blocks.\1.txt_attn_k_norm.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.img_mlp.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.img_mlp.fc_out.\2",
+            r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.txt_mlp.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.txt_mlp.fc_out.\2",
+
+            # 7. Final layers mapping:
+            r"^norm_out\.linear\.(.*)$":
+            r"final_layer.adaLN_modulation.linear.\1",
+            r"^proj_out\.(.*)$":
+            r"final_layer.linear.\1",
+        })
+
+    # Reverse mapping for saving checkpoints: custom -> hf
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    in_channels: int = 65
+    out_channels: int = 32
+    num_attention_heads: int = 16
+    attention_head_dim: int = 128
+    num_layers: int = 54
+    num_refiner_layers: int = 2
+    mlp_ratio: float = 4.0
+    patch_size: int = 1
+    patch_size_t: int = 1
+    qk_norm: str = "rms_norm"
+    text_embed_dim: int = 3584
+    text_embed_2_dim: int = 1472
+    image_embed_dim: int = 1152
+    rope_theta: float = 256.0
+    rope_axes_dim: tuple[int, ...] = (16, 56, 56)
+    target_size: int = 640
+    task_type: str = "i2v"
+    use_meanflow: bool = False
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
+        self.num_channels_latents: int = self.out_channels
+
+
+@dataclass
+class HunyuanVideo15Config(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=HunyuanVideo15ArchConfig)
+
+    prefix: str = "Hunyuan15"

fastvideo/configs/models/encoders/__init__.py

@@ -6,9 +6,11 @@ from fastvideo.configs.models.encoders.clip import (
     CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig)
 from fastvideo.configs.models.encoders.llama import LlamaConfig
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
+from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
     "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig",
-    "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig"
+    "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig",
+    "Qwen2_5_VLConfig"
 ]

fastvideo/configs/models/encoders/base.py

@@ -72,6 +72,7 @@ class EncoderConfig(ModelConfig):
 @dataclass
 class TextEncoderConfig(EncoderConfig):
     arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
+    is_chat_model: bool = False
 
 
 @dataclass

fastvideo/configs/models/encoders/qwen2_5.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.base import (TextEncoderArchConfig,
+                                                    TextEncoderConfig)
+
+
+def _is_transformer_layer(n: str, m) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("embed_tokens")
+
+
+def _is_final_norm(n: str, m) -> bool:
+    return n.endswith("norm")
+
+
+@dataclass
+class Qwen2_5_VLArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 152064
+    hidden_size: int = 8192
+    intermediate_size: int = 29568
+    num_hidden_layers: int = 80
+    num_attention_heads: int = 64
+    num_key_value_heads: int = 8
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 32768
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-05
+    use_cache: bool = True
+    tie_word_embeddings: bool = False
+    rope_theta: float = 1000000.0
+    use_sliding_window: bool = False
+    sliding_window: int | None = 4096
+    max_window_layers: int = 80
+    layer_types: list = field(default_factory=list)
+    attention_dropout: float = 0.0
+    rope_scaling: dict | None = None
+    bos_token_id: int | None = None
+    eos_token_id: int | None = None
+    pad_token_id: int | None = None
+    vision_token_id: int = 151654
+    model_type: str = "qwen2_5_vl_text"
+    dtype: str = "bfloat16"
+
+    stacked_params_mapping: list[tuple[str, str, str
+                                       | int]] = field(default_factory=lambda: [
+                                           (".qkv_proj", ".q_proj", "q"),
+                                           (".qkv_proj", ".k_proj", "k"),
+                                           (".qkv_proj", ".v_proj", "v"),
+                                           (".gate_up_proj", ".gate_proj", 0),
+                                           (".gate_up_proj", ".up_proj", 1),
+                                       ])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda:
+        [_is_transformer_layer, _is_embeddings, _is_final_norm])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.sliding_window = self.sliding_window if self.use_sliding_window else None
+        # for backward compatibility
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.layer_types is None:
+            self.layer_types = [
+                "sliding_attention" if self.sliding_window is not None
+                and i >= self.max_window_layers else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            if self.rope_scaling["type"] == "mrope":
+                self.rope_scaling["type"] = "default"
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+
+        self.tokenizer_kwargs = {
+            "add_generation_prompt": True,
+            "tokenize": True,
+            "return_dict": True,
+            "padding": "max_length",
+            "max_length": 1000 + 108,
+            "truncation": True,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class Qwen2_5_VLConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(
+        default_factory=Qwen2_5_VLArchConfig)
+    prefix: str = "qwen2_5_vl"
+    is_chat_model: bool = True

fastvideo/configs/models/encoders/t5.py

@@ -40,6 +40,8 @@ class T5ArchConfig(TextEncoderArchConfig):
     eos_token_id: int = 1
     classifier_dropout: float = 0.0
     text_len: int = 512
+    dtype: str | None = None
+    gradient_checkpointing: bool = False
     stacked_params_mapping: list[tuple[str, str,
                                        str]] = field(default_factory=lambda: [
                                            # (param_name, shard_name, shard_id)
@@ -68,6 +70,7 @@ class T5ArchConfig(TextEncoderArchConfig):
             "return_attention_mask": True,
             "return_tensors": "pt",
         }
+        self.hidden_size = self.d_model
 
 
 @dataclass

fastvideo/configs/models/vaes/__init__.py

@@ -1,5 +1,6 @@
 from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
+from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
 
@@ -8,4 +9,5 @@ __all__ = [
     "WanVAEConfig",
     "StepVideoVAEConfig",
     "CosmosVAEConfig",
+    "Hunyuan15VAEConfig",
 ]

fastvideo/configs/models/vaes/hunyuan15vae.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class Hunyuan15VAEArchConfig(VAEArchConfig):
+    in_channels: int = 3
+    out_channels: int = 3
+    latent_channels: int = 32
+    block_out_channels: tuple[int, ...] = (128, 256, 512, 1024, 1024)
+    layers_per_block: int = 2
+    spatial_compression_ratio: int = 16
+    temporal_compression_ratio: int = 4
+    downsample_match_channel: bool = True
+    upsample_match_channel: bool = True
+    scaling_factor: float = 1.03682
+
+    def __post_init__(self):
+        self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) -
+                                                  1)
+
+
+@dataclass
+class Hunyuan15VAEConfig(VAEConfig):
+    arch_config: VAEArchConfig = field(default_factory=Hunyuan15VAEArchConfig)

fastvideo/configs/pipelines/__init__.py

@@ -2,6 +2,7 @@ from fastvideo.configs.pipelines.base import (PipelineConfig,
                                               SlidingTileAttnConfig)
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
+from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.registry import (
     get_pipeline_config_cls_from_name)
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
@@ -11,8 +12,8 @@ from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig,
 
 __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
-    "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
-    "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig",
-    "SelfForcingWanT2V480PConfig", "CosmosConfig",
-    "get_pipeline_config_cls_from_name"
+    "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "SlidingTileAttnConfig",
+    "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig",
+    "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig",
+    "CosmosConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/hunyuan15.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+import re
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import HunyuanVideo15Config
+from fastvideo.configs.models.encoders import (BaseEncoderOutput,
+                                               Qwen2_5_VLConfig, T5Config)
+from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+PROMPT_TEMPLATE_TOKEN_LENGTH = 108
+
+PROMPT_TEMPLATE_ENCODE_VIDEO = "You are a helpful assistant. Describe the video by detailing the following aspects: \
+        1. The main content and theme of the video. \
+        2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \
+        3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \
+        4. background environment, light, style and atmosphere. \
+        5. camera angles, movements, and transitions used in the video."
+
+
+def extract_glyph_texts(prompt: str) -> str | None:
+    """
+    Extract glyph texts from prompt using regex pattern.
+
+    Args:
+        prompt: Input prompt string
+
+    Returns:
+        List of extracted glyph texts
+    """
+    pattern = r"\"(.*?)\"|“(.*?)”"
+    matches = re.findall(pattern, prompt)
+    result = [match[0] or match[1] for match in matches]
+    result = list(dict.fromkeys(result)) if len(result) > 1 else result
+
+    if result:
+        formatted_result = ". ".join([f'Text "{text}"'
+                                      for text in result]) + ". "
+    else:
+        formatted_result = None
+
+    return formatted_result
+
+
+def format_text_input(prompt: str, system_message: str) -> list[dict[str, Any]]:
+    """
+    Apply text to template.
+
+    Args:
+        prompt (List[str]): Input text.
+        system_message (str): System message.
+
+    Returns:
+        List[Dict[str, Any]]: List of chat conversation.
+    """
+
+    template = [{
+        "role": "system",
+        "content": system_message
+    }, {
+        "role": "user",
+        "content": prompt if prompt else " "
+    }]
+
+    return template
+
+
+def qwen_preprocess_text(prompt: str) -> list[dict[str, Any]]:
+    output = format_text_input(prompt, PROMPT_TEMPLATE_ENCODE_VIDEO)
+    return output
+
+
+def qwen_postprocess_text(
+        outputs: BaseEncoderOutput,
+        mask: torch.tensor) -> tuple[torch.tensor, torch.tensor]:
+    assert outputs.hidden_states is not None
+    output = outputs.hidden_states[-3]
+    output = output[:, PROMPT_TEMPLATE_TOKEN_LENGTH:]
+    mask = mask[:, PROMPT_TEMPLATE_TOKEN_LENGTH:]
+    return output, mask
+
+
+def byt5_preprocess_text(prompt: str) -> str | None:
+    prompts = [prompt] if isinstance(prompt, str) else prompt
+    glyph_texts = [extract_glyph_texts(p) for p in prompts]
+    return glyph_texts[0]
+
+
+def byt5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
+    return outputs.last_hidden_state
+
+
+@dataclass
+class Hunyuan15T2V480PConfig(PipelineConfig):
+    """Base configuration for HunYuan pipeline architecture."""
+
+    # HunyuanConfig-specific parameters with defaults
+    # DiT
+    dit_config: DiTConfig = field(default_factory=HunyuanVideo15Config)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=Hunyuan15VAEConfig)
+    # Denoising stage
+    flow_shift: int = 5
+
+    # Text encoding stage
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Qwen2_5_VLConfig(), T5Config()))
+    preprocess_text_funcs: tuple[Callable[[Any], Any], ...] = field(
+        default_factory=lambda: (qwen_preprocess_text, byt5_preprocess_text))
+    postprocess_text_funcs: tuple[Callable[..., Any], ...] = field(
+        default_factory=lambda: (qwen_postprocess_text, byt5_postprocess_text))
+
+    # Precision for each component
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp16"
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: ("bf16", "fp32"))
+    text_encoder_crop_start: int = PROMPT_TEMPLATE_TOKEN_LENGTH
+    text_encoder_max_lengths: tuple[int, ...] = field(
+        default_factory=lambda: (1000 + PROMPT_TEMPLATE_TOKEN_LENGTH, 256))
+
+    vae_tiling: bool = True
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig):
+    """Base configuration for HunYuan pipeline architecture."""
+
+    # HunyuanConfig-specific parameters with defaults
+    flow_shift: int = 9

fastvideo/configs/pipelines/registry.py

@@ -7,6 +7,7 @@ from collections.abc import Callable
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
+from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 
 # isort: off
@@ -27,6 +28,10 @@ logger = init_logger(__name__)
 PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
     "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
+    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v":
+    Hunyuan15T2V480PConfig,
+    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v":
+    Hunyuan15T2V720PConfig,
     "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
     "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
     "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": WANV2VConfig,
@@ -56,6 +61,8 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
 PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "hunyuan":
     lambda id: "hunyuan" in id.lower(),
+    "hunyuan15":
+    lambda id: "hunyuan15" in id.lower(),
     "matrixgame":
     lambda id: "matrix-game" in id.lower() or "matrixgame" in id.lower(),
     "wanpipeline":
@@ -78,6 +85,8 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "hunyuan":
     HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
     "matrixgame": MatrixGameI2V480PConfig,
+    "hunyuan15":
+    Hunyuan15T2V480PConfig,  # Base Hunyuan15 config as fallback for any Hunyuan15 variant
     "wanpipeline":
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
@@ -123,6 +132,7 @@ def get_pipeline_config_cls_from_name(
     # First try exact match for specific weights
     if pipeline_name_or_path in PIPE_NAME_TO_CONFIG:
         pipeline_config_cls = PIPE_NAME_TO_CONFIG[pipeline_name_or_path]
+        return pipeline_config_cls
 
     # Try partial matches (for local paths that might include the weight ID)
     for registered_id, config_class in PIPE_NAME_TO_CONFIG.items():

fastvideo/configs/sample/base.py

@@ -50,6 +50,7 @@ class SamplingParam:
     guidance_scale: float = 1.0
     guidance_rescale: float = 0.0
     boundary_ratio: float | None = None
+    sigmas: list[float] | None = None
 
     # TeaCache parameters
     enable_teacache: bool = False

fastvideo/configs/sample/hunyuan15.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+import numpy as np
+from dataclasses import dataclass, field
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+@dataclass
+class Hunyuan15_480P_SamplingParam(SamplingParam):
+    num_inference_steps: int = 50
+
+    num_frames: int = 121
+    height: int = 480
+    width: int = 848
+    fps: int = 24
+
+    guidance_scale: float = 6.0
+    prompt_attention_mask: list = field(default_factory=list)
+    negative_attention_mask: list = field(default_factory=list)
+    sigmas: list[float] | None = field(
+        default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))
+
+    negative_prompt: str = ""
+
+
+@dataclass
+class Hunyuan15_720P_SamplingParam(Hunyuan15_480P_SamplingParam):
+    height: int = 720
+    width: int = 1280

fastvideo/configs/sample/registry.py

@@ -5,6 +5,7 @@ from typing import Any
 
 from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                               HunyuanSamplingParam)
+from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hunyuan15_720P_SamplingParam
 from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
 
 from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam
@@ -35,6 +36,10 @@ logger = init_logger(__name__)
 SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
     "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
+    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v":
+    Hunyuan15_480P_SamplingParam,
+    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v":
+    Hunyuan15_720P_SamplingParam,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
 
     # Wan2.1
@@ -86,6 +91,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
 SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     "hunyuan":
     lambda id: "hunyuan" in id.lower(),
+    "hunyuan15":
+    lambda id: "hunyuan15" in id.lower(),
     "wanpipeline":
     lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo":
@@ -105,6 +112,8 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
 SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     "hunyuan":
     HunyuanSamplingParam,  # Base Hunyuan config as fallback for any Hunyuan variant
+    "hunyuan15":
+    Hunyuan15_480P_SamplingParam,  # Base Hunyuan15 config as fallback for any Hunyuan15 variant
     "wanpipeline":
     WanT2V_1_3B_SamplingParam,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V_14B_480P_SamplingParam,

fastvideo/dataset/dataloader/record_schema.py

@@ -127,7 +127,10 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
 def ode_text_only_record_creator(
         video_name: str, text_embedding: np.ndarray, caption: str,
         trajectory_latents: np.ndarray,
-        trajectory_timesteps: np.ndarray) -> dict[str, Any]:
+        trajectory_timesteps: np.ndarray,
+        text_mask: np.ndarray | None = None,
+        text_embedding_2: np.ndarray | None = None,
+        text_mask_2: np.ndarray | None = None) -> dict[str, Any]:
     """Create a text-only ODE trajectory record matching pyarrow_schema_ode_trajectory_text_only.
 
     Args:
@@ -165,6 +168,25 @@ def ode_text_only_record_creator(
         "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype),
     })
 
+    if text_embedding_2 is not None:
+        record.update({
+            "text_embedding_2_bytes": text_embedding_2.tobytes(),
+            "text_embedding_2_shape": list(text_embedding_2.shape),
+            "text_embedding_2_dtype": str(text_embedding_2.dtype),
+        })
+    if text_mask is not None:
+        record.update({
+            "text_mask_bytes": text_mask.tobytes(),
+            "text_mask_shape": list(text_mask.shape),
+            "text_mask_dtype": str(text_mask.dtype),
+        })
+    if text_mask_2 is not None:
+        record.update({
+            "text_mask_2_bytes": text_mask_2.tobytes(),
+            "text_mask_2_shape": list(text_mask_2.shape),
+            "text_mask_2_dtype": str(text_mask_2.dtype),
+        })
+        
     return record
 
 

fastvideo/dataset/dataloader/schema.py

@@ -90,6 +90,15 @@ pyarrow_schema_ode_trajectory_text_only = pa.schema([
     pa.field("text_embedding_shape", pa.list_(pa.int64())),
     # e.g., 'bfloat16' or 'float32'
     pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_embedding_2_bytes", pa.binary()),
+    pa.field("text_embedding_2_shape", pa.list_(pa.int64())),
+    pa.field("text_embedding_2_dtype", pa.string()),
+    pa.field("text_mask_bytes", pa.binary()),
+    pa.field("text_mask_shape", pa.list_(pa.int64())),
+    pa.field("text_mask_dtype", pa.string()),
+    pa.field("text_mask_2_bytes", pa.binary()),
+    pa.field("text_mask_2_shape", pa.list_(pa.int64())),
+    pa.field("text_mask_2_dtype", pa.string()),
     # --- ODE Trajectory ---
     pa.field("trajectory_latents_bytes", pa.binary()),
     pa.field("trajectory_latents_shape", pa.list_(pa.int64())),

fastvideo/dataset/preprocessing_datasets.py

@@ -17,6 +17,7 @@ from PIL import Image
 from transformers import AutoTokenizer
 
 from fastvideo.logger import init_logger
+from fastvideo.utils import maybe_download_model
 
 logger = init_logger(__name__)
 
@@ -655,7 +656,7 @@ class TextDataset(torch.utils.data.IterableDataset,
         self.seed = seed
 
         # Initialize tokenizer
-        tokenizer_path = os.path.join(args.model_path, "tokenizer")
+        tokenizer_path = os.path.join(maybe_download_model(args.model_path), "tokenizer")
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
                                                   cache_dir=args.cache_dir)
 

fastvideo/layers/activation.py

@@ -87,6 +87,7 @@ _ACTIVATION_REGISTRY = {
     "gelu_pytorch_tanh": lambda: nn.GELU(approximate="tanh"),
     "relu": nn.ReLU,
     "silu": nn.SiLU,
+    "swish": nn.SiLU,
     "quick_gelu": QuickGELU,
 }
 

fastvideo/models/dits/hunyuanvideo15.py

@@ -0,0 +1,853 @@
+# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Optional, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.attention import DistributedAttention, LocalAttention
+from fastvideo.distributed.communication_op import (
+    sequence_model_parallel_all_gather_with_unpad,
+    sequence_model_parallel_shard)
+from fastvideo.configs.models.dits import HunyuanVideo15Config
+from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
+                                        ScaleResidualLayerNormScaleShift)
+from fastvideo.layers.linear import ReplicatedLinear
+# TODO(will-PY-refactor): RMSNorm ....
+from fastvideo.layers.mlp import MLP
+from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
+from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed,
+                                               TimestepEmbedder, unpatchify)
+from fastvideo.models.dits.base import CachableDiT
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.logger import init_logger
+from fastvideo.forward_context import set_forward_context
+from fastvideo.attention.backends.abstract import AttentionMetadata
+
+from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.distributed.utils import create_attention_mask_for_padding
+
+logger = init_logger(__name__)
+
+class HunyuanRMSNorm(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        elementwise_affine=True,
+        eps: float = 1e-6,
+        device=None,
+        dtype=None,
+    ):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.eps = eps
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
+
+    def _norm(self, x) -> torch.Tensor:
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        if hasattr(self, "weight"):
+            output = output * self.weight
+        return output
+
+
+class HunyuanVideo15TimeEmbedding(nn.Module):
+    r"""
+    Time embedding for HunyuanVideo 1.5.
+
+    Supports standard timestep embedding and optional reference timestep embedding for MeanFlow-based super-resolution
+    models.
+
+    Args:
+        embedding_dim (`int`):
+            The dimension of the output embedding.
+    """
+
+    def __init__(self, embedding_dim: int, use_meanflow: bool = False):
+        super().__init__()
+
+        self.timestep_embedder = TimestepEmbedder(hidden_size=embedding_dim)
+
+        self.use_meanflow = use_meanflow
+        self.time_proj_r = None
+        self.timestep_embedder_r = None
+        if use_meanflow:
+            self.timestep_embedder_r = TimestepEmbedder(hidden_size=embedding_dim)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        timestep_r: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        timesteps_emb = self.timestep_embedder(timestep)
+
+        if timestep_r is not None:
+            timesteps_emb_r = self.timestep_embedder_r(timestep_r)
+            timesteps_emb = timesteps_emb + timesteps_emb_r
+
+        return timesteps_emb
+
+
+class HunyuanVideo15ByT5TextProjection(nn.Module):
+    def __init__(self, in_features: int, hidden_size: int, out_features: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(in_features)
+        self.linear_1 = nn.Linear(in_features, hidden_size)
+        self.linear_2 = nn.Linear(hidden_size, hidden_size)
+        self.linear_3 = nn.Linear(hidden_size, out_features)
+        self.act_fn = nn.GELU()
+
+    def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(encoder_hidden_states)
+        hidden_states = self.linear_1(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+        hidden_states = self.linear_3(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideo15ImageProjection(nn.Module):
+    def __init__(self, in_channels: int, hidden_size: int):
+        super().__init__()
+        self.norm_in = nn.LayerNorm(in_channels)
+        self.linear_1 = nn.Linear(in_channels, in_channels)
+        self.act_fn = nn.GELU()
+        self.linear_2 = nn.Linear(in_channels, hidden_size)
+        self.norm_out = nn.LayerNorm(hidden_size)
+
+    def forward(self, image_embeds: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm_in(image_embeds)
+        hidden_states = self.linear_1(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        hidden_states = self.norm_out(hidden_states)
+        return hidden_states
+
+
+class MMDoubleStreamBlock(nn.Module):
+    """
+    A multimodal DiT block with separate modulation for text and image/video,
+    using distributed attention and linear layers.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        mlp_ratio: float,
+        dtype: torch.dtype | None = None,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...]
+        | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.deterministic = False
+        self.num_attention_heads = num_attention_heads
+        head_dim = hidden_size // num_attention_heads
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        # Image modulation components
+        self.img_mod = ModulateProjection(
+            hidden_size,
+            factor=6,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.img_mod",
+        )
+
+        # Fused operations for image stream
+        self.img_attn_norm = LayerNormScaleShift(hidden_size,
+                                                 norm_type="layer",
+                                                 elementwise_affine=False,
+                                                 dtype=dtype)
+        self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            elementwise_affine=False,
+            dtype=dtype)
+        self.img_mlp_residual = ScaleResidual()
+
+        # Image attention components
+        self.img_attn_qkv = ReplicatedLinear(hidden_size,
+                                             hidden_size * 3,
+                                             bias=True,
+                                             params_dtype=dtype,
+                                             prefix=f"{prefix}.img_attn_qkv")
+
+        self.img_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+        self.img_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+
+        self.img_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype,
+                                              prefix=f"{prefix}.img_attn_proj")
+
+        self.img_mlp = MLP(hidden_size,
+                           mlp_hidden_dim,
+                           bias=True,
+                           dtype=dtype,
+                           prefix=f"{prefix}.img_mlp")
+
+        # Text modulation components
+        self.txt_mod = ModulateProjection(
+            hidden_size,
+            factor=6,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.txt_mod",
+        )
+
+        # Fused operations for text stream
+        self.txt_attn_norm = LayerNormScaleShift(hidden_size,
+                                                 norm_type="layer",
+                                                 elementwise_affine=False,
+                                                 dtype=dtype)
+        self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            elementwise_affine=False,
+            dtype=dtype)
+        self.txt_mlp_residual = ScaleResidual()
+
+        # Text attention components
+        self.txt_attn_qkv = ReplicatedLinear(hidden_size,
+                                             hidden_size * 3,
+                                             bias=True,
+                                             params_dtype=dtype)
+
+        # QK norm layers for text
+        self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+        self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+
+        self.txt_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype)
+
+        self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)
+
+        # Distributed attention
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn")
+
+    def forward(
+        self,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        vec: torch.Tensor,
+        freqs_cis: tuple,
+        seq_attention_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Process modulation vectors
+        img_mod_outputs = self.img_mod(vec)
+        (
+            img_attn_shift,
+            img_attn_scale,
+            img_attn_gate,
+            img_mlp_shift,
+            img_mlp_scale,
+            img_mlp_gate,
+        ) = torch.chunk(img_mod_outputs, 6, dim=-1)
+
+        txt_mod_outputs = self.txt_mod(vec)
+        (
+            txt_attn_shift,
+            txt_attn_scale,
+            txt_attn_gate,
+            txt_mlp_shift,
+            txt_mlp_scale,
+            txt_mlp_gate,
+        ) = torch.chunk(txt_mod_outputs, 6, dim=-1)
+
+        # Prepare image for attention using fused operation
+        img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale)
+        # Get QKV for image
+        img_qkv, _ = self.img_attn_qkv(img_attn_input)
+        batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1]
+
+        # Split QKV
+        img_qkv = img_qkv.view(batch_size, image_seq_len, 3,
+                               self.num_attention_heads, -1)
+        img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :,
+                                                                          2]
+
+        # Apply QK-Norm if needed
+        img_q = self.img_attn_q_norm(img_q).to(img_v)
+        img_k = self.img_attn_k_norm(img_k).to(img_v)
+
+        # Prepare text for attention using fused operation
+        txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale)
+
+        # Get QKV for text
+        txt_qkv, _ = self.txt_attn_qkv(txt_attn_input)
+        batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1]
+
+        # Split QKV
+        txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3,
+                               self.num_attention_heads, -1)
+        txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :,
+                                                                          2]
+        # Apply QK-Norm if needed
+        txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
+        txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
+
+        # seq_len = txt_q.shape[1] + img_q.shape[1]
+        # attention_mask = F.pad(encoder_attention_mask, (seq_len - encoder_attention_mask.shape[1], 0), value=True)
+        # attention_mask = attention_mask.bool()
+        # self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
+        # self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
+        # attention_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
+
+        from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+        attn_metadata = FlashAttnMetadataBuilder().build(
+            current_timestep=0,
+            attn_mask=encoder_attention_mask,
+        )
+        # Run distributed attention
+        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
+            img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v, freqs_cis=freqs_cis, attention_mask=seq_attention_mask)
+        
+        img_attn_out, _ = self.img_attn_proj(
+            img_attn.view(batch_size, image_seq_len, -1))
+        # Use fused operation for residual connection, normalization, and modulation
+        img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(
+            img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale)
+
+        # Process image MLP
+        img_mlp_out = self.img_mlp(img_mlp_input)
+        img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)
+
+        # Process text attention output
+        txt_attn_out, _ = self.txt_attn_proj(
+            txt_attn.reshape(batch_size, text_seq_len, -1))
+
+        # Use fused operation for residual connection, normalization, and modulation
+        txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(
+            txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale)
+
+        # Process text MLP
+        txt_mlp_out = self.txt_mlp(txt_mlp_input)
+        txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)
+
+        return img, txt
+
+
+class HunyuanVideo15Transformer3DModel(CachableDiT):
+    r"""
+    A Transformer model for video-like data used in [HunyuanVideo1.5](https://huggingface.co/tencent/HunyuanVideo1.5).
+    """
+
+    # shard single stream, double stream blocks, and refiner_blocks
+    _fsdp_shard_conditions = HunyuanVideo15Config()._fsdp_shard_conditions
+    _compile_conditions = HunyuanVideo15Config()._compile_conditions
+    _supported_attention_backends = HunyuanVideo15Config(
+    )._supported_attention_backends
+    param_names_mapping = HunyuanVideo15Config().param_names_mapping
+    reverse_param_names_mapping = HunyuanVideo15Config(
+    ).reverse_param_names_mapping
+    lora_param_names_mapping = HunyuanVideo15Config().lora_param_names_mapping
+
+    def __init__(
+        self,
+        config: HunyuanVideo15Config,
+        hf_config: dict[str, Any],
+    ) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.num_channels_latents
+        self.out_channels = config.out_channels or config.in_channels
+        self.patch_size = (config.patch_size_t, config.patch_size, config.patch_size)
+
+        # 1. Latent and condition embedders
+        self.img_in = PatchEmbed(self.patch_size, 
+                                    config.in_channels, 
+                                    self.hidden_size,
+                                    prefix=f"{config.prefix}.img_in")
+        self.image_embedder = HunyuanVideo15ImageProjection(config.image_embed_dim, self.hidden_size)
+
+        self.txt_in = SingleTokenRefiner(config.text_embed_dim,
+                                                   self.hidden_size,
+                                                   config.num_attention_heads,
+                                                   depth=config.num_refiner_layers,
+                                                   dtype=None,
+                                                   prefix=f"{config.prefix}.txt_in")
+
+        self.txt_in_2 = HunyuanVideo15ByT5TextProjection(config.text_embed_2_dim, 2048, self.hidden_size)
+
+        self.time_in = HunyuanVideo15TimeEmbedding(self.hidden_size, use_meanflow=config.use_meanflow)
+
+        self.cond_type_embed = nn.Embedding(3, self.hidden_size)
+
+        # 3. Dual stream transformer blocks
+
+        self.double_blocks = nn.ModuleList(
+            [
+                MMDoubleStreamBlock(
+                    hidden_size=self.hidden_size,
+                    num_attention_heads=config.num_attention_heads,
+                    mlp_ratio=config.mlp_ratio,
+                    dtype=None,
+                    supported_attention_backends=self._supported_attention_backends,
+                    prefix=f"{config.prefix}.double_blocks.{i}"
+                )
+                for i in range(config.num_layers)
+            ]
+        )
+
+        # 5. Output projection
+        self.final_layer = FinalLayer(self.hidden_size,
+                                self.patch_size,
+                                self.out_channels,
+                                prefix=f"{config.prefix}.final_layer")
+
+        self.gradient_checkpointing = False
+
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: List[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: List[torch.Tensor],
+        encoder_attention_mask: List[torch.Tensor],
+        guidance: Optional[torch.Tensor] = None,
+        timestep_r: Optional[torch.LongTensor] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        encoder_hidden_states_image = encoder_hidden_states_image[0]
+        encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
+        encoder_attention_mask, encoder_attention_mask_2 = encoder_attention_mask
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # 1. RoPE
+        # Get rotary embeddings
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (post_patch_num_frames, post_patch_height, post_patch_width), self.hidden_size,
+            self.num_attention_heads, self.config.rope_axes_dim, self.config.rope_theta)
+        freqs_cos = freqs_cos.to(hidden_states.device)
+        freqs_sin = freqs_sin.to(hidden_states.device)
+        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
+
+        # 2. Conditional embeddings
+        temb = self.time_in(timestep, timestep_r=timestep_r)
+
+        hidden_states = self.img_in(hidden_states)
+        hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
+
+        current_seq_len = hidden_states.shape[1]
+        sp_world_size = get_sp_world_size()
+        padded_seq_len = current_seq_len * sp_world_size
+        
+        if padded_seq_len > original_seq_len:
+            seq_attention_mask = create_attention_mask_for_padding(
+                seq_len=original_seq_len,
+                padded_seq_len=padded_seq_len,
+                batch_size=batch_size,
+                device=hidden_states.device,
+            )
+        else:
+            seq_attention_mask = None
+
+        # qwen text embedding
+        encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep, encoder_attention_mask)
+
+        encoder_hidden_states_cond_emb = self.cond_type_embed(
+            torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb
+
+        # byt5 text embedding
+        encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2)
+
+        encoder_hidden_states_2_cond_emb = self.cond_type_embed(
+            torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb
+
+        # image embed
+        encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image)
+        is_t2v = torch.all(encoder_hidden_states_image == 0)
+        if is_t2v:
+            encoder_hidden_states_3 = encoder_hidden_states_3 * 0.0
+            encoder_attention_mask_3 = torch.zeros(
+                (batch_size, encoder_hidden_states_3.shape[1]),
+                dtype=encoder_attention_mask.dtype,
+                device=encoder_attention_mask.device,
+            )
+        else:
+            encoder_attention_mask_3 = torch.ones(
+                (batch_size, encoder_hidden_states_3.shape[1]),
+                dtype=encoder_attention_mask.dtype,
+                device=encoder_attention_mask.device,
+            )
+        encoder_hidden_states_3_cond_emb = self.cond_type_embed(
+            2
+            * torch.ones_like(
+                encoder_hidden_states_3[:, :, 0],
+                dtype=torch.long,
+            )
+        )
+        encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb
+
+        # reorder and combine text tokens: combine valid tokens first, then padding
+        encoder_attention_mask = encoder_attention_mask.bool()
+        encoder_attention_mask_2 = encoder_attention_mask_2.bool()
+        encoder_attention_mask_3 = encoder_attention_mask_3.bool()
+        new_encoder_hidden_states = []
+        new_encoder_attention_mask = []
+
+        for text, text_mask, text_2, text_mask_2, image, image_mask in zip(
+            encoder_hidden_states,
+            encoder_attention_mask,
+            encoder_hidden_states_2,
+            encoder_attention_mask_2,
+            encoder_hidden_states_3,
+            encoder_attention_mask_3,
+        ):
+            # Concatenate: [valid_image, valid_byt5, valid_mllm, invalid_image, invalid_byt5, invalid_mllm]
+            new_encoder_hidden_states.append(
+                torch.cat(
+                    [
+                        image[image_mask],  # valid image
+                        text_2[text_mask_2],  # valid byt5
+                        text[text_mask],  # valid mllm
+                        image[~image_mask],  # invalid image (zeroed)
+                        torch.zeros_like(text_2[~text_mask_2]),  # invalid byt5 (zeroed)
+                        torch.zeros_like(text[~text_mask]),  # invalid mllm (zeroed)
+                    ],
+                    dim=0,
+                )
+            )
+            # Apply same reordering to attention masks
+            new_encoder_attention_mask.append(
+                torch.cat(
+                    [
+                        image_mask[image_mask],
+                        text_mask_2[text_mask_2],
+                        text_mask[text_mask],
+                        image_mask[~image_mask],
+                        text_mask_2[~text_mask_2],
+                        text_mask[~text_mask],
+                    ],
+                    dim=0,
+                )
+            )
+
+        encoder_hidden_states = torch.stack(new_encoder_hidden_states)
+        encoder_attention_mask = torch.stack(new_encoder_attention_mask)
+
+        # 4. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for block in self.double_blocks:
+                hidden_states, encoder_hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    temb,
+                    freqs_cis,
+                    seq_attention_mask
+                )
+
+        else:
+            for block in self.double_blocks:
+                hidden_states, encoder_hidden_states = block(
+                    hidden_states,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    temb,
+                    freqs_cis,
+                    seq_attention_mask
+                )
+
+        # Final layer processing
+        hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
+        hidden_states = self.final_layer(hidden_states, temb)
+        # Unpatchify to get original shape
+        hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels)
+
+        return hidden_states
+
+class SingleTokenRefiner(nn.Module):
+    """
+    A token refiner that processes text embeddings with attention to improve
+    their representation for cross-attention with image features.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_size,
+        num_attention_heads,
+        depth=2,
+        qkv_bias=True,
+        dtype=None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        # Input projection
+        # self.input_embedder = ReplicatedLinear(
+        #     in_channels,
+        #     hidden_size,
+        #     bias=True,
+        #     params_dtype=dtype,
+        #     prefix=f"{prefix}.input_embedder")
+        self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True)
+
+        # Timestep embedding
+        self.t_embedder = TimestepEmbedder(hidden_size,
+                                           act_layer="silu",
+                                           dtype=dtype,
+                                           prefix=f"{prefix}.t_embedder")
+
+        # Context embedding
+        self.c_embedder = MLP(in_channels,
+                              hidden_size,
+                              hidden_size,
+                              act_type="silu",
+                              dtype=dtype,
+                              prefix=f"{prefix}.c_embedder")
+
+        # Refiner blocks
+        self.refiner_blocks = nn.ModuleList([
+            IndividualTokenRefinerBlock(
+                hidden_size,
+                num_attention_heads,
+                qkv_bias=qkv_bias,
+                dtype=dtype,
+                prefix=f"{prefix}.refiner_blocks.{i}",
+            ) for i in range(depth)
+        ])
+
+    def forward(self, x, t, mask=None):
+        # Get timestep embeddings
+        timestep_aware_representations = self.t_embedder(t)
+
+        # Get context-aware representations
+        original_dtype = x.dtype
+        if mask is None:
+            context_aware_representations = x.mean(dim=1)
+        else:
+            mask_float = mask.float().unsqueeze(-1)  # [B, L, 1]
+            context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+
+        context_aware_representations = self.c_embedder(
+            context_aware_representations)
+        c = timestep_aware_representations + context_aware_representations
+        # Project input
+        x = self.input_embedder(x)
+        # Process through refiner blocks
+        for block in self.refiner_blocks:
+            x = block(x, c, mask)
+        return x
+
+class IndividualTokenRefinerBlock(nn.Module):
+    """
+    A transformer block for refining individual tokens with self-attention.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        num_attention_heads,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        dtype=None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        # Normalization and attention
+        self.norm1 = nn.LayerNorm(hidden_size,
+                                  eps=1e-6,
+                                  elementwise_affine=True,
+                                  dtype=dtype)
+
+        self.self_attn_qkv = ReplicatedLinear(hidden_size,
+                                              hidden_size * 3,
+                                              bias=qkv_bias,
+                                              params_dtype=dtype,
+                                              prefix=f"{prefix}.self_attn_qkv")
+
+        self.self_attn_proj = ReplicatedLinear(
+            hidden_size,
+            hidden_size,
+            bias=qkv_bias,
+            params_dtype=dtype,
+            prefix=f"{prefix}.self_attn_proj")
+
+        # MLP
+        self.norm2 = nn.LayerNorm(hidden_size,
+                                  eps=1e-6,
+                                  elementwise_affine=True,
+                                  dtype=dtype)
+        self.mlp = MLP(hidden_size,
+                       mlp_hidden_dim,
+                       bias=True,
+                       act_type="silu",
+                       dtype=dtype,
+                       prefix=f"{prefix}.mlp")
+
+        # Modulation
+        self.adaLN_modulation = ModulateProjection(
+            hidden_size,
+            factor=2,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.adaLN_modulation")
+
+        # Scaled dot product attention
+        self.attn = LocalAttention(
+            num_heads=num_attention_heads,
+            head_size=hidden_size // num_attention_heads,
+            # TODO: remove hardcode; remove STA
+            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
+                                          AttentionBackendEnum.TORCH_SDPA),
+        )
+
+    def forward(self, x, c, mask=None):
+        if mask is not None:
+            mask = mask.clone().bool()
+            mask[:, 0] = True  # Prevent attention weights from becoming NaN
+
+        # Get modulation parameters
+        gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1)
+        # Self-attention
+        norm_x = self.norm1(x)
+        qkv, _ = self.self_attn_qkv(norm_x)
+
+        batch_size, seq_len = qkv.shape[0], qkv.shape[1]
+        qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+
+        # Run scaled dot product attention
+        from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+        attn_metadata = FlashAttnMetadataBuilder().build(
+            current_timestep=0,
+            attn_mask=mask,
+        )
+        # Run distributed attention
+        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
+            attn_output = self.attn(q, k, v)  # [B, L, H, D]
+        attn_output = attn_output.reshape(batch_size, seq_len,
+                                          -1)  # [B, L, H*D]
+
+        # Project and apply residual connection with gating
+        attn_out, _ = self.self_attn_proj(attn_output)
+        x = x + attn_out * gate_msa.unsqueeze(1)
+
+        # MLP
+        mlp_out = self.mlp(self.norm2(x))
+        x = x + mlp_out * gate_mlp.unsqueeze(1)
+
+        return x
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT that projects features to pixel space.
+    """
+
+    def __init__(self,
+                 hidden_size,
+                 patch_size,
+                 out_channels,
+                 dtype=None,
+                 prefix: str = "") -> None:
+        super().__init__()
+
+        # Normalization
+        self.norm_final = nn.LayerNorm(hidden_size,
+                                       eps=1e-6,
+                                       elementwise_affine=False,
+                                       dtype=dtype)
+
+        output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels
+
+        self.linear = ReplicatedLinear(hidden_size,
+                                       output_dim,
+                                       bias=True,
+                                       params_dtype=dtype,
+                                       prefix=f"{prefix}.linear")
+
+        # Modulation
+        self.adaLN_modulation = ModulateProjection(
+            hidden_size,
+            factor=2,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.adaLN_modulation")
+
+    def forward(self, x, c):
+        # What the heck HF? Why you change the scale and shift order here???
+        scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        x, _ = self.linear(x)
+        return x

fastvideo/models/encoders/qwen2_5.py

@@ -0,0 +1,387 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from transformers: https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py
+
+import math
+from typing import Any, Optional, Tuple, Union, List, Callable, Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.configs.models.encoders import BaseEncoderOutput, Qwen2_5_VLConfig
+from fastvideo.distributed import get_tp_rank, get_tp_world_size
+from fastvideo.layers.activation import get_act_fn, SiluAndMul
+from fastvideo.layers.layernorm import RMSNorm
+from fastvideo.layers.linear import MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from fastvideo.models.encoders.base import TextEncoder
+from fastvideo.models.loader.weight_utils import default_weight_loader
+from fastvideo.models.mask_utils import sdpa_mask
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+def sdpa_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    dropout: float = 0.0,
+    scaling: Optional[float] = None,
+    is_causal: Optional[bool] = None,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    if kwargs.get("output_attentions", False) or kwargs.get("head_mask") is not None:
+        logger.warning_once(
+            "`sdpa` attention does not support `output_attentions=True` or `head_mask`."
+            " Please set your attention to `eager` if you want any of these features."
+        )
+
+    if hasattr(module, "num_key_value_groups"):
+        key = repeat_kv(key, module.num_key_value_groups)
+        value = repeat_kv(value, module.num_key_value_groups)
+
+    if attention_mask is not None and attention_mask.ndim == 4:
+        attention_mask = attention_mask[:, :, :, : key.shape[-2]]
+    # If attention_mask is not None, convert it to boolean type
+    if attention_mask is not None and attention_mask.dtype != torch.bool:
+        attention_mask = attention_mask.bool()
+
+    attn_output = torch.nn.functional.scaled_dot_product_attention(
+        query,
+        key,
+        value,
+        attn_mask=attention_mask,
+        dropout_p=dropout,
+        scale=scaling,
+        is_causal=is_causal,
+    )
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, None
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):
+    mrope_section = [s * 2 for s in mrope_section]
+    cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
+        unsqueeze_dim
+    )
+    sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
+        unsqueeze_dim
+    )
+
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class Qwen2_5_VLRotaryEmbedding(nn.Module):
+    def __init__(self, config: Qwen2_5_VLConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.config = config
+        
+        self.rope_type = config.rope_scaling.get("rope_type", "default")
+        self.base = config.rope_theta
+        
+        # Simplified initialization
+        head_dim = config.hidden_size // config.num_attention_heads
+        dim = head_dim
+        self.attention_scaling = 1.0
+
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim)
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = inv_freq
+
+    def forward(self, x, position_ids):
+        # In contrast to other models, Qwen2_5_VL has different position ids for the grids
+        # So we expand the inv_freq to shape (3, ...)
+        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class Qwen2_5_VLMLP(nn.Module):
+    def __init__(self, config: Qwen2_5_VLConfig, quant_config: QuantizationConfig | None = None, prefix: str = ""):
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=config.hidden_size,
+            output_sizes=[config.intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=config.intermediate_size,
+            output_size=config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Qwen2_5_VLAttention(nn.Module):
+    def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, quant_config: QuantizationConfig | None = None, prefix: str = ""):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        
+        tp_size = get_tp_world_size()
+        self.total_num_heads = self.num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        
+        self.total_num_kv_heads = self.num_key_value_heads
+        if self.total_num_kv_heads >= tp_size:
+             assert self.total_num_kv_heads % tp_size == 0
+             self.num_kv_heads = self.total_num_kv_heads // tp_size
+        else:
+             assert tp_size % self.total_num_kv_heads == 0
+             self.num_kv_heads = 1 
+
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+        
+        self.layer_type = config.layer_types[layer_idx] if config.layer_types else "full_attention"
+        self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        qkv, _ = self.qkv_proj(hidden_states)
+        query_states, key_states, value_states = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_multimodal_rotary_pos_emb(
+            query_states, key_states, cos, sin, self.config.rope_scaling["mrope_section"]
+        )
+
+        attn_output = sdpa_attention_forward(self, query_states, key_states, value_states, attention_mask, dropout=self.config.attention_dropout, scaling=self.scaling, is_causal=False)[0].reshape(bsz, q_len, -1)
+        attn_output, _ = self.o_proj(attn_output)
+        
+        return attn_output
+
+class Qwen2_5_VLDecoderLayer(nn.Module):
+    def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, quant_config: QuantizationConfig | None = None, prefix: str = ""):
+        super().__init__()
+        self.self_attn = Qwen2_5_VLAttention(config, layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn")
+        self.mlp = Qwen2_5_VLMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp")
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            position_embeddings=position_embeddings,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+        
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        
+        outputs = (hidden_states,)
+            
+        return outputs
+
+class Qwen2_5_VLTextModel(TextEncoder):
+    def __init__(self, config: Qwen2_5_VLConfig):
+        super().__init__(config)
+
+        quant_config = None
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            org_num_embeddings=config.vocab_size
+        )
+        
+        self.layers = nn.ModuleList([
+            Qwen2_5_VLDecoderLayer(config, layer_idx, quant_config=quant_config, prefix=f"{config.prefix}.layers.{layer_idx}")
+            for layer_idx in range(config.num_hidden_layers)
+        ])
+        
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings(input_ids)
+            
+        hidden_states = inputs_embeds
+        
+        if position_ids is None:
+            seq_length = hidden_states.shape[1]
+            cache_position = torch.arange(seq_length, device=hidden_states.device)
+            position_ids = cache_position.view(1, 1, -1).expand(3, hidden_states.shape[0], -1)
+
+        mask_kwargs = {
+            "batch_size": hidden_states.shape[0],
+            "cache_position": cache_position,
+            "kv_length": attention_mask.shape[-1],
+            "kv_offset": 0,
+            "attention_mask": attention_mask
+        }
+
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        
+        all_hidden_states = () if output_hidden_states else None
+        
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+                
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=sdpa_mask(**mask_kwargs),
+                position_ids=position_ids,
+                position_embeddings=position_embeddings,
+                output_attentions=False,
+            )
+            hidden_states = layer_outputs[0]
+            
+        hidden_states = self.norm(hidden_states)
+        
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+            
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            for param_name, weight_name, shard_id in self.config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                # if name.endswith(".bias") and name not in params_dict:
+                #     continue
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                # if name.endswith(".bias") and name not in params_dict:
+                #     continue
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

fastvideo/models/loader/component_loader.py

@@ -412,11 +412,9 @@ class VAELoader(ComponentLoader):
         # Find all safetensors files
         safetensors_list = glob.glob(
             os.path.join(str(model_path), "*.safetensors"))
-        # TODO(PY)
-        assert len(
-            safetensors_list
-        ) == 1, f"Found {len(safetensors_list)} safetensors files in {model_path}"
-        loaded = safetensors_load_file(safetensors_list[0])
+        loaded = {}
+        for sf_file in safetensors_list:
+            loaded.update(safetensors_load_file(sf_file))
         vae.load_state_dict(
             loaded, strict=False)  # We might only load encoder or decoder
 
@@ -478,8 +476,6 @@ class TransformerLoader(ComponentLoader):
             fastvideo_args.pipeline_config.dit_precision]
 
         # Load the model using FSDP loader
-        logger.info("Loading model from %s, default_dtype: %s", cls_name,
-                    default_dtype)
         assert fastvideo_args.hsdp_shard_dim is not None
         model = maybe_load_fsdp_model(
             model_cls=model_cls,

fastvideo/models/mask_utils.py

@@ -0,0 +1,201 @@
+import torch
+from typing import Callable, Optional
+
+def and_masks(*mask_functions: Callable) -> Callable:
+    """Returns a mask function that is the intersection of provided mask functions"""
+    if not all(callable(arg) for arg in mask_functions):
+        raise RuntimeError(f"All inputs should be callable mask_functions: {mask_functions}")
+
+    def and_mask(batch_idx, head_idx, q_idx, kv_idx):
+        result = q_idx.new_ones((), dtype=torch.bool)
+        for mask in mask_functions:
+            result = result & mask(batch_idx, head_idx, q_idx, kv_idx).to(result.device)
+        return result
+
+    return and_mask
+
+def causal_mask_function(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
+    """
+    This creates a basic lower-diagonal causal mask.
+    """
+    return kv_idx <= q_idx
+
+def padding_mask_function(padding_mask: torch.Tensor) -> Callable:
+    """
+    This return the mask_function function corresponding to a 2D padding mask.
+    """
+
+    def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
+        # Note that here the mask should ALWAYS be at least of the max `kv_index` size in the dimension 1. This is because
+        # we cannot pad it here in the mask_function as we don't know the final size, and we cannot try/except, as it is not
+        # vectorizable on accelerator devices
+        return padding_mask[batch_idx, kv_idx]
+
+    return inner_mask
+
+def prepare_padding_mask(
+    attention_mask: Optional[torch.Tensor], kv_length: int, kv_offset: int
+) -> Optional[torch.Tensor]:
+    """
+    From the 2D attention mask, prepare the correct padding mask to use by potentially padding it.
+    """
+    local_padding_mask = attention_mask
+    if attention_mask is not None:
+        # Pad it if necessary
+        if (padding_length := kv_length + kv_offset - attention_mask.shape[-1]) > 0:
+            local_padding_mask = torch.nn.functional.pad(attention_mask, (0, padding_length))
+    return local_padding_mask
+
+def _non_vmap_expansion_sdpa(
+    batch_indices: torch.Tensor, head_indices: torch.Tensor, q_indices: torch.Tensor, kv_indices: torch.Tensor
+):
+    """
+    Used to broadcast our mask_functions over the all 4 dimensions (b_idx, h_idx, q_idx, kv_idx) of the inputs.
+    Allows the usage of any index-based mask function without relying on vmap.
+
+    NOTE: This is limited to index based functions only and is not guaranteed to work otherwise.
+
+    Reference:
+        - https://github.com/huggingface/optimum-onnx/blob/c123e8f4fab61b54a8e0e31ce74462bcacca576e/optimum/exporters/onnx/model_patcher.py#L362-L365
+    """
+    batch_indices = batch_indices[:, None, None, None]
+    head_indices = head_indices[None, :, None, None]
+    q_indices = q_indices[None, None, :, None]
+    kv_indices = kv_indices[None, None, None, :]
+    return batch_indices, head_indices, q_indices, kv_indices
+
+def sdpa_mask(
+    batch_size: int,
+    cache_position: torch.Tensor,
+    kv_length: int,
+    kv_offset: int = 0,
+    mask_function: Callable = causal_mask_function,
+    attention_mask: Optional[torch.Tensor] = None,
+    local_size: Optional[int] = None,
+    allow_is_causal_skip: bool = True,
+    allow_is_bidirectional_skip: bool = False,
+    allow_torch_fix: bool = True,
+    use_vmap: bool = False,
+    **kwargs,
+) -> Optional[torch.Tensor]:
+    """
+    Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that
+    the element should take part in the attention computation, and False that it should not.
+    This function can only be used with torch>=2.5, as the context manager is otherwise not available.
+
+    Args:
+        batch_size (`int`):
+            The batch size of the input sequence.
+        cache_position (`torch.Tensor`):
+            A tensor of shape (query_length,) indicating the current indices of the input sequence elements.
+        kv_length (`int`):
+            The size that the key and value states will have during the attention computation.
+        kv_offset (`int`, optional):
+            An optional offset to indicate at which first position the key and values states will refer to.
+        mask_function (`Callable`):
+            The mask factory function describing the mask pattern.
+        attention_mask (`torch.Tensor`, optional):
+            The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length)
+        local_size (`int`, optional):
+            The size of the local attention, if we do not use full attention. This is used only if `allow_is_causal_skip=True`
+            to try to skip mask creation if possible.
+        allow_is_causal_skip (`bool`, optional):
+            Whether to allow to return `None` for the mask under conditions where we can use the `is_causal` argument in
+            `torch.sdpa` instead. Default to `True`.
+        allow_is_bidirectional_skip (`bool`, optional):
+            Whether to allow to return `None` for the mask under conditions where we do not have to add any bias,
+            i.e. full attention without any padding. Default to `False`.
+        allow_torch_fix (`bool`, optional):
+            Whether to update the mask in case a query is not attending to any tokens, to solve a bug in torch's older
+            versions. We need an arg to skip it when using eager. By default `True`.
+        use_vmap (`bool`, optional):
+            Whether to use `vmap` during the mask construction or not. Allows powerful custom patterns that may not be
+            index-based (for the cost of speed performance). By default `False`.
+
+
+    ## Creating a simple causal mask:
+
+    To create the following causal mask:
+
+        0 ■ ⬚ ⬚ ⬚ ⬚
+        1 ■ ■ ⬚ ⬚ ⬚
+        2 ■ ■ ■ ⬚ ⬚
+        3 ■ ■ ■ ■ ⬚
+        4 ■ ■ ■ ■ ■
+
+    You can do
+
+    ```python
+    >>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5)
+    >>> tensor([[[[ True, False, False, False, False],
+                  [ True,  True, False, False, False],
+                  [ True,  True,  True, False, False],
+                  [ True,  True,  True,  True, False],
+                  [ True,  True,  True,  True,  True]]]])
+    ```
+
+    ## Creating a sliding window mask:
+
+    To create the following sliding window mask (`sliding_window=3`):
+
+        0 ■ ⬚ ⬚ ⬚ ⬚
+        1 ■ ■ ⬚ ⬚ ⬚
+        2 ■ ■ ■ ⬚ ⬚
+        3 ⬚ ■ ■ ■ ⬚
+        4 ⬚ ⬚ ■ ■ ■
+
+    You can do
+
+    ```python
+    >>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5, mask_function=sliding_window_causal_mask_function(3))
+    >>> tensor([[[[ True, False, False, False, False],
+                  [ True,  True, False, False, False],
+                  [ True,  True,  True, False, False],
+                  [False,  True,  True,  True, False],
+                  [False, False,  True,  True,  True]]]])
+    ```
+
+    ## Creating a chunked attention mask
+
+    To create the following chunked attention mask (`chunk_size=3`):
+
+        0 ■ ⬚ ⬚ ⬚ ⬚
+        1 ■ ■ ⬚ ⬚ ⬚
+        2 ■ ■ ■ ⬚ ⬚
+        3 ⬚ ⬚ ⬚ ■ ⬚
+        4 ⬚ ⬚ ⬚ ■ ■
+
+    You can do
+
+    ```python
+    >>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5, mask_function=chunked_causal_mask_function(3, torch.zeros(1, dtype=int)))
+    >>> tensor([[[[ True, False, False, False, False],
+                [ True,  True, False, False, False],
+                [ True,  True,  True, False, False],
+                [False, False, False,  True, False],
+                [False, False, False,  True,  True]]]])
+    ```
+
+    """
+    q_length = cache_position.shape[0]
+
+    # Potentially pad the 2D mask
+    padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset)
+
+    # Potentially add the padding 2D mask
+    if padding_mask is not None:
+        mask_function = and_masks(mask_function, padding_mask_function(padding_mask))
+
+    batch_arange = torch.arange(batch_size, device=cache_position.device)
+    head_arange = torch.arange(1, device=cache_position.device)
+    # Similar to `kv_arange = torch.arange(start=kv_offset, end=kv_offset + kv_length, device=cache_position.device)`
+    # but without data-dependent slicing (i.e. torch.compile friendly)
+    kv_arange = torch.arange(kv_length, device=cache_position.device) + kv_offset
+
+    # Actual mask creation
+    # Apply mask function element-wise through broadcasting
+    attention_mask = mask_function(*_non_vmap_expansion_sdpa(batch_arange, head_arange, cache_position, kv_arange))
+    # Expand the mask to match batch size and query length if they weren't used in the mask function
+    attention_mask = attention_mask.expand(batch_size, -1, q_length, kv_length)
+
+    return attention_mask

fastvideo/models/registry.py

@@ -24,6 +24,8 @@ logger = init_logger(__name__)
 _TEXT_TO_VIDEO_DIT_MODELS = {
     "HunyuanVideoTransformer3DModel":
     ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
+    "HunyuanVideo15Transformer3DModel":
+    ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
@@ -45,6 +47,7 @@ _TEXT_ENCODER_MODELS = {
     "T5EncoderModel": ("encoders", "t5", "T5EncoderModel"),
     "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"),
     "BertModel": ("encoders", "clip", "CLIPTextModel"),
+    "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -56,6 +59,7 @@ _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
 _VAE_MODELS = {
     "AutoencoderKLHunyuanVideo":
     ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
+    "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"),
     "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),
     "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo")
 }

fastvideo/models/vaes/hunyuan15vae.py

@@ -0,0 +1,703 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from diffusers
+
+# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+from fastvideo.layers.activation import get_act_fn
+from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
+from fastvideo.models.vaes.common import ParallelTiledVAE
+
+class HunyuanVideo15CausalConv3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int, int]] = 3,
+        stride: Union[int, Tuple[int, int, int]] = 1,
+        padding: Union[int, Tuple[int, int, int]] = 0,
+        dilation: Union[int, Tuple[int, int, int]] = 1,
+        bias: bool = True,
+        pad_mode: str = "replicate",
+    ) -> None:
+        super().__init__()
+
+        kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+
+        self.pad_mode = pad_mode
+        self.time_causal_padding = (
+            kernel_size[0] // 2,
+            kernel_size[0] // 2,
+            kernel_size[1] // 2,
+            kernel_size[1] // 2,
+            kernel_size[2] - 1,
+            0,
+        )
+
+        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode)
+        return self.conv(hidden_states)
+
+
+class HunyuanVideo15RMS_norm(nn.Module):
+    r"""
+    A custom RMS normalization layer.
+
+    Args:
+        dim (int): The number of dimensions to normalize over.
+        channel_first (bool, optional): Whether the input tensor has channels as the first dimension.
+            Default is True.
+        images (bool, optional): Whether the input represents image data. Default is True.
+        bias (bool, optional): Whether to include a learnable bias term. Default is False.
+    """
+
+    def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
+
+
+class HunyuanVideo15AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = HunyuanVideo15RMS_norm(in_channels, images=False)
+
+        self.to_q = nn.Conv3d(in_channels, in_channels, kernel_size=1)
+        self.to_k = nn.Conv3d(in_channels, in_channels, kernel_size=1)
+        self.to_v = nn.Conv3d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv3d(in_channels, in_channels, kernel_size=1)
+
+    @staticmethod
+    def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None):
+        """Prepare a causal attention mask for 3D videos.
+
+        Args:
+            n_frame (int): Number of frames (temporal length).
+            n_hw (int): Product of height and width.
+            dtype: Desired mask dtype.
+            device: Device for the mask.
+            batch_size (int, optional): If set, expands for batch.
+
+        Returns:
+            torch.Tensor: Causal attention mask.
+        """
+        seq_len = n_frame * n_hw
+        mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
+        for i in range(seq_len):
+            i_frame = i // n_hw
+            mask[i, : (i_frame + 1) * n_hw] = 0
+        if batch_size is not None:
+            mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+        return mask
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        x = self.norm(x)
+
+        query = self.to_q(x)
+        key = self.to_k(x)
+        value = self.to_v(x)
+
+        batch_size, channels, frames, height, width = query.shape
+
+        query = query.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
+        key = key.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
+        value = value.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
+
+        attention_mask = self.prepare_causal_attention_mask(
+            frames, height * width, query.dtype, query.device, batch_size=batch_size
+        )
+
+        x = nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
+
+        # batch_size, 1, frames * height * width, channels
+
+        x = x.squeeze(1).reshape(batch_size, frames, height, width, channels).permute(0, 4, 1, 2, 3)
+        x = self.proj_out(x)
+
+        return x + identity
+
+
+class HunyuanVideo15Upsample(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True):
+        super().__init__()
+        factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2
+        self.conv = HunyuanVideo15CausalConv3d(in_channels, out_channels * factor, kernel_size=3)
+
+        self.add_temporal_upsample = add_temporal_upsample
+        self.repeats = factor * out_channels // in_channels
+
+    @staticmethod
+    def _dcae_upsample_rearrange(tensor, r1=1, r2=2, r3=2):
+        """
+        Convert (b, r1*r2*r3*c, f, h, w) -> (b, c, r1*f, r2*h, r3*w)
+
+        Args:
+            tensor: Input tensor of shape (b, r1*r2*r3*c, f, h, w)
+            r1: temporal upsampling factor
+            r2: height upsampling factor
+            r3: width upsampling factor
+        """
+        b, packed_c, f, h, w = tensor.shape
+        factor = r1 * r2 * r3
+        c = packed_c // factor
+
+        tensor = tensor.view(b, r1, r2, r3, c, f, h, w)
+        tensor = tensor.permute(0, 4, 5, 1, 6, 2, 7, 3)
+        return tensor.reshape(b, c, f * r1, h * r2, w * r3)
+
+    def forward(self, x: torch.Tensor):
+        r1 = 2 if self.add_temporal_upsample else 1
+        h = self.conv(x)
+        if self.add_temporal_upsample:
+            h_first = h[:, :, :1, :, :]
+            h_first = self._dcae_upsample_rearrange(h_first, r1=1, r2=2, r3=2)
+            h_first = h_first[:, : h_first.shape[1] // 2]
+            h_next = h[:, :, 1:, :, :]
+            h_next = self._dcae_upsample_rearrange(h_next, r1=r1, r2=2, r3=2)
+            h = torch.cat([h_first, h_next], dim=2)
+
+            # shortcut computation
+            x_first = x[:, :, :1, :, :]
+            x_first = self._dcae_upsample_rearrange(x_first, r1=1, r2=2, r3=2)
+            x_first = x_first.repeat_interleave(repeats=self.repeats // 2, dim=1)
+
+            x_next = x[:, :, 1:, :, :]
+            x_next = self._dcae_upsample_rearrange(x_next, r1=r1, r2=2, r3=2)
+            x_next = x_next.repeat_interleave(repeats=self.repeats, dim=1)
+            shortcut = torch.cat([x_first, x_next], dim=2)
+
+        else:
+            h = self._dcae_upsample_rearrange(h, r1=r1, r2=2, r3=2)
+            shortcut = x.repeat_interleave(repeats=self.repeats, dim=1)
+            shortcut = self._dcae_upsample_rearrange(shortcut, r1=r1, r2=2, r3=2)
+        return h + shortcut
+
+
+class HunyuanVideo15Downsample(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True):
+        super().__init__()
+        factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2
+        self.conv = HunyuanVideo15CausalConv3d(in_channels, out_channels // factor, kernel_size=3)
+
+        self.add_temporal_downsample = add_temporal_downsample
+        self.group_size = factor * in_channels // out_channels
+
+    @staticmethod
+    def _dcae_downsample_rearrange(tensor, r1=1, r2=2, r3=2):
+        """
+        Convert (b, c, r1*f, r2*h, r3*w) -> (b, r1*r2*r3*c, f, h, w)
+
+        This packs spatial/temporal dimensions into channels (opposite of upsample)
+        """
+        b, c, packed_f, packed_h, packed_w = tensor.shape
+        f, h, w = packed_f // r1, packed_h // r2, packed_w // r3
+
+        tensor = tensor.view(b, c, f, r1, h, r2, w, r3)
+        tensor = tensor.permute(0, 3, 5, 7, 1, 2, 4, 6)
+        return tensor.reshape(b, r1 * r2 * r3 * c, f, h, w)
+
+    def forward(self, x: torch.Tensor):
+        r1 = 2 if self.add_temporal_downsample else 1
+        h = self.conv(x)
+        if self.add_temporal_downsample:
+            h_first = h[:, :, :1, :, :]
+            h_first = self._dcae_downsample_rearrange(h_first, r1=1, r2=2, r3=2)
+            h_first = torch.cat([h_first, h_first], dim=1)
+            h_next = h[:, :, 1:, :, :]
+            h_next = self._dcae_downsample_rearrange(h_next, r1=r1, r2=2, r3=2)
+            h = torch.cat([h_first, h_next], dim=2)
+
+            # shortcut computation
+            x_first = x[:, :, :1, :, :]
+            x_first = self._dcae_downsample_rearrange(x_first, r1=1, r2=2, r3=2)
+            B, C, T, H, W = x_first.shape
+            x_first = x_first.view(B, h.shape[1], self.group_size // 2, T, H, W).mean(dim=2)
+            x_next = x[:, :, 1:, :, :]
+            x_next = self._dcae_downsample_rearrange(x_next, r1=r1, r2=2, r3=2)
+            B, C, T, H, W = x_next.shape
+            x_next = x_next.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2)
+            shortcut = torch.cat([x_first, x_next], dim=2)
+        else:
+            h = self._dcae_downsample_rearrange(h, r1=r1, r2=2, r3=2)
+            shortcut = self._dcae_downsample_rearrange(x, r1=r1, r2=2, r3=2)
+            B, C, T, H, W = shortcut.shape
+            shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2)
+
+        return h + shortcut
+
+
+class HunyuanVideo15ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        non_linearity: str = "swish",
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or in_channels
+
+        self.nonlinearity = get_act_fn(non_linearity)
+
+        self.norm1 = HunyuanVideo15RMS_norm(in_channels, images=False)
+        self.conv1 = HunyuanVideo15CausalConv3d(in_channels, out_channels, kernel_size=3)
+
+        self.norm2 = HunyuanVideo15RMS_norm(out_channels, images=False)
+        self.conv2 = HunyuanVideo15CausalConv3d(out_channels, out_channels, kernel_size=3)
+
+        self.conv_shortcut = None
+        if in_channels != out_channels:
+            self.conv_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv1(hidden_states)
+
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            residual = self.conv_shortcut(residual)
+
+        return hidden_states + residual
+
+
+class HunyuanVideo15MidBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_layers: int = 1,
+        add_attention: bool = True,
+    ) -> None:
+        super().__init__()
+        self.add_attention = add_attention
+
+        # There is always at least one resnet
+        resnets = [
+            HunyuanVideo15ResnetBlock(
+                in_channels=in_channels,
+                out_channels=in_channels,
+            )
+        ]
+        attentions = []
+
+        for _ in range(num_layers):
+            if self.add_attention:
+                attentions.append(HunyuanVideo15AttnBlock(in_channels))
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                HunyuanVideo15ResnetBlock(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                )
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.resnets[0](hidden_states)
+
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            if attn is not None:
+                hidden_states = attn(hidden_states)
+            hidden_states = resnet(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideo15DownBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int = 1,
+        downsample_out_channels: Optional[int] = None,
+        add_temporal_downsample: int = True,
+    ) -> None:
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                HunyuanVideo15ResnetBlock(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if downsample_out_channels is not None:
+            self.downsamplers = nn.ModuleList(
+                [
+                    HunyuanVideo15Downsample(
+                        out_channels,
+                        out_channels=downsample_out_channels,
+                        add_temporal_downsample=add_temporal_downsample,
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideo15UpBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int = 1,
+        upsample_out_channels: Optional[int] = None,
+        add_temporal_upsample: bool = True,
+    ) -> None:
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                HunyuanVideo15ResnetBlock(
+                    in_channels=input_channels,
+                    out_channels=out_channels,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if upsample_out_channels is not None:
+            self.upsamplers = nn.ModuleList(
+                [
+                    HunyuanVideo15Upsample(
+                        out_channels,
+                        out_channels=upsample_out_channels,
+                        add_temporal_upsample=add_temporal_upsample,
+                    )
+                ]
+            )
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for resnet in self.resnets:
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+
+        else:
+            for resnet in self.resnets:
+                hidden_states = resnet(hidden_states)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideo15Encoder3D(nn.Module):
+    r"""
+    3D vae encoder for HunyuanImageRefiner.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 64,
+        block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024),
+        layers_per_block: int = 2,
+        temporal_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 16,
+        downsample_match_channel: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.group_size = block_out_channels[-1] // self.out_channels
+
+        self.conv_in = HunyuanVideo15CausalConv3d(in_channels, block_out_channels[0], kernel_size=3)
+        self.mid_block = None
+        self.down_blocks = nn.ModuleList([])
+
+        input_channel = block_out_channels[0]
+        for i in range(len(block_out_channels)):
+            add_spatial_downsample = i < np.log2(spatial_compression_ratio)
+            output_channel = block_out_channels[i]
+            if not add_spatial_downsample:
+                down_block = HunyuanVideo15DownBlock3D(
+                    num_layers=layers_per_block,
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    downsample_out_channels=None,
+                    add_temporal_downsample=False,
+                )
+                input_channel = output_channel
+            else:
+                add_temporal_downsample = i >= np.log2(spatial_compression_ratio // temporal_compression_ratio)
+                downsample_out_channels = block_out_channels[i + 1] if downsample_match_channel else output_channel
+                down_block = HunyuanVideo15DownBlock3D(
+                    num_layers=layers_per_block,
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    downsample_out_channels=downsample_out_channels,
+                    add_temporal_downsample=add_temporal_downsample,
+                )
+                input_channel = downsample_out_channels
+
+            self.down_blocks.append(down_block)
+
+        self.mid_block = HunyuanVideo15MidBlock(in_channels=block_out_channels[-1])
+
+        self.norm_out = HunyuanVideo15RMS_norm(block_out_channels[-1], images=False)
+        self.conv_act = nn.SiLU()
+        self.conv_out = HunyuanVideo15CausalConv3d(block_out_channels[-1], out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for down_block in self.down_blocks:
+                hidden_states = self._gradient_checkpointing_func(down_block, hidden_states)
+
+            hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)
+        else:
+            for down_block in self.down_blocks:
+                hidden_states = down_block(hidden_states)
+
+            hidden_states = self.mid_block(hidden_states)
+
+        batch_size, _, frame, height, width = hidden_states.shape
+        short_cut = hidden_states.view(batch_size, -1, self.group_size, frame, height, width).mean(dim=2)
+
+        hidden_states = self.norm_out(hidden_states)
+        hidden_states = self.conv_act(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+
+        hidden_states += short_cut
+
+        return hidden_states
+
+
+class HunyuanVideo15Decoder3D(nn.Module):
+    r"""
+    Causal decoder for 3D video-like data used for HunyuanImage-1.5 Refiner.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 32,
+        out_channels: int = 3,
+        block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128),
+        layers_per_block: int = 2,
+        spatial_compression_ratio: int = 16,
+        temporal_compression_ratio: int = 4,
+        upsample_match_channel: bool = True,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.repeat = block_out_channels[0] // self.in_channels
+
+        self.conv_in = HunyuanVideo15CausalConv3d(self.in_channels, block_out_channels[0], kernel_size=3)
+        self.up_blocks = nn.ModuleList([])
+
+        # mid
+        self.mid_block = HunyuanVideo15MidBlock(in_channels=block_out_channels[0])
+
+        # up
+        input_channel = block_out_channels[0]
+        for i in range(len(block_out_channels)):
+            output_channel = block_out_channels[i]
+
+            add_spatial_upsample = i < np.log2(spatial_compression_ratio)
+            add_temporal_upsample = i < np.log2(temporal_compression_ratio)
+            if add_spatial_upsample or add_temporal_upsample:
+                upsample_out_channels = block_out_channels[i + 1] if upsample_match_channel else output_channel
+                up_block = HunyuanVideo15UpBlock3D(
+                    num_layers=self.layers_per_block + 1,
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    upsample_out_channels=upsample_out_channels,
+                    add_temporal_upsample=add_temporal_upsample,
+                )
+                input_channel = upsample_out_channels
+            else:
+                up_block = HunyuanVideo15UpBlock3D(
+                    num_layers=self.layers_per_block + 1,
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    upsample_out_channels=None,
+                    add_temporal_upsample=False,
+                )
+                input_channel = output_channel
+
+            self.up_blocks.append(up_block)
+
+        # out
+        self.norm_out = HunyuanVideo15RMS_norm(block_out_channels[-1], images=False)
+        self.conv_act = nn.SiLU()
+        self.conv_out = HunyuanVideo15CausalConv3d(block_out_channels[-1], out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states) + hidden_states.repeat_interleave(repeats=self.repeat, dim=1)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)
+
+            for up_block in self.up_blocks:
+                hidden_states = self._gradient_checkpointing_func(up_block, hidden_states)
+        else:
+            hidden_states = self.mid_block(hidden_states)
+
+            for up_block in self.up_blocks:
+                hidden_states = up_block(hidden_states)
+
+        # post-process
+        hidden_states = self.norm_out(hidden_states)
+        hidden_states = self.conv_act(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+        return hidden_states
+
+
+class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE):
+    r"""
+    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used for
+    HunyuanVideo-1.5.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+    """
+
+    _supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        config: Hunyuan15VAEConfig,
+    ) -> None:
+        nn.Module.__init__(self)
+        ParallelTiledVAE.__init__(self, config)
+
+        if config.load_encoder:
+            self.encoder = HunyuanVideo15Encoder3D(
+                in_channels=config.in_channels,
+                out_channels=config.latent_channels * 2,
+                block_out_channels=config.block_out_channels,
+                layers_per_block=config.layers_per_block,
+                temporal_compression_ratio=config.temporal_compression_ratio,
+                spatial_compression_ratio=config.spatial_compression_ratio,
+                downsample_match_channel=config.downsample_match_channel,
+            )
+
+        if config.load_decoder:
+            self.decoder = HunyuanVideo15Decoder3D(
+                in_channels=config.latent_channels,
+                out_channels=config.out_channels,
+                block_out_channels=list(reversed(config.block_out_channels)),
+                layers_per_block=config.layers_per_block,
+                temporal_compression_ratio=config.temporal_compression_ratio,
+                spatial_compression_ratio=config.spatial_compression_ratio,
+                upsample_match_channel=config.upsample_match_channel,
+            )
+
+        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+        # intermediate tiles together, the memory requirement can be lowered.
+        self.use_tiling = False
+
+        # The minimal tile height and width for spatial tiling to be used
+        self.tile_sample_min_height = 256
+        self.tile_sample_min_width = 256
+        self.tile_sample_min_num_frames = 2000 # Fill in a random large number, as hy1.5 vae does not use temporal tiling
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.encoder(x)
+        return x
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        dec = self.decoder(z)
+        return dec
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec

fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Hunyuan video diffusion pipeline implementation.
+
+This module contains an implementation of the Hunyuan video diffusion pipeline
+using the modular pipeline architecture.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
+                                        DenoisingStage, InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage,
+                                        Hy15ImageEncodingStage)
+
+# TODO(will): move PRECISION_TO_TYPE to better place
+
+logger = init_logger(__name__)
+
+
+class HunyuanVideo15Pipeline(ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae",
+        "transformer", "scheduler"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage_primary",
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2")
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2")
+                           ],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=Hy15ImageEncodingStage(image_encoder=None,
+                                                    image_processor=None))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = HunyuanVideo15Pipeline

fastvideo/pipelines/pipeline_registry.py

@@ -25,6 +25,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "WanCausalDMDPipeline": "wan",
     "StepVideoPipeline": "stepvideo",
     "HunyuanVideoPipeline": "hunyuan",
+    "HunyuanVideo15Pipeline": "hunyuan15",
     "Cosmos2VideoToWorldPipeline": "cosmos",
     "MatrixGamePipeline": "matrixgame",
     "MatrixGameCausalDMDPipeline": "matrixgame",

fastvideo/pipelines/preprocess/preprocess_pipeline_base.py

@@ -303,7 +303,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
             if data is None:
                 continue
 
-            with torch.inference_mode():
+            with torch.no_grad():
                 # Filter out invalid samples (those with all zeros)
                 valid_indices = []
                 for i, pixel_values in enumerate(data["pixel_values"]):

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -37,7 +37,8 @@ from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage,
                                         InputValidationStage,
                                         LatentPreparationStage,
                                         TextEncodingStage,
-                                        TimestepPreparationStage)
+                                        TimestepPreparationStage,
+                                        Hy15ImageEncodingStage)
 from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict
 
 logger = init_logger(__name__)
@@ -47,7 +48,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
     """ODE Trajectory preprocessing pipeline implementation."""
 
     _required_config_modules = [
-        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+        "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", "transformer", "scheduler"
     ]
     preprocess_dataloader: StatefulDataLoader
     preprocess_loader_iter: Iterator[dict[str, Any]]
@@ -61,19 +62,13 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         """Set up pipeline stages with proper dependency injection."""
         assert fastvideo_args.pipeline_config.flow_shift == 5
-        self.modules["scheduler"] = SelfForcingFlowMatchScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift,
-            sigma_min=0.0,
-            extra_one_step=True)
-        self.modules["scheduler"].set_timesteps(num_inference_steps=48,
-                                                denoising_strength=1.0)
 
         self.add_stage(stage_name="input_validation_stage",
                        stage=InputValidationStage())
         self.add_stage(stage_name="prompt_encoding_stage",
                        stage=TextEncodingStage(
-                           text_encoders=[self.get_module("text_encoder")],
-                           tokenizers=[self.get_module("tokenizer")],
+                           text_encoders=[self.get_module("text_encoder"), self.get_module("text_encoder_2")],
+                           tokenizers=[self.get_module("tokenizer"), self.get_module("tokenizer_2")],
                        ))
         self.add_stage(stage_name="timestep_preparation_stage",
                        stage=TimestepPreparationStage(
@@ -82,6 +77,9 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                        stage=LatentPreparationStage(
                            scheduler=self.get_module("scheduler"),
                            transformer=self.get_module("transformer", None)))
+        self.add_stage(stage_name="image_encoding_stage",
+                        stage=Hy15ImageEncodingStage(image_encoder=None,
+                                                    image_processor=None))
         self.add_stage(stage_name="denoising_stage",
                        stage=DenoisingStage(
                            transformer=self.get_module("transformer"),
@@ -95,11 +93,13 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                                        args):
         """Preprocess text-only data and generate trajectory information."""
 
+        num_encoders = len(self.prompt_encoding_stage.text_encoders)
+
         for batch_idx, data in enumerate(self.pbar):
             if data is None:
                 continue
 
-            with torch.inference_mode():
+            with torch.no_grad():
                 # For text-only processing, we only need text data
                 # Filter out samples without text
                 valid_indices = []
@@ -130,12 +130,9 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                 prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
                     batch_captions,
                     fastvideo_args,
-                    encoder_index=[0],
+                    encoder_index=list(range(num_encoders)),
                     return_attention_mask=True,
                 )
-                prompt_embeds = prompt_embeds_list[0]
-                prompt_attention_masks = prompt_masks_list[0]
-                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
 
                 sampling_params = SamplingParam.from_pretrained(args.model_path)
 
@@ -144,61 +141,52 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                     negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text(
                         sampling_params.negative_prompt,
                         fastvideo_args,
-                        encoder_index=[0],
+                        encoder_index=list(range(num_encoders)),
                         return_attention_mask=True,
                     )
-                    negative_prompt_embed = negative_prompt_embeds_list[0][0]
-                    negative_prompt_attention_mask = negative_prompt_masks_list[
-                        0][0]
                 else:
-                    negative_prompt_embed = None
-                    negative_prompt_attention_mask = None
+                    negative_prompt_embeds_list = []
+                    negative_prompt_masks_list = []
 
                 trajectory_latents = []
                 trajectory_timesteps = []
                 trajectory_decoded = []
 
-                for i, (prompt_embed, prompt_attention_mask) in enumerate(
-                        zip(prompt_embeds, prompt_attention_masks,
-                            strict=False)):
-                    prompt_embed = prompt_embed.unsqueeze(0)
-                    prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
+                # Collect the trajectory data (text-to-video generation)
+                batch = ForwardBatch(**shallow_asdict(sampling_params), )
+                batch.prompt_embeds = prompt_embeds_list
+                batch.prompt_attention_mask = prompt_masks_list
+                batch.negative_prompt_embeds = negative_prompt_embeds_list
+                batch.negative_attention_mask = negative_prompt_masks_list
+                batch.num_inference_steps = 50
+                batch.return_trajectory_latents = True
+                # Enabling this will save the decoded trajectory videos.
+                # Used for debugging.
+                batch.return_trajectory_decoded = False
+                batch.height = args.max_height
+                batch.width = args.max_width
+                batch.num_frames = args.num_frames
+                batch.fps = args.train_fps
+                batch.guidance_scale = 6.0
+                batch.do_classifier_free_guidance = True
 
-                    # Collect the trajectory data (text-to-video generation)
-                    batch = ForwardBatch(**shallow_asdict(sampling_params), )
-                    batch.prompt_embeds = [prompt_embed]
-                    batch.prompt_attention_mask = [prompt_attention_mask]
-                    batch.negative_prompt_embeds = [negative_prompt_embed]
-                    batch.negative_attention_mask = [
-                        negative_prompt_attention_mask
-                    ]
-                    batch.num_inference_steps = 48
-                    batch.return_trajectory_latents = True
-                    # Enabling this will save the decoded trajectory videos.
-                    # Used for debugging.
-                    batch.return_trajectory_decoded = False
-                    batch.height = args.max_height
-                    batch.width = args.max_width
-                    batch.fps = args.train_fps
-                    batch.guidance_scale = 6.0
-                    batch.do_classifier_free_guidance = True
+                result_batch = self.input_validation_stage(
+                    batch, fastvideo_args)
+                result_batch = self.timestep_preparation_stage(
+                    batch, fastvideo_args)
+                result_batch = self.latent_preparation_stage(
+                    result_batch, fastvideo_args)
+                result_batch = self.image_encoding_stage(result_batch, fastvideo_args)
+                result_batch = self.denoising_stage(result_batch,
+                                                    fastvideo_args)
+                result_batch = self.decoding_stage(result_batch,
+                                                    fastvideo_args)
 
-                    result_batch = self.input_validation_stage(
-                        batch, fastvideo_args)
-                    result_batch = self.timestep_preparation_stage(
-                        batch, fastvideo_args)
-                    result_batch = self.latent_preparation_stage(
-                        result_batch, fastvideo_args)
-                    result_batch = self.denoising_stage(result_batch,
-                                                        fastvideo_args)
-                    result_batch = self.decoding_stage(result_batch,
-                                                       fastvideo_args)
-
-                    trajectory_latents.append(
-                        result_batch.trajectory_latents.cpu())
-                    trajectory_timesteps.append(
-                        result_batch.trajectory_timesteps.cpu())
-                    trajectory_decoded.append(result_batch.trajectory_decoded)
+                trajectory_latents.append(
+                    result_batch.trajectory_latents.cpu())
+                trajectory_timesteps.append(
+                    result_batch.trajectory_timesteps.cpu())
+                trajectory_decoded.append(result_batch.trajectory_decoded)
 
                 # Prepare extra features for text-only processing
                 extra_features = {
@@ -209,10 +197,11 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                 if batch.return_trajectory_decoded:
                     for i, decoded_frames in enumerate(trajectory_decoded):
                         for j, decoded_frame in enumerate(decoded_frames):
-                            save_decoded_latents_as_video(
-                                decoded_frame,
-                                f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
-                                args.train_fps)
+                            if j in [40, 44, 49]:
+                                save_decoded_latents_as_video(
+                                    decoded_frame,
+                                    f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
+                                    args.train_fps)
 
                 # Prepare batch data for Parquet dataset
                 batch_data: list[dict[str, Any]] = []
@@ -227,7 +216,10 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                     video_name = os.path.basename(video_path).split(".")[0]
 
                     # Convert tensors to numpy arrays
-                    text_embedding = prompt_embeds[idx].cpu().numpy()
+                    text_embedding = prompt_embeds_list[0].float().cpu().numpy()
+                    text_mask = prompt_masks_list[0].cpu().numpy()
+                    text_embedding_2 = prompt_embeds_list[1].float().cpu().numpy()
+                    text_mask_2 = prompt_masks_list[1].cpu().numpy()
 
                     # Get extra features for this sample
                     sample_extra_features = {}
@@ -253,6 +245,9 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                             "trajectory_latents"],
                         trajectory_timesteps=sample_extra_features[
                             "trajectory_timesteps"],
+                        text_embedding_2=text_embedding_2,
+                        text_mask=text_mask,
+                        text_mask_2=text_mask_2,
                     )
                     batch_data.append(record)
 

fastvideo/pipelines/preprocess/preprocess_pipeline_text.py

@@ -58,7 +58,7 @@ class PreprocessPipeline_Text(BasePreprocessPipeline):
             if data is None:
                 continue
 
-            with torch.inference_mode():
+            with torch.no_grad():
                 # For text-only processing, we only need text data
                 # Filter out samples without text
                 valid_indices = []

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -22,32 +22,33 @@ logger = init_logger(__name__)
 
 
 def main(args) -> None:
-    args.model_path = maybe_download_model(args.model_path)
+    # args.model_path = maybe_download_model(args.model_path)
     maybe_init_distributed_environment_and_model_parallel(1, 1)
     num_gpus = int(os.environ["WORLD_SIZE"])
     assert num_gpus == 1, "Only support 1 GPU"
 
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
+    print(pipeline_config.__class__.__name__)
 
-    kwargs: dict[str, Any] = {}
-    if args.preprocess_task == "text_only":
-        kwargs = {
-            "text_encoder_cpu_offload": False,
-        }
-    else:
-        # Full config for video/image processing
-        kwargs = {
-            "vae_precision": "fp32",
-            "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
-        }
-    pipeline_config.update_config_from_dict(kwargs)
+    # kwargs: dict[str, Any] = {}
+    # if args.preprocess_task == "text_only":
+    #     kwargs = {
+    #         "text_encoder_cpu_offload": False,
+    #     }
+    # else:
+    #     # Full config for video/image processing
+    #     kwargs = {
+    #         "vae_precision": "fp32",
+    #         "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
+    #     }
+    # pipeline_config.update_config_from_dict(kwargs)
 
     fastvideo_args = FastVideoArgs(
         model_path=args.model_path,
         num_gpus=get_world_size(),
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=False,
+        dit_cpu_offload=True,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
         pipeline_config=pipeline_config,
     )
     if args.preprocess_task == "t2v":

fastvideo/pipelines/stages/__init__.py

@@ -16,7 +16,7 @@ from fastvideo.pipelines.stages.denoising import (CosmosDenoisingStage,
 from fastvideo.pipelines.stages.encoding import EncodingStage
 from fastvideo.pipelines.stages.image_encoding import (
     ImageEncodingStage, MatrixGameImageEncodingStage, RefImageEncodingStage,
-    ImageVAEEncodingStage, VideoVAEEncodingStage)
+    ImageVAEEncodingStage, VideoVAEEncodingStage, Hy15ImageEncodingStage)
 from fastvideo.pipelines.stages.input_validation import InputValidationStage
 from fastvideo.pipelines.stages.latent_preparation import (
     CosmosLatentPreparationStage, LatentPreparationStage)
@@ -44,6 +44,7 @@ __all__ = [
     "DecodingStage",
     "ImageEncodingStage",
     "MatrixGameImageEncodingStage",
+    "Hy15ImageEncodingStage",
     "RefImageEncodingStage",
     "ImageVAEEncodingStage",
     "VideoVAEEncodingStage",

fastvideo/pipelines/stages/image_encoding.py

@@ -100,6 +100,33 @@ class ImageEncodingStage(PipelineStage):
         return result
 
 
+class Hy15ImageEncodingStage(ImageEncodingStage):
+    """
+    Stage for encoding image prompts into embeddings for HunyuanVideo1.5 models.
+    """
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify image encoding stage inputs."""
+        return VerificationResult()
+
+    def forward(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        """
+        Encode the prompt into image encoder hidden states.
+        """
+        if batch.pil_image is None:
+            batch.image_embeds = [
+                torch.zeros(1, 729, 1152, device=get_local_torch_device())
+            ]
+
+        raw_latent_shape = list(batch.raw_latent_shape)
+        raw_latent_shape[1] = 1
+        batch.video_latent = torch.zeros(tuple(raw_latent_shape),
+                                         device=get_local_torch_device())
+        return batch
+
+
 class MatrixGameImageEncodingStage(ImageEncodingStage):
     CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
     CLIP_STD = [0.26862954, 0.26130258, 0.27577711]

fastvideo/pipelines/stages/text_encoding.py

@@ -6,6 +6,7 @@ This module contains implementations of prompt encoding stages for diffusion pip
 """
 
 import torch
+from typing import Any
 
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -100,9 +101,9 @@ class TextEncodingStage(PipelineStage):
         """Verify text encoding stage inputs."""
         result = VerificationResult()
         result.add_check("prompt", batch.prompt, V.string_or_list_strings)
-        result.add_check(
-            "negative_prompt", batch.negative_prompt, lambda x: not batch.
-            do_classifier_free_guidance or V.string_not_empty(x))
+        # result.add_check(
+        #     "negative_prompt", batch.negative_prompt, lambda x: not batch.
+        #     do_classifier_free_guidance or V.string_not_empty(x))
         result.add_check("do_classifier_free_guidance",
                          batch.do_classifier_free_guidance, V.bool_value)
         result.add_check("prompt_embeds", batch.prompt_embeds, V.is_list)
@@ -203,20 +204,45 @@ class TextEncodingStage(PipelineStage):
             preprocess_func = preprocess_funcs[i]
             postprocess_func = postprocess_funcs[i]
 
-            processed_texts: list[str] = []
-            for prompt_str in texts:
-                processed_texts.append(preprocess_func(prompt_str))
-
             tok_kwargs = dict(encoder_config.tokenizer_kwargs)
             if max_length is not None:
                 tok_kwargs["max_length"] = max_length
+            elif hasattr(fastvideo_args.pipeline_config,
+                         "text_encoder_max_lengths"):
+                tok_kwargs[
+                    "max_length"] = fastvideo_args.pipeline_config.text_encoder_max_lengths[
+                        i]
+
             if truncation is not None:
                 tok_kwargs["truncation"] = truncation
             if padding is not None:
                 tok_kwargs["padding"] = padding
 
-            text_inputs = tokenizer(processed_texts,
-                                    **tok_kwargs).to(target_device)
+            processed_texts: list[str] = []
+            for prompt_str in texts:
+                processed_text = preprocess_func(prompt_str)
+                if processed_text is not None:
+                    processed_texts.append(processed_text)
+                else:
+                    # Assuming batch_size = 1
+                    prompt_embeds = torch.zeros((1, tok_kwargs["max_length"],
+                                                 encoder_config.hidden_size),
+                                                device=target_device)
+                    attention_mask = torch.zeros((1, tok_kwargs["max_length"]),
+                                                 device=target_device,
+                                                 dtype=torch.int64)
+                    embeds_list.append(prompt_embeds)
+                    attn_masks_list.append(attention_mask)
+                    return self.return_embeds(embeds_list, attn_masks_list,
+                                              return_type,
+                                              return_attention_mask, indices)
+
+            if encoder_config.is_chat_model:
+                text_inputs = tokenizer.apply_chat_template(
+                    processed_texts, **tok_kwargs).to(target_device)
+            else:
+                text_inputs = tokenizer(processed_texts,
+                                        **tok_kwargs).to(target_device)
 
             input_ids = text_inputs["input_ids"]
             attention_mask = text_inputs["attention_mask"]
@@ -228,13 +254,29 @@ class TextEncodingStage(PipelineStage):
                     output_hidden_states=True,
                 )
 
-            prompt_embeds = postprocess_func(outputs)
+            try:
+                prompt_embeds = postprocess_func(outputs)
+            except Exception:
+                prompt_embeds, attention_mask = postprocess_func(
+                    outputs, attention_mask)
+
             if dtype is not None:
                 prompt_embeds = prompt_embeds.to(dtype=dtype)
             embeds_list.append(prompt_embeds)
             if return_attention_mask:
                 attn_masks_list.append(attention_mask)
 
+        return self.return_embeds(embeds_list, attn_masks_list, return_type,
+                                  return_attention_mask, indices)
+
+    def return_embeds(
+        self,
+        embeds_list: list[torch.Tensor],
+        attn_masks_list: list[torch.Tensor],
+        return_type: str = "list",
+        return_attention_mask: bool = False,
+        indices: list[int] | None = None,
+    ) -> Any:
         # Shape results according to return_type
         if return_type == "list":
             if return_attention_mask:

fastvideo/tests/encoders/test_hyt5_encoder.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import numpy as np
+import pytest
+import torch
+from torch.distributed.tensor import DTensor
+from torch.testing import assert_close
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TextEncoderLoader
+from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.configs.models.encoders import T5Config
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+
+@pytest.fixture
+def t5_model_paths():
+    base_model_path = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+    model_path = maybe_download_model(base_model_path)
+    text_encoder_path = os.path.join(model_path, "text_encoder_2")
+    tokenizer_path = os.path.join(model_path, "tokenizer_2")
+    return text_encoder_path, tokenizer_path
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_t5_encoder(t5_model_paths):
+    # Initialize the two model implementations
+    text_encoder_path, tokenizer_path = t5_model_paths
+    hf_config = AutoConfig.from_pretrained(text_encoder_path)
+    print(hf_config)
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision_str = "fp32"
+    precision = PRECISION_TO_TYPE[precision_str]
+    model1 = T5EncoderModel.from_pretrained(text_encoder_path).to(
+        precision).to(device).eval()
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
+
+    args = FastVideoArgs(model_path=text_encoder_path,
+                        pipeline_config=PipelineConfig(text_encoder_configs=(T5Config(),),
+                        text_encoder_precisions=(precision_str,)),
+                        pin_cpu_memory=False)
+    loader = TextEncoderLoader()
+    model2 = loader.load(text_encoder_path, args)
+    model2 = model2.to(precision)
+    model2.eval()
+
+    # Sanity check weights between the two models
+    logger.info("Comparing model weights for sanity check...")
+    params1 = dict(model1.named_parameters())
+    params2 = dict(model2.named_parameters())
+
+    # Check number of parameters
+    logger.info("Model1 has %s parameters", len(params1))
+    logger.info("Model2 has %s parameters", len(params2))
+
+    # check if embed_tokens are the same
+    weights = ["encoder.block.{}.layer.0.layer_norm.weight", \
+               "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\
+                "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \
+                "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight"]
+    
+    for idx in range(hf_config.num_hidden_layers):
+        for w in weights:
+            name1 = w.format(idx)
+            name2 = w.format(idx)
+            p1 = params1[name1]
+            p2 = params2[name2]
+            p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1)
+            assert_close(p1, p2, atol=1e-4, rtol=1e-4)
+    
+
+    # Test with some sample prompts
+    prompts = [
+        "Once upon a time", "The quick brown fox jumps over",
+        "In a galaxy far, far away"
+    ]
+
+    logger.info("Testing T5 encoder with sample prompts")
+
+    with torch.no_grad():
+        for prompt in prompts:
+            logger.info("Testing prompt: %s", prompt)
+
+            # Tokenize the prompt
+            tokens = tokenizer(prompt,
+                               padding="max_length",
+                               max_length=512,
+                               truncation=True,
+                               add_special_tokens=True,
+                               return_tensors="pt").to(device)
+
+            # Get outputs from HuggingFace implementation
+            # filter out padding input_ids
+            # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1]
+            # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1]
+            outputs1 = model1(input_ids=tokens.input_ids,
+                              attention_mask=tokens.attention_mask.float())[0]
+            print("--------------------------------")
+            logger.info("Testing model2")
+
+            # Get outputs from our implementation
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                outputs2 = model2(
+                    input_ids=tokens.input_ids,
+                    attention_mask=tokens.attention_mask,
+                ).last_hidden_state
+
+            # Compare last hidden states
+            last_hidden_state1 = outputs1[tokens.attention_mask == 1]
+            last_hidden_state2 = outputs2[tokens.attention_mask == 1]
+
+            assert last_hidden_state1.shape == last_hidden_state2.shape, \
+                f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}"
+
+            max_diff_hidden = torch.max(
+                torch.abs(last_hidden_state1 - last_hidden_state2))
+            mean_diff_hidden = torch.mean(
+                torch.abs(last_hidden_state1 - last_hidden_state2))
+
+            logger.info("Maximum difference in last hidden states: %s",
+                        max_diff_hidden.item())
+            logger.info("Mean difference in last hidden states: %s",
+                        mean_diff_hidden.item())
+            logger.info("Max memory allocated: %s GB", torch.cuda.max_memory_allocated() / 1024**3)
+            # Check if outputs are similar (allowing for small numerical differences)
+            assert mean_diff_hidden < 1e-4, \
+                f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}"
+            assert max_diff_hidden < 1e-4, \
+                f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}"

fastvideo/tests/encoders/test_qwen2_5_encoder.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import pytest
+import torch
+from torch.distributed.tensor import DTensor
+from torch.testing import assert_close
+from transformers import AutoConfig, AutoTokenizer, Qwen2_5_VLTextModel
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TextEncoderLoader
+from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.configs.models.encoders import Qwen2_5_VLConfig
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29505"
+
+@pytest.fixture
+def qwen_model_path():
+    base_model_path = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+    model_path = maybe_download_model(base_model_path)
+    text_encoder_path = os.path.join(model_path, "text_encoder")
+    tokenizer_path = os.path.join(model_path, "tokenizer")
+    return text_encoder_path, tokenizer_path
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_qwen2_5_encoder(qwen_model_path):
+    text_encoder_path, tokenizer_path = qwen_model_path
+    hf_config = AutoConfig.from_pretrained(text_encoder_path)
+    print(hf_config)
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    # Qwen2.5-VL default dtype is usually bf16
+    precision_str = "fp32"
+    precision = PRECISION_TO_TYPE[precision_str]
+    
+    logger.info(f"Using precision: {precision_str}")
+
+    # Load HF model (Base model)
+    model1 = Qwen2_5_VLTextModel.from_pretrained(text_encoder_path).to(
+        precision).to(device).eval()
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
+    # Load FastVideo model
+    args = FastVideoArgs(model_path=text_encoder_path,
+                        pipeline_config=PipelineConfig(text_encoder_configs=(Qwen2_5_VLConfig(),),
+                        text_encoder_precisions=(precision_str,)),
+                        pin_cpu_memory=False)
+    
+    loader = TextEncoderLoader()
+    model2 = loader.load(text_encoder_path, args)
+    model2 = model2.to(precision)
+    model2.eval()
+
+    # Sanity check weights
+    logger.info("Comparing model weights for sanity check...")
+    params1 = dict(model1.named_parameters())
+    params2 = dict(model2.named_parameters())
+    
+    logger.info("Model1 has %s parameters", len(params1))
+    logger.info("Model2 has %s parameters", len(params2))
+
+    # Check common layers like Norms which are likely not merged/sharded in a way that changes name significantly
+    # or simple linear layers if names match.
+    # Note: FastVideo uses QKVParallelLinear, so q_proj, k_proj, v_proj are merged.
+    # HF Qwen2_5_VL uses separate projections? No, usually they are separate nn.Linear in HF.
+    
+    weights_to_check = [
+        "norm.weight",
+        "layers.{}.self_attn.o_proj.weight",
+        "layers.{}.input_layernorm.weight",
+        "layers.{}.post_attention_layernorm.weight",
+        "layers.{}.mlp.down_proj.weight"
+    ]
+    
+    for idx in range(hf_config.num_hidden_layers):
+        for w in weights_to_check:
+            name1 = w.format(idx)
+            name2 = w.format(idx)
+            p1 = params1[name1]
+            p2 = params2[name2]
+            p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1)
+            
+            # Check shape
+            assert p1.shape == p2.shape, f"Shape mismatch for {w}: {p1.shape} vs {p2.shape}"
+            
+            # Check values
+            assert_close(p1, p2, atol=1e-7, rtol=1e-7, msg=f"Weight mismatch for {w}")
+
+    # Test with sample prompts
+    prompts = [
+        "Hello world",
+        "The quick brown fox jumps over the lazy dog."
+    ]
+    
+    logger.info("Testing with sample prompts")
+    
+    with torch.no_grad():
+        for prompt in prompts:
+            logger.info(f"Prompt: {prompt}")
+            tokens = tokenizer(prompt, return_tensors="pt", padding="max_length", max_length=1000, truncation=True).to(device)
+            
+            # HF Forward
+            # AutoModel for Qwen2.5-VL usually returns BaseModelOutputWithPast
+            outputs1 = model1(
+                input_ids=tokens.input_ids,
+                attention_mask=tokens.attention_mask,
+                output_hidden_states=True
+            ).hidden_states[-3]
+            
+            # FastVideo Forward
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                outputs2 = model2(
+                    input_ids=tokens.input_ids,
+                    attention_mask=tokens.attention_mask,
+                    output_hidden_states=True
+                ).hidden_states[-3]
+                
+            # Compare
+            # Filter padding for comparison if needed, but here we just check raw output matching
+            
+            # Check shapes
+            assert outputs1.shape == outputs2.shape, f"Output shape mismatch: {outputs1.shape} vs {outputs2.shape}"
+            
+            diff = torch.abs(outputs1 - outputs2)
+            max_diff = diff.max().item()
+            mean_diff = diff.mean().item()
+            
+            logger.info(f"Max diff: {max_diff}")
+            logger.info(f"Mean diff: {mean_diff}")
+            
+            # Thresholds
+            # Qwen2.5-VL RoPE is complex, if our implementation is slightly off (e.g. float32 conversion logic in RoPE),
+            # differences might appear. But should be small.
+            if precision_str == "bf16":
+                atol = 5e-2 # relaxed for bf16
+            else:
+                atol = 1e-3
+                
+            if max_diff > atol:
+                logger.warning(f"Max diff {max_diff} > {atol}. Checking if it's acceptable...")
+                # If mean diff is small, maybe just outliers
+                assert mean_diff < atol, f"Mean diff {mean_diff} too high"
+            else:
+                logger.info("Outputs match within tolerance.")
+

fastvideo/tests/vaes/test_hunyuan15_vae.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import os
+
+import pytest
+import torch
+from diffusers import AutoencoderKLHunyuanVideo15
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.logger import init_logger
+# from fastvideo.models.vaes.hunyuanvae import (
+#     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
+from fastvideo.utils import maybe_download_model
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      "data", BASE_MODEL_PATH))
+VAE_PATH = os.path.join(MODEL_PATH, "vae")
+CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_hunyuan_vae():
+    device = torch.device("cuda:0")
+    precision = torch.float32
+    precision_str = "fp32"
+    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=Hunyuan15VAEConfig(), vae_precision=precision_str))
+    args.device = device
+    args.vae_cpu_offload = False
+
+    model1 = AutoencoderKLHunyuanVideo15.from_pretrained(
+        VAE_PATH, torch_dtype=precision).to(device).eval()
+    model1.enable_tiling()
+
+    loader = VAELoader()
+    model2 = loader.load(VAE_PATH, args)
+
+    model2.enable_tiling()
+
+    batch_size = 1
+
+    # Video input [B, C, T, H, W]
+    input_tensor = torch.randn(batch_size,
+                               3,
+                               81,
+                               512,
+                               512,
+                               device=device,
+                               dtype=precision)
+
+    # Disable gradients for inference
+    with torch.no_grad():
+        latent1 = model1.encode(input_tensor, return_dict=False)[0].mode()
+        latent2 = model2.encode(input_tensor).mode()
+
+    assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}"
+    max_diff_encode = torch.max(torch.abs(latent1.float() - latent2.float()))
+    mean_diff_encode = torch.mean(torch.abs(latent1.float() - latent2.float()))
+    logger.info("Maximum difference between encoded latents: %s",
+                max_diff_encode.item())
+    logger.info("Mean difference between encoded latents: %s",
+                mean_diff_encode.item())
+    assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {max_diff_encode.item()}, mean diff = {mean_diff_encode.item()}"
+
+    # Test decoding
+    latent1 = latent1 / model1.config.scaling_factor
+    latent2 = latent2 / model2.config.scaling_factor
+
+    with torch.no_grad():
+        video1 = model1.decode(latent1, return_dict=False)[0]
+        video2 = model2.decode(latent2)
+
+    assert video1.shape == video2.shape, f"Video shapes don't match: {video1.shape} vs {video2.shape}"
+    max_diff_decode = torch.max(torch.abs(video1.float() - video2.float()))
+    mean_diff_decode = torch.mean(torch.abs(video1.float() - video2.float()))
+    logger.info("Maximum difference between decoded videos: %s",
+                max_diff_decode.item())
+    logger.info("Mean difference between decoded videos: %s",
+                mean_diff_decode.item())
+    assert max_diff_decode < 1e-5, f"Decoded videos differ significantly: max diff = {max_diff_decode.item()}, mean diff = {mean_diff_decode.item()}"
+

fastvideo/utils.py

@@ -31,7 +31,7 @@ import filelock
 import imageio
 import numpy as np
 import torch
-import torchvision.utils as make_grid
+from torchvision.utils import make_grid
 import yaml
 from diffusers.loaders.lora_base import (
     _best_guess_weight_name)  # watch out for potetential removal from diffusers