From fdca8fdbaaf338f741f02461df4b025aefe1b97b Mon Sep 17 00:00:00 2001 From: Wei Zhou <69577934+JerryZhou54@users.noreply.github.com> Date: Sat, 20 Dec 2025 22:57:52 -0800 Subject: [PATCH] [New Model] Hunyuan1.5 (#943) --- examples/inference/basic/basic_hy15.py | 42 + fastvideo/attention/backends/flash_attn.py | 55 +- fastvideo/attention/backends/sdpa.py | 33 +- fastvideo/attention/layer.py | 6 +- .../attention/utils/flash_attn_no_pad.py | 99 ++ fastvideo/configs/models/dits/__init__.py | 5 +- fastvideo/configs/models/dits/base.py | 2 + .../configs/models/dits/hunyuanvideo15.py | 157 ++++ fastvideo/configs/models/encoders/__init__.py | 4 +- fastvideo/configs/models/encoders/base.py | 1 + fastvideo/configs/models/encoders/qwen2_5.py | 93 ++ fastvideo/configs/models/encoders/t5.py | 3 + fastvideo/configs/models/vaes/__init__.py | 2 + fastvideo/configs/models/vaes/hunyuan15vae.py | 27 + fastvideo/configs/pipelines/__init__.py | 9 +- fastvideo/configs/pipelines/hunyuan15.py | 139 +++ fastvideo/configs/pipelines/registry.py | 10 + fastvideo/configs/sample/base.py | 1 + fastvideo/configs/sample/hunyuan15.py | 29 + fastvideo/configs/sample/registry.py | 9 + fastvideo/layers/activation.py | 1 + fastvideo/models/dits/hunyuanvideo15.py | 853 ++++++++++++++++++ fastvideo/models/encoders/qwen2_5.py | 387 ++++++++ fastvideo/models/loader/component_loader.py | 10 +- fastvideo/models/mask_utils.py | 201 +++++ fastvideo/models/registry.py | 4 + fastvideo/models/vaes/hunyuan15vae.py | 703 +++++++++++++++ .../pipelines/basic/hunyuan15/__init__.py | 0 .../basic/hunyuan15/hunyuan15_pipeline.py | 74 ++ fastvideo/pipelines/pipeline_registry.py | 1 + fastvideo/pipelines/stages/__init__.py | 3 +- fastvideo/pipelines/stages/image_encoding.py | 27 + fastvideo/pipelines/stages/text_encoding.py | 62 +- fastvideo/tests/encoders/test_hyt5_encoder.py | 140 +++ .../tests/encoders/test_qwen2_5_encoder.py | 150 +++ fastvideo/tests/vaes/test_hunyuan15_vae.py | 91 ++ 36 files changed, 3394 insertions(+), 39 deletions(-) create mode 100644 examples/inference/basic/basic_hy15.py create mode 100644 fastvideo/attention/utils/flash_attn_no_pad.py create mode 100644 fastvideo/configs/models/dits/hunyuanvideo15.py create mode 100644 fastvideo/configs/models/encoders/qwen2_5.py create mode 100644 fastvideo/configs/models/vaes/hunyuan15vae.py create mode 100644 fastvideo/configs/pipelines/hunyuan15.py create mode 100644 fastvideo/configs/sample/hunyuan15.py create mode 100644 fastvideo/models/dits/hunyuanvideo15.py create mode 100644 fastvideo/models/encoders/qwen2_5.py create mode 100644 fastvideo/models/mask_utils.py create mode 100644 fastvideo/models/vaes/hunyuan15vae.py create mode 100644 fastvideo/pipelines/basic/hunyuan15/__init__.py create mode 100644 fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py create mode 100644 fastvideo/tests/encoders/test_hyt5_encoder.py create mode 100644 fastvideo/tests/encoders/test_qwen2_5_encoder.py create mode 100644 fastvideo/tests/vaes/test_hunyuan15_vae.py diff --git a/examples/inference/basic/basic_hy15.py b/examples/inference/basic/basic_hy15.py new file mode 100644 index 000000000..71e491a41 --- /dev/null +++ b/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() diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index f33afd8ed..c1281b8ae 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/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 diff --git a/fastvideo/attention/backends/sdpa.py b/fastvideo/attention/backends/sdpa.py index 782ffd9fc..00a0cee60 100644 --- a/fastvideo/attention/backends/sdpa.py +++ b/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 diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 2ef09da7c..e7985e58e 100644 --- a/fastvideo/attention/layer.py +++ b/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) diff --git a/fastvideo/attention/utils/flash_attn_no_pad.py b/fastvideo/attention/utils/flash_attn_no_pad.py new file mode 100644 index 000000000..1983f8d38 --- /dev/null +++ b/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 diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index f2ab38651..b8691ffed 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/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" ] diff --git a/fastvideo/configs/models/dits/base.py b/fastvideo/configs/models/dits/base.py index 148eb711b..818e0bee3 100644 --- a/fastvideo/configs/models/dits/base.py +++ b/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 diff --git a/fastvideo/configs/models/dits/hunyuanvideo15.py b/fastvideo/configs/models/dits/hunyuanvideo15.py new file mode 100644 index 000000000..6d7f77208 --- /dev/null +++ b/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" diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index 29b5d6113..86a3e3334 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/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" ] diff --git a/fastvideo/configs/models/encoders/base.py b/fastvideo/configs/models/encoders/base.py index 33b2b4c1c..b1a913519 100644 --- a/fastvideo/configs/models/encoders/base.py +++ b/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 diff --git a/fastvideo/configs/models/encoders/qwen2_5.py b/fastvideo/configs/models/encoders/qwen2_5.py new file mode 100644 index 000000000..3af6a6687 --- /dev/null +++ b/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 diff --git a/fastvideo/configs/models/encoders/t5.py b/fastvideo/configs/models/encoders/t5.py index c1de3609c..3fbe75cb0 100644 --- a/fastvideo/configs/models/encoders/t5.py +++ b/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 diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index 12bf5c609..d39254d21 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/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", ] diff --git a/fastvideo/configs/models/vaes/hunyuan15vae.py b/fastvideo/configs/models/vaes/hunyuan15vae.py new file mode 100644 index 000000000..8344986a5 --- /dev/null +++ b/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) diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 8bbb0a60c..e28901889 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/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" ] diff --git a/fastvideo/configs/pipelines/hunyuan15.py b/fastvideo/configs/pipelines/hunyuan15.py new file mode 100644 index 000000000..29e1c5b45 --- /dev/null +++ b/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 diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index d3de0b872..f22fd9a7b 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/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(): diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py index 8d62f95a1..fcf292869 100644 --- a/fastvideo/configs/sample/base.py +++ b/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 diff --git a/fastvideo/configs/sample/hunyuan15.py b/fastvideo/configs/sample/hunyuan15.py new file mode 100644 index 000000000..16871499b --- /dev/null +++ b/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 diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 58969f2f0..1eb9faaa4 100644 --- a/fastvideo/configs/sample/registry.py +++ b/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, diff --git a/fastvideo/layers/activation.py b/fastvideo/layers/activation.py index f20a3b7fa..7e09ebd7b 100644 --- a/fastvideo/layers/activation.py +++ b/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, } diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py new file mode 100644 index 000000000..508ea9648 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/fastvideo/models/encoders/qwen2_5.py b/fastvideo/models/encoders/qwen2_5.py new file mode 100644 index 000000000..bd7c92f72 --- /dev/null +++ b/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 diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 7d19895f3..b0821ee3f 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/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, diff --git a/fastvideo/models/mask_utils.py b/fastvideo/models/mask_utils.py new file mode 100644 index 000000000..ec944c8d9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 04eb6936d..c5c8ad78e 100644 --- a/fastvideo/models/registry.py +++ b/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") } diff --git a/fastvideo/models/vaes/hunyuan15vae.py b/fastvideo/models/vaes/hunyuan15vae.py new file mode 100644 index 000000000..66201bd09 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/fastvideo/pipelines/basic/hunyuan15/__init__.py b/fastvideo/pipelines/basic/hunyuan15/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py new file mode 100644 index 000000000..87a58c692 --- /dev/null +++ b/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 diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index 64b7699a5..0e41765c8 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/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", diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 998d6eb09..cbe24da7b 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/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", diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 66d8c7db6..8bccd29e9 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/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] diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index 1e9454de3..612749bdc 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/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: diff --git a/fastvideo/tests/encoders/test_hyt5_encoder.py b/fastvideo/tests/encoders/test_hyt5_encoder.py new file mode 100644 index 000000000..ef7ebfa27 --- /dev/null +++ b/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()}" diff --git a/fastvideo/tests/encoders/test_qwen2_5_encoder.py b/fastvideo/tests/encoders/test_qwen2_5_encoder.py new file mode 100644 index 000000000..314d5b047 --- /dev/null +++ b/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.") + diff --git a/fastvideo/tests/vaes/test_hunyuan15_vae.py b/fastvideo/tests/vaes/test_hunyuan15_vae.py new file mode 100644 index 000000000..731ebff45 --- /dev/null +++ b/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()}" +