From 23fd3ed3c7dc5deee9f43739fe5bd060bf99be37 Mon Sep 17 00:00:00 2001 From: William Lin Date: Thu, 13 Mar 2025 22:19:41 -0700 Subject: [PATCH 01/12] [Do not merge] V1 encoders and model loading (#261) Co-authored-by: Peiyuan Zhang --- assets/prompt.txt | 2 +- fastvideo/attention/__init__.py | 4 +- .../{distributed_attn.py => flash_attn.py} | 104 +- fastvideo/attention/layer.py | 135 +++ fastvideo/distributed/parallel_state.py | 294 +++++- fastvideo/inference.py | 250 ----- fastvideo/inference_args.py | 20 + fastvideo/inference_engine.py | 9 +- fastvideo/layers/conv.py | 0 fastvideo/layers/pooler.py | 322 ------ fastvideo/layers/resampler.py | 269 ----- fastvideo/layers/rotary_embedding.py | 40 +- fastvideo/layers/vocab_parallel_embedding.py | 8 +- fastvideo/loader/fsdp_load.py | 7 +- fastvideo/loader/loader.py | 323 +++--- fastvideo/loader/utils.py | 184 ++++ fastvideo/loader/weight_utils.py | 206 ++-- fastvideo/model_loader.py | 273 ----- fastvideo/models/__init__.py | 52 +- fastvideo/models/dits/__init__.py | 4 - fastvideo/models/dits/hunyuanvideo.py | 7 +- fastvideo/models/encoders/clip.py | 236 ++++- fastvideo/models/encoders/llama.py | 418 ++++++++ fastvideo/models/encoders/utils.py | 438 ++++++++ fastvideo/models/encoders/vision.py | 6 +- fastvideo/models/flash_attn_no_pad.py | 28 - fastvideo/models/hf_transformer_utils.py | 259 +++++ fastvideo/models/hunyuan/modules/__init__.py | 42 +- .../models/hunyuan_hf/modeling_hunyuan.py | 836 --------------- .../models/hunyuan_hf/pipeline_hunyuan.py | 691 ------------- fastvideo/models/loader.py | 304 ++++++ .../mochi_hf/convert_diffusers_to_mochi.py | 361 ------- .../models/mochi_hf/mochi_latents_utils.py | 45 - fastvideo/models/mochi_hf/modeling_mochi.py | 663 ------------ fastvideo/models/mochi_hf/norm.py | 128 --- fastvideo/models/mochi_hf/pipeline_mochi.py | 757 -------------- fastvideo/models/parameter.py | 2 +- fastvideo/models/{dits => }/registry.py | 38 +- fastvideo/models/schedulers/base.py | 0 .../scheduling_flow_match_euler_discrete.py} | 47 +- fastvideo/models/stepvideo/__init__.py | 7 - fastvideo/models/stepvideo/__version__.py | 1 - fastvideo/models/stepvideo/config.py | 174 ---- .../stepvideo/diffusion/video_pipeline.py | 325 ------ .../models/stepvideo/modules/attentions.py | 96 -- fastvideo/models/stepvideo/modules/blocks.py | 296 ------ fastvideo/models/stepvideo/modules/model.py | 198 ---- .../models/stepvideo/modules/normalization.py | 312 ------ fastvideo/models/stepvideo/modules/rope.py | 90 -- fastvideo/models/stepvideo/parallel.py | 21 - .../models/stepvideo/text_encoder/__init__.py | 12 - .../models/stepvideo/text_encoder/clip.py | 36 - .../stepvideo/text_encoder/flashattention.py | 45 - .../models/stepvideo/text_encoder/stepllm.py | 291 ------ .../stepvideo/text_encoder/tokenizer.py | 209 ---- fastvideo/models/stepvideo/utils/__init__.py | 2 - .../models/stepvideo/utils/quantization.py | 117 --- fastvideo/models/stepvideo/utils/utils.py | 60 -- .../models/stepvideo/utils/video_process.py | 51 - fastvideo/models/stepvideo/vae/vae.py | 975 ------------------ fastvideo/models/text_encoder.py | 255 +++++ fastvideo/models/utils.py | 43 + fastvideo/models/vaes/common.py | 469 +++++++++ fastvideo/models/vaes/hunyuanvae.py | 799 ++++++++++++++ fastvideo/pipelines/__init__.py | 227 +++- .../composed/composed_pipeline_base.py | 8 + fastvideo/pipelines/composed/text_to_video.py | 113 -- fastvideo/pipelines/hunyuan/hunyuan_loader.py | 2 - .../pipelines/implementations/__init__.py | 8 - .../implementations/hunyuan/__init__.py | 4 +- .../implementations/hunyuan/hunyuan_loader.py | 58 +- .../hunyuan/hunyuan_pipeline.py | 80 +- fastvideo/pipelines/loader.py | 10 +- fastvideo/pipelines/pipeline_registry.py | 148 +-- fastvideo/pipelines/stages/decoding.py | 24 +- fastvideo/pipelines/stages/prompt_encoding.py | 2 +- fastvideo/sample/v1_fastvideo_inference.py | 39 +- fastvideo/sample/v2_fastvideo_inference.py | 102 -- fastvideo/v1/test_hunyuan_vae.py | 161 +++ fastvideo/v1/test_hunyuanvideo.py | 4 +- fastvideo/v1/test_hunyuanvideo_load.py | 2 +- scripts/inference/v1_inference_hunyuan.sh | 6 +- scripts/inference/v2_inference_hunyuan.sh | 21 - 83 files changed, 4967 insertions(+), 8748 deletions(-) rename fastvideo/attention/{distributed_attn.py => flash_attn.py} (60%) create mode 100644 fastvideo/attention/layer.py delete mode 100644 fastvideo/inference.py create mode 100644 fastvideo/layers/conv.py delete mode 100644 fastvideo/layers/pooler.py delete mode 100644 fastvideo/layers/resampler.py create mode 100644 fastvideo/loader/utils.py delete mode 100644 fastvideo/model_loader.py delete mode 100644 fastvideo/models/dits/__init__.py create mode 100644 fastvideo/models/encoders/llama.py create mode 100644 fastvideo/models/encoders/utils.py delete mode 100644 fastvideo/models/flash_attn_no_pad.py create mode 100644 fastvideo/models/hf_transformer_utils.py delete mode 100644 fastvideo/models/hunyuan_hf/modeling_hunyuan.py delete mode 100644 fastvideo/models/hunyuan_hf/pipeline_hunyuan.py create mode 100644 fastvideo/models/loader.py delete mode 100644 fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py delete mode 100644 fastvideo/models/mochi_hf/mochi_latents_utils.py delete mode 100644 fastvideo/models/mochi_hf/modeling_mochi.py delete mode 100644 fastvideo/models/mochi_hf/norm.py delete mode 100644 fastvideo/models/mochi_hf/pipeline_mochi.py rename fastvideo/models/{dits => }/registry.py (87%) create mode 100644 fastvideo/models/schedulers/base.py rename fastvideo/models/{stepvideo/diffusion/scheduler.py => schedulers/scheduling_flow_match_euler_discrete.py} (83%) delete mode 100644 fastvideo/models/stepvideo/__init__.py delete mode 100644 fastvideo/models/stepvideo/__version__.py delete mode 100644 fastvideo/models/stepvideo/config.py delete mode 100755 fastvideo/models/stepvideo/diffusion/video_pipeline.py delete mode 100755 fastvideo/models/stepvideo/modules/attentions.py delete mode 100755 fastvideo/models/stepvideo/modules/blocks.py delete mode 100755 fastvideo/models/stepvideo/modules/model.py delete mode 100755 fastvideo/models/stepvideo/modules/normalization.py delete mode 100755 fastvideo/models/stepvideo/modules/rope.py delete mode 100644 fastvideo/models/stepvideo/parallel.py delete mode 100644 fastvideo/models/stepvideo/text_encoder/__init__.py delete mode 100755 fastvideo/models/stepvideo/text_encoder/clip.py delete mode 100755 fastvideo/models/stepvideo/text_encoder/flashattention.py delete mode 100755 fastvideo/models/stepvideo/text_encoder/stepllm.py delete mode 100755 fastvideo/models/stepvideo/text_encoder/tokenizer.py delete mode 100755 fastvideo/models/stepvideo/utils/__init__.py delete mode 100644 fastvideo/models/stepvideo/utils/quantization.py delete mode 100755 fastvideo/models/stepvideo/utils/utils.py delete mode 100644 fastvideo/models/stepvideo/utils/video_process.py delete mode 100755 fastvideo/models/stepvideo/vae/vae.py create mode 100644 fastvideo/models/text_encoder.py create mode 100644 fastvideo/models/vaes/common.py create mode 100644 fastvideo/models/vaes/hunyuanvae.py delete mode 100644 fastvideo/pipelines/composed/text_to_video.py delete mode 100644 fastvideo/pipelines/implementations/__init__.py delete mode 100644 fastvideo/sample/v2_fastvideo_inference.py create mode 100644 fastvideo/v1/test_hunyuan_vae.py delete mode 100644 scripts/inference/v2_inference_hunyuan.sh diff --git a/assets/prompt.txt b/assets/prompt.txt index 65ce027f4..d10865e32 100644 --- a/assets/prompt.txt +++ b/assets/prompt.txt @@ -1,5 +1,5 @@ -A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature. Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting. +A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature. A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere. 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. A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1 diff --git a/fastvideo/attention/__init__.py b/fastvideo/attention/__init__.py index c0000061c..208b46685 100644 --- a/fastvideo/attention/__init__.py +++ b/fastvideo/attention/__init__.py @@ -1 +1,3 @@ -from .distributed_attn import DistributedAttention \ No newline at end of file +from .flash_attn import (DistributedAttention, TextEncoderFlashAttention, LocalAttention) + +__all__ = ["DistributedAttention", "TextEncoderFlashAttention", "LocalAttention"] \ No newline at end of file diff --git a/fastvideo/attention/distributed_attn.py b/fastvideo/attention/flash_attn.py similarity index 60% rename from fastvideo/attention/distributed_attn.py rename to fastvideo/attention/flash_attn.py index 169578e1a..7d570d522 100644 --- a/fastvideo/attention/distributed_attn.py +++ b/fastvideo/attention/flash_attn.py @@ -1,9 +1,11 @@ +from itertools import accumulate +from typing import List, Optional + import torch import torch.nn as nn -from typing import Optional from fastvideo.distributed.communication_op import sequence_model_parallel_all_to_all_4D, sequence_model_parallel_all_gather from fastvideo.distributed.parallel_state import get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size -from flash_attn import flash_attn_func +from flash_attn import flash_attn_func, flash_attn_varlen_func class DistributedAttention(nn.Module): @@ -33,6 +35,7 @@ class DistributedAttention(nn.Module): self.dropout_rate = dropout_rate self.causal = causal self.softmax_scale = softmax_scale + def forward( self, q: torch.Tensor, @@ -132,4 +135,99 @@ class LocalAttention(nn.Module): causal=self.causal ) - return output \ No newline at end of file + return output + +class TextEncoderFlashAttention: + """Simplified Flash Attention implementation for text encoders in diffusion models. + + This implementation removes KV caching and other generation-specific features, + focusing only on the encoder attention pattern needed for text encoders in + diffusion models like DiT. + """ + + def __init__( + self, + softmax_scale: Optional[float] = None, + ) -> None: + """Initialize the text encoder flash attention implementation. + + Args: + num_heads: Number of attention heads + head_size: Size of each attention head + scale: Attention scale factor. If None, defaults to 1/sqrt(head_size) + """ + self.softmax_scale = softmax_scale + + # Set defaults for optional parameters + self.alibi_slopes = None + self.sliding_window = (-1, -1) # No sliding window by default + self.logits_soft_cap = 0 # No soft cap by default + + + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + seq_lens: Optional[List[int]] = None, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass with FlashAttention for text encoder. + + Args: + query: shape = [num_tokens, num_heads, head_size] + key: shape = [num_tokens, num_kv_heads, head_size] + value: shape = [num_tokens, num_kv_heads, head_size] + seq_lens: List of sequence lengths for batched input + output: shape = [num_tokens, num_heads, head_size] + Pre-allocated output tensor (optional) + + Returns: + torch.Tensor: Output tensor with shape [num_tokens, num_heads, head_size] + """ + device = query.device + batch_size = len(seq_lens) if seq_lens is not None else 1 + + if output is None: + output = torch.empty_like(query) + + if seq_lens is None: + # Single sequence case - use standard flash attention + return flash_attn_varlen_func( + q=query, + k=key, + v=value, + softmax_scale=self.softmax_scale, + causal=False, # Non-causal for encoder + # window_size=self.sliding_window, + # alibi_slopes=self.alibi_slopes, + # softcap=self.logits_soft_cap, + out=output, + # fa_version=self.vllm_flash_attn_version, + ) + else: + assert False, "Not implemented" + # Multiple sequences case - use varlen flash attention + max_seq_len = max(seq_lens) + seq_start_loc = list(accumulate(seq_lens, initial=0)) + seq_start_loc_tensor = torch.tensor(seq_start_loc, + dtype=torch.int32, + device=device) + + return flash_attn_varlen_func( + q=query, + k=key, + v=value, + cu_seqlens_q=seq_start_loc_tensor, + cu_seqlens_k=seq_start_loc_tensor, + max_seqlen_q=max_seq_len, + max_seqlen_k=max_seq_len, + softmax_scale=self.scale, + causal=False, # Non-causal for encoder + window_size=self.sliding_window, + alibi_slopes=self.alibi_slopes, + softcap=self.logits_soft_cap, + out=output, + fa_version=self.vllm_flash_attn_version, + ) diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py new file mode 100644 index 000000000..3c1920327 --- /dev/null +++ b/fastvideo/attention/layer.py @@ -0,0 +1,135 @@ +import torch +import torch.nn as nn +from typing import Optional, Dict, Any, Type + +from fastvideo.attention.flash_attn import LocalAttention, DistributedAttention + + +class AttentionBackend: + """Base class for attention backends that provides implementation classes.""" + + @staticmethod + def get_impl_cls() -> Type: + """Returns the implementation class for the attention backend.""" + raise NotImplementedError("Subclasses must implement get_impl_cls") + + +class LocalAttentionBackend(AttentionBackend): + """Backend for local attention implementation.""" + + @staticmethod + def get_impl_cls() -> Type: + """Returns the LocalAttention implementation class.""" + return LocalAttention + + +class DistributedAttentionBackend(AttentionBackend): + """Backend for distributed attention implementation.""" + + @staticmethod + def get_impl_cls() -> Type: + """Returns the DistributedAttention implementation class.""" + return DistributedAttention + + +class AttentionLayer(nn.Module): + """Base attention layer that uses an implementation class from a backend.""" + + def __init__( + self, + attn_backend: AttentionBackend, + dropout_rate: float = 0.0, + causal: bool = False, + softmax_scale: Optional[float] = None, + **extra_impl_args + ): + super().__init__() + impl_cls = attn_backend.get_impl_cls() + self.impl = impl_cls( + dropout_rate=dropout_rate, + causal=causal, + softmax_scale=softmax_scale, + **extra_impl_args + ) + + def forward(self, *args, **kwargs): + """Forward pass delegated to the implementation.""" + return self.impl(*args, **kwargs) + + +class LocalAttentionLayer(AttentionLayer): + """Wrapper for local attention implementation.""" + + def __init__( + self, + dropout_rate: float = 0.0, + causal: bool = False, + softmax_scale: Optional[float] = None, + **extra_impl_args + ): + super().__init__( + LocalAttentionBackend(), + dropout_rate=dropout_rate, + causal=causal, + softmax_scale=softmax_scale, + **extra_impl_args + ) + + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Apply local attention. + + Args: + q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] + k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] + v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] + + Returns: + torch.Tensor: Output tensor after attention + """ + return self.impl(q, k, v) + + +class DistributedAttentionLayer(AttentionLayer): + """Wrapper for distributed attention implementation.""" + + def __init__( + self, + dropout_rate: float = 0.0, + causal: bool = False, + softmax_scale: Optional[float] = None, + **extra_impl_args + ): + super().__init__( + DistributedAttentionBackend(), + dropout_rate=dropout_rate, + causal=causal, + softmax_scale=softmax_scale, + **extra_impl_args + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + replicated_q: Optional[torch.Tensor] = None, + replicated_k: Optional[torch.Tensor] = None, + replicated_v: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Forward pass for distributed attention. + + Args: + q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] + k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] + v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] + replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens + replicated_k (Optional[torch.Tensor]): Replicated key tensor + replicated_v (Optional[torch.Tensor]): Replicated value tensor + + Returns: + Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: + - o (torch.Tensor): Output tensor after attention for the main sequence + - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided + """ + return self.impl(q, k, v, replicated_q, replicated_k, replicated_v) + diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/distributed/parallel_state.py index 467d8917d..56ef861f5 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/distributed/parallel_state.py @@ -850,7 +850,6 @@ def get_sp_group() -> GroupCoordinator: def initialize_model_parallel( tensor_model_parallel_size: int = 1, - pipeline_model_parallel_size: int = 1, sequence_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None: @@ -885,15 +884,15 @@ def initialize_model_parallel( get_world_group().device_group) # Ensure the world size is compatible with the parallelism configuration - assert world_size % (tensor_model_parallel_size * pipeline_model_parallel_size * sequence_model_parallel_size) == 0, \ - f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size}) * " \ - f"pipeline_model_parallel_size ({pipeline_model_parallel_size}) * " \ - f"sequence_model_parallel_size ({sequence_model_parallel_size})" + # assert world_size % (tensor_model_parallel_size * pipeline_model_parallel_size * sequence_model_parallel_size) == 0, \ + # f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size}) * " \ + # f"pipeline_model_parallel_size ({pipeline_model_parallel_size}) * " \ + # f"sequence_model_parallel_size ({sequence_model_parallel_size})" # Check for incompatible parallelism configurations - if sequence_model_parallel_size > 1: - assert tensor_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with tensor parallelism (TP). Please use SP=1 or TP=1." - assert pipeline_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with pipeline parallelism (PP). Please use SP=1 or PP=1." + # if sequence_model_parallel_size > 1: + # assert tensor_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with tensor parallelism (TP). Please use SP=1 or TP=1." + # assert pipeline_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with pipeline parallelism (PP). Please use SP=1 or PP=1." # Build the tensor model-parallel groups. num_tensor_model_parallel_groups: int = (world_size // @@ -933,21 +932,6 @@ def initialize_model_parallel( backend, group_name="sp") - # Build the pipeline model-parallel groups. - num_pipeline_model_parallel_groups: int = (world_size // - pipeline_model_parallel_size) - global _PP - assert _PP is None, ( - "pipeline model parallel group is already initialized") - group_ranks = [] - for i in range(num_pipeline_model_parallel_groups): - ranks = list(range(i, world_size, num_pipeline_model_parallel_groups)) - group_ranks.append(ranks) - _PP = init_model_parallel_group(group_ranks, - get_world_group().local_rank, - backend, - group_name="pp") - def get_sequence_model_parallel_world_size(): """Return world size for the sequence model parallel group.""" @@ -1172,3 +1156,267 @@ def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup], aggregated_data += rank_data return [x == 1 for x in aggregated_data.tolist()] + + +def initialize_tensor_parallel_group( + tensor_model_parallel_size: int = 1, + backend: Optional[str] = None, + group_name_suffix: str = "" +) -> GroupCoordinator: + """Initialize a tensor parallel group for a specific model. + + This function creates a tensor parallel group that can be used with the + patch_tensor_parallel_group context manager. It allows different models + to use different tensor parallelism configurations. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for tensor parallelism that can be used with + the patch_tensor_parallel_group context manager. + + Example usage: + ```python + # Initialize tensor parallel group for model1 + tp_group_model1 = initialize_tensor_parallel_group( + tensor_model_parallel_size=4, + group_name_suffix="model1" + ) + + # Use tensor parallelism for model1 + with patch_tensor_parallel_group(tp_group_model1): + # Run model1 with tensor parallelism + output1 = model1(input1) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group) + + # Ensure the world size is compatible with the parallelism configuration + assert world_size % tensor_model_parallel_size == 0, \ + f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size})" + + # Build the tensor model-parallel groups. + num_tensor_model_parallel_groups: int = (world_size // tensor_model_parallel_size) + tp_group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, + (i + 1) * tensor_model_parallel_size)) + tp_group_ranks.append(ranks) + + # Create TP group coordinator with a unique name + group_name = f"tp_{group_name_suffix}" if group_name_suffix else "tp" + tp_group = init_model_parallel_group(tp_group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name=group_name) + + return tp_group + + +def initialize_sequence_parallel_group( + sequence_model_parallel_size: int = 1, + backend: Optional[str] = None, + group_name_suffix: str = "" +) -> GroupCoordinator: + """Initialize a sequence parallel group for a specific model. + + This function creates a sequence parallel group that can be used with the + patch_sequence_parallel_group context manager. It allows different models + to use different sequence parallelism configurations. + + Arguments: + sequence_model_parallel_size: number of GPUs used for sequence model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for sequence parallelism that can be used with + the patch_sequence_parallel_group context manager. + + Example usage: + ```python + # Initialize sequence parallel group for model2 + sp_group_model2 = initialize_sequence_parallel_group( + sequence_model_parallel_size=2, + group_name_suffix="model2" + ) + + # Use sequence parallelism for model2 + with patch_sequence_parallel_group(sp_group_model2): + # Run model2 with sequence parallelism + output2 = model2(input2) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group) + + # Ensure the world size is compatible with the parallelism configuration + assert world_size % sequence_model_parallel_size == 0, \ + f"World size ({world_size}) must be divisible by sequence_model_parallel_size ({sequence_model_parallel_size})" + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = (world_size // sequence_model_parallel_size) + sp_group_ranks = [] + + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list(range(i * sequence_model_parallel_size, + (i + 1) * sequence_model_parallel_size)) + sp_group_ranks.append(ranks) + + # Create SP group coordinator with a unique name + group_name = f"sp_{group_name_suffix}" if group_name_suffix else "sp" + sp_group = init_model_parallel_group(sp_group_ranks, + get_world_group().local_rank, + backend, + group_name=group_name) + + return sp_group + + +_SP_STATE_PATCHED = False + +@contextmanager +def patch_sequence_parallel_group(sp_group: GroupCoordinator): + """Patch the sp group temporarily until this function ends. + + This method allows running a model with SP while another model uses TP. + + Args: + sp_group (GroupCoordinator): the sp group coordinator + """ + global _SP_STATE_PATCHED + assert not _SP_STATE_PATCHED, "Should not call when it's already patched" + + _SP_STATE_PATCHED = True + old_sp_group = get_sp_group() + global _SP + _SP = sp_group + try: + yield + finally: + # restore the original state + _SP_STATE_PATCHED = False + _SP = old_sp_group + + +_PP_STATE_PATCHED = False + +@contextmanager +def patch_pipeline_parallel_group(pp_group: GroupCoordinator): + """Patch the pp group temporarily until this function ends. + + This method allows running a model with a different PP configuration. + + Args: + pp_group (GroupCoordinator): the pp group coordinator + """ + global _PP_STATE_PATCHED + assert not _PP_STATE_PATCHED, "Should not call when it's already patched" + + _PP_STATE_PATCHED = True + old_pp_group = get_pp_group() + global _PP + _PP = pp_group + try: + yield + finally: + # restore the original state + _PP_STATE_PATCHED = False + _PP = old_pp_group + + +# Example of how to use the independent parallelism functions +""" +Here's a complete example of how to use the independent parallelism functions +for different models: + +```python +import torch +from fastvideo.distributed.parallel_state import ( + init_distributed_environment, + initialize_tensor_parallel_group, + initialize_sequence_parallel_group, + patch_tensor_parallel_group, + patch_sequence_parallel_group, + destroy_model_parallel, + destroy_distributed_environment +) + +# Initialize the distributed environment +init_distributed_environment( + world_size=8, + rank=torch.distributed.get_rank(), + distributed_init_method="tcp://localhost:12345", + local_rank=torch.distributed.get_rank() % torch.cuda.device_count() +) + +try: + # Create a tensor parallel group for model1 with TP=4 + tp_group_model1 = initialize_tensor_parallel_group( + tensor_model_parallel_size=4, + group_name_suffix="model1" + ) + + # Create a sequence parallel group for model2 with SP=2 + sp_group_model2 = initialize_sequence_parallel_group( + sequence_model_parallel_size=2, + group_name_suffix="model2" + ) + + # Create another tensor parallel group for model3 with TP=2 + tp_group_model3 = initialize_tensor_parallel_group( + tensor_model_parallel_size=2, + group_name_suffix="model3" + ) + + # Use model1 with tensor parallelism + with patch_tensor_parallel_group(tp_group_model1): + # Inside this context, get_tp_group() returns tp_group_model1 + # Run model1 with tensor parallelism + output1 = model1(input1) + + # Use model2 with sequence parallelism + with patch_sequence_parallel_group(sp_group_model2): + # Inside this context, get_sp_group() returns sp_group_model2 + # Run model2 with sequence parallelism + output2 = model2(input2) + + # Use model3 with a different tensor parallelism configuration + with patch_tensor_parallel_group(tp_group_model3): + # Inside this context, get_tp_group() returns tp_group_model3 + # Run model3 with tensor parallelism + output3 = model3(input3) + + # You can switch between models as needed + with patch_tensor_parallel_group(tp_group_model1): + # Back to using model1 + more_output1 = model1(more_input1) + +finally: + # Clean up + destroy_model_parallel() + destroy_distributed_environment() +``` + +This approach allows you to: +1. Create separate parallel groups for each model +2. Use different parallelism strategies for different models +3. Switch between models as needed +4. Use unique group names to avoid conflicts + +Note that each model can use its own optimal parallelism strategy without +interfering with other models. +""" diff --git a/fastvideo/inference.py b/fastvideo/inference.py deleted file mode 100644 index 45ee70c06..000000000 --- a/fastvideo/inference.py +++ /dev/null @@ -1,250 +0,0 @@ -import os -import random -import time -from pathlib import Path -from typing import Dict, List, Optional, Union, Any - -import torch -from safetensors.torch import load_file as safetensors_load_file -import json - -from fastvideo.models.hunyuan.constants import NEGATIVE_PROMPT, PRECISION_TO_TYPE, PROMPT_TEMPLATE -# from fastvideo.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline -from fastvideo.pipelines.pipeline_base import DiffusionPipelineBase -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler -# from fastvideo.models.hunyuan.modules import load_model -# from fastvideo.models.hunyuan.text_encoder import TextEncoder -from fastvideo.models.hunyuan.utils.data_utils import align_to -# from fastvideo.models.hunyuan.vae import load_vae -# from fastvideo.utils.parallel_states import nccl_info -from fastvideo.inference_args import InferenceArgs -from fastvideo.platforms import current_platform -from fastvideo.logger import init_logger -from fastvideo.pipelines.loader import PipelineLoader, get_pipeline_loader -from fastvideo.distributed.parallel_state import get_sp_group - -logger = init_logger(__name__) - - -class DiffusionInference: - """ - Unified inference class that works with any diffusion pipeline. - This combines model loading, pipeline creation, and inference in a flexible way. - """ - - def __init__(self, - args: InferenceArgs, - pipeline: DiffusionPipelineBase, - default_negative_prompt: str = NEGATIVE_PROMPT): - """ - Initialize the inference class with a pipeline and args. - - Args: - args: The inference arguments - pipeline: The diffusion pipeline to use - default_negative_prompt: The default negative prompt to use - """ - self.args = args - self.pipeline = pipeline - self.default_negative_prompt = default_negative_prompt - - - - @classmethod - def load_pipeline(cls, inference_args: InferenceArgs): - """ - Create an inference instance from pretrained model components. - - Args: - inference_args: The inference arguments containing model path and other settings - - Returns: - A DiffusionInference instance ready for inference - """ - logger.debug("start loading pipeline") - models_root_path = Path(inference_args.model_path) - if not models_root_path.exists(): - raise ValueError(f"`models_root` not exists: {models_root_path}") - - # Create save folder to save the samples - save_path = inference_args.output_path - os.makedirs(os.path.dirname(save_path), exist_ok=True) - - # Log which platform we're using - print(f"Using platform: {current_platform.__class__.__name__}") - assert current_platform.is_cuda(), "CUDA is not available" - - # init distributed - # _initialize_distributed() - - - # get loader for the pipeline - - # Create pipeline using the components - pipeline_loader = get_pipeline_loader(inference_args) - pipeline = pipeline_loader.load_pipeline( - inference_args=inference_args, - ) - - return cls(inference_args, pipeline, default_negative_prompt=NEGATIVE_PROMPT) - - def predict(self, prompt: str, inference_args: InferenceArgs) -> Dict[str, Any]: - """ - Run inference with the pipeline. - - Returns: - Dictionary with generated samples and metadata - """ - - out_dict = dict() - - batch_size = inference_args.batch_size - num_videos_per_prompt = inference_args.num_videos - seed = inference_args.seed - height = inference_args.height - width = inference_args.width - video_length = inference_args.num_frames - # prompt = inference_args.prompt - negative_prompt = inference_args.neg_prompt - infer_steps = inference_args.num_inference_steps - guidance_scale = inference_args.guidance_scale - flow_shift = inference_args.flow_shift - # mask_strategy = inference_args.mask_strategy - embedded_guidance_scale = inference_args.embedded_cfg_scale - - - # ======================================================================== - # Arguments: seed - # ======================================================================== - if isinstance(seed, torch.Tensor): - seed = seed.tolist() - if seed is None: - seeds = [random.randint(0, 1_000_000) for _ in range(batch_size * num_videos_per_prompt)] - elif isinstance(seed, int): - seeds = [seed + i for _ in range(batch_size) for i in range(num_videos_per_prompt)] - elif isinstance(seed, (list, tuple)): - if len(seed) == batch_size: - seeds = [int(seed[i]) + j for i in range(batch_size) for j in range(num_videos_per_prompt)] - elif len(seed) == batch_size * num_videos_per_prompt: - seeds = [int(s) for s in seed] - else: - raise ValueError( - f"Length of seed must be equal to number of prompt(batch_size) or " - f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}.") - else: - raise ValueError(f"Seed must be an integer, a list of integers, or None, got {seed}.") - # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... - generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds] - out_dict["seeds"] = seeds - - # ======================================================================== - # Arguments: target_width, target_height, target_video_length - # ======================================================================== - if width <= 0 or height <= 0 or video_length <= 0: - raise ValueError( - f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}" - ) - if (video_length - 1) % 4 != 0: - raise ValueError(f"`video_length-1` must be a multiple of 4, got {video_length}") - - logger.info(f"Input (height, width, video_length) = ({height}, {width}, {video_length})") - - target_height = align_to(height, 16) - target_width = align_to(width, 16) - target_video_length = video_length - - out_dict["size"] = (target_height, target_width, target_video_length) - - # ======================================================================== - # Arguments: prompt, new_prompt, negative_prompt - # ======================================================================== - if not isinstance(prompt, str): - raise TypeError(f"`prompt` must be a string, but got {type(prompt)}") - prompt = [prompt.strip()] - - # negative prompt - if negative_prompt is None or negative_prompt == "": - negative_prompt = self.default_negative_prompt - if not isinstance(negative_prompt, str): - raise TypeError(f"`negative_prompt` must be a string, but got {type(negative_prompt)}") - negative_prompt = [negative_prompt.strip()] - - # ======================================================================== - # Scheduler - # ======================================================================== - scheduler = FlowMatchDiscreteScheduler( - shift=flow_shift, - reverse=self.args.flow_reverse, - solver=self.args.flow_solver, - ) - self.pipeline.scheduler = scheduler - - if "884" in self.args.vae: - latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8] - elif "888" in self.args.vae: - latents_size = [(video_length - 1) // 8 + 1, height // 8, width // 8] - n_tokens = latents_size[0] * latents_size[1] * latents_size[2] - - # ======================================================================== - # Print infer args - # ======================================================================== - debug_str = f""" - height: {target_height} - width: {target_width} - video_length: {target_video_length} - prompt: {prompt} - neg_prompt: {negative_prompt} - seed: {seed} - infer_steps: {infer_steps} - num_videos_per_prompt: {num_videos_per_prompt} - guidance_scale: {guidance_scale} - n_tokens: {n_tokens} - flow_shift: {flow_shift} - embedded_guidance_scale: {embedded_guidance_scale}""" - logger.info(debug_str) - # return - # sp_group = get_sp_group() - # local_rank = sp_group.rank - import os - local_rank = os.environ.get("LOCAL_RANK", 0) - device = torch.device(f"cuda:{local_rank}") - batch = ForwardBatch( - prompt=prompt, - negative_prompt=negative_prompt, - num_videos_per_prompt=num_videos_per_prompt, - height=inference_args.height, - width=inference_args.width, - num_frames=inference_args.num_frames, - num_inference_steps=inference_args.num_inference_steps, - guidance_scale=inference_args.guidance_scale, - generator=generator, - eta=0.0, - n_tokens=n_tokens, - data_type="video" if inference_args.num_frames > 1 else "image", - device=device, - extra={}, # Any additional parameters - ) - - print('===============================================') - print(batch) - print('===============================================') - print('===============================================') - print('===============================================') - print(inference_args) - - # ======================================================================== - # Pipeline inference - # ======================================================================== - start_time = time.time() - samples = self.pipeline.forward( - batch=batch, - inference_args=inference_args, - )[0] - out_dict["samples"] = samples - out_dict["prompts"] = prompt - - gen_time = time.time() - start_time - logger.info(f"Success, time: {gen_time}") - - return out_dict \ No newline at end of file diff --git a/fastvideo/inference_args.py b/fastvideo/inference_args.py index 6f0c51022..df8678294 100644 --- a/fastvideo/inference_args.py +++ b/fastvideo/inference_args.py @@ -32,6 +32,10 @@ class InferenceArgs: dit_weight: Optional[str] = None model_dir: Optional[str] = None + # HuggingFace specific parameters + trust_remote_code: bool = False + revision: Optional[str] = None + # Parallelism tp_size: int = 1 sp_size: int = 1 @@ -126,6 +130,8 @@ class InferenceArgs: output_path: str = "outputs/" seed: int = 1024 seeds: Optional[List[int]] = None + device_str: Optional[str] = None + device = None def __post_init__(self): pass @@ -155,6 +161,20 @@ class InferenceArgs: type=str, help="Directory containing StepVideo model", ) + + # HuggingFace specific parameters + parser.add_argument( + "--trust-remote-code", + action="store_true", + default=InferenceArgs.trust_remote_code, + help="Trust remote code when loading HuggingFace models", + ) + parser.add_argument( + "--revision", + type=str, + default=InferenceArgs.revision, + help="The specific model version to use (can be a branch name, tag name, or commit id)", + ) # Parallelism parser.add_argument( diff --git a/fastvideo/inference_engine.py b/fastvideo/inference_engine.py index 6d0315273..2d526d7ec 100644 --- a/fastvideo/inference_engine.py +++ b/fastvideo/inference_engine.py @@ -72,6 +72,11 @@ class InferenceEngine: """ try: logger.info(f"Building pipeline...") + # TODO(will): probably a better place to set device_str? + local_rank = os.environ.get("LOCAL_RANK", 0) + device_str = f"cuda:{local_rank}" + inference_args.device_str = device_str + inference_args.device = torch.device(device_str) # TODO(will): I don't really like this api. # it should be something closer to pipeline_cls.from_pretrained(...) # this way for training we can just do pipeline_cls.from_pretrained( @@ -190,9 +195,7 @@ class InferenceEngine: # return # sp_group = get_sp_group() # local_rank = sp_group.rank - import os - local_rank = os.environ.get("LOCAL_RANK", 0) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(inference_args.device_str) batch = ForwardBatch( prompt=prompt, negative_prompt=negative_prompt, diff --git a/fastvideo/layers/conv.py b/fastvideo/layers/conv.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/layers/pooler.py b/fastvideo/layers/pooler.py deleted file mode 100644 index 0012636ef..000000000 --- a/fastvideo/layers/pooler.py +++ /dev/null @@ -1,322 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from enum import IntEnum -from typing import List, Optional, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F -from transformers import PretrainedConfig -from typing_extensions import assert_never - -from vllm.config import PoolerConfig -from vllm.model_executor.pooling_metadata import (PoolingMetadata, - PoolingTensors) -from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput -from vllm.transformers_utils.config import ( - get_cross_encoder_activation_function) - - -class PoolingType(IntEnum): - """Enumeration for different types of pooling methods.""" - LAST = 0 - ALL = 1 - CLS = 2 - STEP = 3 - MEAN = 4 - - -class SimplePooler(nn.Module): - """A layer that pools specific information from hidden states. - - This layer does the following: - 1. Extracts specific tokens or aggregates data based on pooling method. - 2. Normalizes output if specified. - 3. Returns structured results as `PoolerOutput`. - - Attributes: - pooling_type: The type of pooling to use. - normalize: Whether to normalize the pooled data. - """ - - @staticmethod - def from_pooling_type( - pooling_type: PoolingType, - *, - normalize: bool, - softmax: bool, - step_tag_id: Optional[int] = None, - returned_token_ids: Optional[List[int]] = None, - ) -> "SimplePooler": - if pooling_type == PoolingType.LAST: - assert step_tag_id is None and returned_token_ids is None - return LastPool(normalize=normalize, softmax=softmax) - if pooling_type == PoolingType.ALL: - assert step_tag_id is None and returned_token_ids is None - return AllPool(normalize=normalize, softmax=softmax) - if pooling_type == PoolingType.CLS: - assert step_tag_id is None and returned_token_ids is None - return CLSPool(normalize=normalize, softmax=softmax) - if pooling_type == PoolingType.MEAN: - assert step_tag_id is None and returned_token_ids is None - return MeanPool(normalize=normalize, softmax=softmax) - if pooling_type == PoolingType.STEP: - return StepPool(normalize=normalize, - softmax=softmax, - step_tag_id=step_tag_id, - returned_token_ids=returned_token_ids) - - assert_never(pooling_type) - - def __init__(self, *, normalize: bool, softmax: bool) -> None: - super().__init__() - - self.head = PoolerHead(normalize=normalize, softmax=softmax) - - def get_prompt_lens( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> torch.Tensor: - return PoolingTensors.from_pooling_metadata( - pooling_metadata, hidden_states.device).prompt_lens - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - raise NotImplementedError - - def build_output(self, data: torch.Tensor) -> PoolingSequenceGroupOutput: - return PoolingSequenceGroupOutput(data) - - def forward( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput: - pooled_data = self.extract_states(hidden_states, pooling_metadata) - pooled_data = self.head(pooled_data) - pooled_outputs = [self.build_output(data) for data in pooled_data] - return PoolerOutput(outputs=pooled_outputs) - - -class CLSPool(SimplePooler): - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - first_token_flat_indices = torch.zeros_like(prompt_lens) - first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1] - return hidden_states[first_token_flat_indices] - - -class LastPool(SimplePooler): - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - last_token_flat_indices = torch.cumsum(prompt_lens, dim=0) - 1 - return hidden_states[last_token_flat_indices] - - -class AllPool(SimplePooler): - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - offset = 0 - pooled_data = list[torch.Tensor]() - for prompt_len in prompt_lens: - pooled_data.append(hidden_states[offset:offset + prompt_len]) - offset += prompt_len - - return pooled_data - - -class MeanPool(SimplePooler): - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - cumsum = torch.cumsum(hidden_states, dim=0) - start_indices = torch.cat([ - torch.tensor([0], device=hidden_states.device), - torch.cumsum(prompt_lens[:-1], dim=0) - ]) - end_indices = torch.cumsum(prompt_lens, dim=0) - return (cumsum[end_indices - 1] - cumsum[start_indices] + - hidden_states[start_indices]) / prompt_lens.unsqueeze(1) - - -class StepPool(SimplePooler): - - def __init__( - self, - *, - normalize: bool, - softmax: bool, - step_tag_id: Optional[int] = None, - returned_token_ids: Optional[List[int]] = None, - ): - super().__init__(normalize=normalize, softmax=softmax) - - self.step_tag_id = step_tag_id - self.returned_token_ids = returned_token_ids - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - returned_token_ids = self.returned_token_ids - if returned_token_ids is not None and len(returned_token_ids) > 0: - hidden_states = hidden_states[:, returned_token_ids] - - step_tag_id = self.step_tag_id - - offset = 0 - pooled_data = list[torch.Tensor]() - for prompt_len, seq_data_i in zip(prompt_lens, - pooling_metadata.seq_data.values()): - pooled_data_i = hidden_states[offset:offset + prompt_len] - if step_tag_id is not None: - token_ids = torch.tensor(seq_data_i.prompt_token_ids) - pooled_data_i = pooled_data_i[token_ids == step_tag_id] - - offset += prompt_len - pooled_data.append(pooled_data_i) - - return pooled_data - - -class PoolerHead(nn.Module): - - def __init__(self, *, normalize: bool, softmax: bool) -> None: - super().__init__() - - self.normalize = normalize - self.softmax = softmax - - def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor]): - if self.normalize: - if isinstance(pooled_data, list): - pooled_data = [ - F.normalize(data, p=2, dim=1) for data in pooled_data - ] - else: - pooled_data = F.normalize(pooled_data, p=2, dim=1) - - if self.softmax: - if isinstance(pooled_data, list): - pooled_data = [F.softmax(data, dim=-1) for data in pooled_data] - else: - pooled_data = F.softmax(pooled_data, dim=-1) - - return pooled_data - - -class Pooler(nn.Module): - - @classmethod - def from_config_with_defaults( - cls, - pooler_config: PoolerConfig, - pooling_type: PoolingType, - normalize: bool, - softmax: bool, - step_tag_id: Optional[int] = None, - returned_token_ids: Optional[List[int]] = None, - ) -> SimplePooler: - return SimplePooler.from_pooling_type( - pooling_type=PoolingType[pooler_config.pooling_type] - if pooler_config.pooling_type is not None else pooling_type, - normalize=pooler_config.normalize - if pooler_config.normalize is not None else normalize, - softmax=pooler_config.softmax - if pooler_config.softmax is not None else softmax, - step_tag_id=pooler_config.step_tag_id - if pooler_config.step_tag_id is not None else step_tag_id, - returned_token_ids=pooler_config.returned_token_ids - if pooler_config.returned_token_ids is not None else - returned_token_ids, - ) - - -class CrossEncodingPooler(nn.Module): - """A layer that pools specific information from hidden states. - - This layer does the following: - 1. Extracts specific tokens or aggregates data based on pooling method. - 2. Normalizes output if specified. - 3. Returns structured results as `PoolerOutput`. - - Attributes: - pooling_type: The type of pooling to use. - normalize: Whether to normalize the pooled data. - """ - - def __init__( - self, - config: PretrainedConfig, - classifier: nn.Module, - pooler: Optional[nn.Module] = None, - ): - super().__init__() - self.classifier = classifier - self.pooler = pooler - self.default_activation_function = \ - get_cross_encoder_activation_function(config) - - def forward( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput: - """Pools sentence pair scores from the hidden_states.""" - - prompt_lens = PoolingTensors.from_pooling_metadata( - pooling_metadata, hidden_states.device).prompt_lens - - offset = 0 - pooled_data_lst = [] - for prompt_len in prompt_lens: - pooled_data_i = hidden_states[offset:offset + prompt_len] - - if self.pooler is not None: - final_shape_tensor = self.pooler(pooled_data_i) - else: - final_shape_tensor = self.classifier(pooled_data_i) - - pooled_data_lst.append(final_shape_tensor) - offset += prompt_len - - pooled_output = torch.stack(pooled_data_lst) - - if self.pooler is not None: - # apply classifier once on the full batch if possible - pooled_output = self.classifier(pooled_output) - - scores = self.default_activation_function(pooled_output).squeeze(-1) - - pooled_outputs = [PoolingSequenceGroupOutput(data) for data in scores] - return PoolerOutput(outputs=pooled_outputs) diff --git a/fastvideo/layers/resampler.py b/fastvideo/layers/resampler.py deleted file mode 100644 index 4c9860006..000000000 --- a/fastvideo/layers/resampler.py +++ /dev/null @@ -1,269 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -# Adapted from -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py -# https://huggingface.co/Qwen/Qwen-7B/blob/main/modeling_qwen.py -# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20 -# -# Copyright 2023 The Qwen team. -# Copyright 2023 The vLLM team. -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# 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. -""" -Shared resampler perceiver network used in multimodal models and -related helpers for sincos positional embeddings. - -Example models: Qwen (Qwen-VL), MiniCPM-V 2.0 -""" -import math -from functools import partial -from typing import Callable, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn.functional as F -from torch import nn - -from vllm.model_executor.layers.linear import ReplicatedLinear -from vllm.model_executor.layers.quantization import QuantizationConfig - -DEFAULT_LN = partial(nn.LayerNorm, eps=1e-6) - - -def get_abs_pos(abs_pos: torch.Tensor, tgt_size: Union[torch.Tensor, - int]) -> torch.Tensor: - # abs_pos: L, C - # tgt_size: (H, W) - # return: M, C - src_size = int(math.sqrt(abs_pos.size(0))) - dtype = abs_pos.dtype - if isinstance(tgt_size, int): - tgt_size = (tgt_size, tgt_size) - if (src_size == tgt_size[0] and src_size == tgt_size[1]): - return abs_pos - return (F.interpolate( - abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2), - size=(tgt_size[0], tgt_size[1]), - mode="bicubic", - align_corners=False, - ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype)) - - -# sin/cos positional embedding helpers are adapted from: -# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20 -def get_1d_sincos_pos_embed_from_grid( - embed_dim: int, pos: np.ndarray, - version: Tuple[int, int] = (2, 0)) -> torch.Tensor: - """ - embed_dim: output dimension for each position - pos: a list of positions to be encoded: size (M,) / (H, W) - out: (M, D) / (H, W, D) - """ - assert embed_dim % 2 == 0 - omega = np.arange(embed_dim // 2, dtype=np.float32) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - - if version == (2, 0): - pos = pos.reshape(-1) # (M,) - out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product - emb_sin = np.sin(out) # (M, D/2) - emb_cos = np.cos(out) # (M, D/2) - emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) - else: - out = np.einsum("hw,d->hwd", pos, omega) # (H, W, D/2), outer product - emb_sin = np.sin(out) # (H, W, D/2) - emb_cos = np.cos(out) # (H, W, D/2) - emb = np.concatenate([emb_sin, emb_cos], axis=-1) # (H, W, D) - return emb - - -def get_2d_sincos_pos_embed_from_grid( - embed_dim: int, grid: np.ndarray, - version: Tuple[int, int] = (2, 0)) -> torch.Tensor: - assert embed_dim % 2 == 0 - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid( - embed_dim // 2, grid[0], version) # (H*W, D/2) or (H, W, D/2) - emb_w = get_1d_sincos_pos_embed_from_grid( - embed_dim // 2, grid[1], version) # (H*W, D/2) or (H, W, D/2) - - if version == (2, 0): - emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) - else: - emb = np.concatenate([emb_h, emb_w], axis=-1) # (H, W, D) - return emb - - -def get_2d_sincos_pos_embed( - embed_dim: int, - grid_size: Union[int, Tuple[int, int]], - cls_token: bool = False, - version: Tuple[int, int] = (2, 0), -) -> torch.Tensor: - """ - grid_size: int of the grid height and width - return: - pos_embed: [grid_size*grid_size, embed_dim] or - [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) - """ - if isinstance(grid_size, int): - grid_h_size, grid_w_size = grid_size, grid_size - else: - grid_h_size, grid_w_size = grid_size[0], grid_size[1] - - grid_h = np.arange(grid_h_size, dtype=np.float32) - grid_w = np.arange(grid_w_size, dtype=np.float32) - grid = np.meshgrid(grid_w, grid_h) # here w goes first - grid = np.stack(grid, axis=0) - assert isinstance(grid, np.ndarray) and \ - grid.shape == (2, grid_h_size, grid_w_size) - - if version == (2, 0): - grid = grid.reshape([2, 1, grid_h_size, grid_w_size]) - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, version) - if cls_token: - pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], - axis=0) - else: - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, version) - return pos_embed - - -class BaseResampler(nn.Module): - """ - A 2D perceiver-resampler network with one cross attention layers by - (grid_size**2) learnable queries and 2d sincos pos_emb. - Outputs: - A tensor with the shape of (grid_size**2, embed_dim) - """ - - def __init__(self, - num_queries: int, - embed_dim: int, - num_heads: int, - kv_dim: Optional[int] = None, - norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN, - do_post_projection: bool = True, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "") -> None: - super().__init__() - - self.num_queries = num_queries - self.embed_dim = embed_dim - self.num_heads = num_heads - - self.query = nn.Parameter(torch.empty(self.num_queries, embed_dim)) - - if kv_dim is not None and kv_dim != embed_dim: - self.kv_proj = ReplicatedLinear(kv_dim, - embed_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.kv_proj") - else: - # Maintain the same return value with ReplicatedLinear.forward - self.kv_proj = lambda *args, **kwargs: ( # type: ignore # noqa - nn.Identity()(*args, **kwargs), - None, - ) - self.attn = nn.MultiheadAttention(embed_dim, num_heads) - self.ln_q = norm_layer(embed_dim) - self.ln_kv = norm_layer(embed_dim) - self.do_post_projection = do_post_projection - self.ln_post = norm_layer(embed_dim) if do_post_projection else None - self.proj = nn.Parameter( - (embed_dim**-0.5) * - torch.empty(embed_dim, embed_dim)) if do_post_projection else None - - def _repeat(self, query, N: int): - return query.unsqueeze(1).repeat(1, N, 1) - - -class Resampler2(BaseResampler): - """Resampler-perceiver network to be used for a variety of model types, - e.g., Qwen-vl / Minicpmv 2.0. The main difference is the addition of the - do_post_projection arg, which indicates whether or not there should be - a post layer normalization and projector after the attention. This is - present in minicpmv2.0, but not qwen-vl. - """ - - def __init__(self, - grid_size: int, - embed_dim: int, - num_heads: int, - kv_dim: Optional[int] = None, - norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN, - adaptive: bool = False, - do_post_projection: bool = True, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "") -> None: - super().__init__(grid_size**2, - embed_dim, - num_heads, - kv_dim, - norm_layer, - do_post_projection=do_post_projection, - quant_config=quant_config, - prefix=prefix) - - self.adaptive = adaptive - pos_embed_arr = get_2d_sincos_pos_embed(embed_dim, - grid_size, - version=(2, 0)) - - self.pos_embed = nn.Parameter( - torch.from_numpy(pos_embed_arr).requires_grad_(False)) - - def forward( - self, - x: torch.Tensor, - tgt_sizes: Optional[torch.Tensor] = None, - attn_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if tgt_sizes is None: - tgt_sizes = int(math.sqrt(x.size(1))) - if self.adaptive: - pos_embed_arr = get_2d_sincos_pos_embed(self.embed_dim, - tgt_sizes, - version=(2, 0)) - pos_embed = torch.from_numpy(pos_embed_arr).to(device=x.device, - dtype=x.dtype) - else: - pos_embed = get_abs_pos(self.pos_embed, - tgt_sizes).to(device=x.device, - dtype=x.dtype) - - x, _ = self.kv_proj(x) - x = self.ln_kv(x).permute(1, 0, 2) - - N = x.shape[1] - q = self.ln_q(self.query) - out = self.attn( - self._repeat(q, N) + self.pos_embed.unsqueeze(1), - x + pos_embed.unsqueeze(1), - x, - attn_mask=attn_mask, - )[0] - x = out.permute(1, 0, 2) - if self.do_post_projection: - x = self.ln_post(x) - x = x @ self.proj - return x diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index fe9aca306..974cd23cf 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -500,4 +500,42 @@ def get_rotary_pos_embed( sp_rank=sp_rank, sp_world_size=sp_world_size ) - return freqs_cos, freqs_sin \ No newline at end of file + return freqs_cos, freqs_sin + +_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {} + +def get_rope( + head_size: int, + rotary_dim: int, + max_position: int, + base: int, + is_neox_style: bool = True, + rope_scaling: Optional[Dict[str, Any]] = None, + dtype: Optional[torch.dtype] = None, + partial_rotary_factor: float = 1.0, +) -> RotaryEmbedding: + if dtype is None: + dtype = torch.get_default_dtype() + if rope_scaling is not None: + # Transforms every value that is a list into a tuple for caching calls + rope_scaling_tuple = { + k: tuple(v) if isinstance(v, list) else v + for k, v in rope_scaling.items() + } + rope_scaling_args = tuple(rope_scaling_tuple.items()) + else: + rope_scaling_args = None + if partial_rotary_factor < 1.0: + rotary_dim = int(rotary_dim * partial_rotary_factor) + key = (head_size, rotary_dim, max_position, base, is_neox_style, + rope_scaling_args, dtype) + if key in _ROPE_DICT: + return _ROPE_DICT[key] + + if rope_scaling is None: + rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base, + is_neox_style, dtype) + else: + raise ValueError(f"Unknown RoPE scaling {rope_scaling}") + _ROPE_DICT[key] = rotary_emb + return rotary_emb diff --git a/fastvideo/layers/vocab_parallel_embedding.py b/fastvideo/layers/vocab_parallel_embedding.py index f65dfc3cb..44d7793ab 100644 --- a/fastvideo/layers/vocab_parallel_embedding.py +++ b/fastvideo/layers/vocab_parallel_embedding.py @@ -7,14 +7,14 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter, UninitializedParameter -from vllm.distributed import (divide, get_tensor_model_parallel_rank, +from fastvideo.distributed import (divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce) from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding) -from vllm.model_executor.parameter import BasevLLMParameter -from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform +from fastvideo.models.parameter import BasevLLMParameter +from fastvideo.models.utils import set_weight_attrs +from fastvideo.platforms import current_platform DEFAULT_VOCAB_PADDING_SIZE = 64 diff --git a/fastvideo/loader/fsdp_load.py b/fastvideo/loader/fsdp_load.py index 10e7753d2..5e2a0b3ac 100644 --- a/fastvideo/loader/fsdp_load.py +++ b/fastvideo/loader/fsdp_load.py @@ -3,7 +3,7 @@ # Copyright 2025 The FastVideo Authors. from collections import defaultdict -from typing import Any, Callable, Dict, Generator, List, Optional, Tuple +from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Type from itertools import chain import torch from torch import nn @@ -13,7 +13,7 @@ from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed._tensor import distribute_tensor from torch.nn.modules.module import _IncompatibleKeys from vllm.model_executor.model_loader.weight_utils import safetensors_weights_iterator -from fastvideo.models.dits.registry import ModelRegistry + import contextlib import re @@ -79,14 +79,13 @@ def get_param_names_mapping(mapping_dict: Dict[str, str]) -> Callable[[str], str # TODO(PY): add compile option def load_fsdp_model( - model_name: str, + model_cls: Type[nn.Module], init_params: Dict[str, Any], weight_dir_list: List[str], device: torch.device, cpu_offload: bool = False, default_dtype: Optional[torch.dtype] = torch.bfloat16, ) -> torch.nn.Module: - model_cls, _ = ModelRegistry.resolve_model_cls(model_name) with set_default_dtype(default_dtype), torch.device("meta"): model = model_cls(**init_params) device_mesh = init_device_mesh( diff --git a/fastvideo/loader/loader.py b/fastvideo/loader/loader.py index 1db66e20e..de28e5d63 100644 --- a/fastvideo/loader/loader.py +++ b/fastvideo/loader/loader.py @@ -21,48 +21,22 @@ from contextlib import contextmanager from typing import (Any, Callable, Dict, Generator, Iterable, List, Optional, Tuple, cast) -import gguf import huggingface_hub -import numpy as np import torch -from huggingface_hub import HfApi from torch import nn -from transformers import AutoModelForCausalLM from transformers.utils import SAFE_WEIGHTS_INDEX_NAME -from vllm.attention import Attention -from vllm.config import (LoadConfig, LoadFormat, ModelConfig, ParallelConfig, - VllmConfig, set_current_vllm_config) -from vllm.distributed import (get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size) -from vllm.envs import VLLM_USE_MODELSCOPE -from vllm.logger import init_logger -from vllm.model_executor.layers.linear import (LinearBase, - MergedColumnParallelLinear, - QKVParallelLinear, - ReplicatedLinear, - RowParallelLinear) -from vllm.model_executor.layers.quantization.base_config import ( - QuantizeMethodBase) -from vllm.model_executor.model_loader.tensorizer import ( - TensorizerConfig, is_vllm_tensorized, load_with_tensorizer, - serialize_vllm_model, tensorizer_weights_iterator) -from vllm.model_executor.model_loader.utils import (ParamMapping, - configure_quant_config, - get_model_architecture, - set_default_torch_dtype) -from vllm.model_executor.model_loader.weight_utils import ( +from fastvideo.logger import init_logger +from fastvideo.loader.utils import set_default_torch_dtype +from fastvideo.loader.weight_utils import ( download_safetensors_index_file_from_hf, download_weights_from_hf, filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, - get_gguf_extra_tensor_names, get_lock, gguf_quant_weights_iterator, - initialize_dummy_weights, np_cache_weights_iterator, pt_weights_iterator, - runai_safetensors_weights_iterator, safetensors_weights_iterator) -from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform -from vllm.transformers_utils.s3_utils import glob as s3_glob -from vllm.transformers_utils.utils import is_s3 -from vllm.utils import is_pin_memory_available + pt_weights_iterator, + safetensors_weights_iterator) +from fastvideo.platforms import current_platform +from fastvideo.models.registry import ModelRegistry +from fastvideo.inference_args import InferenceArgs @contextmanager def device_loading_context(module: torch.nn.Module, @@ -86,7 +60,7 @@ def device_loading_context(module: torch.nn.Module, finally: # Restore parameters to their original devices, ignoring new parameters - pin_memory = is_pin_memory_available() + pin_memory = True for name, p in module.named_parameters(): if name in original_device_states: original_device: torch.device = original_device_states[name] @@ -111,90 +85,78 @@ logger = init_logger(__name__) def _initialize_model( - vllm_config: VllmConfig, - *, - prefix: str = "", + model_path: str, + model_config, + inference_args: InferenceArgs, ) -> nn.Module: - """Initialize a model with the given configurations.""" - model_config = vllm_config.model_config - model_class, _ = get_model_architecture(model_config) + """Initialize a model with the given configurations. + + Args: + model_name: The name of the model architecture + model_path: Path to the model directory + inference_args: Inference arguments + component_name: Optional component name (e.g., "text_encoder", "vae") + If provided, will load the config from the component subdirectory + + Returns: + The initialized model + """ + architectures = getattr(model_config, "architectures", []) + if len(architectures) == 0: + raise ValueError("Model config does not contain a valid model architecture") + logger.info("Trying to load model from diffusers format") + class_name = model_config.pop("_class_name") + architectures = [class_name] + if architectures is None: + raise ValueError("Model config does not contain a valid model architecture") + model_cls, _ = ModelRegistry.resolve_model_cls(architectures) + logger.info(f"Loading model_cls: {model_cls}") - if vllm_config.quant_config is not None: - configure_quant_config(vllm_config.quant_config, model_class) - - signatures = inspect.signature(model_class.__init__) - all_params = [param.name for param in signatures.parameters.values()] - if "vllm_config" in all_params and "prefix" in all_params: - # new-style model class - with set_current_vllm_config(vllm_config, check_compile=True): - return model_class(vllm_config=vllm_config, prefix=prefix) - - msg = ("vLLM model class should accept `vllm_config` and `prefix` as " - "input arguments. Possibly you have an old-style model class" - " registered from out of tree and it is used for new vLLM version. " - "Check https://docs.vllm.ai/en/latest/design/arch_overview.html " - "for the design and update the model class accordingly.") - warnings.warn(msg, DeprecationWarning, stacklevel=2) - - logger.warning( - "Trying to guess the arguments for old-style model class %s", - model_class, - ) - # try to be compatible with old-style model class - kwargs = {} - if "prefix" in all_params: - kwargs["prefix"] = prefix - if "config" in all_params: - kwargs["config"] = model_config.hf_config - if "cache_config" in all_params: - kwargs["cache_config"] = vllm_config.cache_config - if "quant_config" in all_params: - kwargs["quant_config"] = vllm_config.quant_config - if "lora_config" in all_params: - kwargs["lora_config"] = vllm_config.lora_config - if "scheduler_config" in all_params: - kwargs["scheduler_config"] = vllm_config.scheduler_config - with set_current_vllm_config(vllm_config, check_compile=True): - return model_class(**kwargs) - - -def _process_weights_after_loading(model: nn.Module, model_config: ModelConfig, - target_device: torch.device) -> None: - for _, module in model.named_modules(): - quant_method = getattr(module, "quant_method", None) - if isinstance(quant_method, QuantizeMethodBase): - # When quant methods need to process weights after loading - # (for repacking, quantizing, etc), they expect parameters - # to be on the global target device. This scope is for the - # case where cpu offloading is used, where we will move the - # parameters onto device for processing and back off after. - with device_loading_context(module, target_device): - quant_method.process_weights_after_loading(module) - - # Currently only used by MLA. - # NOTE: This intentionally happens after other modules so we can easily - # decompress the weights for MLA. - for _, module in model.named_modules(): - if isinstance(module, Attention) and \ - hasattr(module, "process_weights_after_loading"): - # TODO(lucas): see if there is a way to unify the signatures - # of process_weights_after_loading - module.process_weights_after_loading(model_config.dtype) + return model_cls(model_config) + + # Store the config for later use + if not hasattr(inference_args, 'hf_configs'): + inference_args.hf_configs = {} + + if component_name: + inference_args.hf_configs[component_name] = hf_config + else: + inference_args.hf_config = hf_config + + # Initialize the model based on the model_name + if model_name == "CLIPVisionModel": + from fastvideo.models.encoders.clip import CLIPVisionModel + from transformers import CLIPVisionConfig + + # If the config is not a CLIPVisionConfig, convert it + if not isinstance(hf_config, CLIPVisionConfig): + # This might need adjustment based on your specific requirements + clip_config = CLIPVisionConfig.from_dict(hf_config.to_dict()) + else: + clip_config = hf_config + + # Initialize the model with the config + model = CLIPVisionModel(config=clip_config) + else: + # Add more model types as needed + raise ValueError(f"Unsupported model type: {model_name}") + + return model class BaseModelLoader(ABC): """Base class for model loaders.""" - def __init__(self, load_config: LoadConfig): - self.load_config = load_config + def __init__(self, inference_args: InferenceArgs): + self.inference_args = inference_args @abstractmethod - def download_model(self, model_config: ModelConfig) -> None: + def download_model(self, inference_args: InferenceArgs) -> None: """Download a model so that it can be immediately loaded.""" raise NotImplementedError @abstractmethod - def load_model(self, *, vllm_config: VllmConfig) -> nn.Module: + def load_model(self, *, inference_args: InferenceArgs) -> nn.Module: """Load a model with the given configurations.""" raise NotImplementedError @@ -224,40 +186,8 @@ class DefaultModelLoader(BaseModelLoader): counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 - def __init__(self, load_config: LoadConfig): - super().__init__(load_config) - if load_config.model_loader_extra_config: - raise ValueError(f"Model loader extra config is not supported for " - f"load format {load_config.load_format}") - - def _maybe_download_from_modelscope( - self, model: str, revision: Optional[str]) -> Optional[str]: - """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True. - - Returns the path to the downloaded model, or None if the model is not - downloaded from ModelScope.""" - if VLLM_USE_MODELSCOPE: - # download model from ModelScope hub, - # lazy import so that modelscope is not required for normal use. - # pylint: disable=C. - from modelscope.hub.snapshot_download import snapshot_download - - if not os.path.exists(model): - # Use file lock to prevent multiple processes from - # downloading the same model weights at the same time. - with get_lock(model, self.load_config.download_dir): - model_path = snapshot_download( - model_id=model, - cache_dir=self.load_config.download_dir, - local_files_only=huggingface_hub.constants. - HF_HUB_OFFLINE, - revision=revision, - ignore_file_pattern=self.load_config.ignore_patterns, - ) - else: - model_path = model - return model_path - return None + def __init__(self, inference_args: InferenceArgs): + super().__init__(inference_args) def _prepare_weights( self, @@ -269,29 +199,32 @@ class DefaultModelLoader(BaseModelLoader): """Prepare weights for the model. If the model is not local, it will be downloaded.""" - model_name_or_path = (self._maybe_download_from_modelscope( - model_name_or_path, revision) or model_name_or_path) + # model_name_or_path = (self._maybe_download_from_modelscope( + # model_name_or_path, revision) or model_name_or_path) is_local = os.path.isdir(model_name_or_path) - load_format = self.load_config.load_format + assert is_local, "Model path must be a local directory" + # load_format = self.load_config.load_format + # load_format = LoadFormat.AUTO use_safetensors = False index_file = SAFE_WEIGHTS_INDEX_NAME + allow_patterns = ["*.safetensors", "*.bin"] # Some quantized models use .pt files for storing the weights. - if load_format == LoadFormat.AUTO: - allow_patterns = ["*.safetensors", "*.bin"] - elif load_format == LoadFormat.SAFETENSORS: - use_safetensors = True - allow_patterns = ["*.safetensors"] - elif load_format == LoadFormat.MISTRAL: - use_safetensors = True - allow_patterns = ["consolidated*.safetensors"] - index_file = "consolidated.safetensors.index.json" - elif load_format == LoadFormat.PT: - allow_patterns = ["*.pt"] - elif load_format == LoadFormat.NPCACHE: - allow_patterns = ["*.bin"] - else: - raise ValueError(f"Unknown load_format: {load_format}") + # if load_format == LoadFormat.AUTO: + # allow_patterns = ["*.safetensors", "*.bin"] + # elif load_format == LoadFormat.SAFETENSORS: + # use_safetensors = True + # allow_patterns = ["*.safetensors"] + # elif load_format == LoadFormat.MISTRAL: + # use_safetensors = True + # allow_patterns = ["consolidated*.safetensors"] + # index_file = "consolidated.safetensors.index.json" + # elif load_format == LoadFormat.PT: + # allow_patterns = ["*.pt"] + # elif load_format == LoadFormat.NPCACHE: + # allow_patterns = ["*.bin"] + # else: + # raise ValueError(f"Unknown load_format: {load_format}") if fall_back_to_pt: allow_patterns += ["*.pt"] @@ -350,31 +283,22 @@ class DefaultModelLoader(BaseModelLoader): hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( source.model_or_path, source.revision, source.fall_back_to_pt, source.allow_patterns_overrides) - if self.load_config.load_format == LoadFormat.NPCACHE: - # Currently np_cache only support *.bin checkpoints - assert use_safetensors is False - weights_iterator = np_cache_weights_iterator( - source.model_or_path, - self.load_config.download_dir, - hf_folder, - hf_weights_files, - ) - elif use_safetensors: + if use_safetensors: weights_iterator = safetensors_weights_iterator(hf_weights_files) else: weights_iterator = pt_weights_iterator(hf_weights_files) - if current_platform.is_tpu(): - # In PyTorch XLA, we should call `xm.mark_step` frequently so that - # not too many ops are accumulated in the XLA program. - import torch_xla.core.xla_model as xm + # if current_platform.is_tpu(): + # # In PyTorch XLA, we should call `xm.mark_step` frequently so that + # # not too many ops are accumulated in the XLA program. + # import torch_xla.core.xla_model as xm - def _xla_weights_iterator(iterator: Generator): - for weights in iterator: - yield weights - xm.mark_step() + # def _xla_weights_iterator(iterator: Generator): + # for weights in iterator: + # yield weights + # xm.mark_step() - weights_iterator = _xla_weights_iterator(weights_iterator) + # weights_iterator = _xla_weights_iterator(weights_iterator) if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() @@ -384,12 +308,12 @@ class DefaultModelLoader(BaseModelLoader): def _get_all_weights( self, - model_config: ModelConfig, + model_config: Dict[str, Any], model: nn.Module, ) -> Generator[Tuple[str, torch.Tensor], None, None]: primary_weights = DefaultModelLoader.Source( model_config.model, - model_config.revision, + "", prefix="", fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), @@ -405,21 +329,25 @@ class DefaultModelLoader(BaseModelLoader): for source in secondary_weights: yield from self._get_weights_iterator(source) - def download_model(self, model_config: ModelConfig) -> None: - self._prepare_weights(model_config.model, - model_config.revision, - fall_back_to_pt=True, - allow_patterns_overrides=None) + def download_model(self, inference_args: InferenceArgs) -> None: + raise NotImplementedError + # self._prepare_weights(inference_args.model, + # inference_args.revision, + # fall_back_to_pt=True, + # allow_patterns_overrides=None) - def load_model(self, vllm_config: VllmConfig) -> nn.Module: - device_config = vllm_config.device_config - model_config = vllm_config.model_config - target_device = torch.device(device_config.device) - with set_default_torch_dtype(model_config.dtype): + def load_model(self, model_path: str, + model_config: Dict[str, Any], + inference_args: InferenceArgs) -> nn.Module: + logger.info(f"Loading model on device: {inference_args.device_str}") + target_device = torch.device(inference_args.device_str) + # TODO(will): add support for other dtypes + with set_default_torch_dtype(torch.bfloat16): with target_device: - model = _initialize_model(vllm_config=vllm_config) - + model = _initialize_model(model_path, model_config, inference_args) + weights_to_load = {name for name, _ in model.named_parameters()} + model_config.model = model_path loaded_weights = model.load_weights( self._get_all_weights(model_config, model)) self.counter_after_loading_weights = time.perf_counter() @@ -429,24 +357,23 @@ class DefaultModelLoader(BaseModelLoader): self.counter_before_loading_weights) # We only enable strict check for non-quantized models # that have loaded weights tracking currently. - if model_config.quantization is None and loaded_weights is not None: + if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights if weights_not_loaded: raise ValueError( "Following weights were not initialized from " f"checkpoint: {weights_not_loaded}") - _process_weights_after_loading(model, model_config, target_device) - + # TODO(will): add support for training/finetune return model.eval() -def get_model_loader(load_config: LoadConfig) -> BaseModelLoader: +def get_model_loader(inference_args: InferenceArgs) -> BaseModelLoader: """Get a model loader based on the load format.""" - if isinstance(load_config.load_format, type): - return load_config.load_format(load_config) + # if isinstance(load_config.load_format, type): + # return load_config.load_format(load_config) - return DefaultModelLoader(load_config) \ No newline at end of file + return DefaultModelLoader(inference_args) \ No newline at end of file diff --git a/fastvideo/loader/utils.py b/fastvideo/loader/utils.py new file mode 100644 index 000000000..2917fb403 --- /dev/null +++ b/fastvideo/loader/utils.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Utilities for selecting and loading models.""" +import contextlib +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Type + +import torch +import transformers +from torch import nn +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from vllm.config import ModelConfig, ModelImpl +from fastvideo.logger import init_logger +from vllm.model_executor.models import ModelRegistry +from vllm.model_executor.models.adapters import (as_classification_model, + as_embedding_model, + as_reward_model) + +logger = init_logger(__name__) + + +@contextlib.contextmanager +def set_default_torch_dtype(dtype: torch.dtype): + """Sets the default torch dtype to the given dtype.""" + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + yield + torch.set_default_dtype(old_dtype) + + +def load_hf_config_from_subdir(model_path: str, component_name: str, trust_remote_code: bool = False, revision: Optional[str] = None): + """ + Load a HuggingFace config from a component subdirectory. + + Args: + model_path: Path to the model directory + component_name: Name of the component subdirectory (e.g., "text_encoder", "vae") + trust_remote_code: Whether to trust remote code when loading the config + revision: Optional revision to use when loading from HuggingFace Hub + + Returns: + The loaded HuggingFace config + """ + import os + from transformers import AutoConfig + + component_path = os.path.join(model_path, component_name) + + if not os.path.exists(component_path): + raise ValueError(f"Component directory {component_path} does not exist") + + config_path = os.path.join(component_path, "config.json") + if not os.path.exists(config_path): + raise ValueError(f"Config file {config_path} does not exist") + + logger.info(f"Loading config from {config_path}") + config = AutoConfig.from_pretrained( + component_path, + trust_remote_code=trust_remote_code, + revision=revision + ) + + return config + + +def is_transformers_impl_compatible( + arch: str, + module: Optional[transformers.PreTrainedModel] = None) -> bool: + mod = module or getattr(transformers, arch, None) + if mod is None: + return False + if hasattr(mod, "supports_backend"): + return mod.is_backend_compatible() + else: + return mod._supports_flex_attn + + +def resolve_transformers_fallback(model_config: ModelConfig, + architectures: list[str]): + for i, arch in enumerate(architectures): + if arch == "TransformersModel": + continue + auto_map: dict[str, str] = getattr(model_config.hf_config, "auto_map", + None) or dict() + # Make sure that config class is always initialized before model class, + # otherwise the model class won't be able to access the config class, + # the expected auto_map should have correct order like: + # "auto_map": { + # "AutoConfig": "--", + # "AutoModel": "--", + # "AutoModelFor": "--", + # }, + auto_modules = { + name: get_class_from_dynamic_module(module, model_config.model) + for name, module in sorted(auto_map.items(), key=lambda x: x[0]) + } + custom_model_module = auto_modules.get("AutoModel") + # TODO(Isotr0py): Further clean up these raises. + # perhaps handled them in _ModelRegistry._raise_for_unsupported? + if model_config.model_impl == ModelImpl.TRANSFORMERS: + if not is_transformers_impl_compatible(arch, custom_model_module): + raise ValueError( + f"The Transformers implementation of {arch} is not " + "compatible with vLLM.") + architectures[i] = "TransformersModel" + if model_config.model_impl == ModelImpl.AUTO: + if not is_transformers_impl_compatible(arch, custom_model_module): + raise ValueError( + f"{arch} has no vLLM implementation and the Transformers " + "implementation is not compatible with vLLM.") + logger.warning( + "%s has no vLLM implementation, falling back to Transformers " + "implementation. Some features may not be supported and " + "performance may not be optimal.", arch) + architectures[i] = "TransformersModel" + return architectures + + +def get_model_architecture( + model_config: ModelConfig) -> Tuple[Type[nn.Module], str]: + architectures = getattr(model_config.hf_config, "architectures", []) + + # Special handling for quantized Mixtral. + # FIXME(woosuk): This is a temporary hack. + mixtral_supported = [ + "fp8", "compressed-tensors", "gptq_marlin", "awq_marlin" + ] + + if (model_config.quantization is not None + and model_config.quantization not in mixtral_supported + and "MixtralForCausalLM" in architectures): + architectures = ["QuantMixtralForCausalLM"] + + vllm_supported_archs = ModelRegistry.get_supported_archs() + is_vllm_supported = any(arch in vllm_supported_archs + for arch in architectures) + if (not is_vllm_supported + or model_config.model_impl == ModelImpl.TRANSFORMERS): + architectures = resolve_transformers_fallback(model_config, + architectures) + + model_cls, arch = ModelRegistry.resolve_model_cls(architectures) + if model_config.task == "embed": + model_cls = as_embedding_model(model_cls) + elif model_config.task == "classify": + model_cls = as_classification_model(model_cls) + elif model_config.task == "reward": + model_cls = as_reward_model(model_cls) + + return model_cls, arch + + +def get_architecture_class_name(model_config: ModelConfig) -> str: + return get_model_architecture(model_config)[1] + + +@dataclass +class ParamMapping: + """ + A class to handle parameter mapping for model weight loading. + It creates a bidirectional mapping between packed parameters and their + constituent parts. + """ + packed_mapping: Dict[str, List[str]] + inverse_packed_mapping: Dict[str, Tuple[str, + int]] = field(default_factory=dict) + + def __post_init__(self): + for packed_name, sub_params in self.packed_mapping.items(): + # Skip self-contained cases (e.g., {"W_pack": ["W_pack"]}) + if len(sub_params) == 1 and sub_params[0] == packed_name: + continue + for index, param_name in enumerate(sub_params): + self.inverse_packed_mapping[param_name] = ( + packed_name, + index, + ) + + def get_sub_modules(self, + module_name: str) -> Optional[Tuple[str, List[str]]]: + for key, value in self.packed_mapping.items(): + if module_name.endswith(key): + return key, value + return None diff --git a/fastvideo/loader/weight_utils.py b/fastvideo/loader/weight_utils.py index 3e9b88a68..7e514bd3c 100644 --- a/fastvideo/loader/weight_utils.py +++ b/fastvideo/loader/weight_utils.py @@ -25,23 +25,23 @@ from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download from safetensors.torch import load_file, safe_open, save_file from tqdm.auto import tqdm -from vllm.config import LoadConfig, ModelConfig -from vllm.distributed import get_tensor_model_parallel_rank -from vllm.logger import init_logger -from vllm.model_executor.layers.quantization import (QuantizationConfig, - get_quantization_config) -from vllm.platforms import current_platform -from vllm.utils import PlaceholderModule +# from vllm.config import LoadConfig, ModelConfig +from fastvideo.distributed import get_tensor_model_parallel_rank +from fastvideo.logger import init_logger +# from vllm.model_executor.layers.quantization import (QuantizationConfig, +# get_quantization_config) +from fastvideo.platforms import current_platform +# from vllm.utils import PlaceholderModule -try: - from runai_model_streamer import SafetensorsStreamer -except (ImportError, OSError): - # see https://github.com/run-ai/runai-model-streamer/issues/26 - # OSError will be raised on arm64 platform - runai_model_streamer = PlaceholderModule( - "runai_model_streamer") # type: ignore[assignment] - SafetensorsStreamer = runai_model_streamer.placeholder_attr( - "SafetensorsStreamer") +# try: +# from runai_model_streamer import SafetensorsStreamer +# except (ImportError, OSError): +# # see https://github.com/run-ai/runai-model-streamer/issues/26 +# # OSError will be raised on arm64 platform +# runai_model_streamer = PlaceholderModule( +# "runai_model_streamer") # type: ignore[assignment] +# SafetensorsStreamer = runai_model_streamer.placeholder_attr( +# "SafetensorsStreamer") logger = init_logger(__name__) @@ -137,89 +137,89 @@ def convert_bin_to_safetensor_file( # TODO(woosuk): Move this to other place. -def get_quant_config(model_config: ModelConfig, - load_config: LoadConfig) -> QuantizationConfig: +# def get_quant_config(model_config: ModelConfig, +# load_config: LoadConfig) -> QuantizationConfig: - quant_cls = get_quantization_config(model_config.quantization) +# quant_cls = get_quantization_config(model_config.quantization) - # GGUF doesn't have config file - if model_config.quantization == "gguf": - return quant_cls.from_config({}) +# # GGUF doesn't have config file +# if model_config.quantization == "gguf": +# return quant_cls.from_config({}) - # Read the quantization config from the HF model config, if available. - hf_quant_config = getattr(model_config.hf_config, "quantization_config", - None) - # some vision model may keep quantization_config in their text_config - hf_text_config = getattr(model_config.hf_config, "text_config", None) - if hf_quant_config is None and hf_text_config is not None: - hf_quant_config = getattr(hf_text_config, "quantization_config", None) - if hf_quant_config is None: - # compressed-tensors uses a compressions_config - hf_quant_config = getattr(model_config.hf_config, "compression_config", - None) - if hf_quant_config is not None: - return quant_cls.from_config(hf_quant_config) - # In case of bitsandbytes/QLoRA, get quant config from the adapter model. - if model_config.quantization == "bitsandbytes": - if (not load_config.model_loader_extra_config - or "qlora_adapter_name_or_path" - not in load_config.model_loader_extra_config): - return quant_cls.from_config({"adapter_name_or_path": ""}) - model_name_or_path = load_config.model_loader_extra_config[ - "qlora_adapter_name_or_path"] +# # Read the quantization config from the HF model config, if available. +# hf_quant_config = getattr(model_config.hf_config, "quantization_config", +# None) +# # some vision model may keep quantization_config in their text_config +# hf_text_config = getattr(model_config.hf_config, "text_config", None) +# if hf_quant_config is None and hf_text_config is not None: +# hf_quant_config = getattr(hf_text_config, "quantization_config", None) +# if hf_quant_config is None: +# # compressed-tensors uses a compressions_config +# hf_quant_config = getattr(model_config.hf_config, "compression_config", +# None) +# if hf_quant_config is not None: +# return quant_cls.from_config(hf_quant_config) +# # In case of bitsandbytes/QLoRA, get quant config from the adapter model. +# if model_config.quantization == "bitsandbytes": +# if (not load_config.model_loader_extra_config +# or "qlora_adapter_name_or_path" +# not in load_config.model_loader_extra_config): +# return quant_cls.from_config({"adapter_name_or_path": ""}) +# model_name_or_path = load_config.model_loader_extra_config[ +# "qlora_adapter_name_or_path"] - else: - model_name_or_path = model_config.model - is_local = os.path.isdir(model_name_or_path) - if not is_local: - # Download the config files. - with get_lock(model_name_or_path, load_config.download_dir): - hf_folder = snapshot_download( - model_name_or_path, - revision=model_config.revision, - allow_patterns="*.json", - cache_dir=load_config.download_dir, - local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, - tqdm_class=DisabledTqdm, - ) - else: - hf_folder = model_name_or_path +# else: +# model_name_or_path = model_config.model +# is_local = os.path.isdir(model_name_or_path) +# if not is_local: +# # Download the config files. +# with get_lock(model_name_or_path, load_config.download_dir): +# hf_folder = snapshot_download( +# model_name_or_path, +# revision=model_config.revision, +# allow_patterns="*.json", +# cache_dir=load_config.download_dir, +# local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, +# tqdm_class=DisabledTqdm, +# ) +# else: +# hf_folder = model_name_or_path - possible_config_filenames = quant_cls.get_config_filenames() +# possible_config_filenames = quant_cls.get_config_filenames() - # If the quantization config is not found, use the default config. - if not possible_config_filenames: - return quant_cls() +# # If the quantization config is not found, use the default config. +# if not possible_config_filenames: +# return quant_cls() - config_files = glob.glob(os.path.join(hf_folder, "*.json")) +# config_files = glob.glob(os.path.join(hf_folder, "*.json")) - quant_config_files = [ - f for f in config_files if any( - f.endswith(x) for x in possible_config_filenames) - ] - if len(quant_config_files) == 0: - raise ValueError( - f"Cannot find the config file for {model_config.quantization}") - if len(quant_config_files) > 1: - raise ValueError( - f"Found multiple config files for {model_config.quantization}: " - f"{quant_config_files}") +# quant_config_files = [ +# f for f in config_files if any( +# f.endswith(x) for x in possible_config_filenames) +# ] +# if len(quant_config_files) == 0: +# raise ValueError( +# f"Cannot find the config file for {model_config.quantization}") +# if len(quant_config_files) > 1: +# raise ValueError( +# f"Found multiple config files for {model_config.quantization}: " +# f"{quant_config_files}") - quant_config_file = quant_config_files[0] - with open(quant_config_file) as f: - config = json.load(f) +# quant_config_file = quant_config_files[0] +# with open(quant_config_file) as f: +# config = json.load(f) - if model_config.quantization == "bitsandbytes": - config["adapter_name_or_path"] = model_name_or_path - elif model_config.quantization == "modelopt": - if config["producer"]["name"] == "modelopt": - return quant_cls.from_config(config) - else: - raise ValueError( - f"Unsupported quantization config" - f" found for {model_config.quantization} in {f}.") +# if model_config.quantization == "bitsandbytes": +# config["adapter_name_or_path"] = model_name_or_path +# elif model_config.quantization == "modelopt": +# if config["producer"]["name"] == "modelopt": +# return quant_cls.from_config(config) +# else: +# raise ValueError( +# f"Unsupported quantization config" +# f" found for {model_config.quantization} in {f}.") - return quant_cls.from_config(config) +# return quant_cls.from_config(config) def download_weights_from_hf( @@ -436,21 +436,21 @@ def safetensors_weights_iterator( yield name, param -def runai_safetensors_weights_iterator( - hf_weights_files: List[str] -) -> Generator[Tuple[str, torch.Tensor], None, None]: - """Iterate over the weights in the model safetensor files.""" - enable_tqdm = not torch.distributed.is_initialized( - ) or torch.distributed.get_rank() == 0 - with SafetensorsStreamer() as streamer: - for st_file in tqdm( - hf_weights_files, - desc="Loading safetensors using Runai Model Streamer", - disable=not enable_tqdm, - bar_format=_BAR_FORMAT, - ): - streamer.stream_file(st_file) - yield from streamer.get_tensors() +# def runai_safetensors_weights_iterator( +# hf_weights_files: List[str] +# ) -> Generator[Tuple[str, torch.Tensor], None, None]: +# """Iterate over the weights in the model safetensor files.""" +# enable_tqdm = not torch.distributed.is_initialized( +# ) or torch.distributed.get_rank() == 0 +# with SafetensorsStreamer() as streamer: +# for st_file in tqdm( +# hf_weights_files, +# desc="Loading safetensors using Runai Model Streamer", +# disable=not enable_tqdm, +# bar_format=_BAR_FORMAT, +# ): +# streamer.stream_file(st_file) +# yield from streamer.get_tensors() def pt_weights_iterator( diff --git a/fastvideo/model_loader.py b/fastvideo/model_loader.py deleted file mode 100644 index 0b90d183e..000000000 --- a/fastvideo/model_loader.py +++ /dev/null @@ -1,273 +0,0 @@ -""" -Model loader for diffusion models. - -This module provides a base class and implementations for loading model components -needed by diffusion pipelines. -""" - -from abc import ABC, abstractmethod -from typing import Any, Dict, Optional, Type, Union -import os -import torch -import traceback - -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger - -logger = init_logger(__name__) - - -class ModelLoader(ABC): - """ - Abstract base class for model loaders. - - A model loader is responsible for loading the components needed by a diffusion - pipeline, such as the UNet, VAE, text encoder, etc. - """ - - @abstractmethod - def load_components( - self, - inference_args: InferenceArgs, - ) -> Dict[str, Any]: - """ - Load the components needed by a diffusion pipeline. - - Args: - inference_args: The inference arguments. - - Returns: - A dictionary of component names to component objects. - """ - pass - - -class HunYuanModelLoader(ModelLoader): - """ - Model loader for HunYuan models. - - This loader loads the components needed by HunYuan diffusion pipelines. - """ - - def load_components( - self, - inference_args: InferenceArgs, - ) -> Dict[str, Any]: - """ - Load the components needed by a HunYuan diffusion pipeline. - - Args: - inference_args: The inference arguments. - - Returns: - A dictionary of component names to component objects. - - Raises: - FileNotFoundError: If a required model file is not found. - RuntimeError: If there is an error loading a model component. - """ - try: - from fastvideo.models.unet import UNet3DConditionModel - from fastvideo.models.vae import AutoencoderKLCausal3D - from fastvideo.models.text_encoder import TextEncoder - from fastvideo.models.tokenizer import Tokenizer - from fastvideo.models.scheduler import DDIMScheduler - - logger.info(f"Loading model components from {inference_args.model_path}") - - # Check if model path exists - if not os.path.exists(inference_args.model_path): - raise FileNotFoundError(f"Model path not found: {inference_args.model_path}") - - # Check if DiT weight exists - if inference_args.dit_weight and not os.path.exists(inference_args.dit_weight): - raise FileNotFoundError(f"DiT weight not found: {inference_args.dit_weight}") - - # Load UNet - try: - logger.info("Loading UNet...") - unet = UNet3DConditionModel.from_pretrained( - inference_args.dit_weight, - subfolder="", - load_key=inference_args.load_key, - ) - logger.info("UNet loaded successfully") - except Exception as e: - logger.error(f"Error loading UNet: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load UNet: {e}") - - # Load VAE - try: - logger.info("Loading VAE...") - vae = AutoencoderKLCausal3D.from_pretrained( - inference_args.model_path, - subfolder="vae", - use_safetensors=True, - ) - logger.info("VAE loaded successfully") - except Exception as e: - logger.error(f"Error loading VAE: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load VAE: {e}") - - # Configure VAE for sequence parallelism if needed - if inference_args.vae_sp: - try: - logger.info("Configuring VAE for sequence parallelism") - vae.enable_sequence_parallel() - logger.info("VAE configured for sequence parallelism") - except Exception as e: - logger.error(f"Error configuring VAE for sequence parallelism: {e}") - logger.error(traceback.format_exc()) - logger.warning("Continuing without sequence parallelism for VAE") - - # Load text encoders - try: - logger.info("Loading primary text encoder...") - text_encoder = TextEncoder.from_pretrained( - inference_args.model_path, - subfolder="text_encoder", - use_safetensors=True, - ) - logger.info("Primary text encoder loaded successfully") - except Exception as e: - logger.error(f"Error loading primary text encoder: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load primary text encoder: {e}") - - try: - logger.info("Loading secondary text encoder...") - text_encoder_2 = TextEncoder.from_pretrained( - inference_args.model_path, - subfolder="text_encoder_2", - use_safetensors=True, - ) - logger.info("Secondary text encoder loaded successfully") - except Exception as e: - logger.error(f"Error loading secondary text encoder: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load secondary text encoder: {e}") - - # Load tokenizers - try: - logger.info("Loading primary tokenizer...") - tokenizer = Tokenizer.from_pretrained( - inference_args.model_path, - subfolder="tokenizer", - ) - logger.info("Primary tokenizer loaded successfully") - except Exception as e: - logger.error(f"Error loading primary tokenizer: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load primary tokenizer: {e}") - - try: - logger.info("Loading secondary tokenizer...") - tokenizer_2 = Tokenizer.from_pretrained( - inference_args.model_path, - subfolder="tokenizer_2", - ) - logger.info("Secondary tokenizer loaded successfully") - except Exception as e: - logger.error(f"Error loading secondary tokenizer: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load secondary tokenizer: {e}") - - # Load scheduler - try: - logger.info("Loading scheduler...") - scheduler = DDIMScheduler.from_pretrained( - inference_args.model_path, - subfolder="scheduler", - ) - logger.info("Scheduler loaded successfully") - except Exception as e: - logger.error(f"Error loading scheduler: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load scheduler: {e}") - - # Move models to the appropriate device - device = torch.device(f"cuda:{os.environ.get('LOCAL_RANK', '0')}") - logger.info(f"Moving models to device: {device}") - - try: - unet.to(device) - vae.to(device) - text_encoder.to(device) - text_encoder_2.to(device) - logger.info("Models moved to device successfully") - except Exception as e: - logger.error(f"Error moving models to device: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to move models to device: {e}") - - # Set precision - try: - logger.info(f"Setting UNet precision to {inference_args.precision}") - if inference_args.precision == "bf16": - unet = unet.to(torch.bfloat16) - elif inference_args.precision == "fp16": - unet = unet.to(torch.float16) - - logger.info(f"Setting VAE precision to {inference_args.vae_precision}") - if inference_args.vae_precision == "fp16": - vae = vae.to(torch.float16) - - logger.info(f"Setting text encoder precision to {inference_args.text_encoder_precision}") - if inference_args.text_encoder_precision == "fp16": - text_encoder = text_encoder.to(torch.float16) - - logger.info(f"Setting secondary text encoder precision to {inference_args.text_encoder_precision_2}") - if inference_args.text_encoder_precision_2 == "fp16": - text_encoder_2 = text_encoder_2.to(torch.float16) - - logger.info("Precision set successfully") - except Exception as e: - logger.error(f"Error setting precision: {e}") - logger.error(traceback.format_exc()) - logger.warning("Continuing with default precision") - - # Return the components - return { - "unet": unet, - "vae": vae, - "text_encoder": text_encoder, - "text_encoder_2": text_encoder_2, - "tokenizer": tokenizer, - "tokenizer_2": tokenizer_2, - "scheduler": scheduler, - } - except Exception as e: - logger.error(f"Unhandled error in load_components: {e}") - logger.error(traceback.format_exc()) - raise - - -# Factory function to get the appropriate model loader -def get_model_loader(model_type: str) -> Type[ModelLoader]: - """ - Get the appropriate model loader for the given model type. - - Args: - model_type: The type of model to load. - - Returns: - The model loader class. - - Raises: - ValueError: If the model type is not recognized. - """ - model_loaders = { - "hunyuan": HunYuanModelLoader, - "HYVideo-T/2-cfgdistill": HunYuanModelLoader, - } - - if model_type not in model_loaders: - available_models = list(model_loaders.keys()) - raise ValueError( - f"Model type '{model_type}' not recognized. " - f"Available model types: {available_models}" - ) - - return model_loaders[model_type] \ No newline at end of file diff --git a/fastvideo/models/__init__.py b/fastvideo/models/__init__.py index 2c63b0b0d..a6c336936 100644 --- a/fastvideo/models/__init__.py +++ b/fastvideo/models/__init__.py @@ -1,11 +1,57 @@ # SPDX-License-Identifier: Apache-2.0 -from fastvideo.models.parameter import (BasevLLMParameter, - PackedvLLMParameter) -from fastvideo.models.utils import set_random_seed +import torch.nn as nn +from typing import Dict, Any +from fastvideo.loader.loader import get_model_loader +from fastvideo.inference_args import InferenceArgs +from fastvideo.logger import init_logger +import os +from transformers import PretrainedConfig + +logger = init_logger(__name__) + +def get_transformers_model(model_config: PretrainedConfig, model_path: str, inference_args: InferenceArgs) -> nn.Module: + """ + Load a transformers model. + + Args: + architecture: The model architecture name + model_path: Path to the model directory + inference_args: Inference arguments + component_name: Optional component name (e.g., "text_encoder", "vae") + If provided, will load the config from the component subdirectory + + Returns: + The loaded model + """ + model_loader = get_model_loader(inference_args) + return model_loader.load_model(model_config, model_path, inference_args) + +def get_diffusers_model(model_config: Dict[str, Any], model_path: str, inference_args: InferenceArgs) -> nn.Module: + """ + Load a diffusers model. + """ + raise NotImplementedError("Diffusers models are not supported yet.") + +def get_scheduler(module_path: str, architecture: str, inference_args: InferenceArgs) -> Dict: + """Create a scheduler based on the inference args. Can be overridden by subclasses.""" + if hasattr(inference_args, 'denoise_type') and inference_args.denoise_type == "flow": + # TODO(will): add schedulers to register or create a new scheduler registry + # TODO(will): default to config file but allow override through + # inference args. Currently only uses inference args. + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchDiscreteScheduler + return FlowMatchDiscreteScheduler( + shift=inference_args.flow_shift, + reverse=inference_args.flow_reverse, + solver=inference_args.flow_solver, + ) + else: + raise ValueError(f"Invalid denoise type: {inference_args.denoise_type}") __all__ = [ "set_random_seed", "BasevLLMParameter", "PackedvLLMParameter", + "get_model", + "get_scheduler", ] diff --git a/fastvideo/models/dits/__init__.py b/fastvideo/models/dits/__init__.py deleted file mode 100644 index 768f632bd..000000000 --- a/fastvideo/models/dits/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ - -from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT - -__all__ = ["HunyuanVideoDiT"] diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/models/dits/hunyuanvideo.py index f407e645d..49bd3e163 100644 --- a/fastvideo/models/dits/hunyuanvideo.py +++ b/fastvideo/models/dits/hunyuanvideo.py @@ -1,7 +1,7 @@ import torch import torch.nn as nn from typing import Optional, Tuple, List -from fastvideo.attention.distributed_attn import DistributedAttention, LocalAttention +from fastvideo.attention.flash_attn import DistributedAttention, LocalAttention from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.layernorm import LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection, unpatchify @@ -329,7 +329,7 @@ class MMSingleStreamBlock(nn.Module): -class HunyuanVideoDiT(BaseDiT): +class HunyuanVideoTransformer3DModel(BaseDiT): """ HunyuanVideo Transformer backbone adapted for distributed training. @@ -799,9 +799,8 @@ class FinalLayer(nn.Module): ) def forward(self, x, c): - # What the fuck HF? Why you change the scale and shift order here??? + # 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 - diff --git a/fastvideo/models/encoders/clip.py b/fastvideo/models/encoders/clip.py index 907b30f11..67c76163e 100644 --- a/fastvideo/models/encoders/clip.py +++ b/fastvideo/models/encoders/clip.py @@ -5,7 +5,8 @@ from typing import Iterable, Optional, Set, Tuple, Union import torch import torch.nn as nn -from transformers import CLIPVisionConfig +from transformers import CLIPVisionConfig, CLIPTextConfig +from transformers.modeling_outputs import BaseModelOutputWithPooling from vllm.attention.layer import MultiHeadAttention from fastvideo.distributed import divide, get_tensor_model_parallel_world_size @@ -13,12 +14,15 @@ from fastvideo.layers.activation import get_act_fn from fastvideo.layers.linear import (ColumnParallelLinear, QKVParallelLinear, RowParallelLinear) -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.model_loader.weight_utils import default_weight_loader +# TODO: support quantization +# from vllm.model_executor.layers.quantization import QuantizationConfig +from fastvideo.loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import SupportsQuant from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs +class QuantizationConfig: + pass class CLIPEncoderInfo(VisionEncoderInfo[CLIPVisionConfig]): @@ -88,6 +92,48 @@ class CLIPVisionEmbeddings(nn.Module): return embeddings +class CLIPTextEmbeddings(nn.Module): + + def __init__(self, config: CLIPTextConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.Tensor: + seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + max_position_embedding = self.position_embedding.weight.shape[0] + + if seq_length > max_position_embedding: + raise ValueError( + f"Sequence length must be less than max_position_embeddings (got `sequence length`: " + f"{seq_length} and max_position_embeddings: {max_position_embedding}" + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + class CLIPAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -263,6 +309,190 @@ class CLIPEncoder(nn.Module): if return_all_hidden_states: return hidden_states_pool return hidden_states + + + +class CLIPTextTransformer(nn.Module): + + def __init__(self, + config: CLIPTextConfig, + quant_config: Optional[QuantizationConfig] = None, + *, + num_hidden_layers_override: Optional[int] = None, + prefix: str = ""): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPTextEmbeddings(config) + + self.encoder = CLIPEncoder(config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers_override, + prefix=prefix) + + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + # For `pooled_output` computation + self.eos_token_id = config.eos_token_id + + # For attention mask, it differs between `flash_attention_2` and other attention implementations + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) + + # CLIP's text model uses causal mask, prepare it here. + # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 + causal_attention_mask = _create_4d_causal_attention_mask( + input_shape, hidden_states.dtype, device=hidden_states.device + ) + + # expand attention_mask + if attention_mask is not None and not self._use_flash_attention_2: + raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer") + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + causal_attention_mask=causal_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[0] + last_hidden_state = self.final_layer_norm(last_hidden_state) + + if self.eos_token_id == 2: + # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here. + # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added + # ------------------------------------------------------------ + # text_embeds.shape = [batch_size, sequence_length, transformer.width] + # take features from the eot embedding (eot_token is the highest number in each sequence) + # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1), + ] + else: + # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible) + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) + # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) + (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id) + .int() + .argmax(dim=-1), + ] + + if not return_dict: + return (last_hidden_state, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +class CLIPTextModel(nn.Module): + + def __init__( + self, + config: CLIPTextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.text_model = CLIPTextTransformer( + config=config, + quant_config=quant_config, + prefix=prefix) + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + return self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + def load_weights(self, weights: Iterable[Tuple[str, + torch.Tensor]]) -> Set[str]: + + # Define mapping for stacked parameters + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + for name, loaded_weight in weights: + # Handle q_proj, k_proj, v_proj -> qkv_proj mapping + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name in name: + # Replace the weight name with the parameter name + model_param_name = name.replace(weight_name, param_name) + + if model_param_name in params_dict: + param = params_dict[model_param_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(model_param_name) + break + else: + # Use default weight loader for all other parameters + if name in params_dict: + 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 class CLIPVisionTransformer(nn.Module): diff --git a/fastvideo/models/encoders/llama.py b/fastvideo/models/encoders/llama.py new file mode 100644 index 000000000..f1ca25b21 --- /dev/null +++ b/fastvideo/models/encoders/llama.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +"""Inference-only LLaMA model compatible with HuggingFace weights.""" +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union + +import torch +from torch import nn +from transformers import LlamaConfig +from transformers.modeling_outputs import BaseModelOutputWithPast + +from fastvideo.attention import TextEncoderFlashAttention + +from fastvideo.distributed import get_tensor_model_parallel_world_size +from fastvideo.layers.activation import SiluAndMul +from fastvideo.layers.layernorm import RMSNorm +from fastvideo.layers.linear import (MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear) +# from vllm.model_executor.layers.quantization import QuantizationConfig + +from fastvideo.layers.rotary_embedding import get_rope +from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.loader.weight_utils import ( + default_weight_loader, maybe_remap_kv_scale_name) + +from .utils import (extract_layer_index) + +class QuantizationConfig: + pass + + +class LlamaMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + bias: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now.") + 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 LlamaAttention(nn.Module): + + def __init__(self, + config: LlamaConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + bias: bool = False, + bias_o_proj: bool = False, + prefix: str = "") -> None: + super().__init__() + layer_idx = extract_layer_index(prefix) + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr(config, "head_dim", + self.hidden_size // self.total_num_heads) + # Phi models introduced a partial_rotary_factor parameter in the config + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) + self.rotary_dim = int(partial_rotary_factor * self.head_dim) + 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.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias_o_proj, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + is_neox_style = True + is_gguf = quant_config and quant_config.get_name() == "gguf" + if is_gguf and config.model_type == "llama": + is_neox_style = False + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=is_neox_style, + ) + + if hasattr(config, "interleaved_sliding_window"): + interleaved_sliding_window = config.interleaved_sliding_window + if isinstance(interleaved_sliding_window, int): + sliding_window = interleaved_sliding_window + elif isinstance(interleaved_sliding_window, list): + sw_idx = layer_idx % len(interleaved_sliding_window) + sliding_window = interleaved_sliding_window[sw_idx] + else: + raise ValueError( + f"{type(interleaved_sliding_window)} is not supported.") + else: + sliding_window = None + + self.attn = TextEncoderFlashAttention( + self.scaling, + # num_kv_heads=self.num_kv_heads, + # quant_config=quant_config, + # per_layer_sliding_window=sliding_window, + # prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class LlamaDecoderLayer(nn.Module): + + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings) + max_position_embeddings = getattr(config, "max_position_embeddings", + 8192) + # Support abacusai/Smaug-72B-v0.1 with attention_bias + # Support internlm/internlm-7b with bias + attention_bias = getattr(config, "attention_bias", False) or getattr( + config, "bias", False) + bias_o_proj = attention_bias + # support internlm/internlm3-8b with qkv_bias + if hasattr(config, 'qkv_bias'): + attention_bias = config.qkv_bias + + self.self_attn = LlamaAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", + config.num_attention_heads), + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=attention_bias, + bias_o_proj=bias_o_proj, + prefix=f"{prefix}.self_attn", + ) + self.mlp = LlamaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + bias=getattr(config, "mlp_bias", False), + 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, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm( + hidden_states, residual) + hidden_states = self.self_attn(positions=positions, + hidden_states=hidden_states) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class LlamaModel(nn.Module): + + def __init__(self, + config: LlamaConfig, + prefix: str = "", + layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer): + super().__init__() + + quant_config = None + lora_config = None + + self.config = config + self.quant_config = quant_config + lora_vocab = (lora_config.lora_extra_vocab_size * + (lora_config.max_loras or 1)) if lora_config else 0 + self.vocab_size = config.vocab_size + lora_vocab + self.org_vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=quant_config, + ) + + self.layers = nn.ModuleList([ + layer_type( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.layers.{i}" + ) for i in range(config.num_hidden_layers) + ]) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: Optional[torch.Tensor], + positions: torch.Tensor, + inputs_embeds: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + ) -> torch.Tensor: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + residual = None + + all_hidden_states = () if output_hidden_states else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + hidden_states, residual = layer(positions, hidden_states, residual) + + hidden_states, _ = self.norm(hidden_states, residual) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + # TODO(will): maybe unify the output format with other models and use + # our own class + output = BaseModelOutputWithPast( + last_hidden_state=hidden_states, + # past_key_values=past_key_values if use_cache else None, + hidden_states=all_hidden_states, + # attentions=all_self_attns, + ) + + return output + + def load_weights(self, weights: Iterable[Tuple[str, + torch.Tensor]]) -> Set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".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), + ] + 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 + if ("rotary_emb.cos_cached" in name + or "rotary_emb.sin_cached" in name): + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if (self.quant_config is not None and + (scale_name := self.quant_config.get_cache_scale(name))): + # Loading kv cache quantization scales + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else + loaded_weight[0]) + weight_loader(param, loaded_weight) + loaded_params.add(scale_name) + continue + if "scale" in name: + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + for param_name, weight_name, shard_id in 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/encoders/utils.py b/fastvideo/models/encoders/utils.py new file mode 100644 index 000000000..f26828c23 --- /dev/null +++ b/fastvideo/models/encoders/utils.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 + +import itertools +from dataclasses import dataclass, field +from typing import (Callable, Dict, Iterable, List, Literal, Mapping, Optional, + Protocol, Set, Tuple, Union, overload) + +import torch +import torch.nn as nn +from torch.func import functional_call +from transformers import PretrainedConfig + +from fastvideo.logger import init_logger +from fastvideo.loader.weight_utils import default_weight_loader + +logger = init_logger(__name__) + +WeightsMapping = Mapping[str, Optional[str]] +"""If a key maps to a value of `None`, the corresponding weight is ignored.""" + + +@dataclass +class WeightsMapper: + """Maps the name of each weight if they match the following patterns.""" + + orig_to_new_substr: WeightsMapping = field(default_factory=dict) + orig_to_new_prefix: WeightsMapping = field(default_factory=dict) + orig_to_new_suffix: WeightsMapping = field(default_factory=dict) + + def _map_name(self, key: str) -> Optional[str]: + for substr, new_key in self.orig_to_new_substr.items(): + if substr in key: + if new_key is None: + return None + + key = key.replace(substr, new_key, 1) + + for prefix, new_key in self.orig_to_new_prefix.items(): + if key.startswith(prefix): + if new_key is None: + return None + + key = key.replace(prefix, new_key, 1) + + for suffix, new_key in self.orig_to_new_suffix.items(): + if key.endswith(suffix): + if new_key is None: + return None + + key = new_key.join(key.rsplit(suffix, 1)) + + return key + + def apply( + self, weights: Iterable[Tuple[str, torch.Tensor]] + ) -> Iterable[Tuple[str, torch.Tensor]]: + return ((out_name, data) for name, data in weights + if (out_name := self._map_name(name)) is not None) + + +class AutoWeightsLoader: + """ + Helper class to load weights into a :class:`torch.nn.Module`. It is able + to automatically detect child modules and parameters while iterating over + the weights only once. + + The weight loading logic for individual modules can be overridden + by defining a ``load_weights`` method. + + Similarly, the weight loading logic for individual parameters can be + overridden by defining a ``weight_loader`` method. + + Detailed weight loading information can be viewed by setting the + environment variable ``VLLM_LOGGING_LEVEL=DEBUG``. + """ + + def __init__( + self, + module: nn.Module, + *, + skip_prefixes: Optional[List[str]] = None, + ignore_unexpected_prefixes: Optional[List[str]] = None, + ) -> None: + super().__init__() + + self.module = module + self.skip_prefixes = skip_prefixes or [] + self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or [] + + def _groupby_prefix( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + ) -> Iterable[Tuple[str, Iterable[Tuple[str, torch.Tensor]]]]: + weights_by_parts = ((weight_name.split(".", 1), weight_data) + for weight_name, weight_data in weights) + + for prefix, group in itertools.groupby(weights_by_parts, + key=lambda x: x[0][0]): + yield ( + prefix, + # Because maxsplit=1 in weight_name.split(...), + # the length of `parts` must either be 1 or 2 + (("" if len(parts) == 1 else parts[1], weights_data) + for parts, weights_data in group), + ) + + def _get_qualname(self, prefix: str, rest: str) -> str: + if prefix == "": + return rest + if rest == "": + return prefix + + return ".".join((prefix, rest)) + + def _can_skip(self, qualname: str) -> bool: + return any(qualname.startswith(p) for p in self.skip_prefixes) + + def _can_ignore_unexpected(self, qualname: str) -> bool: + return any( + qualname.startswith(p) for p in self.ignore_unexpected_prefixes) + + def _load_param( + self, + base_prefix: str, + param: nn.Parameter, + weights: Iterable[Tuple[str, torch.Tensor]], + ) -> Iterable[str]: + for weight_name, weight_data in weights: + weight_qualname = self._get_qualname(base_prefix, weight_name) + + if self._can_skip(weight_qualname): + logger.debug("Skipping weight %s", weight_qualname) + + continue + + if weight_name != "": + if self._can_ignore_unexpected(weight_qualname): + logger.debug("Ignoring weight %s", weight_qualname) + + continue + + raise ValueError( + f"Attempted to load nested weight '{weight_qualname}' " + f"into a single parameter '{base_prefix}'") + + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, weight_data) + + logger.debug("Loaded weight %s with shape %s", weight_qualname, + param.shape) + + yield weight_qualname + + def _load_module( + self, + base_prefix: str, + module: nn.Module, + weights: Iterable[Tuple[str, torch.Tensor]], + ) -> Iterable[str]: + if isinstance(module, PPMissingLayer): + return + + # Avoid infinite recursion since this function is typically + # called inside load_weights of the module itself + if module != self.module: + module_load_weights = getattr(module, "load_weights", None) + if callable(module_load_weights): + loaded_params = module_load_weights(weights) + if loaded_params is None: + logger.warning( + "Unable to collect loaded parameters " + "for module %s", module) + else: + yield from map( + lambda x: self._get_qualname(base_prefix, x), + loaded_params, + ) + + child_modules = dict(module.named_children()) + child_params = dict(module.named_parameters(recurse=False)) + + for child_prefix, child_weights in self._groupby_prefix(weights): + prefix = self._get_qualname(base_prefix, child_prefix) + + if child_prefix in child_modules: + if self._can_skip(prefix + "."): + logger.debug("Skipping module %s", prefix) + + continue + + yield from self._load_module(prefix, + child_modules[child_prefix], + child_weights) + elif child_prefix in child_params: + if self._can_skip(prefix): + logger.debug("Skipping param %s", prefix) + + continue + + yield from self._load_param(prefix, child_params[child_prefix], + child_weights) + else: + can_skip_module = self._can_skip(prefix + ".") + can_skip_param = self._can_skip(prefix) + if can_skip_module or can_skip_param: + logger.debug("Skipping missing %s", prefix) + + continue + + can_ignore_module = self._can_ignore_unexpected(prefix + ".") + can_ignore_param = self._can_ignore_unexpected(prefix) + if can_ignore_module or can_ignore_param: + logger.debug("Ignoring missing %s", prefix) + + continue + + msg = (f"There is no module or parameter named '{prefix}' " + f"in {type(self.module).__name__}") + raise ValueError(msg) + + def load_weights( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + *, + mapper: Optional[WeightsMapper] = None, + ) -> Set[str]: + if mapper is not None: + weights = mapper.apply(weights) + + autoloaded_weights = set(self._load_module("", self.module, weights)) + return autoloaded_weights + + +@overload +def flatten_bn(x: torch.Tensor) -> torch.Tensor: + ... + + +@overload +def flatten_bn(x: List[torch.Tensor]) -> List[torch.Tensor]: + ... + + +@overload +def flatten_bn( + x: Union[List[torch.Tensor], torch.Tensor], + *, + concat: Literal[True], +) -> torch.Tensor: + ... + + +@overload +def flatten_bn( + x: Union[List[torch.Tensor], torch.Tensor], + *, + concat: bool = False, +) -> Union[List[torch.Tensor], torch.Tensor]: + ... + + +def flatten_bn( + x: Union[List[torch.Tensor], torch.Tensor], + *, + concat: bool = False, +) -> Union[List[torch.Tensor], torch.Tensor]: + """ + Flatten the ``B`` and ``N`` dimensions of batched multimodal inputs. + + The input tensor should have shape ``(B, N, ...)```. + """ + if isinstance(x, torch.Tensor): + return x.flatten(0, 1) + + if concat: + return torch.cat(x) + + return [x_n for x_b in x for x_n in x_b] + + +class LayerFn(Protocol): + + def __call__(self, prefix: str) -> torch.nn.Module: + ... + + +class PPMissingLayer(torch.nn.Identity): + """ + A placeholder layer for missing layers in a pipeline parallel model. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + + +_CPU_OFFLOAD_BYTES = 0 +_CPU_OFFLOAD_MAX_BYTES = 0 + + +def set_cpu_offload_max_bytes(max_bytes: int) -> None: + global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES + _CPU_OFFLOAD_BYTES = 0 + _CPU_OFFLOAD_MAX_BYTES = max_bytes + + +def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module: + device = next(module.parameters()).device + + if device == torch.device("cpu"): + return module + + global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES + if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + return module + + pin_memory = True + + # offload parameters to CPU + # use pin_memory if possible, which helps cudagraph capture speed + offloaded_parameters = False + for p in module.parameters(): + if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + # we use per-parameter offloading + # one module might have some parameters offloaded and some not + break + + # `torch.empty_like` does not support `pin_memory` argument + cpu_data = torch.empty_strided(size=p.data.size(), + stride=p.data.stride(), + dtype=p.data.dtype, + layout=p.data.layout, + device='cpu', + pin_memory=pin_memory) + cpu_data.copy_(p.data) + p.data = cpu_data + _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size() + offloaded_parameters = True + + if offloaded_parameters: + original_forward = module.forward + + def forward(*args, **kwargs): + module.forward = original_forward + device_state = { + # here we blindly call `to(device)` + # if the parameter is already on the device, it will be a no-op + k: v.to(device, non_blocking=True) + for k, v in module.state_dict().items() + } + output = functional_call(module, + device_state, + args=args, + kwargs=kwargs) + module.forward = forward + return output + + module.forward = forward + + return module + + +def maybe_prefix(prefix: str, name: str) -> str: + """Add a prefix to a name if the prefix is non-empty. + + Args: + prefix: The prefix to add. If empty, no prefix will be added. + name: The name to potentially prefix. + + Returns: + The string "prefix.name" if prefix was non-empty, otherwise just "name". + """ + return name if not prefix else f"{prefix}.{name}" + + +def extract_layer_index(layer_name: str) -> int: + """ + Extract the layer index from the module name. + Examples: + - "encoder.layers.0" -> 0 + - "encoder.layers.1.self_attn" -> 1 + - "2.self_attn" -> 2 + - "model.encoder.layers.0.sub.1" -> ValueError + """ + subnames = layer_name.split(".") + int_vals: List[int] = [] + for subname in subnames: + try: + int_vals.append(int(subname)) + except ValueError: + continue + assert len(int_vals) == 1, (f"layer name {layer_name} should" + " only contain one integer") + return int_vals[0] + + +def cast_overflow_tensors( + tensors: torch.Tensor, + offset: float = 1000, +) -> torch.Tensor: + if tensors.isinf().any() or tensors.isnan().any(): + clamp_value = torch.finfo(tensors.dtype).max - offset + tensors = torch.clamp(tensors, min=-clamp_value, max=clamp_value) + return tensors + +def make_layers( + num_hidden_layers: int, + layer_fn: LayerFn, + prefix: str, +) -> torch.nn.ModuleList: + """Make a list of layers with the given layer function. + + Args: + num_hidden_layers: Number of layers to create. + layer_fn: Function that creates a layer given a prefix. + prefix: Prefix for layer names. + + Returns: + A ModuleList containing all layers. + """ + modules = torch.nn.ModuleList([ + maybe_offload_to_cpu(layer_fn(prefix=f"{prefix}.{idx}")) + for idx in range(num_hidden_layers) + ]) + return modules + + +def add_prefix(name: str, prefix: str) -> str: + """Add a weight path prefix to a module name. + + Args: + name: base module name. + prefix: weight prefix str to added to the front of `name` concatenated with `.`. + + Returns: + The string `prefix.name` if prefix is non-empty, otherwise just `name`. + """ + return name if not prefix else f"{prefix}.{name}" diff --git a/fastvideo/models/encoders/vision.py b/fastvideo/models/encoders/vision.py index 778e5640b..c9ff3e3a6 100644 --- a/fastvideo/models/encoders/vision.py +++ b/fastvideo/models/encoders/vision.py @@ -6,11 +6,11 @@ from typing import Final, Generic, Optional, Protocol, TypeVar, Union import torch from transformers import PretrainedConfig -import vllm.envs as envs +import fastvideo.envs as envs from vllm.attention.selector import (backend_name_to_enum, get_global_forced_attn_backend) -from vllm.logger import init_logger -from vllm.platforms import _Backend, current_platform +from fastvideo.logger import init_logger +from fastvideo.platforms import _Backend, current_platform logger = init_logger(__name__) diff --git a/fastvideo/models/flash_attn_no_pad.py b/fastvideo/models/flash_attn_no_pad.py deleted file mode 100644 index f686e70a6..000000000 --- a/fastvideo/models/flash_attn_no_pad.py +++ /dev/null @@ -1,28 +0,0 @@ -from einops import rearrange -from flash_attn import flash_attn_varlen_qkvpacked_func -from flash_attn.bert_padding import pad_input, unpad_input - - -def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None): - # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27 - batch_size = qkv.shape[0] - seqlen = qkv.shape[1] - nheads = qkv.shape[-2] - x = rearrange(qkv, "b s three h d -> b s (three h d)") - x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask) - - x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) - output_unpad = flash_attn_varlen_qkvpacked_func( - x_unpad, - cu_seqlens, - max_s, - dropout_p, - softmax_scale=softmax_scale, - causal=causal, - ) - 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/models/hf_transformer_utils.py b/fastvideo/models/hf_transformer_utils.py new file mode 100644 index 000000000..b4a7ed042 --- /dev/null +++ b/fastvideo/models/hf_transformer_utils.py @@ -0,0 +1,259 @@ +# Copyright 2023-2024 SGLang Team +# 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. +# ============================================================================== +"""Utilities for Huggingface Transformers.""" + +import contextlib +import os +import warnings +from pathlib import Path +from typing import Dict, Optional, Type, Union, Any +import json + +from huggingface_hub import snapshot_download +from transformers import ( + AutoConfig, + AutoProcessor, + AutoTokenizer, + PretrainedConfig, + PreTrainedTokenizer, + PreTrainedTokenizerFast, +) +from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES + +# from fastvideo.models.configs import ChatGLMConfig, DbrxConfig, ExaoneConfig, Qwen2_5_VLConfig + +_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { + # ChatGLMConfig.model_type: ChatGLMConfig, + # DbrxConfig.model_type: DbrxConfig, + # ExaoneConfig.model_type: ExaoneConfig, + # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig, +} + +for name, cls in _CONFIG_REGISTRY.items(): + with contextlib.suppress(ValueError): + AutoConfig.register(name, cls) + + +def download_from_hf(model_path: str): + if os.path.exists(model_path): + return model_path + + return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"]) + + +def get_hf_config( + model: str, + trust_remote_code: bool, + revision: Optional[str] = None, + model_override_args: Optional[dict] = None, + inference_args: Optional[dict] = None, + **kwargs, +): + is_gguf = check_gguf_file(model) + if is_gguf: + kwargs["gguf_file"] = model + model = Path(model).parent + + config = AutoConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + if config.model_type in _CONFIG_REGISTRY: + config_class = _CONFIG_REGISTRY[config.model_type] + config = config_class.from_pretrained(model, revision=revision) + # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`. + setattr(config, "_name_or_path", model) + if model_override_args: + config.update(model_override_args) + + # Special architecture mapping check for GGUF models + if is_gguf: + if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: + raise RuntimeError(f"Can't get gguf config for {config.model_type}.") + model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type] + config.update({"architectures": [model_type]}) + + return config + + +def get_diffusers_config( + model: str, + inference_args: Optional[dict] = None, +) -> Dict[str, Any]: + """Gets a configuration for the given diffusers model. + + Args: + model: The model name or path. + inference_args: Optional inference arguments to override in the config. + + Returns: + The loaded configuration. + """ + # Check if the model path exists + if os.path.exists(model): + config_file = os.path.join(model, "config.json") + if os.path.exists(config_file): + try: + # Load the config directly from the file + with open(config_file, "r") as f: + config_dict = json.load(f) + + # TODO(will): apply any overrides from inference args + return config_dict + except Exception as e: + raise RuntimeError(f"Failed to load diffusers config from {config_file}: {e}") + else: + raise RuntimeError(f"Diffusers config file not found at {model}") + + +# Models don't use the same configuration key for determining the maximum +# context length. Store them here so we can sanely check them. +# NOTE: The ordering here is important. Some models have two of these and we +# have a preference for which value gets used. +CONTEXT_LENGTH_KEYS = [ + "max_sequence_length", + "seq_length", + "max_seq_len", + "model_max_length", + "max_position_embeddings", +] + + +def get_context_length(config): + """Get the context length of a model from a huggingface model configs.""" + text_config = config + rope_scaling = getattr(text_config, "rope_scaling", None) + if rope_scaling: + rope_scaling_factor = rope_scaling.get("factor", 1) + if "original_max_position_embeddings" in rope_scaling: + rope_scaling_factor = 1 + if rope_scaling.get("rope_type", None) == "llama3": + rope_scaling_factor = 1 + else: + rope_scaling_factor = 1 + + for key in CONTEXT_LENGTH_KEYS: + val = getattr(text_config, key, None) + if val is not None: + return int(rope_scaling_factor * val) + return 2048 + + +# A fast LLaMA tokenizer with the pre-processed `tokenizer.json` file. +_FAST_LLAMA_TOKENIZER = "hf-internal-testing/llama-tokenizer" + + +def get_tokenizer( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: Optional[str] = None, + **kwargs, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + """Gets a tokenizer for the given model name via Huggingface.""" + if tokenizer_mode == "slow": + if kwargs.get("use_fast", False): + raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.") + kwargs["use_fast"] = False + + is_gguf = check_gguf_file(tokenizer_name) + if is_gguf: + kwargs["gguf_file"] = tokenizer_name + tokenizer_name = Path(tokenizer_name).parent + + try: + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + tokenizer_revision=tokenizer_revision, + clean_up_tokenization_spaces=False, + **kwargs, + ) + except TypeError as e: + # The LLaMA tokenizer causes a protobuf error in some environments. + err_msg = ( + "Failed to load the tokenizer. If you are using a LLaMA V1 model " + f"consider using '{_FAST_LLAMA_TOKENIZER}' instead of the " + "original tokenizer." + ) + raise RuntimeError(err_msg) from e + except ValueError as e: + # If the error pertains to the tokenizer class not existing or not + # currently being imported, suggest using the --trust-remote-code flag. + if not trust_remote_code and ( + "does not exist or is not currently imported." in str(e) + or "requires you to execute the tokenizer file" in str(e) + ): + err_msg = ( + "Failed to load the tokenizer. If the tokenizer is a custom " + "tokenizer not yet available in the HuggingFace transformers " + "library, consider setting `trust_remote_code=True` in LLM " + "or using the `--trust-remote-code` flag in the CLI." + ) + raise RuntimeError(err_msg) from e + else: + raise e + + if not isinstance(tokenizer, PreTrainedTokenizerFast): + warnings.warn( + "Using a slow tokenizer. This might cause a significant " + "slowdown. Consider using a fast tokenizer instead." + ) + + attach_additional_stop_token_ids(tokenizer) + return tokenizer + + +def get_processor( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: Optional[str] = None, + **kwargs, +): + processor = AutoProcessor.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + tokenizer_revision=tokenizer_revision, + **kwargs, + ) + + attach_additional_stop_token_ids(processor.tokenizer) + return processor + + +def attach_additional_stop_token_ids(tokenizer): + # Special handling for stop token <|eom_id|> generated by llama 3 tool use. + if "<|eom_id|>" in tokenizer.get_added_vocab(): + tokenizer.additional_stop_token_ids = set( + [tokenizer.get_added_vocab()["<|eom_id|>"]] + ) + else: + tokenizer.additional_stop_token_ids = None + + +def check_gguf_file(model: Union[str, os.PathLike]) -> bool: + """Check if the file is a GGUF model.""" + model = Path(model) + if not model.is_file(): + return False + elif model.suffix == ".gguf": + return True + + with open(model, "rb") as f: + header = f.read(4) + return header == b"GGUF" diff --git a/fastvideo/models/hunyuan/modules/__init__.py b/fastvideo/models/hunyuan/modules/__init__.py index c85b51a4c..3bb947c4e 100644 --- a/fastvideo/models/hunyuan/modules/__init__.py +++ b/fastvideo/models/hunyuan/modules/__init__.py @@ -1,25 +1,25 @@ -from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer +# from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer -def load_model(args, in_channels, out_channels, factor_kwargs): - """load hunyuan video model +# def load_model(args, in_channels, out_channels, factor_kwargs): +# """load hunyuan video model - Args: - args (dict): model args - in_channels (int): input channels number - out_channels (int): output channels number - factor_kwargs (dict): factor kwargs +# Args: +# args (dict): model args +# in_channels (int): input channels number +# out_channels (int): output channels number +# factor_kwargs (dict): factor kwargs - Returns: - model (nn.Module): The hunyuan video model - """ - if args.model in HUNYUAN_VIDEO_CONFIG.keys(): - model = HYVideoDiffusionTransformer( - in_channels=in_channels, - out_channels=out_channels, - **HUNYUAN_VIDEO_CONFIG[args.model], - **factor_kwargs, - ) - return model - else: - raise NotImplementedError() +# Returns: +# model (nn.Module): The hunyuan video model +# """ +# if args.model in HUNYUAN_VIDEO_CONFIG.keys(): +# model = HYVideoDiffusionTransformer( +# in_channels=in_channels, +# out_channels=out_channels, +# **HUNYUAN_VIDEO_CONFIG[args.model], +# **factor_kwargs, +# ) +# return model +# else: +# raise NotImplementedError() diff --git a/fastvideo/models/hunyuan_hf/modeling_hunyuan.py b/fastvideo/models/hunyuan_hf/modeling_hunyuan.py deleted file mode 100644 index 39b1a815a..000000000 --- a/fastvideo/models/hunyuan_hf/modeling_hunyuan.py +++ /dev/null @@ -1,836 +0,0 @@ -# Copyright 2024 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, List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin -from diffusers.models.attention import FeedForward -from diffusers.models.attention_processor import Attention, AttentionProcessor -from diffusers.models.embeddings import (CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, - get_1d_rotary_pos_embed) -from diffusers.models.modeling_outputs import Transformer2DModelOutput -from diffusers.models.modeling_utils import ModelMixin -from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle -from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers - -from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad -from fastvideo.utils.communications import all_gather, all_to_all_4D -from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def shrink_head(encoder_state, dim): - local_heads = encoder_state.shape[dim] // nccl_info.sp_size - return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) - - -class HunyuanVideoAttnProcessor2_0: - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "HunyuanVideoAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.") - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - - sequence_length = hidden_states.size(1) - encoder_sequence_length = encoder_hidden_states.size(1) - if attn.add_q_proj is None and encoder_hidden_states is not None: - hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) - - # 1. QKV projections - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2) - key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2) - value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2) - - # 2. QK normalization - if attn.norm_q is not None: - query = attn.norm_q(query).to(value) - if attn.norm_k is not None: - key = attn.norm_k(key).to(value) - - image_rotary_emb = ( - shrink_head(image_rotary_emb[0], dim=0), - shrink_head(image_rotary_emb[1], dim=0), - ) - - # 3. Rotational positional embeddings applied to latent stream - if image_rotary_emb is not None: - from diffusers.models.embeddings import apply_rotary_emb - - if attn.add_q_proj is None and encoder_hidden_states is not None: - query = torch.cat( - [ - apply_rotary_emb(query[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb), - query[:, :, -encoder_hidden_states.shape[1]:], - ], - dim=2, - ) - key = torch.cat( - [ - apply_rotary_emb(key[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb), - key[:, :, -encoder_hidden_states.shape[1]:], - ], - dim=2, - ) - else: - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - # 4. Encoder condition QKV projection and normalization - if attn.add_q_proj is not None and encoder_hidden_states is not None: - encoder_query = attn.add_q_proj(encoder_hidden_states) - encoder_key = attn.add_k_proj(encoder_hidden_states) - encoder_value = attn.add_v_proj(encoder_hidden_states) - - encoder_query = encoder_query.unflatten(2, (attn.heads, -1)).transpose(1, 2) - encoder_key = encoder_key.unflatten(2, (attn.heads, -1)).transpose(1, 2) - encoder_value = encoder_value.unflatten(2, (attn.heads, -1)).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_query = attn.norm_added_q(encoder_query).to(encoder_value) - if attn.norm_added_k is not None: - encoder_key = attn.norm_added_k(encoder_key).to(encoder_value) - - query = torch.cat([query, encoder_query], dim=2) - key = torch.cat([key, encoder_key], dim=2) - value = torch.cat([value, encoder_value], dim=2) - - if get_sequence_parallel_state(): - query_img, query_txt = query[:, :, :sequence_length, :], query[:, :, sequence_length:, :] - key_img, key_txt = key[:, :, :sequence_length, :], key[:, :, sequence_length:, :] - value_img, value_txt = value[:, :, :sequence_length, :], value[:, :, sequence_length:, :] - query_img = all_to_all_4D(query_img, scatter_dim=1, gather_dim=2) # - key_img = all_to_all_4D(key_img, scatter_dim=1, gather_dim=2) - value_img = all_to_all_4D(value_img, scatter_dim=1, gather_dim=2) - - query_txt = shrink_head(query_txt, dim=1) - key_txt = shrink_head(key_txt, dim=1) - value_txt = shrink_head(value_txt, dim=1) - query = torch.cat([query_img, query_txt], dim=2) - key = torch.cat([key_img, key_txt], dim=2) - value = torch.cat([value_img, value_txt], dim=2) - - query = query.unsqueeze(2) - key = key.unsqueeze(2) - value = value.unsqueeze(2) - qkv = torch.cat([query, key, value], dim=2) - qkv = qkv.transpose(1, 3) - - # 5. Attention - attention_mask = attention_mask[:, 0, :] - seq_len = qkv.shape[1] - attn_len = attention_mask.shape[1] - attention_mask = F.pad(attention_mask, (seq_len - attn_len, 0), value=True) - - hidden_states = flash_attn_no_pad(qkv, attention_mask, causal=False, dropout_p=0.0, softmax_scale=None) - - if get_sequence_parallel_state(): - hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( - (sequence_length * nccl_info.sp_size, encoder_sequence_length), dim=1) - hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) - encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous() - hidden_states = hidden_states.flatten(2, 3) - hidden_states = hidden_states.to(query.dtype) - encoder_hidden_states = encoder_hidden_states.flatten(2, 3) - encoder_hidden_states = encoder_hidden_states.to(query.dtype) - else: - hidden_states = hidden_states.flatten(2, 3) - hidden_states = hidden_states.to(query.dtype) - - # 6. Output projection - if encoder_hidden_states is not None: - hidden_states, encoder_hidden_states = ( - hidden_states[:, :-encoder_hidden_states.shape[1]], - hidden_states[:, -encoder_hidden_states.shape[1]:], - ) - - if encoder_hidden_states is not None: - if getattr(attn, "to_out", None) is not None: - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - if getattr(attn, "to_add_out", None) is not None: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - - -class HunyuanVideoPatchEmbed(nn.Module): - - def __init__( - self, - patch_size: Union[int, Tuple[int, int, int]] = 16, - in_chans: int = 3, - embed_dim: int = 768, - ) -> None: - super().__init__() - - patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size - self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.proj(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) # BCFHW -> BNC - return hidden_states - - -class HunyuanVideoAdaNorm(nn.Module): - - def __init__(self, in_features: int, out_features: Optional[int] = None) -> None: - super().__init__() - - out_features = out_features or 2 * in_features - self.linear = nn.Linear(in_features, out_features) - self.nonlinearity = nn.SiLU() - - def forward(self, - temb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - temb = self.linear(self.nonlinearity(temb)) - gate_msa, gate_mlp = temb.chunk(2, dim=1) - gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1) - return gate_msa, gate_mlp - - -class HunyuanVideoIndividualTokenRefinerBlock(nn.Module): - - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - mlp_width_ratio: str = 4.0, - mlp_drop_rate: float = 0.0, - attention_bias: bool = True, - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - - self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) - self.attn = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - heads=num_attention_heads, - dim_head=attention_head_dim, - bias=attention_bias, - ) - - self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) - self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate) - - self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size) - - def forward( - self, - hidden_states: torch.Tensor, - temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - norm_hidden_states = self.norm1(hidden_states) - - attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=None, - attention_mask=attention_mask, - ) - - gate_msa, gate_mlp = self.norm_out(temb) - hidden_states = hidden_states + attn_output * gate_msa - - ff_output = self.ff(self.norm2(hidden_states)) - hidden_states = hidden_states + ff_output * gate_mlp - - return hidden_states - - -class HunyuanVideoIndividualTokenRefiner(nn.Module): - - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - num_layers: int, - mlp_width_ratio: float = 4.0, - mlp_drop_rate: float = 0.0, - attention_bias: bool = True, - ) -> None: - super().__init__() - - self.refiner_blocks = nn.ModuleList([ - HunyuanVideoIndividualTokenRefinerBlock( - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - mlp_width_ratio=mlp_width_ratio, - mlp_drop_rate=mlp_drop_rate, - attention_bias=attention_bias, - ) for _ in range(num_layers) - ]) - - def forward( - self, - hidden_states: torch.Tensor, - temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - ) -> None: - self_attn_mask = None - if attention_mask is not None: - batch_size = attention_mask.shape[0] - seq_len = attention_mask.shape[1] - attention_mask = attention_mask.to(hidden_states.device).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) - self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() - self_attn_mask[:, :, :, 0] = True - - for block in self.refiner_blocks: - hidden_states = block(hidden_states, temb, self_attn_mask) - - return hidden_states - - -class HunyuanVideoTokenRefiner(nn.Module): - - def __init__( - self, - in_channels: int, - num_attention_heads: int, - attention_head_dim: int, - num_layers: int, - mlp_ratio: float = 4.0, - mlp_drop_rate: float = 0.0, - attention_bias: bool = True, - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - - self.time_text_embed = CombinedTimestepTextProjEmbeddings(embedding_dim=hidden_size, - pooled_projection_dim=in_channels) - self.proj_in = nn.Linear(in_channels, hidden_size, bias=True) - self.token_refiner = HunyuanVideoIndividualTokenRefiner( - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - num_layers=num_layers, - mlp_width_ratio=mlp_ratio, - mlp_drop_rate=mlp_drop_rate, - attention_bias=attention_bias, - ) - - def forward( - self, - hidden_states: torch.Tensor, - timestep: torch.LongTensor, - attention_mask: Optional[torch.LongTensor] = None, - ) -> torch.Tensor: - if attention_mask is None: - pooled_projections = hidden_states.mean(dim=1) - else: - original_dtype = hidden_states.dtype - mask_float = attention_mask.float().unsqueeze(-1) - pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1) - pooled_projections = pooled_projections.to(original_dtype) - - temb = self.time_text_embed(timestep, pooled_projections) - hidden_states = self.proj_in(hidden_states) - hidden_states = self.token_refiner(hidden_states, temb, attention_mask) - - return hidden_states - - -class HunyuanVideoRotaryPosEmbed(nn.Module): - - def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None: - super().__init__() - - self.patch_size = patch_size - self.patch_size_t = patch_size_t - self.rope_dim = rope_dim - self.theta = theta - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - rope_sizes = [ - num_frames * nccl_info.sp_size // self.patch_size_t, height // self.patch_size, width // self.patch_size - ] - - axes_grids = [] - for i in range(3): - # Note: The following line diverges from original behaviour. We create the grid on the device, whereas - # original implementation creates it on CPU and then moves it to device. This results in numerical - # differences in layerwise debugging outputs, but visually it is the same. - grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32) - axes_grids.append(grid) - grid = torch.meshgrid(*axes_grids, indexing="ij") # [W, H, T] - grid = torch.stack(grid, dim=0) # [3, W, H, T] - - freqs = [] - for i in range(3): - freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True) - freqs.append(freq) - - freqs_cos = torch.cat([f[0] for f in freqs], dim=1) # (W * H * T, D / 2) - freqs_sin = torch.cat([f[1] for f in freqs], dim=1) # (W * H * T, D / 2) - return freqs_cos, freqs_sin - - -class HunyuanVideoSingleTransformerBlock(nn.Module): - - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - mlp_ratio: float = 4.0, - qk_norm: str = "rms_norm", - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - mlp_dim = int(hidden_size * mlp_ratio) - - self.attn = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=hidden_size, - bias=True, - processor=HunyuanVideoAttnProcessor2_0(), - qk_norm=qk_norm, - eps=1e-6, - pre_only=True, - ) - - self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm") - self.proj_mlp = nn.Linear(hidden_size, mlp_dim) - self.act_mlp = nn.GELU(approximate="tanh") - self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ) -> torch.Tensor: - text_seq_length = encoder_hidden_states.shape[1] - hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) - - residual = hidden_states - - # 1. Input normalization - norm_hidden_states, gate = self.norm(hidden_states, emb=temb) - mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) - - norm_hidden_states, norm_encoder_hidden_states = ( - norm_hidden_states[:, :-text_seq_length, :], - norm_hidden_states[:, -text_seq_length:, :], - ) - - # 2. Attention - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - attn_output = torch.cat([attn_output, context_attn_output], dim=1) - - # 3. Modulation and residual connection - hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) - hidden_states = gate.unsqueeze(1) * self.proj_out(hidden_states) - hidden_states = hidden_states + residual - - hidden_states, encoder_hidden_states = ( - hidden_states[:, :-text_seq_length, :], - hidden_states[:, -text_seq_length:, :], - ) - return hidden_states, encoder_hidden_states - - -class HunyuanVideoTransformerBlock(nn.Module): - - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - mlp_ratio: float, - qk_norm: str = "rms_norm", - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - - self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm") - self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm") - - self.attn = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - added_kv_proj_dim=hidden_size, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=hidden_size, - context_pre_only=False, - bias=True, - processor=HunyuanVideoAttnProcessor2_0(), - qk_norm=qk_norm, - eps=1e-6, - ) - - self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") - - self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - # 1. Input normalization - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) - norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( - encoder_hidden_states, emb=temb) - - # 2. Joint attention - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=freqs_cis, - ) - - # 3. Modulation and residual connection - hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1) - encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1) - - norm_hidden_states = self.norm2(hidden_states) - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - - # 4. Feed-forward - ff_output = self.ff(norm_hidden_states) - context_ff_output = self.ff_context(norm_encoder_hidden_states) - - hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output - encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output - - return hidden_states, encoder_hidden_states - - -class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): - r""" - A Transformer model for video-like data used in [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo). - - Args: - in_channels (`int`, defaults to `16`): - The number of channels in the input. - out_channels (`int`, defaults to `16`): - The number of channels in the output. - num_attention_heads (`int`, defaults to `24`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `128`): - The number of channels in each head. - num_layers (`int`, defaults to `20`): - The number of layers of dual-stream blocks to use. - num_single_layers (`int`, defaults to `40`): - The number of layers of single-stream blocks to use. - num_refiner_layers (`int`, defaults to `2`): - The number of layers of refiner blocks to use. - mlp_ratio (`float`, defaults to `4.0`): - The ratio of the hidden layer size to the input size in the feedforward network. - patch_size (`int`, defaults to `2`): - The size of the spatial patches to use in the patch embedding layer. - patch_size_t (`int`, defaults to `1`): - The size of the tmeporal patches to use in the patch embedding layer. - qk_norm (`str`, defaults to `rms_norm`): - The normalization to use for the query and key projections in the attention layers. - guidance_embeds (`bool`, defaults to `True`): - Whether to use guidance embeddings in the model. - text_embed_dim (`int`, defaults to `4096`): - Input dimension of text embeddings from the text encoder. - pooled_projection_dim (`int`, defaults to `768`): - The dimension of the pooled projection of the text embeddings. - rope_theta (`float`, defaults to `256.0`): - The value of theta to use in the RoPE layer. - rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`): - The dimensions of the axes to use in the RoPE layer. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - in_channels: int = 16, - out_channels: int = 16, - num_attention_heads: int = 24, - attention_head_dim: int = 128, - num_layers: int = 20, - num_single_layers: int = 40, - num_refiner_layers: int = 2, - mlp_ratio: float = 4.0, - patch_size: int = 2, - patch_size_t: int = 1, - qk_norm: str = "rms_norm", - guidance_embeds: bool = True, - text_embed_dim: int = 4096, - pooled_projection_dim: int = 768, - rope_theta: float = 256.0, - rope_axes_dim: Tuple[int] = (16, 56, 56), - ) -> None: - super().__init__() - - inner_dim = num_attention_heads * attention_head_dim - out_channels = out_channels or in_channels - - # 1. Latent and condition embedders - self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim) - self.context_embedder = HunyuanVideoTokenRefiner(text_embed_dim, - num_attention_heads, - attention_head_dim, - num_layers=num_refiner_layers) - self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim) - - # 2. RoPE - self.rope = HunyuanVideoRotaryPosEmbed(patch_size, patch_size_t, rope_axes_dim, rope_theta) - - # 3. Dual stream transformer blocks - self.transformer_blocks = nn.ModuleList([ - HunyuanVideoTransformerBlock(num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm) - for _ in range(num_layers) - ]) - - # 4. Single stream transformer blocks - self.single_transformer_blocks = nn.ModuleList([ - HunyuanVideoSingleTransformerBlock(num_attention_heads, - attention_head_dim, - mlp_ratio=mlp_ratio, - qk_norm=qk_norm) for _ in range(num_single_layers) - ]) - - # 5. Output projection - self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): - if hasattr(module, "get_processor"): - processors[f"{name}.processor"] = module.get_processor() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes.") - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def _set_gradient_checkpointing(self, module, value=False): - if hasattr(module, "gradient_checkpointing"): - module.gradient_checkpointing = value - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_attention_mask: torch.Tensor, - guidance: torch.Tensor = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - return_dict: bool = True, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: - if guidance is None: - guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16) - - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.") - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p, p_t = self.config.patch_size, self.config.patch_size_t - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p - post_patch_width = width // p - - pooled_projections = encoder_hidden_states[:, 0, :self.config.pooled_projection_dim] - encoder_hidden_states = encoder_hidden_states[:, 1:] - - # 1. RoPE - image_rotary_emb = self.rope(hidden_states) - - # 2. Conditional embeddings - temb = self.time_text_embed(timestep, guidance, pooled_projections) - hidden_states = self.x_embedder(hidden_states) - encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) - - # 3. Attention mask preparation - latent_sequence_length = hidden_states.shape[1] - condition_sequence_length = encoder_hidden_states.shape[1] - sequence_length = latent_sequence_length + condition_sequence_length - attention_mask = torch.zeros(batch_size, - sequence_length, - sequence_length, - device=hidden_states.device, - dtype=torch.bool) # [B, N, N] - - effective_condition_sequence_length = encoder_attention_mask.sum(dim=1, dtype=torch.int) - effective_sequence_length = latent_sequence_length + effective_condition_sequence_length - - for i in range(batch_size): - attention_mask[i, :effective_sequence_length[i], :effective_sequence_length[i]] = True - - # 4. Transformer blocks - if torch.is_grad_enabled() and self.gradient_checkpointing: - - def create_custom_forward(module, return_dict=None): - - def custom_forward(*inputs): - if return_dict is not None: - return module(*inputs, return_dict=return_dict) - else: - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - - for block in self.transformer_blocks: - hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - hidden_states, - encoder_hidden_states, - temb, - attention_mask, - image_rotary_emb, - **ckpt_kwargs, - ) - - for block in self.single_transformer_blocks: - hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - hidden_states, - encoder_hidden_states, - temb, - attention_mask, - image_rotary_emb, - **ckpt_kwargs, - ) - - else: - for block in self.transformer_blocks: - hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask, - image_rotary_emb) - - for block in self.single_transformer_blocks: - hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask, - image_rotary_emb) - - # 5. Output projection - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width, - -1, p_t, p, p) - hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7) - hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - - if not return_dict: - return (hidden_states, ) - - return Transformer2DModelOutput(sample=hidden_states) diff --git a/fastvideo/models/hunyuan_hf/pipeline_hunyuan.py b/fastvideo/models/hunyuan_hf/pipeline_hunyuan.py deleted file mode 100644 index a30f083e1..000000000 --- a/fastvideo/models/hunyuan_hf/pipeline_hunyuan.py +++ /dev/null @@ -1,691 +0,0 @@ -# Copyright 2024 The HunyuanVideo 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. - -import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn.functional as F -from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback -from diffusers.loaders import HunyuanVideoLoraLoaderMixin -from diffusers.models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel -from diffusers.pipelines.hunyuan_video.pipeline_output import HunyuanVideoPipelineOutput -from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from diffusers.schedulers import FlowMatchEulerDiscreteScheduler -from diffusers.utils import logging, replace_example_docstring -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from einops import rearrange -from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast - -from fastvideo.utils.communications import all_gather -from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -EXAMPLE_DOC_STRING = """ - Examples: - ```python - >>> import torch - >>> from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel - >>> from diffusers.utils import export_to_video - - >>> model_id = "tencent/HunyuanVideo" - >>> transformer = HunyuanVideoTransformer3DModel.from_pretrained( - ... model_id, subfolder="transformer", torch_dtype=torch.bfloat16 - ... ) - >>> pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.float16) - >>> pipe.vae.enable_tiling() - >>> pipe.to("cuda") - - >>> output = pipe( - ... prompt="A cat walks on the grass, realistic", - ... height=320, - ... width=512, - ... num_frames=61, - ... num_inference_steps=30, - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=15) - ``` -""" - -DEFAULT_PROMPT_TEMPLATE = { - "template": ("<|start_header_id|>system<|end_header_id|>\n\nDescribe 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:<|eot_id|>" - "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"), - "crop_start": - 95, -} - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps -def retrieve_timesteps( - scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, - **kwargs, -): - r""" - Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles - custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. - - Args: - scheduler (`SchedulerMixin`): - The scheduler to get timesteps from. - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` - must be `None`. - device (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): - Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, - `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): - Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, - `num_inference_steps` and `timesteps` must be `None`. - - Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the - second element is the number of inference steps. - """ - if timesteps is not None and sigmas is not None: - raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") - if timesteps is not None: - accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accepts_timesteps: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" timestep schedules. Please check whether you are using the correct scheduler.") - scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - elif sigmas is not None: - accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accept_sigmas: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" sigmas schedules. Please check whether you are using the correct scheduler.") - scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - else: - scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) - timesteps = scheduler.timesteps - return timesteps, num_inference_steps - - -class HunyuanVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin): - r""" - Pipeline for text-to-video generation using HunyuanVideo. - - This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - text_encoder ([`LlamaModel`]): - [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - tokenizer_2 (`LlamaTokenizer`): - Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - transformer ([`HunyuanVideoTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKLHunyuanVideo`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder_2 ([`CLIPTextModel`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically - the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. - tokenizer_2 (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - text_encoder: LlamaModel, - tokenizer: LlamaTokenizerFast, - transformer: HunyuanVideoTransformer3DModel, - vae: AutoencoderKLHunyuanVideo, - scheduler: FlowMatchEulerDiscreteScheduler, - text_encoder_2: CLIPTextModel, - tokenizer_2: CLIPTokenizer, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - ) - - self.vae_scale_factor_temporal = (self.vae.temporal_compression_ratio - if hasattr(self, "vae") and self.vae is not None else 4) - self.vae_scale_factor_spatial = (self.vae.spatial_compression_ratio - if hasattr(self, "vae") and self.vae is not None else 8) - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - def _get_llama_prompt_embeds( - self, - prompt: Union[str, List[str]], - prompt_template: Dict[str, Any], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 256, - num_hidden_layers_to_skip: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: - device = device or self._execution_device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - prompt = [prompt_template["template"].format(p) for p in prompt] - - crop_start = prompt_template.get("crop_start", None) - if crop_start is None: - prompt_template_input = self.tokenizer( - prompt_template["template"], - padding="max_length", - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=False, - ) - crop_start = prompt_template_input["input_ids"].shape[-1] - # Remove <|eot_id|> token and placeholder {} - crop_start -= 2 - - max_sequence_length += crop_start - text_inputs = self.tokenizer( - prompt, - max_length=max_sequence_length, - padding="max_length", - truncation=True, - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=True, - ) - text_input_ids = text_inputs.input_ids.to(device=device) - prompt_attention_mask = text_inputs.attention_mask.to(device=device) - - prompt_embeds = self.text_encoder( - input_ids=text_input_ids, - attention_mask=prompt_attention_mask, - output_hidden_states=True, - ).hidden_states[-(num_hidden_layers_to_skip + 1)] - prompt_embeds = prompt_embeds.to(dtype=dtype) - - if crop_start is not None and crop_start > 0: - prompt_embeds = prompt_embeds[:, crop_start:] - prompt_attention_mask = prompt_attention_mask[:, crop_start:] - - # duplicate text embeddings for each generation per prompt, using mps friendly method - _, seq_len, _ = prompt_embeds.shape - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.repeat(1, num_videos_per_prompt) - prompt_attention_mask = prompt_attention_mask.view(batch_size * num_videos_per_prompt, seq_len) - - return prompt_embeds, prompt_attention_mask - - def _get_clip_prompt_embeds( - self, - prompt: Union[str, List[str]], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 77, - ) -> torch.Tensor: - device = device or self._execution_device - dtype = dtype or self.text_encoder_2.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = self.tokenizer_2( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ) - - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1]) - logger.warning("The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_sequence_length} tokens: {removed_text}") - - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, -1) - - return prompt_embeds - - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]] = None, - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, - num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 256, - ): - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds( - prompt, - prompt_template, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=max_sequence_length, - ) - - if pooled_prompt_embeds is None: - if prompt_2 is None and pooled_prompt_embeds is None: - prompt_2 = prompt - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=77, - ) - - return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - prompt_template=None, - ): - if height % 16 != 0 or width % 16 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") - - if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs - for k in callback_on_step_end_tensor_inputs): - raise ValueError( - f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" - ) - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two.") - elif prompt_2 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two.") - elif prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.") - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): - raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") - - if prompt_template is not None: - if not isinstance(prompt_template, dict): - raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}") - if "template" not in prompt_template: - raise ValueError( - f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}") - - def prepare_latents( - self, - batch_size: int, - num_channels_latents: 32, - height: int = 720, - width: int = 1280, - num_frames: int = 129, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if latents is not None: - return latents.to(device=device, dtype=dtype) - - shape = ( - batch_size, - num_channels_latents, - num_frames, - int(height) // self.vae_scale_factor_spatial, - int(width) // self.vae_scale_factor_spatial, - ) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators.") - - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - return latents - - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def interrupt(self): - return self._interrupt - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - prompt_2: Union[str, List[str]] = None, - height: int = 720, - width: int = 1280, - num_frames: int = 129, - num_inference_steps: int = 50, - sigmas: List[float] = None, - guidance_scale: float = 6.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, - MultiPipelineCallbacks]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, - max_sequence_length: int = 256, - ): - r""" - The call function to the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead. - height (`int`, defaults to `720`): - The height in pixels of the generated image. - width (`int`, defaults to `1280`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `129`): - The number of frames in the generated video. - num_inference_steps (`int`, defaults to `50`): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - sigmas (`List[float]`, *optional*): - Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in - their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed - will be used. - guidance_scale (`float`, defaults to `6.0`): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. Note that the only available HunyuanVideo model is - CFG-distilled, which means that traditional guidance between unconditional and conditional latent is - not applied. - num_videos_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor is generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not - provided, text embeddings are generated from the `prompt` input argument. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a plain tuple. - attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in - [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - clip_skip (`int`, *optional*): - Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that - the output of the pre-final layer will be used for computing the prompt embeddings. - callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): - A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of - each denoising step during the inference. with the following arguments: `callback_on_step_end(self: - DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a - list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. - - Examples: - - Returns: - [`~HunyuanVideoPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned, otherwise a `tuple` is returned - where the first element is a list with the generated images and the second element is a list of `bool`s - indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. - """ - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - height, - width, - prompt_embeds, - callback_on_step_end_tensor_inputs, - prompt_template, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - device = self._execution_device - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - # 3. Encode input prompt - prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt( - prompt=prompt, - prompt_2=prompt, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - ) - - transformer_dtype = self.transformer.dtype - prompt_embeds = prompt_embeds.to(transformer_dtype) - prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) - if pooled_prompt_embeds is not None: - pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype) - - # 4. Prepare timesteps - sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - sigmas=sigmas, - ) - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_latent_frames, - torch.float32, - device, - generator, - latents, - ) - # check sequence_parallel - world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group - if get_sequence_parallel_state(): - latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous() - latents = latents[:, :, rank, :, :, :] - - # 6. Prepare guidance condition - guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0 - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - self._num_timesteps = len(timesteps) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - latent_model_input = latents.to(transformer_dtype) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latents.shape[0]).to(latents.dtype) - if pooled_prompt_embeds.shape[-1] != prompt_embeds.shape[-1]: - pooled_prompt_embeds_padding = F.pad( - pooled_prompt_embeds, - (0, prompt_embeds.shape[2] - pooled_prompt_embeds.shape[1]), - value=0, - ).unsqueeze(1) - encoder_hidden_states = torch.cat([pooled_prompt_embeds_padding, prompt_embeds], dim=1) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=encoder_hidden_states, # [1, 257, 4096] - timestep=timestep, - encoder_attention_mask=prompt_attention_mask, - guidance=guidance, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - - if callback_on_step_end is not None: - callback_kwargs = {} - for k in callback_on_step_end_tensor_inputs: - callback_kwargs[k] = locals()[k] - callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) - - latents = callback_outputs.pop("latents", latents) - prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) - - # call the callback, if provided - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() - - if get_sequence_parallel_state(): - latents = all_gather(latents, dim=2) - - if not output_type == "latent": - latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video, ) - - return HunyuanVideoPipelineOutput(frames=video) diff --git a/fastvideo/models/loader.py b/fastvideo/models/loader.py new file mode 100644 index 000000000..ee3be6425 --- /dev/null +++ b/fastvideo/models/loader.py @@ -0,0 +1,304 @@ +from abc import ABC, abstractmethod + +import torch +from fastvideo.inference_args import InferenceArgs +from fastvideo.logger import init_logger +import os +from fastvideo.distributed.parallel_state import initialize_sequence_parallel_group +import glob +from fastvideo.loader.fsdp_load import load_fsdp_model +from fastvideo.loader.loader import get_model_loader +from transformers import PretrainedConfig, AutoTokenizer +from fastvideo.models.hf_transformer_utils import get_hf_config, get_diffusers_config +from fastvideo.models import get_scheduler +from fastvideo.models.registry import ModelRegistry +from safetensors.torch import load_file as safetensors_load_file + +import json + + +logger = init_logger(__name__) + + +class ComponentLoader(ABC): + """Base class for loading a specific type of model component.""" + + def __init__(self, device=None): + self.device = device + + @abstractmethod + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """ + Load the component based on the model path, architecture, and inference args. + + Args: + model_path: Path to the component model + architecture: Architecture of the component model + inference_args: Inference arguments + + Returns: + The loaded component + """ + raise NotImplementedError + + @classmethod + def for_module_type(cls, module_type: str, transformers_or_diffusers: str) -> 'ComponentLoader': + """ + Factory method to create a component loader for a specific module type. + + Args: + module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler") + transformers_or_diffusers: Whether the module is from transformers or diffusers + + Returns: + A component loader for the specified module type + """ + # Map of module types to their loader classes and expected library + module_loaders = { + "scheduler": (SchedulerLoader, "diffusers"), + "transformer": (TransformerLoader, "diffusers"), + "vae": (VAELoader, "diffusers"), + "text_encoder": (TextEncoderLoader, "transformers"), + "text_encoder_2": (TextEncoderLoader, "transformers"), + "tokenizer": (TokenizerLoader, "transformers"), + "tokenizer_2": (TokenizerLoader, "transformers"), + } + + if module_type in module_loaders: + loader_cls, expected_library = module_loaders[module_type] + # Assert that the library matches what's expected for this module type + assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + return loader_cls() + + # For unknown module types, use a generic loader + logger.warning(f"No specific loader found for module type: {module_type}. Using generic loader.") + return GenericComponentLoader(transformers_or_diffusers) + +class TextEncoderLoader(ComponentLoader): + """Loader for text encoders.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the text encoders based on the model path, architecture, and inference args.""" + use_v1 = True + if not use_v1: + return self.load_v0(model_path, architecture, inference_args) + else: + return self.load_v1(model_path, architecture, inference_args) + + def load_v1(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the text encoders based on the model path, architecture, and inference args.""" + model_config: PretrainedConfig = get_hf_config( + model=model_path, + trust_remote_code=inference_args.trust_remote_code, + revision=inference_args.revision, + model_override_args=None, + inference_args=inference_args, + ) + logger.info(f"HF Model config: {model_config}") + model_loader = get_model_loader(inference_args) + model = model_loader.load_model(model_path, model_config, inference_args) + return model + + def load_v0(self, model_path: str, architecture: str, inference_args: InferenceArgs): + from fastvideo.pipelines.hunyuan.constants import PROMPT_TEMPLATE + from fastvideo.models.hunyuan.text_encoder import TextEncoder + # Text encoder + if inference_args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0) + elif inference_args.prompt_template is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template].get("crop_start", 0) + else: + crop_start = 0 + max_length = inference_args.text_len + crop_start + + # prompt_template + prompt_template = (PROMPT_TEMPLATE[inference_args.prompt_template] if inference_args.prompt_template is not None else None) + + # prompt_template_video + prompt_template_video = (PROMPT_TEMPLATE[inference_args.prompt_template_video] + if inference_args.prompt_template_video is not None else None) + + text_encoder = TextEncoder( + text_encoder_type=inference_args.text_encoder, + text_encoder_path=model_path, + max_length=max_length, + text_encoder_precision=inference_args.text_encoder_precision, + tokenizer_type=inference_args.tokenizer, + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=inference_args.hidden_state_skip_layer, + apply_final_norm=inference_args.apply_final_norm, + reproduce=inference_args.reproduce, + logger=logger, + device=self.device if not inference_args.use_cpu_offload else "cpu", + ) + return text_encoder + +class TokenizerLoader(ComponentLoader): + """Loader for tokenizers.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the tokenizer based on the model path, architecture, and inference args.""" + logger.info(f"Loading tokenizer from {model_path}") + + tokenizer = AutoTokenizer.from_pretrained( + model_path, + # TODO(will): pass these tokenizer kwargs from inference args? Maybe + # other method of config? + padding_size='right', + ) + logger.info(f"Loaded tokenizer: {tokenizer.__class__.__name__}") + return tokenizer + +class VAELoader(ComponentLoader): + """Loader for VAE.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the VAE based on the model path, architecture, and inference args.""" + # TODO(will): move this to a constants file + from fastvideo.pipelines.hunyuan.constants import PRECISION_TO_TYPE + + config = get_diffusers_config(model=model_path) + + class_name = config.pop("_class_name") + assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported." + config.pop("_diffusers_version") + + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) + + vae = vae_cls(**config).to(inference_args.device) + + # 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 {path}" + loaded = safetensors_load_file(safetensors_list[0]) + vae.load_state_dict(loaded) + dtype = PRECISION_TO_TYPE[inference_args.vae_precision] + vae = vae.eval().to(dtype) + + # TODO(will): should we define hunyuan vae config class? + vae_kwargs = { + "s_ratio": config["spatial_compression_ratio"], + "t_ratio": config["temporal_compression_ratio"], + } + + vae.kwargs = vae_kwargs + + return vae + +class TransformerLoader(ComponentLoader): + """Loader for transformer.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the transformer based on the model path, architecture, and inference args.""" + + model_config = get_diffusers_config(model=model_path) + cls_name = model_config.pop("_class_name") + if cls_name is None: + raise ValueError(f"Model config does not contain a _class_name attribute. " + "Only diffusers format is supported.") + model_config.pop("_diffusers_version") + + model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) + + # Find all safetensors files + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) + if not safetensors_list: + raise ValueError(f"No safetensors files found in {model_path}") + + logger.info(f"Loading model from {len(safetensors_list)} safetensors files in {model_path}") + + # initialize_sequence_parallel_group(inference_args.sp_size) + + # Load the model using FSDP loader + logger.info(f"Loading model from {cls_name}") + model = load_fsdp_model( + model_cls=model_cls, + init_params=model_config, + weight_dir_list=safetensors_list, + device=inference_args.device, + cpu_offload=inference_args.use_cpu_offload + ) + + total_params = sum(p.numel() for p in model.parameters()) + logger.info(f"Loaded model with {total_params / 1e9:.2f}B parameters") + + model.eval() + return model + +class SchedulerLoader(ComponentLoader): + """Loader for scheduler.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the scheduler based on the model path, architecture, and inference args.""" + + scheduler = get_scheduler( + module_path=model_path, + architecture=architecture, + inference_args=inference_args, + ) + logger.info(f"Scheduler loaded: {scheduler}") + return scheduler + +class GenericComponentLoader(ComponentLoader): + """Generic loader for components that don't have a specific loader.""" + + def __init__(self, library="transformers"): + super().__init__() + self.library = library + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load a generic component based on the model path, architecture, and inference args.""" + logger.warning(f"Using generic loader for {model_path} with library {self.library}") + + if self.library == "transformers": + from transformers import AutoModel + + model = AutoModel.from_pretrained( + model_path, + trust_remote_code=inference_args.trust_remote_code, + revision=inference_args.revision, + ) + logger.info(f"Loaded generic transformers model: {model.__class__.__name__}") + return model + elif self.library == "diffusers": + logger.warning(f"Generic loading for diffusers components is not fully implemented") + from fastvideo.models.hf_transformer_utils import get_diffusers_config + + model_config = get_diffusers_config(model=model_path) + logger.info(f"Diffusers Model config: {model_config}") + # This is a placeholder - in a real implementation, you'd need to handle this properly + return None + else: + raise ValueError(f"Unsupported library: {self.library}") + +class PipelineComponentLoader: + """ + Utility class for loading pipeline components. + This replaces the chain of if-else statements in load_pipeline_module. + """ + + @staticmethod + def load_module(module_name: str, component_model_path: str, transformers_or_diffusers: str, + architecture: str, inference_args: InferenceArgs): + """ + Load a pipeline module. + + Args: + module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler") + component_model_path: Path to the component model + transformers_or_diffusers: Whether the module is from transformers or diffusers + architecture: Architecture of the component model + inference_args: Inference arguments + + Returns: + The loaded module + """ + logger.info(f"Loading {module_name} using {transformers_or_diffusers} from {component_model_path}") + + # Get the appropriate loader for this module type + loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers) + + # Load the module + return loader.load(component_model_path, architecture, inference_args) diff --git a/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py b/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py deleted file mode 100644 index 4249ba971..000000000 --- a/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py +++ /dev/null @@ -1,361 +0,0 @@ -import argparse -import os - -import torch -from safetensors.torch import save_file - -parser = argparse.ArgumentParser() -parser.add_argument("--diffusers_path", required=True, type=str) -parser.add_argument("--transformer_path", type=str, default=None, help="Path to save transformer model") -parser.add_argument("--vae_encoder_path", type=str, default=None, help="Path to save VAE encoder model") -parser.add_argument("--vae_decoder_path", type=str, default=None, help="Path to save VAE decoder model") - -args = parser.parse_args() - - -def reverse_scale_shift(weight, dim): - scale, shift = weight.chunk(2, dim=0) - new_weight = torch.cat([shift, scale], dim=0) - return new_weight - - -def reverse_proj_gate(weight): - gate, proj = weight.chunk(2, dim=0) - new_weight = torch.cat([proj, gate], dim=0) - return new_weight - - -def convert_diffusers_transformer_to_mochi(state_dict): - original_state_dict = state_dict.copy() - new_state_dict = {} - - # Convert patch_embed - new_state_dict["x_embedder.proj.weight"] = original_state_dict.pop("patch_embed.proj.weight") - new_state_dict["x_embedder.proj.bias"] = original_state_dict.pop("patch_embed.proj.bias") - - # Convert time_embed - new_state_dict["t_embedder.mlp.0.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.weight") - new_state_dict["t_embedder.mlp.0.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.bias") - new_state_dict["t_embedder.mlp.2.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.weight") - new_state_dict["t_embedder.mlp.2.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.bias") - new_state_dict["t5_y_embedder.to_kv.weight"] = original_state_dict.pop("time_embed.pooler.to_kv.weight") - new_state_dict["t5_y_embedder.to_kv.bias"] = original_state_dict.pop("time_embed.pooler.to_kv.bias") - new_state_dict["t5_y_embedder.to_q.weight"] = original_state_dict.pop("time_embed.pooler.to_q.weight") - new_state_dict["t5_y_embedder.to_q.bias"] = original_state_dict.pop("time_embed.pooler.to_q.bias") - new_state_dict["t5_y_embedder.to_out.weight"] = original_state_dict.pop("time_embed.pooler.to_out.weight") - new_state_dict["t5_y_embedder.to_out.bias"] = original_state_dict.pop("time_embed.pooler.to_out.bias") - new_state_dict["t5_yproj.weight"] = original_state_dict.pop("time_embed.caption_proj.weight") - new_state_dict["t5_yproj.bias"] = original_state_dict.pop("time_embed.caption_proj.bias") - - # Convert transformer blocks - num_layers = 48 - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - new_prefix = f"blocks.{i}." - - # norm1 - new_state_dict[new_prefix + "mod_x.weight"] = original_state_dict.pop(block_prefix + "norm1.linear.weight") - new_state_dict[new_prefix + "mod_x.bias"] = original_state_dict.pop(block_prefix + "norm1.linear.bias") - - if i < num_layers - 1: - new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(block_prefix + - "norm1_context.linear.weight") - new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(block_prefix + - "norm1_context.linear.bias") - else: - new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(block_prefix + - "norm1_context.linear_1.weight") - new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(block_prefix + - "norm1_context.linear_1.bias") - - # Visual attention - q = original_state_dict.pop(block_prefix + "attn1.to_q.weight") - k = original_state_dict.pop(block_prefix + "attn1.to_k.weight") - v = original_state_dict.pop(block_prefix + "attn1.to_v.weight") - qkv_weight = torch.cat([q, k, v], dim=0) - new_state_dict[new_prefix + "attn.qkv_x.weight"] = qkv_weight - - new_state_dict[new_prefix + "attn.q_norm_x.weight"] = original_state_dict.pop(block_prefix + - "attn1.norm_q.weight") - new_state_dict[new_prefix + "attn.k_norm_x.weight"] = original_state_dict.pop(block_prefix + - "attn1.norm_k.weight") - new_state_dict[new_prefix + "attn.proj_x.weight"] = original_state_dict.pop(block_prefix + - "attn1.to_out.0.weight") - new_state_dict[new_prefix + "attn.proj_x.bias"] = original_state_dict.pop(block_prefix + "attn1.to_out.0.bias") - - # Context attention - q = original_state_dict.pop(block_prefix + "attn1.add_q_proj.weight") - k = original_state_dict.pop(block_prefix + "attn1.add_k_proj.weight") - v = original_state_dict.pop(block_prefix + "attn1.add_v_proj.weight") - qkv_weight = torch.cat([q, k, v], dim=0) - new_state_dict[new_prefix + "attn.qkv_y.weight"] = qkv_weight - - new_state_dict[new_prefix + "attn.q_norm_y.weight"] = original_state_dict.pop(block_prefix + - "attn1.norm_added_q.weight") - new_state_dict[new_prefix + "attn.k_norm_y.weight"] = original_state_dict.pop(block_prefix + - "attn1.norm_added_k.weight") - if i < num_layers - 1: - new_state_dict[new_prefix + "attn.proj_y.weight"] = original_state_dict.pop(block_prefix + - "attn1.to_add_out.weight") - new_state_dict[new_prefix + "attn.proj_y.bias"] = original_state_dict.pop(block_prefix + - "attn1.to_add_out.bias") - - # MLP - new_state_dict[new_prefix + "mlp_x.w1.weight"] = reverse_proj_gate( - original_state_dict.pop(block_prefix + "ff.net.0.proj.weight")) - new_state_dict[new_prefix + "mlp_x.w2.weight"] = original_state_dict.pop(block_prefix + "ff.net.2.weight") - if i < num_layers - 1: - new_state_dict[new_prefix + "mlp_y.w1.weight"] = reverse_proj_gate( - original_state_dict.pop(block_prefix + "ff_context.net.0.proj.weight")) - new_state_dict[new_prefix + "mlp_y.w2.weight"] = original_state_dict.pop(block_prefix + - "ff_context.net.2.weight") - - # Output layers - new_state_dict["final_layer.mod.weight"] = reverse_scale_shift(original_state_dict.pop("norm_out.linear.weight"), - dim=0) - new_state_dict["final_layer.mod.bias"] = reverse_scale_shift(original_state_dict.pop("norm_out.linear.bias"), dim=0) - new_state_dict["final_layer.linear.weight"] = original_state_dict.pop("proj_out.weight") - new_state_dict["final_layer.linear.bias"] = original_state_dict.pop("proj_out.bias") - - new_state_dict["pos_frequencies"] = original_state_dict.pop("pos_frequencies") - - print("Remaining Keys:", original_state_dict.keys()) - - return new_state_dict - - -def convert_diffusers_vae_to_mochi(state_dict): - original_state_dict = state_dict.copy() - encoder_state_dict = {} - decoder_state_dict = {} - - # Convert encoder - prefix = "encoder." - - encoder_state_dict["layers.0.weight"] = original_state_dict.pop(f"{prefix}proj_in.weight") - encoder_state_dict["layers.0.bias"] = original_state_dict.pop(f"{prefix}proj_in.bias") - - # Convert block_in - for i in range(3): - encoder_state_dict[f"layers.{i+1}.stack.0.weight"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight") - encoder_state_dict[f"layers.{i+1}.stack.0.bias"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias") - encoder_state_dict[f"layers.{i+1}.stack.2.weight"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.conv1.conv.weight") - encoder_state_dict[f"layers.{i+1}.stack.2.bias"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.conv1.conv.bias") - encoder_state_dict[f"layers.{i+1}.stack.3.weight"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight") - encoder_state_dict[f"layers.{i+1}.stack.3.bias"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias") - encoder_state_dict[f"layers.{i+1}.stack.5.weight"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.conv2.conv.weight") - encoder_state_dict[f"layers.{i+1}.stack.5.bias"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.conv2.conv.bias") - - # Convert down_blocks - down_block_layers = [3, 4, 6] - for block in range(3): - encoder_state_dict[f"layers.{block+4}.layers.0.weight"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.conv_in.conv.weight") - encoder_state_dict[f"layers.{block+4}.layers.0.bias"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.conv_in.conv.bias") - - for i in range(down_block_layers[block]): - # Convert resnets - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.weight"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.bias"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.weight"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.bias"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.weight"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.bias"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.weight"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.bias"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias") - - # Convert attentions - q = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_q.weight") - k = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_k.weight") - v = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_v.weight") - qkv_weight = torch.cat([q, k, v], dim=0) - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.qkv.weight"] = qkv_weight - - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.weight"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.weight") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.bias"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.bias") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.weight"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.weight") - encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.bias"] = original_state_dict.pop( - f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.bias") - - # Convert block_out - for i in range(3): - encoder_state_dict[f"layers.{i+7}.stack.0.weight"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight") - encoder_state_dict[f"layers.{i+7}.stack.0.bias"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias") - encoder_state_dict[f"layers.{i+7}.stack.2.weight"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.conv1.conv.weight") - encoder_state_dict[f"layers.{i+7}.stack.2.bias"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.conv1.conv.bias") - encoder_state_dict[f"layers.{i+7}.stack.3.weight"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight") - encoder_state_dict[f"layers.{i+7}.stack.3.bias"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias") - encoder_state_dict[f"layers.{i+7}.stack.5.weight"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.conv2.conv.weight") - encoder_state_dict[f"layers.{i+7}.stack.5.bias"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.conv2.conv.bias") - - q = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_q.weight") - k = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_k.weight") - v = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_v.weight") - qkv_weight = torch.cat([q, k, v], dim=0) - encoder_state_dict[f"layers.{i+7}.attn_block.attn.qkv.weight"] = qkv_weight - - encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.weight"] = original_state_dict.pop( - f"{prefix}block_out.attentions.{i}.to_out.0.weight") - encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.bias"] = original_state_dict.pop( - f"{prefix}block_out.attentions.{i}.to_out.0.bias") - encoder_state_dict[f"layers.{i+7}.attn_block.norm.weight"] = original_state_dict.pop( - f"{prefix}block_out.norms.{i}.norm_layer.weight") - encoder_state_dict[f"layers.{i+7}.attn_block.norm.bias"] = original_state_dict.pop( - f"{prefix}block_out.norms.{i}.norm_layer.bias") - - # Convert output layers - encoder_state_dict["output_norm.weight"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.weight") - encoder_state_dict["output_norm.bias"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.bias") - encoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight") - - # Convert decoder - prefix = "decoder." - - decoder_state_dict["blocks.0.0.weight"] = original_state_dict.pop(f"{prefix}conv_in.weight") - decoder_state_dict["blocks.0.0.bias"] = original_state_dict.pop(f"{prefix}conv_in.bias") - - # Convert block_in - for i in range(3): - decoder_state_dict[f"blocks.0.{i+1}.stack.0.weight"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight") - decoder_state_dict[f"blocks.0.{i+1}.stack.0.bias"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias") - decoder_state_dict[f"blocks.0.{i+1}.stack.2.weight"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.conv1.conv.weight") - decoder_state_dict[f"blocks.0.{i+1}.stack.2.bias"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.conv1.conv.bias") - decoder_state_dict[f"blocks.0.{i+1}.stack.3.weight"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight") - decoder_state_dict[f"blocks.0.{i+1}.stack.3.bias"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias") - decoder_state_dict[f"blocks.0.{i+1}.stack.5.weight"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.conv2.conv.weight") - decoder_state_dict[f"blocks.0.{i+1}.stack.5.bias"] = original_state_dict.pop( - f"{prefix}block_in.resnets.{i}.conv2.conv.bias") - - # Convert up_blocks - up_block_layers = [6, 4, 3] - for block in range(3): - for i in range(up_block_layers[block]): - decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.weight"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight") - decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.bias"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias") - decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.weight"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight") - decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.bias"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias") - decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.weight"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight") - decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.bias"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias") - decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.weight"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight") - decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.bias"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias") - decoder_state_dict[f"blocks.{block+1}.proj.weight"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.proj.weight") - decoder_state_dict[f"blocks.{block+1}.proj.bias"] = original_state_dict.pop( - f"{prefix}up_blocks.{block}.proj.bias") - - # Convert block_out - for i in range(3): - decoder_state_dict[f"blocks.4.{i}.stack.0.weight"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight") - decoder_state_dict[f"blocks.4.{i}.stack.0.bias"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias") - decoder_state_dict[f"blocks.4.{i}.stack.2.weight"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.conv1.conv.weight") - decoder_state_dict[f"blocks.4.{i}.stack.2.bias"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.conv1.conv.bias") - decoder_state_dict[f"blocks.4.{i}.stack.3.weight"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight") - decoder_state_dict[f"blocks.4.{i}.stack.3.bias"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias") - decoder_state_dict[f"blocks.4.{i}.stack.5.weight"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.conv2.conv.weight") - decoder_state_dict[f"blocks.4.{i}.stack.5.bias"] = original_state_dict.pop( - f"{prefix}block_out.resnets.{i}.conv2.conv.bias") - - # Convert output layers - decoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight") - decoder_state_dict["output_proj.bias"] = original_state_dict.pop(f"{prefix}proj_out.bias") - - return encoder_state_dict, decoder_state_dict - - -def ensure_safetensors_extension(path): - if not path.endswith(".safetensors"): - path = path + ".safetensors" - return path - - -def ensure_directory_exists(path): - directory = os.path.dirname(path) - if directory: - os.makedirs(directory, exist_ok=True) - - -def main(args): - from diffusers import MochiPipeline - - pipe = MochiPipeline.from_pretrained(args.diffusers_path) - - if args.transformer_path: - transformer_path = ensure_safetensors_extension(args.transformer_path) - ensure_directory_exists(transformer_path) - - print("Converting transformer model...") - transformer_state_dict = convert_diffusers_transformer_to_mochi(pipe.transformer.state_dict()) - save_file(transformer_state_dict, transformer_path) - print(f"Saved transformer to {transformer_path}") - - if args.vae_encoder_path and args.vae_decoder_path: - encoder_path = ensure_safetensors_extension(args.vae_encoder_path) - decoder_path = ensure_safetensors_extension(args.vae_decoder_path) - - ensure_directory_exists(encoder_path) - ensure_directory_exists(decoder_path) - - print("Converting VAE models...") - encoder_state_dict, decoder_state_dict = convert_diffusers_vae_to_mochi(pipe.vae.state_dict()) - - save_file(encoder_state_dict, encoder_path) - print(f"Saved VAE encoder to {encoder_path}") - - save_file(decoder_state_dict, decoder_path) - print(f"Saved VAE decoder to {decoder_path}") - elif args.vae_encoder_path or args.vae_decoder_path: - print("Warning: Both VAE encoder and decoder paths must be specified to convert VAE models.") - - -if __name__ == "__main__": - main(args) diff --git a/fastvideo/models/mochi_hf/mochi_latents_utils.py b/fastvideo/models/mochi_hf/mochi_latents_utils.py deleted file mode 100644 index c554d56c7..000000000 --- a/fastvideo/models/mochi_hf/mochi_latents_utils.py +++ /dev/null @@ -1,45 +0,0 @@ -import torch - -mochi_latents_mean = torch.tensor([ - -0.06730895953510081, - -0.038011381506090416, - -0.07477820912866141, - -0.05565264470995561, - 0.012767231469026969, - -0.04703542746246419, - 0.043896967884726704, - -0.09346305707025976, - -0.09918314763016893, - -0.008729793427399178, - -0.011931556316503654, - -0.0321993391887285, -]).view(1, 12, 1, 1, 1) -mochi_latents_std = torch.tensor([ - 0.9263795028493863, - 0.9248894543193766, - 0.9393059390890617, - 0.959253732819592, - 0.8244560132752793, - 0.917259975397747, - 0.9294154431013696, - 1.3720942357788521, - 0.881393668867029, - 0.9168315692124348, - 0.9185249279345552, - 0.9274757570805041, -]).view(1, 12, 1, 1, 1) -mochi_scaling_factor = 1.0 - - -def normalize_dit_input(model_type, latents): - if model_type == "mochi": - latents_mean = mochi_latents_mean.to(latents.device, latents.dtype) - latents_std = mochi_latents_std.to(latents.device, latents.dtype) - latents = (latents - latents_mean) / latents_std - return latents - elif model_type == "hunyuan_hf": - return latents * 0.476986 - elif model_type == "hunyuan": - return latents * 0.476986 - else: - raise NotImplementedError(f"model_type {model_type} not supported") diff --git a/fastvideo/models/mochi_hf/modeling_mochi.py b/fastvideo/models/mochi_hf/modeling_mochi.py deleted file mode 100644 index b3da2ec8c..000000000 --- a/fastvideo/models/mochi_hf/modeling_mochi.py +++ /dev/null @@ -1,663 +0,0 @@ -# Copyright 2024 The Genmo 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, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.loaders import PeftAdapterMixin -from diffusers.models.attention import FeedForward as HF_FeedForward -from diffusers.models.attention_processor import Attention -from diffusers.models.embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed -from diffusers.models.modeling_utils import ModelMixin -from diffusers.models.normalization import AdaLayerNormContinuous -from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers -from diffusers.utils.torch_utils import maybe_allow_in_graph -from liger_kernel.ops.swiglu import LigerSiLUMulFunction - -from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad -from fastvideo.models.mochi_hf.norm import (MochiLayerNormContinuous, MochiModulatedRMSNorm, MochiRMSNorm, - MochiRMSNormZero) -from fastvideo.utils.communications import all_gather, all_to_all_4D -from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class FeedForward(HF_FeedForward): - - def __init__( - self, - dim: int, - dim_out: Optional[int] = None, - mult: int = 4, - dropout: float = 0.0, - activation_fn: str = "geglu", - final_dropout: bool = False, - inner_dim=None, - bias: bool = True, - ): - super().__init__(dim, dim_out, mult, dropout, activation_fn, final_dropout, inner_dim, bias) - assert activation_fn == "swiglu" - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.net[0].proj(hidden_states) - hidden_states, gate = hidden_states.chunk(2, dim=-1) - - return self.net[2](LigerSiLUMulFunction.apply(gate, hidden_states)) - - -class MochiAttention(nn.Module): - - def __init__( - self, - query_dim: int, - processor: "MochiAttnProcessor2_0", - heads: int = 8, - dim_head: int = 64, - dropout: float = 0.0, - bias: bool = False, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, - out_dim: int = None, - out_context_dim: int = None, - out_bias: bool = True, - context_pre_only: bool = False, - eps: float = 1e-5, - ): - super().__init__() - self.inner_dim = out_dim if out_dim is not None else dim_head * heads - self.out_dim = out_dim if out_dim is not None else query_dim - self.out_context_dim = out_context_dim if out_context_dim else query_dim - self.context_pre_only = context_pre_only - - self.heads = out_dim // dim_head if out_dim is not None else heads - - self.norm_q = MochiRMSNorm(dim_head, eps) - self.norm_k = MochiRMSNorm(dim_head, eps) - self.norm_added_q = MochiRMSNorm(dim_head, eps) - self.norm_added_k = MochiRMSNorm(dim_head, eps) - - self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) - self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias) - self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias) - - self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) - self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) - if self.context_pre_only is not None: - self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) - - self.to_out = nn.ModuleList([]) - self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) - self.to_out.append(nn.Dropout(dropout)) - - if not self.context_pre_only: - self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias) - - self.processor = processor - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - **kwargs, - ): - return self.processor( - self, - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - **kwargs, - ) - - -class MochiAttnProcessor2_0: - """Attention processor used in Mochi.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.") - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - # [b, s, h * d] - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - # [b, s, h=24, d=128] - query = query.unflatten(2, (attn.heads, -1)) - key = key.unflatten(2, (attn.heads, -1)) - value = value.unflatten(2, (attn.heads, -1)) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - # [b, 256, h * d] - encoder_query = attn.add_q_proj(encoder_hidden_states) - encoder_key = attn.add_k_proj(encoder_hidden_states) - encoder_value = attn.add_v_proj(encoder_hidden_states) - - # [b, 256, h=24, d=128] - encoder_query = encoder_query.unflatten(2, (attn.heads, -1)) - encoder_key = encoder_key.unflatten(2, (attn.heads, -1)) - encoder_value = encoder_value.unflatten(2, (attn.heads, -1)) - - if attn.norm_added_q is not None: - encoder_query = attn.norm_added_q(encoder_query) - if attn.norm_added_k is not None: - encoder_key = attn.norm_added_k(encoder_key) - - if image_rotary_emb is not None: - freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1] - # shard the head dimension - if get_sequence_parallel_state(): - # B, S, H, D to (S, B,) H, D - # batch_size, seq_len, attn_heads, head_dim - query = all_to_all_4D(query, scatter_dim=2, gather_dim=1) - key = all_to_all_4D(key, scatter_dim=2, gather_dim=1) - value = all_to_all_4D(value, scatter_dim=2, gather_dim=1) - - def shrink_head(encoder_state, dim): - local_heads = encoder_state.shape[dim] // nccl_info.sp_size - return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) - - encoder_query = shrink_head(encoder_query, dim=2) - encoder_key = shrink_head(encoder_key, dim=2) - encoder_value = shrink_head(encoder_value, dim=2) - if image_rotary_emb is not None: - freqs_cos = shrink_head(freqs_cos, dim=1) - freqs_sin = shrink_head(freqs_sin, dim=1) - - if image_rotary_emb is not None: - - def apply_rotary_emb(x, freqs_cos, freqs_sin): - x_even = x[..., 0::2].float() - x_odd = x[..., 1::2].float() - cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype) - sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype) - - return torch.stack([cos, sin], dim=-1).flatten(-2) - - query = apply_rotary_emb(query, freqs_cos, freqs_sin) - key = apply_rotary_emb(key, freqs_cos, freqs_sin) - - # query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) - # encoder_query, encoder_key, encoder_value = ( - # encoder_query.transpose(1, 2), - # encoder_key.transpose(1, 2), - # encoder_value.transpose(1, 2), - # ) - # [b, s, h, d] - sequence_length = query.size(1) - encoder_sequence_length = encoder_query.size(1) - - # H - query = torch.cat([query, encoder_query], dim=1).unsqueeze(2) - key = torch.cat([key, encoder_key], dim=1).unsqueeze(2) - value = torch.cat([value, encoder_value], dim=1).unsqueeze(2) - # B, S, 3, H, D - qkv = torch.cat([query, key, value], dim=2) - - attn_mask = encoder_attention_mask[:, :].bool() - attn_mask = F.pad(attn_mask, (sequence_length, 0), value=True) - hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) - - # hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask = None, dropout_p=0.0, is_causal=False) - - # valid_lengths = encoder_attention_mask.sum(dim=1) + sequence_length - # def no_padding_mask(score, b, h, q_idx, kv_idx): - # return torch.where(kv_idx < valid_lengths[b],score, -float("inf")) - - # hidden_states = flex_attention(query, key, value, score_mod=no_padding_mask) - if get_sequence_parallel_state(): - hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( - (sequence_length, encoder_sequence_length), dim=1) - # B, S, H, D - hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) - encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous() - hidden_states = hidden_states.flatten(2, 3) - hidden_states = hidden_states.to(query.dtype) - encoder_hidden_states = encoder_hidden_states.flatten(2, 3) - encoder_hidden_states = encoder_hidden_states.to(query.dtype) - else: - hidden_states = hidden_states.flatten(2, 3) - hidden_states = hidden_states.to(query.dtype) - - hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( - (sequence_length, encoder_sequence_length), dim=1) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if hasattr(attn, "to_add_out"): - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - - -@maybe_allow_in_graph -class MochiTransformerBlock(nn.Module): - r""" - Transformer block used in [Mochi](https://huggingface.co/genmo/mochi-1-preview). - - Args: - dim (`int`): - The number of channels in the input and output. - num_attention_heads (`int`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`): - The number of channels in each head. - qk_norm (`str`, defaults to `"rms_norm"`): - The normalization layer to use. - activation_fn (`str`, defaults to `"swiglu"`): - Activation function to use in feed-forward. - context_pre_only (`bool`, defaults to `False`): - Whether or not to process context-related conditions with additional layers. - eps (`float`, defaults to `1e-6`): - Epsilon value for normalization layers. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - pooled_projection_dim: int, - qk_norm: str = "rms_norm", - activation_fn: str = "swiglu", - context_pre_only: bool = False, - eps: float = 1e-6, - ) -> None: - super().__init__() - - self.context_pre_only = context_pre_only - self.ff_inner_dim = (4 * dim * 2) // 3 - self.ff_context_inner_dim = (4 * pooled_projection_dim * 2) // 3 - - self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False) - - if not context_pre_only: - self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False) - else: - self.norm1_context = MochiLayerNormContinuous( - embedding_dim=pooled_projection_dim, - conditioning_embedding_dim=dim, - eps=eps, - ) - - self.attn1 = MochiAttention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - bias=False, - added_kv_proj_dim=pooled_projection_dim, - added_proj_bias=False, - out_dim=dim, - out_context_dim=pooled_projection_dim, - context_pre_only=context_pre_only, - processor=MochiAttnProcessor2_0(), - eps=1e-5, - ) - - # TODO(aryan): norm_context layers are not needed when `context_pre_only` is True - self.norm2 = MochiModulatedRMSNorm(eps=eps) - self.norm2_context = (MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None) - - self.norm3 = MochiModulatedRMSNorm(eps) - self.norm3_context = (MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None) - - self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False) - self.ff_context = None - if not context_pre_only: - self.ff_context = FeedForward( - pooled_projection_dim, - inner_dim=self.ff_context_inner_dim, - activation_fn=activation_fn, - bias=False, - ) - - self.norm4 = MochiModulatedRMSNorm(eps=eps) - self.norm4_context = MochiModulatedRMSNorm(eps=eps) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor, - temb: torch.Tensor, - image_rotary_emb: Optional[torch.Tensor] = None, - output_attn=False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) - - if not self.context_pre_only: - ( - norm_encoder_hidden_states, - enc_gate_msa, - enc_scale_mlp, - enc_gate_mlp, - ) = self.norm1_context(encoder_hidden_states, temb) - else: - norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb) - - attn_hidden_states, context_attn_hidden_states = self.attn1( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - ) - - hidden_states = hidden_states + self.norm2(attn_hidden_states, torch.tanh(gate_msa).unsqueeze(1)) - norm_hidden_states = self.norm3(hidden_states, (1 + scale_mlp.unsqueeze(1).to(torch.float32))) - ff_output = self.ff(norm_hidden_states) - hidden_states = hidden_states + self.norm4(ff_output, torch.tanh(gate_mlp).unsqueeze(1)) - - if not self.context_pre_only: - encoder_hidden_states = encoder_hidden_states + self.norm2_context(context_attn_hidden_states, - torch.tanh(enc_gate_msa).unsqueeze(1)) - norm_encoder_hidden_states = self.norm3_context( - encoder_hidden_states, - (1 + enc_scale_mlp.unsqueeze(1).to(torch.float32)), - ) - context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + self.norm4_context(context_ff_output, - torch.tanh(enc_gate_mlp).unsqueeze(1)) - - if not output_attn: - attn_hidden_states = None - return hidden_states, encoder_hidden_states, attn_hidden_states - - -class MochiRoPE(nn.Module): - r""" - RoPE implementation used in [Mochi](https://huggingface.co/genmo/mochi-1-preview). - - Args: - base_height (`int`, defaults to `192`): - Base height used to compute interpolation scale for rotary positional embeddings. - base_width (`int`, defaults to `192`): - Base width used to compute interpolation scale for rotary positional embeddings. - """ - - def __init__(self, base_height: int = 192, base_width: int = 192) -> None: - super().__init__() - - self.target_area = base_height * base_width - - def _centers(self, start, stop, num, device, dtype) -> torch.Tensor: - edges = torch.linspace(start, stop, num + 1, device=device, dtype=dtype) - return (edges[:-1] + edges[1:]) / 2 - - def _get_positions( - self, - num_frames: int, - height: int, - width: int, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ) -> torch.Tensor: - scale = (self.target_area / (height * width))**0.5 - t = torch.arange(num_frames * nccl_info.sp_size, device=device, dtype=dtype) - h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype) - w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype) - - grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij") - - positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3) - return positions - - def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: - with torch.autocast(freqs.device.type, enabled=False): - # Always run ROPE freqs computation in FP32 - freqs = torch.einsum( - "nd,dhf->nhf", # codespell:ignore - pos.to(torch.float32), # codespell:ignore - freqs.to(torch.float32)) - freqs_cos = torch.cos(freqs) - freqs_sin = torch.sin(freqs) - return freqs_cos, freqs_sin - - def forward( - self, - pos_frequencies: torch.Tensor, - num_frames: int, - height: int, - width: int, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - pos = self._get_positions(num_frames, height, width, device, dtype) - rope_cos, rope_sin = self._create_rope(pos_frequencies, pos) - return rope_cos, rope_sin - - -@maybe_allow_in_graph -class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): - r""" - A Transformer model for video-like data introduced in [Mochi](https://huggingface.co/genmo/mochi-1-preview). - - Args: - patch_size (`int`, defaults to `2`): - The size of the patches to use in the patch embedding layer. - num_attention_heads (`int`, defaults to `24`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `128`): - The number of channels in each head. - num_layers (`int`, defaults to `48`): - The number of layers of Transformer blocks to use. - in_channels (`int`, defaults to `12`): - The number of channels in the input. - out_channels (`int`, *optional*, defaults to `None`): - The number of channels in the output. - qk_norm (`str`, defaults to `"rms_norm"`): - The normalization layer to use. - text_embed_dim (`int`, defaults to `4096`): - Input dimension of text embeddings from the text encoder. - time_embed_dim (`int`, defaults to `256`): - Output dimension of timestep embeddings. - activation_fn (`str`, defaults to `"swiglu"`): - Activation function to use in feed-forward. - max_sequence_length (`int`, defaults to `256`): - The maximum sequence length of text embeddings supported. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - patch_size: int = 2, - num_attention_heads: int = 24, - attention_head_dim: int = 128, - num_layers: int = 48, - pooled_projection_dim: int = 1536, - in_channels: int = 12, - out_channels: Optional[int] = None, - qk_norm: str = "rms_norm", - text_embed_dim: int = 4096, - time_embed_dim: int = 256, - activation_fn: str = "swiglu", - max_sequence_length: int = 256, - ) -> None: - super().__init__() - - inner_dim = num_attention_heads * attention_head_dim - out_channels = out_channels or in_channels - - self.patch_embed = PatchEmbed( - patch_size=patch_size, - in_channels=in_channels, - embed_dim=inner_dim, - pos_embed_type=None, - ) - - self.time_embed = MochiCombinedTimestepCaptionEmbedding( - embedding_dim=inner_dim, - pooled_projection_dim=pooled_projection_dim, - text_embed_dim=text_embed_dim, - time_embed_dim=time_embed_dim, - num_attention_heads=8, - ) - - self.pos_frequencies = nn.Parameter(torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0)) - self.rope = MochiRoPE() - - self.transformer_blocks = nn.ModuleList([ - MochiTransformerBlock( - dim=inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - pooled_projection_dim=pooled_projection_dim, - qk_norm=qk_norm, - activation_fn=activation_fn, - context_pre_only=i == num_layers - 1, - ) for i in range(num_layers) - ]) - - self.norm_out = AdaLayerNormContinuous( - inner_dim, - inner_dim, - elementwise_affine=False, - eps=1e-6, - norm_type="layer_norm", - ) - self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels) - - self.gradient_checkpointing = False - - def _set_gradient_checkpointing(self, module, value=False): - if hasattr(module, "gradient_checkpointing"): - module.gradient_checkpointing = value - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_attention_mask: torch.Tensor, - output_features=False, - output_features_stride=8, - attention_kwargs: Optional[Dict[str, Any]] = None, - return_dict: bool = False, - ) -> torch.Tensor: - assert (return_dict is False), "return_dict is not supported in MochiTransformer3DModel" - - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if (attention_kwargs is not None and attention_kwargs.get("scale", None) is not None): - logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.") - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p = self.config.patch_size - - post_patch_height = height // p - post_patch_width = width // p - # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux - timestep = 1000 - timestep - temb, encoder_hidden_states = self.time_embed( - timestep, - encoder_hidden_states, - encoder_attention_mask, - hidden_dtype=hidden_states.dtype, - ) - - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.patch_embed(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2) - - image_rotary_emb = self.rope( - self.pos_frequencies, - num_frames, - post_patch_height, - post_patch_width, - device=hidden_states.device, - dtype=torch.float32, - ) - attn_outputs_list = [] - for i, block in enumerate(self.transformer_blocks): - if self.gradient_checkpointing: - - def create_custom_forward(module): - - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = ({"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}) - ( - hidden_states, - encoder_hidden_states, - attn_outputs, - ) = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - hidden_states, - encoder_hidden_states, - encoder_attention_mask, - temb, - image_rotary_emb, - output_features, - **ckpt_kwargs, - ) - else: - hidden_states, encoder_hidden_states, attn_outputs = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - temb=temb, - image_rotary_emb=image_rotary_emb, - output_attn=output_features, - ) - if i % output_features_stride == 0: - attn_outputs_list.append(attn_outputs) - - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1) - hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5) - output = hidden_states.reshape(batch_size, -1, num_frames, height, width) - - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - - if not output_features: - attn_outputs_list = None - else: - attn_outputs_list = torch.stack(attn_outputs_list, dim=0) - # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux - return (-output, attn_outputs_list) diff --git a/fastvideo/models/mochi_hf/norm.py b/fastvideo/models/mochi_hf/norm.py deleted file mode 100644 index d91bb95f1..000000000 --- a/fastvideo/models/mochi_hf/norm.py +++ /dev/null @@ -1,128 +0,0 @@ -# Copyright 2024 The Genmo 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 Tuple - -import torch -import torch.nn as nn - - -class MochiModulatedRMSNorm(nn.Module): - - def __init__(self, eps: float): - super().__init__() - - self.eps = eps - - def forward(self, hidden_states, scale=None): - hidden_states_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.eps) - if scale is not None: - hidden_states = hidden_states * scale - - hidden_states = hidden_states.to(hidden_states_dtype) - - return hidden_states - - -class MochiRMSNorm(nn.Module): - - def __init__(self, dim, eps: float, elementwise_affine=True): - super().__init__() - - self.eps = eps - if elementwise_affine: - self.weight = nn.Parameter(torch.ones(dim)) - else: - self.weight = None - - def forward(self, hidden_states): - hidden_states_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.eps) - if self.weight is not None: - # convert into half-precision if necessary - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - hidden_states = hidden_states * self.weight - hidden_states = hidden_states.to(hidden_states_dtype) - - return hidden_states - - -class MochiLayerNormContinuous(nn.Module): - - def __init__( - self, - embedding_dim: int, - conditioning_embedding_dim: int, - eps=1e-5, - bias=True, - ): - super().__init__() - - # AdaLN - self.silu = nn.SiLU() - self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) - self.norm = MochiModulatedRMSNorm(eps=eps) - - def forward( - self, - x: torch.Tensor, - conditioning_embedding: torch.Tensor, - ) -> torch.Tensor: - input_dtype = x.dtype - - # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) - scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) - x = self.norm(x, (1 + scale.unsqueeze(1).to(torch.float32))) - - return x.to(input_dtype) - - -class MochiRMSNormZero(nn.Module): - r""" - Adaptive RMS Norm used in Mochi. - Parameters: - embedding_dim (`int`): The size of each embedding vector. - """ - - def __init__( - self, - embedding_dim: int, - hidden_dim: int, - eps: float = 1e-5, - elementwise_affine: bool = False, - ) -> None: - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, hidden_dim) - self.norm = MochiModulatedRMSNorm(eps=eps) - - def forward(self, hidden_states: torch.Tensor, - emb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - hidden_states_dtype = hidden_states.dtype - - emb = self.linear(self.silu(emb)) - scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) - - hidden_states = self.norm(hidden_states, (1 + scale_msa[:, None].to(torch.float32))) - hidden_states = hidden_states.to(hidden_states_dtype) - - return hidden_states, gate_msa, scale_mlp, gate_mlp diff --git a/fastvideo/models/mochi_hf/pipeline_mochi.py b/fastvideo/models/mochi_hf/pipeline_mochi.py deleted file mode 100644 index f42192407..000000000 --- a/fastvideo/models/mochi_hf/pipeline_mochi.py +++ /dev/null @@ -1,757 +0,0 @@ -# Copyright 2024 Black Forest Labs 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. - -import copy -import inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback -from diffusers.loaders import Mochi1LoraLoaderMixin -from diffusers.models.autoencoders import AutoencoderKL -from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput -from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from diffusers.schedulers import FlowMatchEulerDiscreteScheduler -from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from einops import rearrange -from transformers import T5EncoderModel, T5TokenizerFast - -from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel -from fastvideo.utils.communications import all_gather -from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import MochiPipeline - >>> from diffusers.utils import export_to_video - - >>> pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - >>> prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k." - >>> frames = pipe(prompt, num_inference_steps=28, guidance_scale=3.5).frames[0] - >>> export_to_video(frames, "mochi.mp4") - ``` -""" - - -def calculate_shift( - image_seq_len, - base_seq_len: int = 256, - max_seq_len: int = 4096, - base_shift: float = 0.5, - max_shift: float = 1.16, -): - m = (max_shift - base_shift) / (max_seq_len - base_seq_len) - b = base_shift - m * base_seq_len - mu = image_seq_len * m + b - return mu - - -# from: https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77 -def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None): - if linear_steps is None: - linear_steps = num_steps // 2 - linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)] - threshold_noise_step_diff = linear_steps - threshold_noise * num_steps - quadratic_steps = num_steps - linear_steps - quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2) - linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2) - const = quadratic_coef * (linear_steps**2) - quadratic_sigma_schedule = [ - quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps) - ] - sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule - sigma_schedule = [1.0 - x for x in sigma_schedule] - return sigma_schedule - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps -def retrieve_timesteps( - scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, - **kwargs, -): - r""" - Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles - custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. - - Args: - scheduler (`SchedulerMixin`): - The scheduler to get timesteps from. - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` - must be `None`. - device (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): - Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, - `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): - Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, - `num_inference_steps` and `timesteps` must be `None`. - - Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the - second element is the number of inference steps. - """ - if timesteps is not None and sigmas is not None: - raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") - if timesteps is not None: - accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accepts_timesteps: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" timestep schedules. Please check whether you are using the correct scheduler.") - scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - elif sigmas is not None: - accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accept_sigmas: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" sigmas schedules. Please check whether you are using the correct scheduler.") - scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - else: - scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) - timesteps = scheduler.timesteps - return timesteps, num_inference_steps - - -class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): - r""" - The mochi pipeline for text-to-video generation. - - Reference: https://github.com/genmoai/models - - Args: - transformer ([`MochiTransformer3DModel`]): - Conditional Transformer architecture to denoise the encoded video latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`T5EncoderModel`]): - [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically - the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). - tokenizer (`T5TokenizerFast`): - Second Tokenizer of class - [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). - """ - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: T5EncoderModel, - tokenizer: T5TokenizerFast, - transformer: MochiTransformer3DModel, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_spatial_scale_factor = 8 - self.vae_temporal_scale_factor = 6 - self.patch_size = 2 - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_scale_factor) - self.tokenizer_max_length = (self.tokenizer.model_max_length - if hasattr(self, "tokenizer") and self.tokenizer is not None else 77) - self.default_height = 480 - self.default_width = 848 - - # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds - def _get_t5_prompt_embeds( - self, - prompt: Union[str, List[str]] = None, - num_videos_per_prompt: int = 1, - max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ): - device = device or self._execution_device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.attention_mask - prompt_attention_mask = prompt_attention_mask.bool().to(device) - - untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids - - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1]) - logger.warning("The following part of your input was truncated because `max_sequence_length` is set to " - f" {max_sequence_length} tokens: {removed_text}") - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0] - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - _, seq_len, _ = prompt_embeds.shape - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - - prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) - prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) - - return prompt_embeds, prompt_attention_mask - - # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - do_classifier_free_guidance: bool = True, - num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than `1`). - do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): - Whether to use classifier free guidance or not. - num_videos_per_prompt (`int`, *optional*, defaults to 1): - Number of videos that should be generated per prompt. torch device to place the resulting embeddings on - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - device: (`torch.device`, *optional*): - torch device - dtype: (`torch.dtype`, *optional*): - torch dtype - """ - device = device or self._execution_device - - prompt = [prompt] if isinstance(prompt, str) else prompt - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds( - prompt=prompt, - num_videos_per_prompt=num_videos_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - dtype=dtype, - ) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt = (batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt) - - if prompt is not None and type(prompt) is not type(negative_prompt): - raise TypeError( - f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}.") - elif batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`.") - - ( - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self._get_t5_prompt_embeds( - prompt=negative_prompt, - num_videos_per_prompt=num_videos_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - dtype=dtype, - ) - - return ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) - - def check_inputs( - self, - prompt, - height, - width, - callback_on_step_end_tensor_inputs=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - ): - if height % 8 != 0 or width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") - - if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs - for k in callback_on_step_end_tensor_inputs): - raise ValueError( - f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" - ) - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two.") - elif prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.") - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if prompt_embeds is not None and prompt_attention_mask is None: - raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") - - if (negative_prompt_embeds is not None and negative_prompt_attention_mask is None): - raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") - - if prompt_embeds is not None and negative_prompt_embeds is not None: - if prompt_embeds.shape != negative_prompt_embeds.shape: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}.") - if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: - raise ValueError( - "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" - f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" - f" {negative_prompt_attention_mask.shape}.") - - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - num_frames, - dtype, - device, - generator, - latents=None, - ): - height = height // self.vae_spatial_scale_factor - width = width // self.vae_spatial_scale_factor - num_frames = (num_frames - 1) // self.vae_temporal_scale_factor + 1 - - shape = (batch_size, num_channels_latents, num_frames, height, width) - - if latents is not None: - return latents.to(device=device, dtype=dtype) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators.") - - latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32) - latents = latents.to(dtype) - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1.0 - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def interrupt(self): - return self._interrupt - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_frames: int = 19, - num_inference_steps: int = 64, - timesteps: List[int] = None, - guidance_scale: float = 4.5, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - max_sequence_length: int = 256, - return_all_states=False, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. This is set to 1024 by default for the best results. - width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. This is set to 1024 by default for the best results. - num_frames (`int`, defaults to 16): - The number of video frames to generate - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - timesteps (`List[int]`, *optional*): - Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument - in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is - passed will be used. Must be in descending order. - guidance_scale (`float`, defaults to `4.5`): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - num_videos_per_prompt (`int`, *optional*, defaults to 1): - The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will ge generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for text embeddings. - negative_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not - provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. - negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.mochi.MochiPipelineOutput`] instead of a plain tuple. - attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in - [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising steps during the inference. The function is called - with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, - callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by - `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. - max_sequence_length (`int` defaults to `256`): - Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.mochi.MochiPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.mochi.MochiPipelineOutput`] is returned, otherwise a `tuple` - is returned where the first element is a list with the generated images. - """ - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - height = height or self.default_height - width = width or self.default_width - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt=prompt, - height=height, - width=width, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self._execution_device - - # 3. Prepare text embeddings - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - max_sequence_length=max_sequence_length, - device=device, - ) - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) - - # 4. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - prompt_embeds.dtype, - device, - generator, - latents, - ) - world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group - if get_sequence_parallel_state(): - latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous() - latents = latents[:, :, rank, :, :, :] - - original_noise = copy.deepcopy(latents) - # 5. Prepare timestep - # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77 - threshold_noise = 0.025 - sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise) - sigmas = np.array(sigmas) - # check if of type FlowMatchEulerDiscreteScheduler - if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - timesteps, - sigmas, - ) - else: - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - ) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 6. Denoising loop - self._progress_bar_config = {"disable": nccl_info.rank_within_group != 0} - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - latent_model_input = (torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - encoder_attention_mask=prompt_attention_mask, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - # Mochi CFG + Sampling runs in FP32 - noise_pred = noise_pred.to(torch.float32) - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0] - latents = latents.to(latents_dtype) - - if latents.dtype != latents_dtype: - if torch.backends.mps.is_available(): - # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 - latents = latents.to(latents_dtype) - - if callback_on_step_end is not None: - callback_kwargs = {} - for k in callback_on_step_end_tensor_inputs: - callback_kwargs[k] = locals()[k] - callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) - - latents = callback_outputs.pop("latents", latents) - prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) - - # call the callback, if provided - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() - - if XLA_AVAILABLE: - xm.mark_step() - - if get_sequence_parallel_state(): - latents = all_gather(latents, dim=2) - # latents_shape = list(latents.shape) - # full_shape = [latents_shape[0] * world_size] + latents_shape[1:] - # all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device) - # torch.distributed.all_gather_into_tensor(all_latents, latents) - # latents_list = list(all_latents.chunk(world_size, dim=0)) - # latents = torch.cat(latents_list, dim=2) - - if output_type == "latent": - video = latents - else: - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = (hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None) - has_latents_std = (hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None) - if has_latents_mean and has_latents_std: - latents_mean = (torch.tensor(self.vae.config.latents_mean).view(1, 12, 1, 1, - 1).to(latents.device, latents.dtype)) - latents_std = (torch.tensor(self.vae.config.latents_std).view(1, 12, 1, 1, - 1).to(latents.device, latents.dtype)) - latents = (latents * latents_std / self.vae.config.scaling_factor + latents_mean) - else: - latents = latents / self.vae.config.scaling_factor - - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - if return_all_states: - # Pay extra attention here: - # prompt_embeds with shape torch.Size([2, 256]), where prompt_embeds[1] is the prompt_embeds for the actual prompt - # prompt_embeds[0] is for negative prompt - return original_noise, video, latents, prompt_embeds, prompt_attention_mask - - if not return_dict: - return (video, ) - - return MochiPipelineOutput(frames=video) diff --git a/fastvideo/models/parameter.py b/fastvideo/models/parameter.py index 30e901e98..f4f8b148a 100644 --- a/fastvideo/models/parameter.py +++ b/fastvideo/models/parameter.py @@ -431,4 +431,4 @@ def _adjust_shard_indexes_for_packing(shard_size, shard_offset, packed_factor, shard_size=shard_size, shard_offset=shard_offset, marlin_tile_size=marlin_tile_size) - return shard_size, shard_offset + return shard_size, shard_offset \ No newline at end of file diff --git a/fastvideo/models/dits/registry.py b/fastvideo/models/registry.py similarity index 87% rename from fastvideo/models/dits/registry.py rename to fastvideo/models/registry.py index 8a1a2b248..cc02f9e6d 100644 --- a/fastvideo/models/dits/registry.py +++ b/fastvideo/models/registry.py @@ -13,18 +13,34 @@ import cloudpickle from torch import nn from fastvideo.logger import logger -# huggingface class name: (fastvideo module name, fastvideo class name) -_TEXT_TO_VIDEO_MODELS = { - "HunyuanVideoTransformer3DModel": ("hunyuanvideo", "HunyuanVideoDiT"), +# huggingface class name: (component_name, fastvideo module name, fastvideo class name) +_TEXT_TO_VIDEO_DIT_MODELS = { + "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"), } -_IMAGE_TO_VIDEO_MODELS = { - +_IMAGE_TO_VIDEO_DIT_MODELS = { + # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), +} + +_TEXT_ENCODER_MODELS = { + "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"), + "LlamaModel": ("encoders", "llama", "LlamaModel"), +} + +_IMAGE_ENCODER_MODELS = { + # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"), +} + +_VAE_MODELS = { + "AutoencoderKLHunyuanVideo": ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), } _FAST_VIDEO_MODELS = { - **_TEXT_TO_VIDEO_MODELS, - **_IMAGE_TO_VIDEO_MODELS, + **_TEXT_TO_VIDEO_DIT_MODELS, + **_IMAGE_TO_VIDEO_DIT_MODELS, + **_TEXT_ENCODER_MODELS, + **_IMAGE_ENCODER_MODELS, + **_VAE_MODELS, } _SUBPROCESS_COMMAND = [ @@ -111,6 +127,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel): Represents a model that has not been imported in the main process. """ module_name: str + component_name: str class_name: str # Performed in another process to avoid initializing CUDA @@ -128,7 +145,7 @@ def _try_load_model_cls( model_arch: str, model: _BaseRegisteredModel, ) -> Optional[Type[nn.Module]]: - from vllm.platforms import current_platform + from fastvideo.platforms import current_platform current_platform.verify_model_arch(model_arch) try: return model.load_model_cls() @@ -266,8 +283,9 @@ class _ModelRegistry: ModelRegistry = _ModelRegistry({ model_arch: _LazyRegisteredModel( - module_name=f"fastvideo.models.dits.{mod_relname}", + module_name=f"fastvideo.models.{component_name}.{mod_relname}", + component_name=component_name, class_name=cls_name, ) - for model_arch, (mod_relname, cls_name) in _FAST_VIDEO_MODELS.items() + for model_arch, (component_name, mod_relname, cls_name) in _FAST_VIDEO_MODELS.items() }) \ No newline at end of file diff --git a/fastvideo/models/schedulers/base.py b/fastvideo/models/schedulers/base.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/models/stepvideo/diffusion/scheduler.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py similarity index 83% rename from fastvideo/models/stepvideo/diffusion/scheduler.py rename to fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index 1246bbf97..e9e96edb1 100644 --- a/fastvideo/models/stepvideo/diffusion/scheduler.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -1,3 +1,22 @@ +# Copyright 2024 Stability AI, Katherine Crowson 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. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== + from dataclasses import dataclass from typing import Optional, Tuple, Union @@ -36,6 +55,8 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): timestep_spacing (`str`, defaults to `"linspace"`): The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. reverse (`bool`, defaults to `True`): Whether to reverse the timestep schedule. """ @@ -47,9 +68,10 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): def __init__( self, num_train_timesteps: int = 1000, - reverse: bool = False, + shift: float = 1.0, + reverse: bool = True, solver: str = "euler", - device: Union[str, torch.device] = None, + n_tokens: Optional[int] = None, ): sigmas = torch.linspace(1, 0, num_train_timesteps + 1) @@ -63,8 +85,6 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): self._step_index = None self._begin_index = None - self.device = device - self.supported_solver = ["euler"] if solver not in self.supported_solver: raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}") @@ -100,8 +120,8 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, num_inference_steps: int, - time_shift: float = 13.0, device: Union[str, torch.device] = None, + n_tokens: int = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -114,17 +134,16 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): n_tokens (`int`, *optional*): Number of tokens in the input sequence. """ - device = device or self.device self.num_inference_steps = num_inference_steps - sigmas = torch.linspace(1, 0, num_inference_steps + 1, device=device) - sigmas = self.sd3_time_shift(sigmas, time_shift) + sigmas = torch.linspace(1, 0, num_inference_steps + 1) + sigmas = self.sd3_time_shift(sigmas) if not self.config.reverse: sigmas = 1 - sigmas self.sigmas = sigmas - self.timesteps = sigmas[:-1] + self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device) # Reset step index self._step_index = None @@ -154,15 +173,15 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: return sample - def sd3_time_shift(self, t: torch.Tensor, time_shift: float = 13.0): - return (time_shift * t) / (1 + (time_shift - 1) * t) + def sd3_time_shift(self, t: torch.Tensor): + return (self.config.shift * t) / (1 + (self.config.shift - 1) * t) def step( self, model_output: torch.FloatTensor, timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, - return_dict: bool = False, + return_dict: bool = True, ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion @@ -212,9 +231,9 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): self._step_index += 1 if not return_dict: - return prev_sample + return (prev_sample, ) return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) def __len__(self): - return self.config.num_train_timesteps + return self.config.num_train_timesteps \ No newline at end of file diff --git a/fastvideo/models/stepvideo/__init__.py b/fastvideo/models/stepvideo/__init__.py deleted file mode 100644 index e919871df..000000000 --- a/fastvideo/models/stepvideo/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -import os - -os.environ["NCCL_DEBUG"] = "ERROR" - -from .diffusion.scheduler import * -from .diffusion.video_pipeline import * -from .modules.model import * \ No newline at end of file diff --git a/fastvideo/models/stepvideo/__version__.py b/fastvideo/models/stepvideo/__version__.py deleted file mode 100644 index 3f5c4a7d6..000000000 --- a/fastvideo/models/stepvideo/__version__.py +++ /dev/null @@ -1 +0,0 @@ -__version__ = "0.1.0" diff --git a/fastvideo/models/stepvideo/config.py b/fastvideo/models/stepvideo/config.py deleted file mode 100644 index d7adb50bd..000000000 --- a/fastvideo/models/stepvideo/config.py +++ /dev/null @@ -1,174 +0,0 @@ -import argparse - - -def parse_args(namespace=None): - parser = argparse.ArgumentParser(description="StepVideo inference script") - - parser = add_extra_models_args(parser) - parser = add_denoise_schedule_args(parser) - parser = add_inference_args(parser) - parser = add_parallel_args(parser) - - args = parser.parse_args(namespace=namespace) - - return args - - -def add_extra_models_args(parser: argparse.ArgumentParser): - group = parser.add_argument_group(title="Extra models args, including vae, text encoders and tokenizers)") - - group.add_argument( - "--vae_url", - type=str, - default='127.0.0.1', - help="vae url.", - ) - group.add_argument( - "--caption_url", - type=str, - default='127.0.0.1', - help="caption url.", - ) - - return parser - - -def add_denoise_schedule_args(parser: argparse.ArgumentParser): - group = parser.add_argument_group(title="Denoise schedule args") - - # Flow Matching - group.add_argument( - "--time_shift", - type=float, - default=7.0, - help="Shift factor for flow matching schedulers.", - ) - group.add_argument( - "--flow_reverse", - action="store_true", - help="If reverse, learning/sampling from t=1 -> t=0.", - ) - group.add_argument( - "--flow_solver", - type=str, - default="euler", - help="Solver for flow matching.", - ) - - return parser - - -def add_inference_args(parser: argparse.ArgumentParser): - group = parser.add_argument_group(title="Inference args") - - # ======================== Model loads ======================== - group.add_argument( - "--model_dir", - type=str, - default="./ckpts", - help="Root path of all the models, including t2v models and extra models.", - ) - group.add_argument( - "--model_resolution", - type=str, - default="540p", - choices=["540p"], - help="Root path of all the models, including t2v models and extra models.", - ) - group.add_argument( - "--use-cpu-offload", - action="store_true", - help="Use CPU offload for the model load.", - ) - - # ======================== Inference general setting ======================== - group.add_argument( - "--batch_size", - type=int, - default=1, - help="Batch size for inference and evaluation.", - ) - group.add_argument( - "--infer_steps", - type=int, - default=50, - help="Number of denoising steps for inference.", - ) - group.add_argument( - "--save_path", - type=str, - default="./results", - help="Path to save the generated samples.", - ) - group.add_argument( - "--name_suffix", - type=str, - default="", - help="Suffix for the names of saved samples.", - ) - group.add_argument( - "--num_videos", - type=int, - default=1, - help="Number of videos to generate for each prompt.", - ) - # ---sample size--- - group.add_argument( - "--num_frames", - type=int, - default=204, - help="How many frames to sample from a video. ", - ) - group.add_argument( - "--height", - type=int, - default=544, - help="The height of video sample", - ) - group.add_argument( - "--width", - type=int, - default=992, - help="The width of video sample", - ) - # --- prompt --- - group.add_argument( - "--prompt", - type=str, - default=None, - help="Prompt for sampling during evaluation.", - ) - group.add_argument("--seed", type=int, default=1234, help="Seed for evaluation.") - - # Classifier-Free Guidance - group.add_argument("--pos_magic", - type=str, - default="超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。", - help="Positive magic prompt for sampling.") - group.add_argument("--neg_magic", - type=str, - default="画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。", - help="Negative magic prompt for sampling.") - group.add_argument("--cfg_scale", type=float, default=9.0, help="Classifier free guidance scale.") - - return parser - - -def add_parallel_args(parser: argparse.ArgumentParser): - group = parser.add_argument_group(title="Parallel args") - - # ======================== Model loads ======================== - group.add_argument( - "--ulysses_degree", - type=int, - default=8, - help="Ulysses degree.", - ) - group.add_argument( - "--ring_degree", - type=int, - default=1, - help="Ulysses degree.", - ) - - return parser diff --git a/fastvideo/models/stepvideo/diffusion/video_pipeline.py b/fastvideo/models/stepvideo/diffusion/video_pipeline.py deleted file mode 100755 index 264c2c356..000000000 --- a/fastvideo/models/stepvideo/diffusion/video_pipeline.py +++ /dev/null @@ -1,325 +0,0 @@ -# Copyright 2025 StepFun Inc. All Rights Reserved. - -import asyncio -import pickle -from dataclasses import dataclass -from typing import Dict, List, Optional, Union - -import numpy as np -import torch -from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from diffusers.utils import BaseOutput - -from fastvideo.models.stepvideo.diffusion.scheduler import FlowMatchDiscreteScheduler -from fastvideo.models.stepvideo.modules.model import StepVideoModel -from fastvideo.models.stepvideo.utils import VideoProcessor - - -def call_api_gen(url, api, port=8080): - url = f"http://{url}:{port}/{api}-api" - import aiohttp - - async def _fn(samples, *args, **kwargs): - if api == 'vae': - data = { - "samples": samples, - } - elif api == 'caption': - data = { - "prompts": samples, - } - else: - raise Exception(f"Not supported api: {api}...") - - async with aiohttp.ClientSession() as sess: - data_bytes = pickle.dumps(data) - async with sess.get(url, data=data_bytes, timeout=12000) as response: - result = bytearray() - while not response.content.at_eof(): - chunk = await response.content.read(1024) - result += chunk - response_data = pickle.loads(result) - return response_data - - return _fn - - -@dataclass -class StepVideoPipelineOutput(BaseOutput): - video: Union[torch.Tensor, np.ndarray] - - -class StepVideoPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-video generation using StepVideo. - - This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - transformer ([`StepVideoModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae_url: - remote vae server's url. - caption_url: - remote caption (stepllm and clip) server's url. - """ - - def __init__( - self, - transformer: StepVideoModel, - scheduler: FlowMatchDiscreteScheduler, - vae_url: str = '127.0.0.1', - caption_url: str = '127.0.0.1', - save_path: str = './results', - name_suffix: str = '', - ): - super().__init__() - - self.register_modules( - transformer=transformer, - scheduler=scheduler, - ) - - self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 8 - self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 16 - self.video_processor = VideoProcessor(save_path, name_suffix) - - self.vae_url = vae_url - self.caption_url = caption_url - self.setup_api(self.vae_url, self.caption_url) - - def setup_api(self, vae_url, caption_url): - self.vae_url = vae_url - self.caption_url = caption_url - self.caption = call_api_gen(caption_url, 'caption') - self.vae = call_api_gen(vae_url, 'vae') - return self - - def encode_prompt( - self, - prompt: str, - neg_magic: str = '', - pos_magic: str = '', - ): - device = self._execution_device - prompts = [prompt + pos_magic] - bs = len(prompts) - prompts += [neg_magic] * bs - - data = asyncio.run(self.caption(prompts)) - prompt_embeds, prompt_attention_mask, clip_embedding = data['y'].to(device), data['y_mask'].to( - device), data['clip_embedding'].to(device) - - return prompt_embeds, clip_embedding, prompt_attention_mask - - def decode_vae(self, samples): - samples = asyncio.run(self.vae(samples.cpu())) - return samples - - def check_inputs(self, num_frames, width, height): - num_frames = max(num_frames // 17 * 17, 1) - width = max(width // 16 * 16, 16) - height = max(height // 16 * 16, 16) - return num_frames, width, height - - def prepare_latents( - self, - batch_size: int, - num_channels_latents: 64, - height: int = 544, - width: int = 992, - num_frames: int = 204, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if latents is not None: - return latents.to(device=device, dtype=dtype) - - num_frames, width, height = self.check_inputs(num_frames, width, height) - shape = ( - batch_size, - max(num_frames // 17 * 3, 1), - num_channels_latents, - int(height) // self.vae_scale_factor_spatial, - int(width) // self.vae_scale_factor_spatial, - ) # b,f,c,h,w - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators.") - - if generator is None: - generator = torch.Generator(device=self._execution_device) - - latents = torch.randn(shape, generator=generator, device=device, dtype=dtype) - return latents - - @torch.inference_mode() - def __call__( - self, - prompt: Union[str, List[str]] = None, - height: int = 544, - width: int = 992, - num_frames: int = 204, - num_inference_steps: int = 50, - guidance_scale: float = 9.0, - time_shift: float = 13.0, - neg_magic: str = "", - pos_magic: str = "", - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "mp4", - output_file_name: Optional[str] = "", - return_dict: bool = True, - mask_strategy: Optional[Dict[str, list]] = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - height (`int`, defaults to `544`): - The height in pixels of the generated image. - width (`int`, defaults to `992`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `204`): - The number of frames in the generated video. - num_inference_steps (`int`, defaults to `50`): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - guidance_scale (`float`, defaults to `9.0`): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - num_videos_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor is generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between `PIL.Image` or `np.array`. - output_file_name(`str`, *optional*`): - The output mp4 file name. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`StepVideoPipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~StepVideoPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`StepVideoPipelineOutput`] is returned, otherwise a `tuple` is returned - where the first element is a list with the generated images and the second element is a list of `bool`s - indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. - """ - - # 1. Check inputs. Raise error if not correct - device = self._execution_device - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - do_classifier_free_guidance = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, prompt_embeds_2, prompt_attention_mask = self.encode_prompt( - prompt=prompt, - neg_magic=neg_magic, - pos_magic=pos_magic, - ) - - transformer_dtype = self.transformer.dtype - prompt_embeds = prompt_embeds.to(transformer_dtype) - prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) - prompt_embeds_2 = prompt_embeds_2.to(transformer_dtype) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps=num_inference_steps, time_shift=time_shift, device=device) - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - torch.bfloat16, - device, - generator, - latents, - ) - - def dict_to_3d_list(best_masks, t_max=50, l_max=48, h_max=48): - result = [[[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max)] - if best_masks is None: - return result - for key, value in best_masks.items(): - timestep, layer, head = map(int, key.split('_')) - result[timestep][layer][head] = value - return result - - mask_strategy = dict_to_3d_list(mask_strategy) - - #best_mask_selections = None - # 7. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(self.scheduler.timesteps): - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = latent_model_input.to(transformer_dtype) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - encoder_hidden_states_2=prompt_embeds_2, - return_dict=False, - mask_strategy=mask_strategy[i], - ) - # perform guidance - if do_classifier_free_guidance: - noise_pred_text, noise_pred_uncond = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(model_output=noise_pred, timestep=t, sample=latents) - - progress_bar.update() - - if not torch.distributed.is_initialized() or int(torch.distributed.get_rank()) == 0: - if not output_type == "latent": - video = self.decode_vae(latents) - video = self.video_processor.postprocess_video(video, - output_file_name=output_file_name, - output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video, ) - - return StepVideoPipelineOutput(video=video) diff --git a/fastvideo/models/stepvideo/modules/attentions.py b/fastvideo/models/stepvideo/modules/attentions.py deleted file mode 100755 index 1011fae5c..000000000 --- a/fastvideo/models/stepvideo/modules/attentions.py +++ /dev/null @@ -1,96 +0,0 @@ -import torch -import torch.nn as nn -from einops import rearrange -from flash_attn import flash_attn_func - -try: - from st_attn import sliding_tile_attention -except ImportError: - print("Could not load Sliding Tile Attention.") - sliding_tile_attention = None - -from fastvideo.utils.communications import all_to_all_4D -from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info - - -class Attention(nn.Module): - - def __init__(self): - super().__init__() - - def attn_processor(self, attn_type): - if attn_type == 'torch': - return self.torch_attn_func - elif attn_type == 'parallel': - return self.parallel_attn_func - else: - raise Exception('Not supported attention type...') - - def tile(self, x, sp_size): - x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=36 // sp_size, h=48, w=48) - return rearrange(x, - "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", - n_t=6, - n_h=6, - n_w=6, - ts_t=6, - ts_h=8, - ts_w=8) - - def untile(self, x, sp_size): - x = rearrange(x, - "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d", - n_t=6, - n_h=6, - n_w=6, - ts_t=6, - ts_h=8, - ts_w=8) - return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=36 // sp_size, h=48, w=48) - - def torch_attn_func(self, q, k, v, attn_mask=None, causal=False, drop_rate=0.0, **kwargs): - - if attn_mask is not None and attn_mask.dtype != torch.bool: - attn_mask = attn_mask.to(q.dtype) - - if attn_mask is not None and attn_mask.ndim == 3: ## no head - n_heads = q.shape[2] - attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) - - q, k, v = map(lambda x: rearrange(x, 'b s h d -> b h s d'), (q, k, v)) - x = torch.nn.functional.scaled_dot_product_attention(q, - k, - v, - attn_mask=attn_mask, - dropout_p=drop_rate, - is_causal=causal) - x = rearrange(x, 'b h s d -> b s h d') - return x - - def parallel_attn_func(self, q, k, v, causal=False, mask_strategy=None, **kwargs): - if get_sequence_parallel_state(): - q = all_to_all_4D(q, scatter_dim=2, gather_dim=1) - k = all_to_all_4D(k, scatter_dim=2, gather_dim=1) - v = all_to_all_4D(v, scatter_dim=2, gather_dim=1) - - if mask_strategy[0] is not None: - q = self.tile(q, nccl_info.sp_size).transpose(1, 2).contiguous() - k = self.tile(k, nccl_info.sp_size).transpose(1, 2).contiguous() - v = self.tile(v, nccl_info.sp_size).transpose(1, 2).contiguous() - - head_num = q.size(1) # 48 // sp_size - current_rank = nccl_info.rank_within_group - - start_head = current_rank * head_num - windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)] - - x = sliding_tile_attention(q, k, v, windows, 0, False).transpose(1, 2).contiguous() - x = self.untile(x, nccl_info.sp_size) - else: - x = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False) - - if get_sequence_parallel_state(): - x = all_to_all_4D(x, scatter_dim=1, gather_dim=2) - - x = x.to(q.dtype) - return x diff --git a/fastvideo/models/stepvideo/modules/blocks.py b/fastvideo/models/stepvideo/modules/blocks.py deleted file mode 100755 index befcfec7e..000000000 --- a/fastvideo/models/stepvideo/modules/blocks.py +++ /dev/null @@ -1,296 +0,0 @@ -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -from typing import Optional - -import torch -import torch.nn as nn -from einops import rearrange - -from fastvideo.models.stepvideo.modules.attentions import Attention -from fastvideo.models.stepvideo.modules.normalization import RMSNorm -from fastvideo.models.stepvideo.modules.rope import RoPE3D - - -class SelfAttention(Attention): - - def __init__(self, hidden_dim, head_dim, bias=False, with_rope=True, with_qk_norm=True, attn_type='torch'): - super().__init__() - self.head_dim = head_dim - self.n_heads = hidden_dim // head_dim - - self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=bias) - self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias) - - self.with_rope = with_rope - self.with_qk_norm = with_qk_norm - if self.with_qk_norm: - self.q_norm = RMSNorm(head_dim, elementwise_affine=True) - self.k_norm = RMSNorm(head_dim, elementwise_affine=True) - - if self.with_rope: - self.rope_3d = RoPE3D(freq=1e4, F0=1.0, scaling_factor=1.0) - self.rope_ch_split = [64, 32, 32] - - self.core_attention = self.attn_processor(attn_type=attn_type) - self.parallel = attn_type == 'parallel' - - def apply_rope3d(self, x, fhw_positions, rope_ch_split, parallel=True): - x = self.rope_3d(x, fhw_positions, rope_ch_split, parallel) - return x - - def forward(self, x, cu_seqlens=None, max_seqlen=None, rope_positions=None, attn_mask=None, mask_strategy=None): - xqkv = self.wqkv(x) - xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim) - - xq, xk, xv = torch.split(xqkv, [self.head_dim] * 3, dim=-1) ## seq_len, n, dim - - if self.with_qk_norm: - xq = self.q_norm(xq) - xk = self.k_norm(xk) - - if self.with_rope: - xq = self.apply_rope3d(xq, rope_positions, self.rope_ch_split, parallel=self.parallel) - xk = self.apply_rope3d(xk, rope_positions, self.rope_ch_split, parallel=self.parallel) - - output = self.core_attention(xq, - xk, - xv, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - attn_mask=attn_mask, - mask_strategy=mask_strategy) - output = rearrange(output, 'b s h d -> b s (h d)') - output = self.wo(output) - - return output - - -class CrossAttention(Attention): - - def __init__(self, hidden_dim, head_dim, bias=False, with_qk_norm=True, attn_type='torch'): - super().__init__() - self.head_dim = head_dim - self.n_heads = hidden_dim // head_dim - - self.wq = nn.Linear(hidden_dim, hidden_dim, bias=bias) - self.wkv = nn.Linear(hidden_dim, hidden_dim * 2, bias=bias) - self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias) - - self.with_qk_norm = with_qk_norm - if self.with_qk_norm: - self.q_norm = RMSNorm(head_dim, elementwise_affine=True) - self.k_norm = RMSNorm(head_dim, elementwise_affine=True) - - self.core_attention = self.attn_processor(attn_type=attn_type) - - def forward(self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, attn_mask=None): - xq = self.wq(x) - xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim) - - xkv = self.wkv(encoder_hidden_states) - xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim) - - xk, xv = torch.split(xkv, [self.head_dim] * 2, dim=-1) ## seq_len, n, dim - - if self.with_qk_norm: - xq = self.q_norm(xq) - xk = self.k_norm(xk) - - output = self.core_attention(xq, xk, xv, attn_mask=attn_mask) - - output = rearrange(output, 'b s h d -> b s (h d)') - output = self.wo(output) - - return output - - -class GELU(nn.Module): - r""" - GELU activation function with tanh approximation support with `approximate="tanh"`. - - Parameters: - dim_in (`int`): The number of channels in the input. - dim_out (`int`): The number of channels in the output. - approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation. - bias (`bool`, defaults to True): Whether to use a bias in the linear layer. - """ - - def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True): - super().__init__() - self.proj = nn.Linear(dim_in, dim_out, bias=bias) - self.approximate = approximate - - def gelu(self, gate: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.gelu(gate, approximate=self.approximate) - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states = self.gelu(hidden_states) - return hidden_states - - -class FeedForward(nn.Module): - - def __init__( - self, - dim: int, - inner_dim: Optional[int] = None, - dim_out: Optional[int] = None, - mult: int = 4, - bias: bool = False, - ): - super().__init__() - inner_dim = dim * mult if inner_dim is None else inner_dim - dim_out = dim if dim_out is None else dim_out - self.net = nn.ModuleList([ - GELU(dim, inner_dim, approximate="tanh", bias=bias), - nn.Identity(), - nn.Linear(inner_dim, dim_out, bias=bias) - ]) - - def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: - for module in self.net: - hidden_states = module(hidden_states) - return hidden_states - - -def modulate(x, scale, shift): - x = x * (1 + scale) + shift - return x - - -def gate(x, gate): - x = gate * x - return x - - -class StepVideoTransformerBlock(nn.Module): - r""" - A basic Transformer block. - - Parameters: - dim (`int`): The number of channels in the input and output. - num_attention_heads (`int`): The number of heads to use for multi-head attention. - attention_head_dim (`int`): The number of channels in each head. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - num_embeds_ada_norm (: - obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. - attention_bias (: - obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used. - double_self_attention (`bool`, *optional*): - Whether to use two self-attention layers. In this case no cross attention layers are used. - upcast_attention (`bool`, *optional*): - Whether to upcast the attention computation to float32. This is useful for mixed precision training. - norm_elementwise_affine (`bool`, *optional*, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_type (`str`, *optional*, defaults to `"layer_norm"`): - The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. - final_dropout (`bool` *optional*, defaults to False): - Whether to apply a final dropout after the last feed-forward layer. - attention_type (`str`, *optional*, defaults to `"default"`): - The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. - positional_embeddings (`str`, *optional*, defaults to `None`): - The type of positional embeddings to apply to. - num_positional_embeddings (`int`, *optional*, defaults to `None`): - The maximum number of positional embeddings to apply. - """ - - def __init__(self, - dim: int, - attention_head_dim: int, - norm_eps: float = 1e-5, - ff_inner_dim: Optional[int] = None, - ff_bias: bool = False, - attention_type: str = 'parallel'): - super().__init__() - self.dim = dim - self.norm1 = nn.LayerNorm(dim, eps=norm_eps) - self.attn1 = SelfAttention(dim, - attention_head_dim, - bias=False, - with_rope=True, - with_qk_norm=True, - attn_type=attention_type) - - self.norm2 = nn.LayerNorm(dim, eps=norm_eps) - self.attn2 = CrossAttention(dim, attention_head_dim, bias=False, with_qk_norm=True, attn_type='torch') - - self.ff = FeedForward(dim=dim, inner_dim=ff_inner_dim, dim_out=dim, bias=ff_bias) - - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) - - @torch.no_grad() - def forward(self, - q: torch.Tensor, - kv: Optional[torch.Tensor] = None, - timestep: Optional[torch.LongTensor] = None, - attn_mask=None, - rope_positions: list = None, - mask_strategy=None) -> torch.Tensor: - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (torch.clone(chunk) for chunk in ( - self.scale_shift_table[None] + timestep.reshape(-1, 6, self.dim)).chunk(6, dim=1)) - - scale_shift_q = modulate(self.norm1(q), scale_msa, shift_msa) - - attn_q = self.attn1(scale_shift_q, rope_positions=rope_positions, mask_strategy=mask_strategy) - - q = gate(attn_q, gate_msa) + q - - attn_q = self.attn2(q, kv, attn_mask) - - q = attn_q + q - - scale_shift_q = modulate(self.norm2(q), scale_mlp, shift_mlp) - - ff_output = self.ff(scale_shift_q) - - q = gate(ff_output, gate_mlp) + q - - return q - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - patch_size=64, - in_channels=3, - embed_dim=768, - layer_norm=False, - flatten=True, - bias=True, - ): - super().__init__() - - self.flatten = flatten - self.layer_norm = layer_norm - - self.proj = nn.Conv2d(in_channels, - embed_dim, - kernel_size=(patch_size, patch_size), - stride=patch_size, - bias=bias) - - def forward(self, latent): - latent = self.proj(latent).to(latent.dtype) - if self.flatten: - latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC - if self.layer_norm: - latent = self.norm(latent) - - return latent diff --git a/fastvideo/models/stepvideo/modules/model.py b/fastvideo/models/stepvideo/modules/model.py deleted file mode 100755 index 2864403ad..000000000 --- a/fastvideo/models/stepvideo/modules/model.py +++ /dev/null @@ -1,198 +0,0 @@ -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -from typing import Dict, Optional - -import torch -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.models.modeling_utils import ModelMixin -from einops import rearrange, repeat -from torch import nn - -from fastvideo.models.stepvideo.modules.blocks import PatchEmbed, StepVideoTransformerBlock -from fastvideo.models.stepvideo.modules.normalization import AdaLayerNormSingle, PixArtAlphaTextProjection -from fastvideo.models.stepvideo.parallel import parallel_forward -from fastvideo.models.stepvideo.utils import with_empty_init - - -class StepVideoModel(ModelMixin, ConfigMixin): - _no_split_modules = ["StepVideoTransformerBlock", "PatchEmbed"] - - @with_empty_init - @register_to_config - def __init__( - self, - num_attention_heads: int = 48, - attention_head_dim: int = 128, - in_channels: int = 64, - out_channels: Optional[int] = 64, - num_layers: int = 48, - dropout: float = 0.0, - patch_size: int = 1, - norm_type: str = "ada_norm_single", - norm_elementwise_affine: bool = False, - norm_eps: float = 1e-6, - use_additional_conditions: Optional[bool] = False, - caption_channels: Optional[int] | list | tuple = [6144, 1024], - attention_type: Optional[str] = "parallel", - ): - super().__init__() - - # Set some common variables used across the board. - self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim - self.out_channels = in_channels if out_channels is None else out_channels - - self.use_additional_conditions = use_additional_conditions - - self.pos_embed = PatchEmbed( - patch_size=patch_size, - in_channels=self.config.in_channels, - embed_dim=self.inner_dim, - ) - - self.transformer_blocks = nn.ModuleList([ - StepVideoTransformerBlock(dim=self.inner_dim, - attention_head_dim=self.config.attention_head_dim, - attention_type=attention_type) for _ in range(self.config.num_layers) - ]) - - # 3. Output blocks. - self.norm_out = nn.LayerNorm(self.inner_dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) - self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels) - self.patch_size = patch_size - - self.adaln_single = AdaLayerNormSingle(self.inner_dim, use_additional_conditions=self.use_additional_conditions) - - if isinstance(self.config.caption_channels, int): - caption_channel = self.config.caption_channels - else: - caption_channel, clip_channel = self.config.caption_channels - self.clip_projection = nn.Linear(clip_channel, self.inner_dim) - - self.caption_norm = nn.LayerNorm(caption_channel, eps=norm_eps, elementwise_affine=norm_elementwise_affine) - - self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channel, hidden_size=self.inner_dim) - - self.parallel = attention_type == 'parallel' - - def patchfy(self, hidden_states): - hidden_states = rearrange(hidden_states, 'b f c h w -> (b f) c h w') - hidden_states = self.pos_embed(hidden_states) - return hidden_states - - def prepare_attn_mask(self, encoder_attention_mask, encoder_hidden_states, q_seqlen): - kv_seqlens = encoder_attention_mask.sum(dim=1).int() - mask = torch.zeros([len(kv_seqlens), q_seqlen, max(kv_seqlens)], - dtype=torch.bool, - device=encoder_attention_mask.device) - encoder_hidden_states = encoder_hidden_states[:, :max(kv_seqlens)] - for i, kv_len in enumerate(kv_seqlens): - mask[i, :, :kv_len] = 1 - return encoder_hidden_states, mask - - @parallel_forward - def block_forward(self, - hidden_states, - encoder_hidden_states=None, - timestep=None, - rope_positions=None, - attn_mask=None, - parallel=True, - mask_strategy=None): - - for i, block in enumerate(self.transformer_blocks): - hidden_states = block(hidden_states, - encoder_hidden_states, - timestep=timestep, - attn_mask=attn_mask, - rope_positions=rope_positions, - mask_strategy=mask_strategy[i]) - - return hidden_states - - @torch.inference_mode() - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_hidden_states_2: Optional[torch.Tensor] = None, - timestep: Optional[torch.LongTensor] = None, - added_cond_kwargs: Dict[str, torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - fps: torch.Tensor = None, - return_dict: bool = True, - mask_strategy=None, - ): - assert hidden_states.ndim == 5 - "hidden_states's shape should be (bsz, f, ch, h ,w)" - - bsz, frame, _, height, width = hidden_states.shape - height, width = height // self.patch_size, width // self.patch_size - - hidden_states = self.patchfy(hidden_states) - len_frame = hidden_states.shape[1] - - if self.use_additional_conditions: - added_cond_kwargs = { - "resolution": torch.tensor([(height, width)] * bsz, - device=hidden_states.device, - dtype=hidden_states.dtype), - "nframe": torch.tensor([frame] * bsz, device=hidden_states.device, dtype=hidden_states.dtype), - "fps": fps - } - else: - added_cond_kwargs = {} - - timestep, embedded_timestep = self.adaln_single(timestep, added_cond_kwargs=added_cond_kwargs) - - encoder_hidden_states = self.caption_projection(self.caption_norm(encoder_hidden_states)) - - if encoder_hidden_states_2 is not None and hasattr(self, 'clip_projection'): - clip_embedding = self.clip_projection(encoder_hidden_states_2) - encoder_hidden_states = torch.cat([clip_embedding, encoder_hidden_states], dim=1) - - hidden_states = rearrange(hidden_states, '(b f) l d-> b (f l) d', b=bsz, f=frame, l=len_frame).contiguous() - encoder_hidden_states, attn_mask = self.prepare_attn_mask(encoder_attention_mask, - encoder_hidden_states, - q_seqlen=frame * len_frame) - - hidden_states = self.block_forward(hidden_states, - encoder_hidden_states, - timestep=timestep, - rope_positions=[frame, height, width], - attn_mask=attn_mask, - parallel=self.parallel, - mask_strategy=mask_strategy) - - hidden_states = rearrange(hidden_states, 'b (f l) d -> (b f) l d', b=bsz, f=frame, l=len_frame) - - embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d', f=frame).contiguous() - - shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states) - # Modulation - hidden_states = hidden_states * (1 + scale) + shift - hidden_states = self.proj_out(hidden_states) - - # unpatchify - hidden_states = hidden_states.reshape(shape=(-1, height, width, self.patch_size, self.patch_size, - self.out_channels)) - - hidden_states = rearrange(hidden_states, 'n h w p q c -> n c h p w q') - output = hidden_states.reshape(shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)) - - output = rearrange(output, '(b f) c h w -> b f c h w', f=frame) - - if return_dict: - return {'x': output} - return output diff --git a/fastvideo/models/stepvideo/modules/normalization.py b/fastvideo/models/stepvideo/modules/normalization.py deleted file mode 100755 index 18950be82..000000000 --- a/fastvideo/models/stepvideo/modules/normalization.py +++ /dev/null @@ -1,312 +0,0 @@ -import math -from typing import Dict, Optional, Tuple - -import torch -import torch.nn as nn - - -class RMSNorm(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): - """ - 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 - - -ACTIVATION_FUNCTIONS = { - "swish": nn.SiLU(), - "silu": nn.SiLU(), - "mish": nn.Mish(), - "gelu": nn.GELU(), - "relu": nn.ReLU(), -} - - -def get_activation(act_fn: str) -> nn.Module: - """Helper function to get activation function from string. - - Args: - act_fn (str): Name of activation function. - - Returns: - nn.Module: Activation function. - """ - - act_fn = act_fn.lower() - if act_fn in ACTIVATION_FUNCTIONS: - return ACTIVATION_FUNCTIONS[act_fn] - else: - raise ValueError(f"Unsupported activation function: {act_fn}") - - -def get_timestep_embedding( - timesteps: torch.Tensor, - embedding_dim: int, - flip_sin_to_cos: bool = False, - downscale_freq_shift: float = 1, - scale: float = 1, - max_period: int = 10000, -): - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - - :param timesteps: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the - embeddings. :return: an [N x dim] Tensor of positional embeddings. - """ - assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) - exponent = exponent / (half_dim - downscale_freq_shift) - - emb = torch.exp(exponent) - emb = timesteps[:, None].float() * emb[None, :] - - # scale embeddings - emb = scale * emb - - # concat sine and cosine embeddings - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) - - # flip sine and cosine embeddings - if flip_sin_to_cos: - emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) - - # zero pad - if embedding_dim % 2 == 1: - emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) - return emb - - -class Timesteps(nn.Module): - - def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - - def forward(self, timesteps): - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - ) - return t_emb - - -class TimestepEmbedding(nn.Module): - - def __init__(self, - in_channels: int, - time_embed_dim: int, - act_fn: str = "silu", - out_dim: int = None, - post_act_fn: Optional[str] = None, - cond_proj_dim=None, - sample_proj_bias=True): - super().__init__() - linear_cls = nn.Linear - - self.linear_1 = linear_cls( - in_channels, - time_embed_dim, - bias=sample_proj_bias, - ) - - if cond_proj_dim is not None: - self.cond_proj = linear_cls( - cond_proj_dim, - in_channels, - bias=False, - ) - else: - self.cond_proj = None - - self.act = get_activation(act_fn) - - if out_dim is not None: - time_embed_dim_out = out_dim - else: - time_embed_dim_out = time_embed_dim - - self.linear_2 = linear_cls( - time_embed_dim, - time_embed_dim_out, - bias=sample_proj_bias, - ) - - if post_act_fn is None: - self.post_act = None - else: - self.post_act = get_activation(post_act_fn) - - def forward(self, sample, condition=None): - if condition is not None: - sample = sample + self.cond_proj(condition) - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - - if self.post_act is not None: - sample = self.post_act(sample) - return sample - - -class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): - - def __init__(self, embedding_dim, size_emb_dim, use_additional_conditions: bool = False): - super().__init__() - - self.outdim = size_emb_dim - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - - self.use_additional_conditions = use_additional_conditions - if self.use_additional_conditions: - self.additional_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.resolution_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim) - self.nframe_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - self.fps_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - - def forward(self, timestep, resolution=None, nframe=None, fps=None): - hidden_dtype = next(self.timestep_embedder.parameters()).dtype - - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) - - if self.use_additional_conditions: - batch_size = timestep.shape[0] - resolution_emb = self.additional_condition_proj(resolution.flatten()).to(hidden_dtype) - resolution_emb = self.resolution_embedder(resolution_emb).reshape(batch_size, -1) - nframe_emb = self.additional_condition_proj(nframe.flatten()).to(hidden_dtype) - nframe_emb = self.nframe_embedder(nframe_emb).reshape(batch_size, -1) - conditioning = timesteps_emb + resolution_emb + nframe_emb - - if fps is not None: - fps_emb = self.additional_condition_proj(fps.flatten()).to(hidden_dtype) - fps_emb = self.fps_embedder(fps_emb).reshape(batch_size, -1) - conditioning = conditioning + fps_emb - else: - conditioning = timesteps_emb - - return conditioning - - -class AdaLayerNormSingle(nn.Module): - r""" - Norm layer adaptive layer norm single (adaLN-single). - - As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - use_additional_conditions (`bool`): To use additional conditions for normalization or not. - """ - - def __init__(self, embedding_dim: int, use_additional_conditions: bool = False, time_step_rescale=1000): - super().__init__() - - self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim, - size_emb_dim=embedding_dim // 2, - use_additional_conditions=use_additional_conditions) - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) - - self.time_step_rescale = time_step_rescale ## timestep usually in [0, 1], we rescale it to [0,1000] for stability - - def forward( - self, - timestep: torch.Tensor, - added_cond_kwargs: Dict[str, torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - embedded_timestep = self.emb(timestep * self.time_step_rescale, **added_cond_kwargs) - - out = self.linear(self.silu(embedded_timestep)) - - return out, embedded_timestep - - -class PixArtAlphaTextProjection(nn.Module): - """ - Projects caption embeddings. Also handles dropout for classifier-free guidance. - - Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py - """ - - def __init__(self, in_features, hidden_size): - super().__init__() - self.linear_1 = nn.Linear( - in_features, - hidden_size, - bias=True, - ) - self.act_1 = nn.GELU(approximate="tanh") - self.linear_2 = nn.Linear( - hidden_size, - hidden_size, - bias=True, - ) - - def forward(self, caption): - hidden_states = self.linear_1(caption) - hidden_states = self.act_1(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states diff --git a/fastvideo/models/stepvideo/modules/rope.py b/fastvideo/models/stepvideo/modules/rope.py deleted file mode 100755 index c9b094d66..000000000 --- a/fastvideo/models/stepvideo/modules/rope.py +++ /dev/null @@ -1,90 +0,0 @@ -import torch - -from fastvideo.utils.parallel_states import nccl_info - - -class RoPE1D: - - def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0): - self.base = freq - self.F0 = F0 - self.scaling_factor = scaling_factor - self.cache = {} - - def get_cos_sin(self, D, seq_len, device, dtype): - if (D, seq_len, device, dtype) not in self.cache: - inv_freq = 1.0 / (self.base**(torch.arange(0, D, 2).float().to(device) / D)) - t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) - freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype) - freqs = torch.cat((freqs, freqs), dim=-1) - cos = freqs.cos() # (Seq, Dim) - sin = freqs.sin() - self.cache[D, seq_len, device, dtype] = (cos, sin) - return self.cache[D, seq_len, device, dtype] - - @staticmethod - def rotate_half(x): - x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:] - return torch.cat((-x2, x1), dim=-1) - - def apply_rope1d(self, tokens, pos1d, cos, sin): - assert pos1d.ndim == 2 - cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :] - sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :] - return (tokens * cos) + (self.rotate_half(tokens) * sin) - - def __call__(self, tokens, positions): - """ - input: - * tokens: batch_size x ntokens x nheads x dim - * positions: batch_size x ntokens (t position of each token) - output: - * tokens after applying RoPE2D (batch_size x ntokens x nheads x dim) - """ - D = tokens.size(3) - assert positions.ndim == 2 # Batch, Seq - cos, sin = self.get_cos_sin(D, int(positions.max()) + 1, tokens.device, tokens.dtype) - tokens = self.apply_rope1d(tokens, positions, cos, sin) - return tokens - - -class RoPE3D(RoPE1D): - - def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0): - super(RoPE3D, self).__init__(freq, F0, scaling_factor) - self.position_cache = {} - - def get_mesh_3d(self, rope_positions, bsz): - f, h, w = rope_positions - - if f"{f}-{h}-{w}" not in self.position_cache: - x = torch.arange(f, device='cpu') - y = torch.arange(h, device='cpu') - z = torch.arange(w, device='cpu') - self.position_cache[f"{f}-{h}-{w}"] = torch.cartesian_prod(x, y, z).view(1, f * h * w, 3).expand(bsz, -1, 3) - return self.position_cache[f"{f}-{h}-{w}"] - - def __call__(self, tokens, rope_positions, ch_split, parallel=False): - """ - input: - * tokens: batch_size x ntokens x nheads x dim - * rope_positions: list of (f, h, w) - output: - * tokens after applying RoPE2D (batch_size x ntokens x nheads x dim) - """ - assert sum(ch_split) == tokens.size(-1) - - mesh_grid = self.get_mesh_3d(rope_positions, bsz=tokens.shape[0]) - out = [] - for i, (D, x) in enumerate(zip(ch_split, torch.split(tokens, ch_split, dim=-1))): - cos, sin = self.get_cos_sin(D, int(mesh_grid.max()) + 1, tokens.device, tokens.dtype) - - if parallel: - mesh = torch.chunk(mesh_grid[:, :, i], nccl_info.sp_size, dim=1)[nccl_info.rank_within_group].clone() - else: - mesh = mesh_grid[:, :, i].clone() - x = self.apply_rope1d(x, mesh.to(tokens.device), cos, sin) - out.append(x) - - tokens = torch.cat(out, dim=-1) - return tokens diff --git a/fastvideo/models/stepvideo/parallel.py b/fastvideo/models/stepvideo/parallel.py deleted file mode 100644 index 9ce6ee3ce..000000000 --- a/fastvideo/models/stepvideo/parallel.py +++ /dev/null @@ -1,21 +0,0 @@ -import torch - -from fastvideo.utils.communications import all_gather -from fastvideo.utils.parallel_states import nccl_info - - -def parallel_forward(fn_): - - def wrapTheFunction(_, hidden_states, *args, **kwargs): - if kwargs['parallel']: - hidden_states = torch.chunk(hidden_states, nccl_info.sp_size, dim=-2)[nccl_info.rank_within_group] - kwargs['attn_mask'] = torch.chunk(kwargs['attn_mask'], nccl_info.sp_size, - dim=-2)[nccl_info.rank_within_group] - output = fn_(_, hidden_states, *args, **kwargs) - - if kwargs['parallel']: - output = all_gather(output.contiguous(), dim=-2) - - return output - - return wrapTheFunction diff --git a/fastvideo/models/stepvideo/text_encoder/__init__.py b/fastvideo/models/stepvideo/text_encoder/__init__.py deleted file mode 100644 index bb56aad46..000000000 --- a/fastvideo/models/stepvideo/text_encoder/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -import os - -import torch - -from fastvideo.models.stepvideo.config import parse_args - -try: - args = parse_args() - torch.ops.load_library( - os.path.join(args.model_dir, 'lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so')) -except Exception as err: - print(err) diff --git a/fastvideo/models/stepvideo/text_encoder/clip.py b/fastvideo/models/stepvideo/text_encoder/clip.py deleted file mode 100755 index 8e0386eef..000000000 --- a/fastvideo/models/stepvideo/text_encoder/clip.py +++ /dev/null @@ -1,36 +0,0 @@ -import os - -import torch -import torch.nn as nn -from transformers import BertModel, BertTokenizer - - -class HunyuanClip(nn.Module): - """ - Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py - hunyuan's clip used BertModel and BertTokenizer, so we copy it. - """ - - def __init__(self, model_dir, max_length=77): - super(HunyuanClip, self).__init__() - - self.max_length = max_length - self.tokenizer = BertTokenizer.from_pretrained(os.path.join(model_dir, 'tokenizer')) - self.text_encoder = BertModel.from_pretrained(os.path.join(model_dir, 'clip_text_encoder')) - - @torch.no_grad - def forward(self, prompts, with_mask=True): - self.device = next(self.text_encoder.parameters()).device - text_inputs = self.tokenizer( - prompts, - padding="max_length", - max_length=self.max_length, - truncation=True, - return_attention_mask=True, - return_tensors="pt", - ) - prompt_embeds = self.text_encoder( - text_inputs.input_ids.to(self.device), - attention_mask=text_inputs.attention_mask.to(self.device) if with_mask else None, - ) - return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output diff --git a/fastvideo/models/stepvideo/text_encoder/flashattention.py b/fastvideo/models/stepvideo/text_encoder/flashattention.py deleted file mode 100755 index 647698f4b..000000000 --- a/fastvideo/models/stepvideo/text_encoder/flashattention.py +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -import torch - - -def flash_attn_func(q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=True, - return_attn_probs=False, - tp_group_rank=0, - tp_group_size=1): - softmax_scale = q.size(-1)**(-0.5) if softmax_scale is None else softmax_scale - return torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale, causal, return_attn_probs, None, - tp_group_rank, tp_group_size)[0] - - -class FlashSelfAttention(torch.nn.Module): - - def __init__( - self, - attention_dropout=0.0, - ): - super().__init__() - self.dropout_p = attention_dropout - - def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None): - if cu_seqlens is None: - output = flash_attn_func(q, k, v, dropout_p=self.dropout_p) - else: - raise ValueError('cu_seqlens is not supported!') - - return output diff --git a/fastvideo/models/stepvideo/text_encoder/stepllm.py b/fastvideo/models/stepvideo/text_encoder/stepllm.py deleted file mode 100755 index 7e971477a..000000000 --- a/fastvideo/models/stepvideo/text_encoder/stepllm.py +++ /dev/null @@ -1,291 +0,0 @@ -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -import os -from typing import Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange -from transformers.modeling_utils import PretrainedConfig, PreTrainedModel - -from fastvideo.models.stepvideo.modules.normalization import RMSNorm -from fastvideo.models.stepvideo.text_encoder.flashattention import FlashSelfAttention -from fastvideo.models.stepvideo.text_encoder.tokenizer import LLaMaEmbedding, Wrapped_StepChatTokenizer -from fastvideo.models.stepvideo.utils import with_empty_init - - -def safediv(n, d): - q, r = divmod(n, d) - assert r == 0 - return q - - -class MultiQueryAttention(nn.Module): - - def __init__(self, cfg, layer_id=None): - super().__init__() - - self.head_dim = cfg.hidden_size // cfg.num_attention_heads - self.max_seq_len = cfg.seq_length - self.use_flash_attention = cfg.use_flash_attn - assert self.use_flash_attention, 'FlashAttention is required!' - - self.n_groups = cfg.num_attention_groups - self.tp_size = 1 - self.n_local_heads = cfg.num_attention_heads - self.n_local_groups = self.n_groups - - self.wqkv = nn.Linear( - cfg.hidden_size, - cfg.hidden_size + self.head_dim * 2 * self.n_groups, - bias=False, - ) - self.wo = nn.Linear( - cfg.hidden_size, - cfg.hidden_size, - bias=False, - ) - - assert self.use_flash_attention, 'non-Flash attention not supported yet.' - self.core_attention = FlashSelfAttention(attention_dropout=cfg.attention_dropout) - - self.layer_id = layer_id - - def forward( - self, - x: torch.Tensor, - mask: Optional[torch.Tensor], - cu_seqlens: Optional[torch.Tensor], - max_seq_len: Optional[torch.Tensor], - ): - seqlen, bsz, dim = x.shape - xqkv = self.wqkv(x) - - xq, xkv = torch.split( - xqkv, - (dim // self.tp_size, self.head_dim * 2 * self.n_groups // self.tp_size), - dim=-1, - ) - - # gather on 1st dimension - xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim) - xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim) - xk, xv = xkv.chunk(2, -1) - - # rotary embedding + flash attn - xq = rearrange(xq, "s b h d -> b s h d") - xk = rearrange(xk, "s b h d -> b s h d") - xv = rearrange(xv, "s b h d -> b s h d") - - q_per_kv = self.n_local_heads // self.n_local_groups - if q_per_kv > 1: - b, s, h, d = xk.size() - if h == 1: - xk = xk.expand(b, s, q_per_kv, d) - xv = xv.expand(b, s, q_per_kv, d) - else: - ''' To cover the cases where h > 1, we have - the following implementation, which is equivalent to: - xk = xk.repeat_interleave(q_per_kv, dim=-2) - xv = xv.repeat_interleave(q_per_kv, dim=-2) - but can avoid calling aten::item() that involves cpu. - ''' - idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten() - xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() - xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() - - if self.use_flash_attention: - output = self.core_attention(xq, xk, xv, cu_seqlens=cu_seqlens, max_seq_len=max_seq_len) - # reduce-scatter only support first dimension now - output = rearrange(output, "b s h d -> s b (h d)").contiguous() - else: - xq, xk, xv = [rearrange(x, "b s ... -> s b ...").contiguous() for x in (xq, xk, xv)] - output = self.core_attention(xq, xk, xv, mask) - output = self.wo(output) - return output - - -class FeedForward(nn.Module): - - def __init__( - self, - cfg, - dim: int, - hidden_dim: int, - layer_id: int, - multiple_of: int = 256, - ): - super().__init__() - - hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) - - def swiglu(x): - x = torch.chunk(x, 2, dim=-1) - return F.silu(x[0]) * x[1] - - self.swiglu = swiglu - - self.w1 = nn.Linear( - dim, - 2 * hidden_dim, - bias=False, - ) - self.w2 = nn.Linear( - hidden_dim, - dim, - bias=False, - ) - - def forward(self, x): - x = self.swiglu(self.w1(x)) - output = self.w2(x) - return output - - -class TransformerBlock(nn.Module): - - def __init__(self, cfg, layer_id: int): - super().__init__() - - self.n_heads = cfg.num_attention_heads - self.dim = cfg.hidden_size - self.head_dim = cfg.hidden_size // cfg.num_attention_heads - self.attention = MultiQueryAttention( - cfg, - layer_id=layer_id, - ) - - self.feed_forward = FeedForward( - cfg, - dim=cfg.hidden_size, - hidden_dim=cfg.ffn_hidden_size, - layer_id=layer_id, - ) - self.layer_id = layer_id - self.attention_norm = RMSNorm( - cfg.hidden_size, - eps=cfg.layernorm_epsilon, - ) - self.ffn_norm = RMSNorm( - cfg.hidden_size, - eps=cfg.layernorm_epsilon, - ) - - def forward( - self, - x: torch.Tensor, - mask: Optional[torch.Tensor], - cu_seqlens: Optional[torch.Tensor], - max_seq_len: Optional[torch.Tensor], - ): - residual = self.attention.forward(self.attention_norm(x), mask, cu_seqlens, max_seq_len) - h = x + residual - ffn_res = self.feed_forward.forward(self.ffn_norm(h)) - out = h + ffn_res - return out - - -class Transformer(nn.Module): - - def __init__( - self, - config, - max_seq_size=8192, - ): - super().__init__() - self.num_layers = config.num_layers - self.layers = self._build_layers(config) - - def _build_layers(self, config): - layers = torch.nn.ModuleList() - for layer_id in range(self.num_layers): - layers.append(TransformerBlock( - config, - layer_id=layer_id + 1, - )) - return layers - - def forward( - self, - hidden_states, - attention_mask, - cu_seqlens=None, - max_seq_len=None, - ): - - if max_seq_len is not None and not isinstance(max_seq_len, torch.Tensor): - max_seq_len = torch.tensor(max_seq_len, dtype=torch.int32, device="cpu") - - for lid, layer in enumerate(self.layers): - hidden_states = layer( - hidden_states, - attention_mask, - cu_seqlens, - max_seq_len, - ) - return hidden_states - - -class Step1Model(PreTrainedModel): - config_class = PretrainedConfig - - @with_empty_init - def __init__( - self, - config, - ): - super().__init__(config) - self.tok_embeddings = LLaMaEmbedding(config) - self.transformer = Transformer(config) - - def forward( - self, - input_ids=None, - attention_mask=None, - ): - - hidden_states = self.tok_embeddings(input_ids) - - hidden_states = self.transformer( - hidden_states, - attention_mask, - ) - return hidden_states - - -class STEP1TextEncoder(torch.nn.Module): - - def __init__(self, model_dir, max_length=320): - super(STEP1TextEncoder, self).__init__() - self.max_length = max_length - self.text_tokenizer = Wrapped_StepChatTokenizer(os.path.join(model_dir, 'step1_chat_tokenizer.model')) - text_encoder = Step1Model.from_pretrained(model_dir) - self.text_encoder = text_encoder.eval().to(torch.bfloat16) - - @torch.no_grad - def forward(self, prompts, with_mask=True, max_length=None): - self.device = next(self.text_encoder.parameters()).device - with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16): - if type(prompts) is str: - prompts = [prompts] - - txt_tokens = self.text_tokenizer(prompts, - max_length=max_length or self.max_length, - padding="max_length", - truncation=True, - return_tensors="pt") - y = self.text_encoder(txt_tokens.input_ids.to(self.device), - attention_mask=txt_tokens.attention_mask.to(self.device) if with_mask else None) - y_mask = txt_tokens.attention_mask - return y.transpose(0, 1), y_mask diff --git a/fastvideo/models/stepvideo/text_encoder/tokenizer.py b/fastvideo/models/stepvideo/text_encoder/tokenizer.py deleted file mode 100755 index 62ca0d608..000000000 --- a/fastvideo/models/stepvideo/text_encoder/tokenizer.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -from typing import List - -import torch -import torch.nn as nn - - -class LLaMaEmbedding(nn.Module): - """Language model embeddings. - - Arguments: - hidden_size: hidden size - vocab_size: vocabulary size - max_sequence_length: maximum size of sequence. This - is used for positional embedding - embedding_dropout_prob: dropout probability for embeddings - init_method: weight initialization method - num_tokentypes: size of the token-type embeddings. 0 value - will ignore this embedding - """ - - def __init__( - self, - cfg, - ): - super().__init__() - self.hidden_size = cfg.hidden_size - self.params_dtype = cfg.params_dtype - self.fp32_residual_connection = cfg.fp32_residual_connection - self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32 - self.word_embeddings = torch.nn.Embedding( - cfg.padded_vocab_size, - self.hidden_size, - ) - self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) - - def forward(self, input_ids): - # Embeddings. - if self.embedding_weights_in_fp32: - self.word_embeddings = self.word_embeddings.to(torch.float32) - embeddings = self.word_embeddings(input_ids) - if self.embedding_weights_in_fp32: - embeddings = embeddings.to(self.params_dtype) - self.word_embeddings = self.word_embeddings.to(self.params_dtype) - - # Data format change to avoid explicit transposes : [b s h] --> [s b h]. - embeddings = embeddings.transpose(0, 1).contiguous() - - # If the input flag for fp32 residual connection is set, convert for float. - if self.fp32_residual_connection: - embeddings = embeddings.float() - - # Dropout. - embeddings = self.embedding_dropout(embeddings) - - return embeddings - - -class StepChatTokenizer: - """Step Chat Tokenizer""" - - def __init__( - self, - model_file, - name="StepChatTokenizer", - bot_token="<|BOT|>", # Begin of Turn - eot_token="<|EOT|>", # End of Turn - call_start_token="<|CALL_START|>", # Call Start - call_end_token="<|CALL_END|>", # Call End - think_start_token="<|THINK_START|>", # Think Start - think_end_token="<|THINK_END|>", # Think End - mask_start_token="<|MASK_1e69f|>", # Mask start - mask_end_token="<|UNMASK_1e69f|>", # Mask end - ): - import sentencepiece - - self._tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file) - - self._vocab = {} - self._inv_vocab = {} - - self._special_tokens = {} - self._inv_special_tokens = {} - - self._t5_tokens = [] - - for idx in range(self._tokenizer.get_piece_size()): - text = self._tokenizer.id_to_piece(idx) - self._inv_vocab[idx] = text - self._vocab[text] = idx - - if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown(idx): - self._special_tokens[text] = idx - self._inv_special_tokens[idx] = text - - self._unk_id = self._tokenizer.unk_id() - self._bos_id = self._tokenizer.bos_id() - self._eos_id = self._tokenizer.eos_id() - - for token in [bot_token, eot_token, call_start_token, call_end_token, think_start_token, think_end_token]: - assert token in self._vocab, f"Token '{token}' not found in tokenizer" - assert token in self._special_tokens, f"Token '{token}' is not a special token" - - for token in [mask_start_token, mask_end_token]: - assert token in self._vocab, f"Token '{token}' not found in tokenizer" - - self._bot_id = self._tokenizer.piece_to_id(bot_token) - self._eot_id = self._tokenizer.piece_to_id(eot_token) - self._call_start_id = self._tokenizer.piece_to_id(call_start_token) - self._call_end_id = self._tokenizer.piece_to_id(call_end_token) - self._think_start_id = self._tokenizer.piece_to_id(think_start_token) - self._think_end_id = self._tokenizer.piece_to_id(think_end_token) - self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token) - self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token) - - self._underline_id = self._tokenizer.piece_to_id("\u2581") - - @property - def vocab(self): - return self._vocab - - @property - def inv_vocab(self): - return self._inv_vocab - - @property - def vocab_size(self): - return self._tokenizer.vocab_size() - - def tokenize(self, text: str) -> List[int]: - return self._tokenizer.encode_as_ids(text) - - def detokenize(self, token_ids: List[int]) -> str: - return self._tokenizer.decode_ids(token_ids) - - -class Tokens: - - def __init__(self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len) -> None: - self.input_ids = input_ids - self.attention_mask = attention_mask - self.cu_input_ids = cu_input_ids - self.cu_seqlens = cu_seqlens - self.max_seq_len = max_seq_len - - def to(self, device): - self.input_ids = self.input_ids.to(device) - self.attention_mask = self.attention_mask.to(device) - self.cu_input_ids = self.cu_input_ids.to(device) - self.cu_seqlens = self.cu_seqlens.to(device) - return self - - -class Wrapped_StepChatTokenizer(StepChatTokenizer): - - def __call__(self, text, max_length=320, padding="max_length", truncation=True, return_tensors="pt"): - # [bos, ..., eos, pad, pad, ..., pad] - self.BOS = 1 - self.EOS = 2 - self.PAD = 2 - out_tokens = [] - attn_mask = [] - if len(text) == 0: - part_tokens = [self.BOS] + [self.EOS] - valid_size = len(part_tokens) - if len(part_tokens) < max_length: - part_tokens += [self.PAD] * (max_length - valid_size) - out_tokens.append(part_tokens) - attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) - else: - for part in text: - part_tokens = self.tokenize(part) - part_tokens = part_tokens[:(max_length - 2)] # leave 2 space for bos and eos - part_tokens = [self.BOS] + part_tokens + [self.EOS] - valid_size = len(part_tokens) - if len(part_tokens) < max_length: - part_tokens += [self.PAD] * (max_length - valid_size) - out_tokens.append(part_tokens) - attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) - - out_tokens = torch.tensor(out_tokens, dtype=torch.long) - attn_mask = torch.tensor(attn_mask, dtype=torch.long) - - # padding y based on tp size - padded_len = 0 - padded_flag = True if padded_len > 0 else False - if padded_flag: - pad_tokens = torch.tensor([[self.PAD] * max_length], device=out_tokens.device) - pad_attn_mask = torch.tensor([[1] * padded_len + [0] * (max_length - padded_len)], device=attn_mask.device) - out_tokens = torch.cat([out_tokens, pad_tokens], dim=0) - attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0) - - # cu_seqlens - cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0) - seqlen = attn_mask.sum(dim=1).tolist() - cu_seqlens = torch.cumsum(torch.tensor([0] + seqlen), 0).to(device=out_tokens.device, dtype=torch.int32) - max_seq_len = max(seqlen) - return Tokens(out_tokens, cu_out_tokens, attn_mask, cu_seqlens, max_seq_len) diff --git a/fastvideo/models/stepvideo/utils/__init__.py b/fastvideo/models/stepvideo/utils/__init__.py deleted file mode 100755 index a77226b13..000000000 --- a/fastvideo/models/stepvideo/utils/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .utils import * -from .video_process import * \ No newline at end of file diff --git a/fastvideo/models/stepvideo/utils/quantization.py b/fastvideo/models/stepvideo/utils/quantization.py deleted file mode 100644 index e87e38670..000000000 --- a/fastvideo/models/stepvideo/utils/quantization.py +++ /dev/null @@ -1,117 +0,0 @@ -# from stepvideo.diffusion.video_pipeline import StepVideoPipeline -import torch -import torch.nn as nn -from torch.nn import functional as F - - -def get_fp_maxval(bits=8, mantissa_bit=3, sign_bits=1): - _bits = torch.tensor(bits) - _mantissa_bit = torch.tensor(mantissa_bit) - _sign_bits = torch.tensor(sign_bits) - M = torch.clamp(torch.round(_mantissa_bit), 1, _bits - _sign_bits) - E = _bits - _sign_bits - M - bias = 2**(E - 1) - 1 - mantissa = 1 - for i in range(mantissa_bit - 1): - mantissa += 1 / (2**(i + 1)) - maxval = mantissa * 2**(2**E - 1 - bias) - return maxval - - -def quantize_to_fp8(x, bits=8, mantissa_bit=3, sign_bits=1): - """ - Default is E4M3. - """ - bits = torch.tensor(bits) - mantissa_bit = torch.tensor(mantissa_bit) - sign_bits = torch.tensor(sign_bits) - M = torch.clamp(torch.round(mantissa_bit), 1, bits - sign_bits) - E = bits - sign_bits - M - bias = 2**(E - 1) - 1 - mantissa = 1 - for i in range(mantissa_bit - 1): - mantissa += 1 / (2**(i + 1)) - maxval = mantissa * 2**(2**E - 1 - bias) - minval = -maxval - minval = -maxval if sign_bits == 1 else torch.zeros_like(maxval) - input_clamp = torch.min(torch.max(x, minval), maxval) - log_scales = torch.clamp((torch.floor(torch.log2(torch.abs(input_clamp)) + bias)).detach(), 1.0) - log_scales = 2.0**(log_scales - M - bias.type(x.dtype)) - # dequant - qdq_out = torch.round(input_clamp / log_scales) * log_scales - return qdq_out, log_scales - - -def fp8_tensor_quant(x, scale, bits=8, mantissa_bit=3, sign_bits=1): - for i in range(len(x.shape) - 1): - scale = scale.unsqueeze(-1) - new_x = x / scale - quant_dequant_x, log_scales = quantize_to_fp8(new_x, bits=bits, mantissa_bit=mantissa_bit, sign_bits=sign_bits) - return quant_dequant_x, scale, log_scales - - -def fp8_activation_dequant(qdq_out, scale, dtype): - qdq_out = qdq_out.type(dtype) - quant_dequant_x = qdq_out * scale.to(dtype) - return quant_dequant_x - - -def fp8_linear_forward(cls, original_dtype, input): - weight_dtype = cls.weight.dtype - ##### - if cls.weight.dtype != torch.float8_e4m3fn: - assert False - maxval = get_fp_maxval() - scale = torch.max(torch.abs(cls.weight.flatten())) / maxval - linear_weight, scale, log_scales = fp8_tensor_quant(cls.weight, scale) - linear_weight = linear_weight.to(torch.float8_e4m3fn) - weight_dtype = linear_weight.dtype - else: - scale = cls.fp8_scale.to(cls.weight.device) - linear_weight = cls.weight - ##### - - if weight_dtype == torch.float8_e4m3fn: - if True or len(input.shape) == 3: - cls_dequant = fp8_activation_dequant(linear_weight, scale, original_dtype) - if cls.bias is not None: - print(f"input dtype: {input.dtype}") - print(f"cls_dequant dtype: {cls_dequant.dtype}") - print(f"cls.bias dtype: {cls.bias.dtype}") - - output = F.linear(input, cls_dequant, cls.bias) - else: - output = F.linear(input, cls_dequant) - return output - else: - return cls.original_forward(input.to(original_dtype)) - else: - return cls.original_forward(input) - - -def convert_fp8_linear(module, original_dtype, params_to_keep={}): - setattr(module, "fp8_matmul_enabled", True) - fp8_layers = [] - scale_dict = {} - counter = 0 - for key, layer in module.named_modules(): - if isinstance(layer, nn.Linear) and 'transformer_blocks' in key: - print(f"Converting {key} to FP8") - fp8_layers.append(key) - original_forward = layer.forward - maxval = get_fp_maxval() - scale = torch.max(torch.abs(layer.weight.flatten())) / maxval - - original_weight = layer.weight.data # Store a reference to the original weights - quantized_weight, scale, _ = fp8_tensor_quant(original_weight, scale) - scale_dict[key] = scale - layer.weight = torch.nn.Parameter(quantized_weight.to(torch.float8_e4m3fn)) - del original_weight # Delete the reference to the original weights - torch.cuda.empty_cache() - - # print(f"layer weight dtype: {layer.weight.dtype} for layer {key}") - setattr(layer, "fp8_scale", scale.to(dtype=original_dtype)) - setattr(layer, "original_forward", original_forward) - setattr(layer, "forward", lambda input, m=layer: fp8_linear_forward(m, original_dtype, input)) - counter += 1 - return scale_dict diff --git a/fastvideo/models/stepvideo/utils/utils.py b/fastvideo/models/stepvideo/utils/utils.py deleted file mode 100755 index 835393277..000000000 --- a/fastvideo/models/stepvideo/utils/utils.py +++ /dev/null @@ -1,60 +0,0 @@ -import random -from functools import wraps - -import numpy as np -import torch -import torch.utils._device - - -def setup_seed(seed): - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - - -class EmptyInitOnDevice(torch.overrides.TorchFunctionMode): - - def __init__(self, device=None): - self.device = device - - def __torch_function__(self, func, types, args=(), kwargs=None): - kwargs = kwargs or {} - if getattr(func, '__module__', None) == 'torch.nn.init': - if 'tensor' in kwargs: - return kwargs['tensor'] - else: - return args[0] - if self.device is not None and func in torch.utils._device._device_constructors( - ) and kwargs.get('device') is None: - kwargs['device'] = self.device - return func(*args, **kwargs) - - -def with_empty_init(func): - - @wraps(func) - def wrapper(*args, **kwargs): - with EmptyInitOnDevice('cpu'): - return func(*args, **kwargs) - - return wrapper - - -def culens2mask(cu_seqlens=None, cu_seqlens_kv=None, max_seqlen=None, max_seqlen_kv=None, is_causal=False): - assert len(cu_seqlens) == len(cu_seqlens_kv) - "q k v should have same bsz..." - bsz = len(cu_seqlens) - 1 - seqlens = cu_seqlens[1:] - cu_seqlens[:-1] - seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - - attn_mask = torch.zeros(bsz, max_seqlen, max_seqlen_kv, dtype=torch.bool) - for i, (seq_len, seq_len_kv) in enumerate(zip(seqlens, seqlens_kv)): - if is_causal: - attn_mask[i, :seq_len, :seq_len_kv] = torch.triu(torch.ones(seq_len, seq_len_kv), diagonal=1).bool() - else: - attn_mask[i, :seq_len, :seq_len_kv] = torch.ones([seq_len, seq_len_kv], dtype=torch.bool) - - return attn_mask diff --git a/fastvideo/models/stepvideo/utils/video_process.py b/fastvideo/models/stepvideo/utils/video_process.py deleted file mode 100644 index 93e6af546..000000000 --- a/fastvideo/models/stepvideo/utils/video_process.py +++ /dev/null @@ -1,51 +0,0 @@ -import os - -import imageio -import numpy as np -import torch - - -class VideoProcessor: - - def __init__(self, save_path: str = './results', name_suffix: str = ''): - self.save_path = save_path - os.makedirs(self.save_path, exist_ok=True) - self.name_suffix = name_suffix - - def crop2standard540p(self, vid_array): - _, height, width, _ = vid_array.shape - height_center = height // 2 - width_center = width // 2 - if width_center > height_center: ## horizon mode - return vid_array[:, height_center - 270:height_center + 270, width_center - 480:width_center + 480] - elif width_center < height_center: ## portrait mode - return vid_array[:, height_center - 480:height_center + 480, width_center - 270:width_center + 270] - else: - return vid_array - - def save_imageio_video(self, video_array: np.array, output_filename: str, fps=25, codec='libx264'): - - ffmpeg_params = [ - "-vf", - "atadenoise=0a=0.1:0b=0.1:1a=0.1:1b=0.1", # denoise - ] - - with imageio.get_writer(output_filename, fps=fps, codec=codec, ffmpeg_params=ffmpeg_params) as vid_writer: - for img_array in video_array: - vid_writer.append_data(img_array) - - def postprocess_video(self, video_tensor, output_file_name='', output_type="mp4", crop2standard540p=True): - if len(self.name_suffix) == 0: - video_path = os.path.join(self.save_path, f"{output_file_name}.{output_type}") - else: - video_path = os.path.join(self.save_path, f"{output_file_name}-{self.name_suffix}.{output_type}") - - video_tensor = torch.cat([t for t in video_tensor], dim=-2) - video_tensor = (video_tensor.cpu().clamp(-1, 1) + 1) * 127.5 - video_array = video_tensor.clamp(0, 255).to(torch.uint8).numpy().transpose(0, 2, 3, 1) - - if crop2standard540p: - video_array = self.crop2standard540p(video_array) - - self.save_imageio_video(video_array, video_path) - print(f"Saved the generated video in {video_path}") diff --git a/fastvideo/models/stepvideo/vae/vae.py b/fastvideo/models/stepvideo/vae/vae.py deleted file mode 100755 index 1dbf532c0..000000000 --- a/fastvideo/models/stepvideo/vae/vae.py +++ /dev/null @@ -1,975 +0,0 @@ -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -import torch -from einops import rearrange -from torch import nn -from torch.nn import functional as F - -from fastvideo.models.stepvideo.utils import with_empty_init - - -def base_group_norm(x, norm_layer, act_silu=False, channel_last=False): - if hasattr(base_group_norm, 'spatial') and base_group_norm.spatial: - assert channel_last - x_shape = x.shape - x = x.flatten(0, 1) - if channel_last: - # Permute to NCHW format - x = x.permute(0, 3, 1, 2) - - out = F.group_norm(x.contiguous(), norm_layer.num_groups, norm_layer.weight, norm_layer.bias, norm_layer.eps) - if act_silu: - out = F.silu(out) - - if channel_last: - # Permute back to NHWC format - out = out.permute(0, 2, 3, 1) - - out = out.view(x_shape) - else: - if channel_last: - # Permute to NCHW format - x = x.permute(0, 3, 1, 2) - out = F.group_norm(x.contiguous(), norm_layer.num_groups, norm_layer.weight, norm_layer.bias, norm_layer.eps) - if act_silu: - out = F.silu(out) - if channel_last: - # Permute back to NHWC format - out = out.permute(0, 2, 3, 1) - return out - - -def base_conv2d(x, conv_layer, channel_last=False, residual=None): - if channel_last: - x = x.permute(0, 3, 1, 2) # NHWC to NCHW - out = F.conv2d(x, conv_layer.weight, conv_layer.bias, stride=conv_layer.stride, padding=conv_layer.padding) - if residual is not None: - if channel_last: - residual = residual.permute(0, 3, 1, 2) # NHWC to NCHW - out += residual - if channel_last: - out = out.permute(0, 2, 3, 1) # NCHW to NHWC - return out - - -def base_conv3d(x, conv_layer, channel_last=False, residual=None, only_return_output=False): - if only_return_output: - size = cal_outsize(x.shape, conv_layer.weight.shape, conv_layer.stride, conv_layer.padding) - return torch.empty(size, device=x.device, dtype=x.dtype) - if channel_last: - x = x.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW - out = F.conv3d(x, conv_layer.weight, conv_layer.bias, stride=conv_layer.stride, padding=conv_layer.padding) - if residual is not None: - if channel_last: - residual = residual.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW - out += residual - if channel_last: - out = out.permute(0, 2, 3, 4, 1) # NCDHW to NDHWC - return out - - -def cal_outsize(input_sizes, kernel_sizes, stride, padding): - stride_d, stride_h, stride_w = stride - padding_d, padding_h, padding_w = padding - dilation_d, dilation_h, dilation_w = 1, 1, 1 - - in_d = input_sizes[1] - in_h = input_sizes[2] - in_w = input_sizes[3] - - kernel_d = kernel_sizes[2] - kernel_h = kernel_sizes[3] - kernel_w = kernel_sizes[4] - out_channels = kernel_sizes[0] - - out_d = calc_out_(in_d, padding_d, dilation_d, kernel_d, stride_d) - out_h = calc_out_(in_h, padding_h, dilation_h, kernel_h, stride_h) - out_w = calc_out_(in_w, padding_w, dilation_w, kernel_w, stride_w) - size = [input_sizes[0], out_d, out_h, out_w, out_channels] - return size - - -def calc_out_(in_size, padding, dilation, kernel, stride): - return (in_size + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1 - - -def base_conv3d_channel_last(x, conv_layer, residual=None): - in_numel = x.numel() - out_numel = int(x.numel() * conv_layer.out_channels / conv_layer.in_channels) - if (in_numel >= 2**30) or (out_numel >= 2**30): - assert conv_layer.stride[0] == 1, "time split asks time stride = 1" - - B, T, H, W, C = x.shape - K = conv_layer.kernel_size[0] - - chunks = 4 - chunk_size = T // chunks - - if residual is None: - out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual, only_return_output=True) - else: - out_nhwc = residual - - assert B == 1 - for i in range(chunks): - if i == chunks - 1: - xi = x[:1, chunk_size * i:] - out_nhwci = out_nhwc[:1, chunk_size * i:] - else: - xi = x[:1, chunk_size * i:chunk_size * (i + 1) + K - 1] - out_nhwci = out_nhwc[:1, chunk_size * i:chunk_size * (i + 1)] - if residual is not None: - if i == chunks - 1: - ri = residual[:1, chunk_size * i:] - else: - ri = residual[:1, chunk_size * i:chunk_size * (i + 1)] - else: - ri = None - out_nhwci.copy_(base_conv3d(xi, conv_layer, channel_last=True, residual=ri)) - else: - out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual) - return out_nhwc - - -class Upsample2D(nn.Module): - - def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.use_conv_transpose = use_conv_transpose - - if use_conv: - self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1) - else: - assert "Not Supported" - self.conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1) - - def forward(self, x, output_size=None): - assert x.shape[-1] == self.channels - - if self.use_conv_transpose: - return self.conv(x) - - if output_size is None: - x = F.interpolate(x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last), - scale_factor=2.0, - mode='nearest').permute(0, 2, 3, 1).contiguous() - else: - x = F.interpolate(x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last), - size=output_size, - mode='nearest').permute(0, 2, 3, 1).contiguous() - - # x = self.conv(x) - x = base_conv2d(x, self.conv, channel_last=True) - return x - - -class Downsample2D(nn.Module): - - def __init__(self, channels, use_conv=False, out_channels=None, padding=1): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.padding = padding - stride = 2 - - if use_conv: - self.conv = nn.Conv2d(self.channels, self.out_channels, 3, stride=stride, padding=padding) - else: - assert self.channels == self.out_channels - self.conv = nn.AvgPool2d(kernel_size=stride, stride=stride) - - def forward(self, x): - assert x.shape[-1] == self.channels - if self.use_conv and self.padding == 0: - pad = (0, 0, 0, 1, 0, 1) - x = F.pad(x, pad, mode="constant", value=0) - - assert x.shape[-1] == self.channels - # x = self.conv(x) - x = base_conv2d(x, self.conv, channel_last=True) - return x - - -class CausalConv(nn.Module): - - def __init__(self, chan_in, chan_out, kernel_size, **kwargs): - super().__init__() - - if isinstance(kernel_size, int): - kernel_size = kernel_size if isinstance(kernel_size, tuple) else ((kernel_size, ) * 3) - time_kernel_size, height_kernel_size, width_kernel_size = kernel_size - - self.dilation = kwargs.pop('dilation', 1) - self.stride = kwargs.pop('stride', 1) - if isinstance(self.stride, int): - self.stride = (self.stride, 1, 1) - time_pad = self.dilation * (time_kernel_size - 1) + max((1 - self.stride[0]), 0) - height_pad = height_kernel_size // 2 - width_pad = width_kernel_size // 2 - self.time_causal_padding = (width_pad, width_pad, height_pad, height_pad, time_pad, 0) - self.time_uncausal_padding = (width_pad, width_pad, height_pad, height_pad, 0, 0) - - self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs) - self.is_first_run = True - - def forward(self, x, is_init=True, residual=None): - x = nn.functional.pad(x, self.time_causal_padding if is_init else self.time_uncausal_padding) - - x = self.conv(x) - if residual is not None: - x.add_(residual) - return x - - -class ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - factor: int, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.factor = factor - assert out_channels * factor**3 % in_channels == 0 - self.repeats = out_channels * factor**3 // in_channels - - def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: - x = x.repeat_interleave(self.repeats, dim=1) - x = x.view(x.size(0), self.out_channels, self.factor, self.factor, self.factor, x.size(2), x.size(3), x.size(4)) - x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() - x = x.view(x.size(0), self.out_channels, - x.size(2) * self.factor, - x.size(4) * self.factor, - x.size(6) * self.factor) - x = x[:, :, self.factor - 1:, :, :] - return x - - -class ConvPixelShuffleUpSampleLayer3D(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - factor: int, - ): - super().__init__() - self.factor = factor - out_ratio = factor**3 - self.conv = CausalConv(in_channels, out_channels * out_ratio, kernel_size=kernel_size) - - def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: - x = self.conv(x, is_init) - x = self.pixel_shuffle_3d(x, self.factor) - return x - - @staticmethod - def pixel_shuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: - batch_size, channels, depth, height, width = x.size() - new_channels = channels // (factor**3) - new_depth = depth * factor - new_height = height * factor - new_width = width * factor - - x = x.view(batch_size, new_channels, factor, factor, factor, depth, height, width) - x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() - x = x.view(batch_size, new_channels, new_depth, new_height, new_width) - x = x[:, :, factor - 1:, :, :] - return x - - -class ConvPixelUnshuffleDownSampleLayer3D(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - factor: int, - ): - super().__init__() - self.factor = factor - out_ratio = factor**3 - assert out_channels % out_ratio == 0 - self.conv = CausalConv(in_channels, out_channels // out_ratio, kernel_size=kernel_size) - - def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: - x = self.conv(x, is_init) - x = self.pixel_unshuffle_3d(x, self.factor) - return x - - @staticmethod - def pixel_unshuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: - pad = (0, 0, 0, 0, factor - 1, 0) # (left, right, top, bottom, front, back) - x = F.pad(x, pad) - B, C, D, H, W = x.shape - x = x.view(B, C, D // factor, factor, H // factor, factor, W // factor, factor) - x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() - x = x.view(B, C * factor**3, D // factor, H // factor, W // factor) - return x - - -class PixelUnshuffleChannelAveragingDownSampleLayer3D(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - factor: int, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.factor = factor - assert in_channels * factor**3 % out_channels == 0 - self.group_size = in_channels * factor**3 // out_channels - - def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: - pad = (0, 0, 0, 0, self.factor - 1, 0) # (left, right, top, bottom, front, back) - x = F.pad(x, pad) - B, C, D, H, W = x.shape - x = x.view(B, C, D // self.factor, self.factor, H // self.factor, self.factor, W // self.factor, self.factor) - x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() - x = x.view(B, C * self.factor**3, D // self.factor, H // self.factor, W // self.factor) - x = x.view(B, self.out_channels, self.group_size, D // self.factor, H // self.factor, W // self.factor) - x = x.mean(dim=2) - return x - - -def base_group_norm_with_zero_pad(x, norm_layer, act_silu=True, pad_size=2): - out_shape = list(x.shape) - out_shape[1] += pad_size - out = torch.empty(out_shape, dtype=x.dtype, device=x.device) - out[:, pad_size:] = base_group_norm(x, norm_layer, act_silu=act_silu, channel_last=True) - out[:, :pad_size] = 0 - return out - - -class CausalConvChannelLast(CausalConv): - - def __init__(self, chan_in, chan_out, kernel_size, **kwargs): - super().__init__(chan_in, chan_out, kernel_size, **kwargs) - - self.time_causal_padding = (0, 0) + self.time_causal_padding - self.time_uncausal_padding = (0, 0) + self.time_uncausal_padding - - def forward(self, x, is_init=True, residual=None): - if self.is_first_run: - self.is_first_run = False - # self.conv.weight = nn.Parameter(self.conv.weight.permute(0,2,3,4,1).contiguous()) - - x = nn.functional.pad(x, self.time_causal_padding if is_init else self.time_uncausal_padding) - - x = base_conv3d_channel_last(x, self.conv, residual=residual) - return x - - -class CausalConvAfterNorm(CausalConv): - - def __init__(self, chan_in, chan_out, kernel_size, **kwargs): - super().__init__(chan_in, chan_out, kernel_size, **kwargs) - - if self.time_causal_padding == (1, 1, 1, 1, 2, 0): - self.conv = nn.Conv3d(chan_in, - chan_out, - kernel_size, - stride=self.stride, - dilation=self.dilation, - padding=(0, 1, 1), - **kwargs) - else: - self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs) - self.is_first_run = True - - def forward(self, x, is_init=True, residual=None): - if self.is_first_run: - self.is_first_run = False - - if self.time_causal_padding == (1, 1, 1, 1, 2, 0): - pass - else: - x = nn.functional.pad(x, self.time_causal_padding).contiguous() - - x = base_conv3d_channel_last(x, self.conv, residual=residual) - return x - - -class AttnBlock(nn.Module): - - def __init__(self, in_channels): - super().__init__() - - self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels) - self.q = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) - self.k = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) - self.v = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) - self.proj_out = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) - - def attention(self, x, is_init=True): - x = base_group_norm(x, self.norm, act_silu=False, channel_last=True) - q = self.q(x, is_init) - k = self.k(x, is_init) - v = self.v(x, is_init) - - b, t, h, w, c = q.shape - q, k, v = map(lambda x: rearrange(x, "b t h w c -> b 1 (t h w) c"), (q, k, v)) - x = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) - x = rearrange(x, "b 1 (t h w) c -> b t h w c", t=t, h=h, w=w) - - return x - - def forward(self, x): - x = x.permute(0, 2, 3, 4, 1).contiguous() - h = self.attention(x) - x = self.proj_out(h, residual=x) - x = x.permute(0, 4, 1, 2, 3) - return x - - -class Resnet3DBlock(nn.Module): - - def __init__( - self, - in_channels, - out_channels=None, - temb_channels=512, - conv_shortcut=False, - ): - super().__init__() - - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - - self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels) - self.conv1 = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) - if temb_channels > 0: - self.temb_proj = nn.Linear(temb_channels, out_channels) - - self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels) - self.conv2 = CausalConvAfterNorm(out_channels, out_channels, kernel_size=3) - - assert conv_shortcut is False - self.use_conv_shortcut = conv_shortcut - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - self.conv_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) - else: - self.nin_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=1) - - def forward(self, x, temb=None, is_init=True): - x = x.permute(0, 2, 3, 4, 1).contiguous() - - h = base_group_norm_with_zero_pad(x, self.norm1, act_silu=True, pad_size=2) - h = self.conv1(h) - if temb is not None: - h = h + self.temb_proj(nn.functional.silu(temb))[:, :, None, None] - - x = self.nin_shortcut(x) if self.in_channels != self.out_channels else x - - h = base_group_norm_with_zero_pad(h, self.norm2, act_silu=True, pad_size=2) - x = self.conv2(h, residual=x) - - x = x.permute(0, 4, 1, 2, 3) - return x - - -class Downsample3D(nn.Module): - - def __init__(self, in_channels, with_conv, stride): - super().__init__() - - self.with_conv = with_conv - if with_conv: - self.conv = CausalConv(in_channels, in_channels, kernel_size=3, stride=stride) - - def forward(self, x, is_init=True): - if self.with_conv: - x = self.conv(x, is_init) - else: - x = nn.functional.avg_pool3d(x, kernel_size=2, stride=2) - return x - - -class VideoEncoder(nn.Module): - - def __init__( - self, - ch=32, - ch_mult=(4, 8, 16, 16), - num_res_blocks=2, - in_channels=3, - z_channels=16, - double_z=True, - down_sampling_layer=[1, 2], - resamp_with_conv=True, - version=1, - ): - super().__init__() - - temb_ch = 0 - - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - - # downsampling - self.conv_in = CausalConv(in_channels, ch, kernel_size=3) - self.down_sampling_layer = down_sampling_layer - - in_ch_mult = (1, ) + tuple(ch_mult) - self.down = nn.ModuleList() - for i_level in range(self.num_resolutions): - block = nn.ModuleList() - attn = nn.ModuleList() - block_in = ch * in_ch_mult[i_level] - block_out = ch * ch_mult[i_level] - for i_block in range(self.num_res_blocks): - block.append(Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch)) - block_in = block_out - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions - 1: - if i_level in self.down_sampling_layer: - down.downsample = Downsample3D(block_in, resamp_with_conv, stride=(2, 2, 2)) - else: - down.downsample = Downsample2D(block_in, resamp_with_conv, padding=0) #DIFF - self.down.append(down) - - # middle - self.mid = nn.Module() - self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) - self.mid.attn_1 = AttnBlock(block_in) - self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) - - # end - self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) - self.version = version - if version == 2: - channels = 4 * z_channels * 2**3 - self.conv_patchify = ConvPixelUnshuffleDownSampleLayer3D(block_in, channels, kernel_size=3, factor=2) - self.shortcut_pathify = PixelUnshuffleChannelAveragingDownSampleLayer3D(block_in, channels, 2) - self.shortcut_out = PixelUnshuffleChannelAveragingDownSampleLayer3D( - channels, 2 * z_channels if double_z else z_channels, 1) - self.conv_out = CausalConvChannelLast(channels, 2 * z_channels if double_z else z_channels, kernel_size=3) - else: - self.conv_out = CausalConvAfterNorm(block_in, 2 * z_channels if double_z else z_channels, kernel_size=3) - - @torch.inference_mode() - def forward(self, x, video_frame_num, is_init=True): - # timestep embedding - temb = None - - t = video_frame_num - - # downsampling - h = self.conv_in(x, is_init) - - # make it real channel last, but behave like normal layout - h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) - - for i_level in range(self.num_resolutions): - for i_block in range(self.num_res_blocks): - h = self.down[i_level].block[i_block](h, temb, is_init) - if len(self.down[i_level].attn) > 0: - h = self.down[i_level].attn[i_block](h) - - if i_level != self.num_resolutions - 1: - if isinstance(self.down[i_level].downsample, Downsample2D): - _, _, t, _, _ = h.shape - h = rearrange(h, "b c t h w -> (b t) h w c", t=t) - h = self.down[i_level].downsample(h) - h = rearrange(h, "(b t) h w c -> b c t h w", t=t) - else: - h = self.down[i_level].downsample(h, is_init) - - h = self.mid.block_1(h, temb, is_init) - h = self.mid.attn_1(h) - h = self.mid.block_2(h, temb, is_init) - - h = h.permute(0, 2, 3, 4, 1).contiguous() # b c l h w -> b l h w c - if self.version == 2: - h = base_group_norm(h, self.norm_out, act_silu=True, channel_last=True) - h = h.permute(0, 4, 1, 2, 3).contiguous() - shortcut = self.shortcut_pathify(h, is_init) - h = self.conv_patchify(h, is_init) - h = h.add_(shortcut) - shortcut = self.shortcut_out(h, is_init).permute(0, 2, 3, 4, 1) - h = self.conv_out(h.permute(0, 2, 3, 4, 1).contiguous(), is_init) - h = h.add_(shortcut) - else: - h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2) - h = self.conv_out(h, is_init) - h = h.permute(0, 4, 1, 2, 3) # b l h w c -> b c l h w - - h = rearrange(h, "b c t h w -> b t c h w") - return h - - -class Res3DBlockUpsample(nn.Module): - - def __init__(self, input_filters, num_filters, down_sampling_stride, down_sampling=False): - super().__init__() - - self.input_filters = input_filters - self.num_filters = num_filters - - self.act_ = nn.SiLU(inplace=True) - - self.conv1 = CausalConvChannelLast(num_filters, num_filters, kernel_size=[3, 3, 3]) - self.norm1 = nn.GroupNorm(32, num_filters) - - self.conv2 = CausalConvChannelLast(num_filters, num_filters, kernel_size=[3, 3, 3]) - self.norm2 = nn.GroupNorm(32, num_filters) - - self.down_sampling = down_sampling - if down_sampling: - self.down_sampling_stride = down_sampling_stride - else: - self.down_sampling_stride = [1, 1, 1] - - if num_filters != input_filters or down_sampling: - self.conv3 = CausalConvChannelLast(input_filters, - num_filters, - kernel_size=[1, 1, 1], - stride=self.down_sampling_stride) - self.norm3 = nn.GroupNorm(32, num_filters) - - def forward(self, x, is_init=False): - x = x.permute(0, 2, 3, 4, 1).contiguous() - - residual = x - - h = self.conv1(x, is_init) - h = base_group_norm(h, self.norm1, act_silu=True, channel_last=True) - - h = self.conv2(h, is_init) - h = base_group_norm(h, self.norm2, act_silu=False, channel_last=True) - - if self.down_sampling or self.num_filters != self.input_filters: - x = self.conv3(x, is_init) - x = base_group_norm(x, self.norm3, act_silu=False, channel_last=True) - - h.add_(x) - h = self.act_(h) - if residual is not None: - h.add_(residual) - - h = h.permute(0, 4, 1, 2, 3) - return h - - -class Upsample3D(nn.Module): - - def __init__(self, in_channels, scale_factor=2): - super().__init__() - - self.scale_factor = scale_factor - self.conv3d = Res3DBlockUpsample(input_filters=in_channels, - num_filters=in_channels, - down_sampling_stride=(1, 1, 1), - down_sampling=False) - - def forward(self, x, is_init=True, is_split=True): - b, c, t, h, w = x.shape - - # x = x.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3).to(memory_format=torch.channels_last_3d) - if is_split: - split_size = c // 8 - x_slices = torch.split(x, split_size, dim=1) - x = [nn.functional.interpolate(x, scale_factor=self.scale_factor) for x in x_slices] - x = torch.cat(x, dim=1) - else: - x = nn.functional.interpolate(x, scale_factor=self.scale_factor) - - x = self.conv3d(x, is_init) - return x - - -class VideoDecoder(nn.Module): - - def __init__( - self, - ch=128, - z_channels=16, - out_channels=3, - ch_mult=(1, 2, 4, 4), - num_res_blocks=2, - temporal_up_layers=[2, 3], - temporal_downsample=4, - resamp_with_conv=True, - version=1, - ): - super().__init__() - - temb_ch = 0 - - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.temporal_downsample = temporal_downsample - - block_in = ch * ch_mult[self.num_resolutions - 1] - self.version = version - if version == 2: - channels = 4 * z_channels * 2**3 - self.conv_in = CausalConv(z_channels, channels, kernel_size=3) - self.shortcut_in = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(z_channels, channels, 1) - self.conv_unpatchify = ConvPixelShuffleUpSampleLayer3D(channels, block_in, kernel_size=3, factor=2) - self.shortcut_unpathify = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(channels, block_in, 2) - else: - self.conv_in = CausalConv(z_channels, block_in, kernel_size=3) - - # middle - self.mid = nn.Module() - self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) - self.mid.attn_1 = AttnBlock(block_in) - self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) - - # upsampling - self.up_id = len(temporal_up_layers) - self.video_frame_num = 1 - self.cur_video_frame_num = self.video_frame_num // 2**self.up_id + 1 - self.up = nn.ModuleList() - for i_level in reversed(range(self.num_resolutions)): - block = nn.ModuleList() - attn = nn.ModuleList() - block_out = ch * ch_mult[i_level] - for i_block in range(self.num_res_blocks + 1): - block.append(Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch)) - block_in = block_out - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - if i_level in temporal_up_layers: - up.upsample = Upsample3D(block_in) - self.cur_video_frame_num = self.cur_video_frame_num * 2 - else: - up.upsample = Upsample2D(block_in, resamp_with_conv) - self.up.insert(0, up) # prepend to get consistent order - - # end - self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) - self.conv_out = CausalConvAfterNorm(block_in, out_channels, kernel_size=3) - - @torch.inference_mode() - def forward(self, z, is_init=True): - z = rearrange(z, "b t c h w -> b c t h w") - - h = self.conv_in(z, is_init=is_init) - if self.version == 2: - shortcut = self.shortcut_in(z, is_init=is_init) - h = h.add_(shortcut) - shortcut = self.shortcut_unpathify(h, is_init=is_init) - h = self.conv_unpatchify(h, is_init=is_init) - h = h.add_(shortcut) - - temb = None - - h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) - h = self.mid.block_1(h, temb, is_init=is_init) - h = self.mid.attn_1(h) - h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) - h = self.mid.block_2(h, temb, is_init=is_init) - - # upsampling - for i_level in reversed(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks + 1): - h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) - h = self.up[i_level].block[i_block](h, temb, is_init=is_init) - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h) - if i_level != 0: - if isinstance(self.up[i_level].upsample, Upsample2D): - B = h.size(0) - h = h.permute(0, 2, 3, 4, 1).flatten(0, 1) - h = self.up[i_level].upsample(h) - h = h.unflatten(0, (B, -1)).permute(0, 4, 1, 2, 3) - else: - h = self.up[i_level].upsample(h, is_init=is_init) - - # end - h = h.permute(0, 2, 3, 4, 1) # b c l h w -> b l h w c - h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2) - h = self.conv_out(h) - h = h.permute(0, 4, 1, 2, 3) - - if is_init: - h = h[:, :, (self.temporal_downsample - 1):] - return h - - -def rms_norm(input, normalized_shape, eps=1e-6): - dtype = input.dtype - input = input.to(torch.float32) - variance = input.pow(2).flatten(-len(normalized_shape)).mean(-1)[(..., ) + (None, ) * len(normalized_shape)] - input = input * torch.rsqrt(variance + eps) - return input.to(dtype) - - -class DiagonalGaussianDistribution(object): - - def __init__(self, parameters, deterministic=False, rms_norm_mean=False, only_return_mean=False): - self.parameters = parameters - self.mean, self.logvar = torch.chunk(parameters, 2, dim=-3) #N,[X],C,H,W - self.logvar = torch.clamp(self.logvar, -30.0, 20.0) - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - self.deterministic = deterministic - if self.deterministic: - self.var = self.std = torch.zeros_like(self.mean, - device=self.parameters.device, - dtype=self.parameters.dtype) - if rms_norm_mean: - self.mean = rms_norm(self.mean, self.mean.size()[1:]) - self.only_return_mean = only_return_mean - - def sample(self, generator=None): - # make sure sample is on the same device - # as the parameters and has same dtype - sample = torch.randn(self.mean.shape, generator=generator, device=self.parameters.device) - sample = sample.to(dtype=self.parameters.dtype) - x = self.mean + self.std * sample - if self.only_return_mean: - return self.mean - else: - return x - - -class AutoencoderKL(nn.Module): - - @with_empty_init - def __init__( - self, - in_channels=3, - out_channels=3, - z_channels=16, - num_res_blocks=2, - model_path=None, - weight_dict={}, - world_size=1, - version=1, - ): - super().__init__() - - self.frame_len = 17 - self.latent_len = 3 if version == 2 else 5 - - base_group_norm.spatial = True if version == 2 else False - - self.encoder = VideoEncoder( - in_channels=in_channels, - z_channels=z_channels, - num_res_blocks=num_res_blocks, - version=version, - ) - - self.decoder = VideoDecoder( - z_channels=z_channels, - out_channels=out_channels, - num_res_blocks=num_res_blocks, - version=version, - ) - - if model_path is not None: - weight_dict = self.init_from_ckpt(model_path) - if len(weight_dict) != 0: - self.load_from_dict(weight_dict) - self.convert_channel_last() - - self.world_size = world_size - - def init_from_ckpt(self, model_path): - from safetensors import safe_open - p = {} - with safe_open(model_path, framework="pt", device="cpu") as f: - for k in f.keys(): - tensor = f.get_tensor(k) - if k.startswith("decoder.conv_out."): - k = k.replace("decoder.conv_out.", "decoder.conv_out.conv.") - p[k] = tensor - return p - - def load_from_dict(self, p): - self.load_state_dict(p) - - def convert_channel_last(self): - #Conv2d NCHW->NHWC - pass - - def naive_encode(self, x, is_init_image=True): - b, len, c, h, w = x.size() - x = rearrange(x, 'b l c h w -> b c l h w').contiguous() - z = self.encoder(x, len, True) # 下采样[1, 4, 8, 16, 16] - return z - - @torch.inference_mode() - def encode(self, x): - # b (nc cf) c h w -> (b nc) cf c h w -> encode -> (b nc) cf c h w -> b (nc cf) c h w - chunks = list(x.split(self.frame_len, dim=1)) - for i in range(len(chunks)): - chunks[i] = self.naive_encode(chunks[i], True) - z = torch.cat(chunks, dim=1) - - posterior = DiagonalGaussianDistribution(z) - return posterior.sample() - - def decode_naive(self, z, is_init=True): - z = z.to(next(self.decoder.parameters()).dtype) - dec = self.decoder(z, is_init) - return dec - - @torch.inference_mode() - def decode(self, z): - # b (nc cf) c h w -> (b nc) cf c h w -> decode -> (b nc) c cf h w -> b (nc cf) c h w - chunks = list(z.split(self.latent_len, dim=1)) - - if self.world_size > 1: - chunks_total_num = len(chunks) - max_num_per_rank = (chunks_total_num + self.world_size - 1) // self.world_size - rank = torch.distributed.get_rank() - chunks_ = chunks[max_num_per_rank * rank:max_num_per_rank * (rank + 1)] - if len(chunks_) < max_num_per_rank: - chunks_.extend(chunks[:max_num_per_rank - len(chunks_)]) - chunks = chunks_ - - for i in range(len(chunks)): - chunks[i] = self.decode_naive(chunks[i], True).permute(0, 2, 1, 3, 4) - x = torch.cat(chunks, dim=1) - - if self.world_size > 1: - x_ = torch.empty([x.size(0), (self.world_size * max_num_per_rank) * self.frame_len, *x.shape[2:]], - dtype=x.dtype, - device=x.device) - torch.distributed.all_gather_into_tensor(x_, x) - x = x_[:, :chunks_total_num * self.frame_len] - - x = self.mix(x) - return x - - def mix(self, x): - remain_scale = 0.6 - mix_scale = 1. - remain_scale - front = slice(self.frame_len - 1, x.size(1) - 1, self.frame_len) - back = slice(self.frame_len, x.size(1), self.frame_len) - x[:, back] = x[:, back] * remain_scale + x[:, front] * mix_scale - x[:, front] = x[:, front] * remain_scale + x[:, back] * mix_scale - return x diff --git a/fastvideo/models/text_encoder.py b/fastvideo/models/text_encoder.py new file mode 100644 index 000000000..fa582d2c8 --- /dev/null +++ b/fastvideo/models/text_encoder.py @@ -0,0 +1,255 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn + +from transformers.utils import ModelOutput +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +def use_default(value, default): + return value if value is not None else default + +@dataclass +class TextEncoderModelOutput(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: + hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + text_outputs (`list`, *optional*, returned when `return_texts=True` is passed): + List of decoded texts. + """ + + hidden_state: torch.FloatTensor = None + attention_mask: Optional[torch.LongTensor] = None + hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None + text_outputs: Optional[list] = None + + +class TextEncoder(nn.Module): + + def __init__( + self, + text_encoder, + tokenizer, + max_length: int, + text_encoder_precision: Optional[str] = None, + text_encoder_path: Optional[str] = None, + output_key: Optional[str] = None, + use_attention_mask: bool = True, + prompt_template: Optional[dict] = None, + prompt_template_video: Optional[dict] = None, + hidden_state_skip_layer: Optional[int] = None, + apply_final_norm: bool = False, + reproduce: bool = False, + device=None, + ): + super().__init__() + # TODO(will): check if there's a cleaner way to do this + self.text_encoder_type = text_encoder.config.architectures[0] + self.max_length = max_length + self.precision = text_encoder_precision + self.model_path = text_encoder_path + self.use_attention_mask = use_attention_mask + if prompt_template_video is not None: + assert (use_attention_mask is True), "Attention mask is True required when training videos." + self.prompt_template = prompt_template + self.prompt_template_video = prompt_template_video + self.hidden_state_skip_layer = hidden_state_skip_layer + self.apply_final_norm = apply_final_norm + self.reproduce = reproduce + + self.use_template = self.prompt_template is not None + if self.use_template: + assert (isinstance(self.prompt_template, dict) and "template" in self.prompt_template + ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}" + assert "{}" in str(self.prompt_template["template"]), ( + "`prompt_template['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template['template']}") + + self.use_video_template = self.prompt_template_video is not None + if self.use_video_template: + if self.prompt_template_video is not None: + assert ( + isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video + ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}" + assert "{}" in str(self.prompt_template_video["template"]), ( + "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template_video['template']}") + + if "T5" in self.text_encoder_type: + self.output_key = output_key or "last_hidden_state" + elif "CLIPTextModel" in self.text_encoder_type: + self.output_key = output_key or "pooler_output" + elif "LlamaModel" in self.text_encoder_type or "glm" in self.text_encoder_type: + self.output_key = output_key or "last_hidden_state" + else: + raise ValueError(f"Unsupported text encoder type: {self.text_encoder_type}") + + self.model = text_encoder + # self.dtype = self.model.dtype + self.device = device + + self.tokenizer = tokenizer + + def __repr__(self): + return f"{self.text_encoder_type} ({self.precision} - {self.model_path})" + + @staticmethod + def apply_text_to_template(text, template, prevent_empty_text=True): + """ + Apply text to template. + + Args: + text (str): Input text. + template (str or list): Template string or list of chat conversation. + prevent_empty_text (bool): If True, we will prevent the user text from being empty + by adding a space. Defaults to True. + """ + if isinstance(template, str): + # Will send string to tokenizer. Used for llm + return template.format(text) + else: + raise TypeError(f"Unsupported template type: {type(template)}") + + def text2tokens(self, text, data_type="image"): + """ + Tokenize the input text. + + Args: + text (str or list): Input text. + """ + tokenize_input_type = "str" + if self.use_template: + if data_type == "image": + prompt_template = self.prompt_template["template"] + elif data_type == "video": + prompt_template = self.prompt_template_video["template"] + else: + raise ValueError(f"Unsupported data type: {data_type}") + if isinstance(text, (list, tuple)): + text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text] + if isinstance(text[0], list): + tokenize_input_type = "list" + elif isinstance(text, str): + text = self.apply_text_to_template(text, prompt_template) + if isinstance(text, list): + tokenize_input_type = "list" + else: + raise TypeError(f"Unsupported text type: {type(text)}") + + kwargs = dict( + truncation=True, + max_length=self.max_length, + padding="max_length", + return_tensors="pt", + ) + if tokenize_input_type == "str": + return self.tokenizer( + text, + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + **kwargs, + ) + elif tokenize_input_type == "list": + return self.tokenizer.apply_chat_template( + text, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **kwargs, + ) + else: + raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}") + + def encode( + self, + batch_encoding, + use_attention_mask=None, + output_hidden_states=False, + do_sample=None, + hidden_state_skip_layer=None, + return_texts=False, + data_type="image", + device=None, + ): + """ + Args: + batch_encoding (dict): Batch encoding from tokenizer. + use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask. + Defaults to None. + output_hidden_states (bool): Whether to output hidden states. If False, return the value of + self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer, + output_hidden_states will be set True. Defaults to False. + do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None. + When self.produce is False, do_sample is set to True by default. + hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer. + If None, self.output_key will be used. Defaults to None. + return_texts (bool): Whether to return the decoded texts. Defaults to False. + """ + device = self.model.device if device is None else device + use_attention_mask = use_default(use_attention_mask, self.use_attention_mask) + hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer) + do_sample = use_default(do_sample, not self.reproduce) + attention_mask = (batch_encoding["attention_mask"].to(device) if use_attention_mask else None) + print(f"batch_encoding: {batch_encoding}") + # logger.info(f"attention_mask: {attention_mask}") + outputs = self.model( + input_ids=batch_encoding["input_ids"].to(device), + # attention_mask=attention_mask, + output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None, + ) + if hidden_state_skip_layer is not None: + last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)] + # Real last hidden state already has layer norm applied. So here we only apply it + # for intermediate layers. + if hidden_state_skip_layer > 0 and self.apply_final_norm: + last_hidden_state = self.model.final_layer_norm(last_hidden_state) + else: + last_hidden_state = outputs[self.output_key] + + # Remove hidden states of instruction tokens, only keep prompt tokens. + if self.use_template: + if data_type == "image": + crop_start = self.prompt_template.get("crop_start", -1) + elif data_type == "video": + crop_start = self.prompt_template_video.get("crop_start", -1) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if crop_start > 0: + last_hidden_state = last_hidden_state[:, crop_start:] + attention_mask = (attention_mask[:, crop_start:] if use_attention_mask else None) + + if output_hidden_states: + return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states) + return TextEncoderModelOutput(last_hidden_state, attention_mask) + + def forward( + self, + text, + use_attention_mask=None, + output_hidden_states=False, + do_sample=False, + hidden_state_skip_layer=None, + return_texts=False, + ): + batch_encoding = self.text2tokens(text) + return self.encode( + batch_encoding, + use_attention_mask=use_attention_mask, + output_hidden_states=output_hidden_states, + do_sample=do_sample, + hidden_state_skip_layer=hidden_state_skip_layer, + return_texts=return_texts, + ) diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 31d230b2e..8f639c018 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -5,6 +5,49 @@ from typing import Any, Dict, Optional import torch +# TODO(PY): move it elsewhere +def auto_attributes(init_func): + """ + Decorator that automatically adds all initialization arguments as object attributes. + + Example: + @auto_attributes + def __init__(self, a=1, b=2): + pass + + # This will automatically set: + # - self.a = 1 and self.b = 2 + # - self.config.a = 1 and self.config.b = 2 + """ + def wrapper(self, *args, **kwargs): + # Get the function signature + import inspect + signature = inspect.signature(init_func) + parameters = signature.parameters + + # Get parameter names (excluding 'self') + param_names = list(parameters.keys())[1:] + + # Bind arguments to parameters + bound_args = signature.bind(self, *args, **kwargs) + bound_args.apply_defaults() + + # Create config object if it doesn't exist + if not hasattr(self, 'config'): + self.config = type('Config', (), {})() + + # Set attributes on self and self.config + for name in param_names: + if name in bound_args.arguments: + value = bound_args.arguments[name] + setattr(self, name, value) + setattr(self.config, name, value) + + # Call the original __init__ function + return init_func(self, *args, **kwargs) + + return wrapper + def set_random_seed(seed: int) -> None: from fastvideo.platforms import current_platform diff --git a/fastvideo/models/vaes/common.py b/fastvideo/models/vaes/common.py new file mode 100644 index 000000000..125b21e01 --- /dev/null +++ b/fastvideo/models/vaes/common.py @@ -0,0 +1,469 @@ +import torch +from typing import Optional, Tuple +import numpy as np +import torch.nn as nn +from abc import abstractmethod, ABC +import torch.distributed as dist +from fastvideo.distributed import tensor_model_parallel_all_gather, get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size +from math import prod +class ParallelTiledVAE(ABC): + def __init__(self, *args, **kwargs): + # Check if subclass has defined all required properties + required_attributes = [ + 'tile_sample_min_height', + 'tile_sample_min_width', + 'tile_sample_min_num_frames', + 'tile_sample_stride_height', + 'tile_sample_stride_width', + 'tile_sample_stride_num_frames', + 'use_tiling' + ] + + for attr in required_attributes: + if not hasattr(self, attr): + raise AttributeError(f"Subclasses of ParallelVAE must define '{attr}' property") + + @abstractmethod + def _encode(self, *args, **kwargs): + pass + + @abstractmethod + def _decode(self, *args, **kwargs): + pass + + + def encode(self, x: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = x.shape + + if self.use_tiling and num_frames > self.tile_sample_min_num_frames: + latents = self.tiled_encode(x) + elif self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): + latents = self.spatial_tiled_encode(x) + else: + latents = self._encode(x) + return DiagonalGaussianDistribution(latents) + + + + def decode(self, z: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + + if self.use_tiling and get_sequence_model_parallel_world_size() > 1: + return self.parallel_tiled_decode(z) + if self.use_tiling and num_frames > tile_latent_min_num_frames: + return self.tiled_decode(z) + + if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): + return self.spatial_tiled_decode(z) + + return self._decode(z) + + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + batch_size, num_channels, num_frames, height, width = x.shape + latent_height = height // self.spatial_compression_ratio + latent_width = width // self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + # Split x into overlapping tiles and encode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, self.tile_sample_stride_height): + row = [] + for j in range(0, width, self.tile_sample_stride_width): + tile = x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] + tile = self._encode(tile) + row.append(tile) + rows.append(row) + + return self._merge_spatial_tiles(rows, blend_height, blend_width,tile_latent_stride_height, tile_latent_stride_width) + + def _parallel_data_generator(self, gathered_results, + gathered_dim_metadata): + global_idx = 0 + for i, per_rank_metadata in enumerate(gathered_dim_metadata): + _start_shape = 0 + for shape in per_rank_metadata: + mul_shape = prod(shape) + yield (gathered_results[i, _start_shape:_start_shape + + mul_shape].reshape(shape), global_idx) + _start_shape += mul_shape + global_idx += 1 + + def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor: + """ + Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs + """ + world_size, rank = get_sequence_model_parallel_world_size(), get_sequence_model_parallel_rank() + B, C, T, H, W = z.shape + + # Calculate parameters + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + + # Calculate tile dimensions + num_t_tiles = (T + tile_latent_stride_num_frames - 1) // tile_latent_stride_num_frames + num_h_tiles = (H + tile_latent_stride_height - 1) // tile_latent_stride_height + num_w_tiles = (W + tile_latent_stride_width - 1) // tile_latent_stride_width + total_spatial_tiles = num_h_tiles * num_w_tiles + total_tiles = num_t_tiles * total_spatial_tiles + + # Calculate tiles per rank and padding + tiles_per_rank = (total_tiles + world_size - 1) // world_size + start_tile_idx = rank * tiles_per_rank + end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles) + + local_results = [] + local_dim_metadata = [] + # Process assigned tiles + for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + + # Calculate positions + t_start = t_idx * tile_latent_stride_num_frames + h_start = h_idx * tile_latent_stride_height + w_start = w_idx * tile_latent_stride_width + + # Extract and process tile + tile = z[:, :, t_start:t_start + tile_latent_min_num_frames + 1, + h_start:h_start + tile_latent_min_height, + w_start:w_start + tile_latent_min_width] + + # Process tile + tile = self._decode(tile) + + if t_start > 0: + tile = tile[:, :, 1:, :, :] + + # Store metadata + shape = tile.shape + # Store decoded data (flattened) + decoded_flat = tile.reshape(-1) + local_results.append(decoded_flat) + local_dim_metadata.append(shape) + + results = torch.cat(local_results, dim=0).contiguous() + del local_results + torch.cuda.empty_cache() + # first gather size to pad the results + local_size = torch.tensor([results.size(0)], + device=results.device, + dtype=torch.int64) + all_sizes = [ + torch.zeros(1, device=results.device, dtype=torch.int64) + for _ in range(world_size) + ] + dist.all_gather(all_sizes, local_size) + max_size = max(size.item() for size in all_sizes) + padded_results = torch.zeros(max_size, device=results.device) + padded_results[:results.size(0)] = results + del results + torch.cuda.empty_cache() + # Gather all results + gathered_dim_metadata = [None] * world_size + gathered_results = torch.zeros_like(padded_results).repeat( + world_size, *[1] * len(padded_results.shape) + ).contiguous() # use contiguous to make sure it won't copy data in the following operations + # TODO (PY): use fastvideo distributed methods + dist.all_gather_into_tensor(gathered_results, padded_results) + dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) + # Process gathered results + data = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] + for _ in range(num_t_tiles)] + for current_data, global_idx in self._parallel_data_generator( + gathered_results, gathered_dim_metadata): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + data[t_idx][h_idx][w_idx] = current_data + # Merge results + result_slices = [] + last_slice_data = None + for i, tem_data in enumerate(data): + slice_data = self._merge_spatial_tiles(tem_data, blend_height, blend_width, + self.tile_sample_stride_height, self.tile_sample_stride_width) + if i > 0: + slice_data = self.blend_t(last_slice_data, slice_data, blend_num_frames) + result_slices.append(slice_data[:, :, :self.tile_sample_stride_num_frames, :, :]) + else: + result_slices.append(slice_data[:, :, :self.tile_sample_stride_num_frames + 1, :, :]) + last_slice_data = slice_data + dec = torch.cat(result_slices, dim=2) + + return dec + + def _merge_spatial_tiles(self, tiles, blend_height, blend_width, stride_height, stride_width): + """Helper function to merge spatial tiles with blending""" + result_rows = [] + for i, row in enumerate(tiles): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(tiles[i - 1][j], tile, + blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :stride_height, :stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + + Returns: + `torch.Tensor`: + The decoded images. + """ + + batch_size, num_channels, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, tile_latent_stride_height): + row = [] + for j in range(0, width, tile_latent_stride_width): + tile = z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + decoded = self._decode(tile) + row.append(decoded) + rows.append(row) + return self._merge_spatial_tiles(rows, blend_height, blend_width,self.tile_sample_stride_height, self.tile_sample_stride_width) + + def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = x.shape + latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1 + + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames + + row = [] + for i in range(0, num_frames, self.tile_sample_stride_num_frames): + tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :] + if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width): + tile = self.spatial_tiled_encode(tile) + else: + tile = self._encode(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_num_frames) + result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :]) + else: + result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :]) + + enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames] + return enc + + def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = z.shape + num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + + row = [] + for i in range(0, num_frames, tile_latent_stride_num_frames): + tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :] + if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height): + decoded = self.spatial_tiled_decode(tile) + else: + decoded = self._decode(tile) + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + result_row = [] + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_num_frames) + result_row.append(tile[:, :, : self.tile_sample_stride_num_frames, :, :]) + else: + result_row.append(tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :]) + + dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames] + return dec + + + + def enable_tiling( + self, + tile_sample_min_height: Optional[int] = None, + tile_sample_min_width: Optional[int] = None, + tile_sample_min_num_frames: Optional[int] = None, + tile_sample_stride_height: Optional[float] = None, + tile_sample_stride_width: Optional[float] = None, + tile_sample_stride_num_frames: Optional[float] = None, + ) -> None: + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + + Args: + tile_sample_min_height (`int`, *optional*): + The minimum height required for a sample to be separated into tiles across the height dimension. + tile_sample_min_width (`int`, *optional*): + The minimum width required for a sample to be separated into tiles across the width dimension. + tile_sample_min_num_frames (`int`, *optional*): + The minimum number of frames required for a sample to be separated into tiles across the frame + dimension. + tile_sample_stride_height (`int`, *optional*): + The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are + no tiling artifacts produced across the height dimension. + tile_sample_stride_width (`int`, *optional*): + The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling + artifacts produced across the width dimension. + tile_sample_stride_num_frames (`int`, *optional*): + The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts + produced across the frame dimension. + """ + self.use_tiling = True + self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames + self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height + self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width + self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames + + + + def disable_tiling(self) -> None: + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_tiling = False + + + +# adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691 +class DiagonalGaussianDistribution(object): + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, device=self.parameters.device, dtype=self.parameters.dtype + ) + + def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = torch.randn( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + else: + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=[1, 2, 3], + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar, + dim=[1, 2, 3], + ) + + def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + + def mode(self) -> torch.Tensor: + return self.mean \ No newline at end of file diff --git a/fastvideo/models/vaes/hunyuanvae.py b/fastvideo/models/vaes/hunyuanvae.py new file mode 100644 index 000000000..f62d89dc5 --- /dev/null +++ b/fastvideo/models/vaes/hunyuanvae.py @@ -0,0 +1,799 @@ +# Copyright 2024 The Hunyuan Team, The HuggingFace Team and The FastVideo 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.models.vaes.common import DiagonalGaussianDistribution, ParallelTiledVAE +from fastvideo.layers.activation import get_act_fn + +from fastvideo.models.utils import auto_attributes + +def prepare_causal_attention_mask( + num_frames: int, height_width: int, dtype: torch.dtype, device: torch.device, batch_size: int = None +) -> torch.Tensor: + indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device) + indices_blocks = indices.repeat_interleave(height_width) + x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy") + mask = torch.where(x <= y, 0, -float("inf")).to(dtype=dtype) + + if batch_size is not None: + mask = mask.unsqueeze(0).expand(batch_size, -1, -1) + return mask + +class HunyuanVAEAttention(nn.Module): + def __init__(self, in_channels, heads, dim_head, eps, norm_num_groups, bias): + super().__init__() + self.in_channels = in_channels + self.heads = heads + self.dim_head = dim_head + self.eps = eps + self.norm_num_groups = norm_num_groups + self.bias = bias + + inner_dim = heads * dim_head + + # Define the projection layers + self.to_q = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_k = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_v = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_out = nn.Sequential( + nn.Linear(inner_dim, in_channels, bias=bias) + ) + + # Optional normalization layers + self.group_norm = nn.GroupNorm(norm_num_groups, in_channels, eps=eps, affine=True) + + def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + residual = hidden_states + + batch_size, sequence_length, _ = hidden_states.shape + + hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + # Project to query, key, value + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + # Reshape for multi-head attention + head_dim = self.dim_head + + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + # Perform scaled dot-product attention + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + # Reshape back + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # Linear projection + hidden_states = self.to_out(hidden_states) + + # Residual connection and rescale + hidden_states = hidden_states + residual + + return hidden_states + +class HunyuanVideoCausalConv3d(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 HunyuanVideoUpsampleCausal3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + kernel_size: int = 3, + stride: int = 1, + bias: bool = True, + upsample_factor: Tuple[float, float, float] = (2, 2, 2), + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + self.upsample_factor = upsample_factor + + self.conv = HunyuanVideoCausalConv3d(in_channels, out_channels, kernel_size, stride, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_frames = hidden_states.size(2) + + first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2) + first_frame = F.interpolate( + first_frame.squeeze(2), scale_factor=self.upsample_factor[1:], mode="nearest" + ).unsqueeze(2) + + if num_frames > 1: + # See: https://github.com/pytorch/pytorch/issues/81665 + # Unless you have a version of pytorch where non-contiguous implementation of F.interpolate + # is fixed, this will raise either a runtime error, or fail silently with bad outputs. + # If you are encountering an error here, make sure to try running encoding/decoding with + # `vae.enable_tiling()` first. If that doesn't work, open an issue at: + # https://github.com/huggingface/diffusers/issues + other_frames = other_frames.contiguous() + other_frames = F.interpolate(other_frames, scale_factor=self.upsample_factor, mode="nearest") + hidden_states = torch.cat((first_frame, other_frames), dim=2) + else: + hidden_states = first_frame + + hidden_states = self.conv(hidden_states) + return hidden_states + + +class HunyuanVideoDownsampleCausal3D(nn.Module): + def __init__( + self, + channels: int, + out_channels: Optional[int] = None, + padding: int = 1, + kernel_size: int = 3, + bias: bool = True, + stride=2, + ) -> None: + super().__init__() + out_channels = out_channels or channels + + self.conv = HunyuanVideoCausalConv3d(channels, out_channels, kernel_size, stride, padding, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv(hidden_states) + return hidden_states + + +class HunyuanVideoResnetBlockCausal3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + dropout: float = 0.0, + groups: int = 32, + eps: float = 1e-6, + non_linearity: str = "silu", + ) -> None: + super().__init__() + out_channels = out_channels or in_channels + + self.nonlinearity = get_act_fn(non_linearity) + + self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True) + self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0) + + self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True) + self.dropout = nn.Dropout(dropout) + self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0) + + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = HunyuanVideoCausalConv3d(in_channels, out_channels, 1, 1, 0) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.contiguous() + 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.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + + hidden_states = hidden_states + residual + return hidden_states + + +class HunyuanVideoMidBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_attention: bool = True, + attention_head_dim: int = 1, + ) -> None: + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + + # There is always at least one resnet + resnets = [ + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ] + attentions = [] + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + HunyuanVAEAttention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + eps=resnet_eps, + norm_num_groups=resnet_groups, + bias=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(self.resnets[0], hidden_states) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) + attention_mask = prepare_causal_attention_mask( + num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size + ) + hidden_states = attn(hidden_states, attention_mask=attention_mask) + hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) + + else: + hidden_states = self.resnets[0](hidden_states) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) + attention_mask = prepare_causal_attention_mask( + num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size + ) + hidden_states = attn(hidden_states, attention_mask=attention_mask) + hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + + hidden_states = resnet(hidden_states) + + return hidden_states + + +class HunyuanVideoDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_downsample: bool = True, + downsample_stride: int = 2, + downsample_padding: int = 1, + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=out_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + HunyuanVideoDownsampleCausal3D( + out_channels, + out_channels=out_channels, + padding=downsample_padding, + stride=downsample_stride, + ) + ] + ) + else: + self.downsamplers = 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.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class HunyuanVideoUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_upsample: bool = True, + upsample_scale_factor: Tuple[int, int, int] = (2, 2, 2), + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=input_channels, + out_channels=out_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + HunyuanVideoUpsampleCausal3D( + out_channels, + out_channels=out_channels, + upsample_factor=upsample_scale_factor, + ) + ] + ) + 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 HunyuanVideoEncoder3D(nn.Module): + r""" + Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ), + block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + temporal_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ) -> None: + super().__init__() + + self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + if down_block_type != "HunyuanVideoDownBlock3D": + raise ValueError(f"Unsupported down_block_type: {down_block_type}") + + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int(np.log2(temporal_compression_ratio)) + + if temporal_compression_ratio == 4: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool( + i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block + ) + elif temporal_compression_ratio == 8: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i < num_time_downsample_layers) + else: + raise ValueError(f"Unsupported time_compression_ratio: {temporal_compression_ratio}") + + downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) + downsample_stride_T = (2,) if add_time_downsample else (1,) + downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) + + down_block = HunyuanVideoDownBlock3D( + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=bool(add_spatial_downsample or add_time_downsample), + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + downsample_stride=downsample_stride, + downsample_padding=0, + ) + + self.down_blocks.append(down_block) + + self.mid_block = HunyuanVideoMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + add_attention=mid_block_add_attention, + ) + + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], conv_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) + + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class HunyuanVideoDecoder3D(nn.Module): + r""" + Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ), + block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1) + self.up_blocks = nn.ModuleList([]) + + # mid + self.mid_block = HunyuanVideoMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + add_attention=mid_block_add_attention, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + if up_block_type != "HunyuanVideoUpBlock3D": + raise ValueError(f"Unsupported up_block_type: {up_block_type}") + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_upsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_upsample = bool(i < num_spatial_upsample_layers) + add_time_upsample = bool( + i >= len(block_out_channels) - 1 - num_time_upsample_layers and not is_final_block + ) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") + + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) + upsample_scale_factor_T = (2,) if add_time_upsample else (1,) + upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW) + + up_block = HunyuanVideoUpBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + add_upsample=bool(add_spatial_upsample or add_time_upsample), + upsample_scale_factor=upsample_scale_factor, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + ) + + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], 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: + 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.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): + r""" + A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. + Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). + + 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 + @auto_attributes + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 16, + down_block_types: Tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ), + up_block_types: Tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ), + block_out_channels: Tuple[int] = (128, 256, 512, 512), + layers_per_block: int = 2, + act_fn: str = "silu", + norm_num_groups: int = 32, + scaling_factor: float = 0.476986, + spatial_compression_ratio: int = 8, + temporal_compression_ratio: int = 4, + mid_block_add_attention: bool = True, + load_encoder: bool = True, + load_decoder: bool = True, + ) -> None: + super().__init__() + + # TODO(will): only pass in config. We do this by manually defining a + # config for hunyuan vae + self.block_out_channels = block_out_channels + + if load_encoder: + self.encoder = HunyuanVideoEncoder3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + double_z=True, + mid_block_add_attention=mid_block_add_attention, + temporal_compression_ratio=temporal_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + ) + self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) + + if load_decoder: + self.decoder = HunyuanVideoDecoder3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + time_compression_ratio=temporal_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) + + + + # 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 = True + + # 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 = 16 + + # The minimal distance between two spatial tiles + self.tile_sample_stride_height = 192 + self.tile_sample_stride_width = 192 + self.tile_sample_stride_num_frames = 12 + + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + x = self.encoder(x) + enc = self.quant_conv(x) + return enc + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + 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. + """ + 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/__init__.py b/fastvideo/pipelines/__init__.py index 18b53bbc4..a1919ea27 100644 --- a/fastvideo/pipelines/__init__.py +++ b/fastvideo/pipelines/__init__.py @@ -3,12 +3,25 @@ Diffusion pipelines for FastVideo. This package contains diffusion pipelines for generating videos and images. """ +import os +import json +from copy import deepcopy from typing import Dict, Optional, Type, Any -# First, import the registry -from fastvideo.pipelines.pipeline_registry import PipelineRegistry, register_pipeline +from fastvideo.pipelines.pipeline_registry import PipelineRegistry from fastvideo.inference_args import InferenceArgs +from fastvideo.logger import init_logger +from huggingface_hub import snapshot_download +from transformers import PretrainedConfig +from fastvideo.models.hf_transformer_utils import get_hf_config, get_diffusers_config +from fastvideo.models import get_scheduler +from fastvideo.loader.loader import get_model_loader +import glob +from fastvideo.loader.fsdp_load import load_fsdp_model +from fastvideo.models.loader import PipelineComponentLoader + +logger = init_logger(__name__) # Then import the base classes from fastvideo.pipelines.composed.composed_pipeline_base import ( @@ -20,30 +33,199 @@ def get_pipeline_type(inference_args: InferenceArgs) -> str: # hardcode for now return "hunyuan_video" +def maybe_download_model(model_path: str) -> str: + """ + Check if the model path is a Hugging Face Hub model ID and download it if needed. + + Args: + model_path: Local path or Hugging Face Hub model ID + + Returns: + Local path to the model + """ + + # If the path exists locally, return it + if os.path.exists(model_path): + logger.info(f"Model already exists locally at {model_path}") + return model_path + + # Otherwise, assume it's a HF Hub model ID and try to download it + try: + logger.info(f"Downloading model snapshot from HF Hub for {model_path}...") + local_path = snapshot_download( + repo_id=model_path, + # allow_patterns=["*.json", "*.bin", "*.safetensors", "*.pt", "*.pth", "*.ckpt"], + ignore_patterns=["*.onnx", "*.msgpack"], + ) + logger.info(f"Downloaded model to {local_path}") + return local_path + except Exception as e: + raise ValueError(f"Could not find model at {model_path} and failed to download from HF Hub: {e}") + +def verify_model_config_and_directory(model_path: str) -> dict: + """ + Verify that the model directory contains a valid diffusers configuration. + + Args: + model_path: Path to the model directory + + Returns: + The loaded model configuration as a dictionary + """ + + # Check for model_index.json which is required for diffusers models + config_path = os.path.join(model_path, "model_index.json") + if not os.path.exists(config_path): + raise ValueError( + f"Model directory {model_path} does not contain model_index.json. " + "Only Hugging Face diffusers format is supported." + ) + + # Check for transformer and vae directories + transformer_dir = os.path.join(model_path, "transformer") + vae_dir = os.path.join(model_path, "vae") + + if not os.path.exists(transformer_dir): + raise ValueError(f"Model directory {model_path} does not contain a transformer/ directory.") + + if not os.path.exists(vae_dir): + raise ValueError(f"Model directory {model_path} does not contain a vae/ directory.") + + # Load the config + try: + with open(config_path, "r") as f: + config = json.load(f) + except Exception as e: + raise ValueError(f"Failed to load model configuration from {config_path}: {e}") + + # Verify diffusers version exists + if "_diffusers_version" not in config: + raise ValueError(f"model_index.json does not contain _diffusers_version") + + logger.info(f"Diffusers version: {config['_diffusers_version']}") + return config + +def load_pipeline_module(module_name: str, component_model_path: str, transformers_or_diffusers: str, architecture: str, inference_args: InferenceArgs) -> Any: + """ + Load a pipeline module using the appropriate loader. + + Args: + module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler") + component_model_path: Path to the component model + transformers_or_diffusers: Whether the module is from transformers or diffusers + architecture: Architecture of the component model + inference_args: Inference arguments + + Returns: + The loaded module + """ + return PipelineComponentLoader.load_module( + module_name=module_name, + component_model_path=component_model_path, + transformers_or_diffusers=transformers_or_diffusers, + architecture=architecture, + inference_args=inference_args + ) + +def load_pipeline_modules(model_path: str, config: Dict, inference_args: InferenceArgs) -> dict[str, Any]: + """ + Load the pipeline modules from the config. + + Args: + config: The model_index.json config + inference_args: Inference arguments + + Returns: + Dictionary mapping module names to loaded modules + """ + logger.info(f"Loading pipeline modules from config: {config}") + modules_config = deepcopy(config) + + # remove keys that are not pipeline modules + modules_config.pop("_class_name") + modules_config.pop("_diffusers_version") + + # some sanity checks + assert len(modules_config) > 1, "model_index.json must contain at least one pipeline module" + + required_modules = ["vae", "text_encoder", "transformer", "scheduler", "tokenizer"] + for module_name in required_modules: + if module_name not in modules_config: + raise ValueError(f"model_index.json must contain a {module_name} module") + logger.info(f"Diffusers config passed sanity checks") + + + # all the component models used by the pipeline + pipeline_modules = {} + for module_name, (transformers_or_diffusers, architecture) in modules_config.items(): + component_model_path = os.path.join(model_path, module_name) + module = load_pipeline_module( + module_name, + component_model_path, + transformers_or_diffusers, + architecture, + inference_args, + ) + + pipeline_modules[module_name] = module + + + # Check if all required modules were loaded + for module_name in required_modules: + if module_name not in pipeline_modules or pipeline_modules[module_name] is None: + logger.warning(f"Required module {module_name} was not loaded properly") + + return pipeline_modules + def build_pipeline(inference_args: InferenceArgs) -> ComposedPipelineBase: """ Only works with valid hf diffusers configs. (model_index.json) - We want to build a pipeline based on the inference args. - 1. load the correct hf config from disk or download from hub - 2. based on the config, determine the pipeline class and pipeline loader class - 3. parse the config to get the model component (vae, text_encoders, etc...) - names and paths - 4. the pipeline loader class will use the model component names and paths to load - the models. - 5. the pipeline class will be composed of the models returned by the pipeline loader + We want to build a pipeline based on the inference args mode_path: + 1. download the model from the hub if it's not already downloaded + 2. verify the model config and directory + 3. based on the config, determine the pipeline class + 4. parse the config to get the model components (vae, text_encoders, etc...) + 5. the pipeline loader class will use the model component names and paths to load + 6. the pipeline class will be composed of the models returned by the pipeline loader """ # Get pipeline type - pipeline_type = get_pipeline_type(inference_args) + model_path = inference_args.model_path + model_path = maybe_download_model(model_path) + # inference_args.downloaded_model_path = model_path + logger.info(f"Model path: {model_path}") + config = verify_model_config_and_directory(model_path) + + pipeline_architecture = config.get("_class_name") + if pipeline_architecture is None: + raise ValueError("Model config does not contain a _class_name attribute. " + "Only diffusers format is supported.") - if not hasattr(inference_args, 'model_type') or not inference_args.model_type: - inference_args.model_type = pipeline_type + pipeline_cls, pipeline_architecture = PipelineRegistry.resolve_pipeline_cls(pipeline_architecture) + + # instantiate the pipeline + pipeline = pipeline_cls() - # TODO(will): Get pipeline loader (hardcoded for now) - from fastvideo.pipelines.loader import get_pipeline_loader - pipeline_loader = get_pipeline_loader(inference_args) + try: + pipeline_modules = load_pipeline_modules(model_path, config, inference_args) + except Exception as e: + logger.error(f"Failed to load pipeline modules: {e}") + raise ValueError(f"Failed to load pipeline modules: {e}") + + logger.info(f"Initializing encoders") + pipeline.initialize_encoders(pipeline_modules, inference_args) + + logger.info(f"Registering modules") + pipeline.register_modules(pipeline_modules) - # Load and return pipeline - return pipeline_loader.load_pipeline_v2(inference_args) + logger.info(f"Setting up pipeline") + pipeline.setup_pipeline(inference_args) + + logger.info(f"Initializing pipeline") + pipeline.initialize_pipeline(inference_args) + + # pipeline is now initialized and ready to use + return pipeline + def list_available_pipelines() -> Dict[str, Type[Any]]: @@ -55,16 +237,9 @@ def list_available_pipelines() -> Dict[str, Type[Any]]: """ return PipelineRegistry.list() - -# Import all pipeline implementations to register them -# These imports should be at the end to avoid circular imports -from fastvideo.pipelines.implementations.hunyuan import HunyuanVideoPipeline - __all__ = [ - "create_pipeline", + "build_pipeline", "list_available_pipelines", "ComposedPipelineBase", "DiffusionPipelineOutput", - "register_pipeline", - "HunyuanVideoPipeline", ] \ No newline at end of file diff --git a/fastvideo/pipelines/composed/composed_pipeline_base.py b/fastvideo/pipelines/composed/composed_pipeline_base.py index fd04d81e0..284eddd0c 100644 --- a/fastvideo/pipelines/composed/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed/composed_pipeline_base.py @@ -63,6 +63,14 @@ class ComposedPipelineBase(ABC): """ ... + @abstractmethod + def initialize_encoders(self, modules: Dict[str, Any], inference_args: InferenceArgs): + """ + Initialize the encoders. Will remove the encoders/tokenizers modules from the + modules. Will add the TextEncoder or ImageEncoder to the modules. + """ + ... + def register_modules(self, modules: Dict[str, Any]): """ Register modules with the pipeline and its stages. diff --git a/fastvideo/pipelines/composed/text_to_video.py b/fastvideo/pipelines/composed/text_to_video.py deleted file mode 100644 index fbc18bb4f..000000000 --- a/fastvideo/pipelines/composed/text_to_video.py +++ /dev/null @@ -1,113 +0,0 @@ -""" -Text-to-video pipeline implementation. - -This module contains a standard text-to-video pipeline implementation that can be -used as a base for more specific implementations. -""" - -from typing import Any, Dict, List, Optional, Union, Tuple -import torch - -from fastvideo.pipelines.composed.composed_pipeline_base import ComposedPipelineBase -from fastvideo.pipelines.stages import ( - InputValidationStage, - PromptEncodingStage, - TimestepPreparationStage, - LatentPreparationStage, - ConditioningStage, - DenoisingStage, - DecodingStage, - PostProcessingStage, -) -# Import from the main pipelines package to avoid circular imports -from fastvideo.pipelines import register_pipeline - - -@register_pipeline("text-to-video") -class TextToVideoPipeline(ComposedPipelineBase): - """ - Standard text-to-video pipeline. - - This pipeline implements a standard text-to-video diffusion process with the - following stages: - 1. Input validation - 2. Prompt encoding - 3. Timestep preparation - 4. Latent preparation - 5. Conditioning (classifier-free guidance) - 6. Denoising - 7. Decoding - 8. Post-processing - - This pipeline supports models with either a single text encoder or dual text encoders. - Different model architectures may use secondary encoders for various purposes such as - additional conditioning information, temporal control, stylistic aspects, or other - model-specific features. - """ - - is_video_pipeline = True - - def __init__( - self, - input_validation_stage: InputValidationStage, - prompt_encoding_stage: PromptEncodingStage, - timestep_preparation_stage: TimestepPreparationStage, - latent_preparation_stage: LatentPreparationStage, - conditioning_stage: ConditioningStage, - denoising_stage: DenoisingStage, - decoding_stage: DecodingStage, - post_processing_stage: Optional[PostProcessingStage] = None, - secondary_prompt_encoding_stage: Optional[PromptEncodingStage] = None, - ): - """ - Initialize the text-to-video pipeline. - - Args: - input_validation_stage: The input validation stage. - prompt_encoding_stage: The primary prompt encoding stage. - timestep_preparation_stage: The timestep preparation stage. - latent_preparation_stage: The latent preparation stage. - conditioning_stage: The conditioning stage. - denoising_stage: The denoising stage. - decoding_stage: The decoding stage. - post_processing_stage: The post-processing stage (optional). - secondary_prompt_encoding_stage: Optional secondary prompt encoding stage for models - with dual text encoders. The purpose of this encoder varies by model architecture - and may handle additional conditioning aspects such as temporal control, style, - or other model-specific features. - """ - super().__init__() - - # Add the stages in order - self.add_stage(input_validation_stage) - self.add_stage(prompt_encoding_stage) - - # Add optional secondary prompt encoding stage if provided - if secondary_prompt_encoding_stage is not None: - self.add_stage(secondary_prompt_encoding_stage) - - self.add_stage(timestep_preparation_stage) - self.add_stage(latent_preparation_stage) - self.add_stage(conditioning_stage) - self.add_stage(denoising_stage) - self.add_stage(decoding_stage) - - if post_processing_stage is not None: - self.add_stage(post_processing_stage) - - @property - def components(self) -> Dict[str, Any]: - """Get all components used by this pipeline.""" - components = {} - - # Add components from each stage - for stage in self._stages: - for attr_name in dir(stage): - if attr_name.startswith("__") or attr_name.startswith("_"): - continue - - attr = getattr(stage, attr_name) - if isinstance(attr, torch.nn.Module): - components[attr_name] = attr - - return components \ No newline at end of file diff --git a/fastvideo/pipelines/hunyuan/hunyuan_loader.py b/fastvideo/pipelines/hunyuan/hunyuan_loader.py index 6ee639f3f..b2d3d76e6 100644 --- a/fastvideo/pipelines/hunyuan/hunyuan_loader.py +++ b/fastvideo/pipelines/hunyuan/hunyuan_loader.py @@ -221,7 +221,6 @@ class HunyuanPipelineLoader(PipelineLoader): hidden_state_skip_layer=inference_args.hidden_state_skip_layer, apply_final_norm=inference_args.apply_final_norm, reproduce=inference_args.reproduce, - logger=logger, device=self.device if not inference_args.use_cpu_offload else "cpu", ) text_encoder_2 = None @@ -232,7 +231,6 @@ class HunyuanPipelineLoader(PipelineLoader): text_encoder_precision=inference_args.text_encoder_precision_2, tokenizer_type=inference_args.tokenizer_2, reproduce=inference_args.reproduce, - logger=logger, device=self.device if not inference_args.use_cpu_offload else "cpu", ) diff --git a/fastvideo/pipelines/implementations/__init__.py b/fastvideo/pipelines/implementations/__init__.py deleted file mode 100644 index 616c6d670..000000000 --- a/fastvideo/pipelines/implementations/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -""" -Concrete implementations of diffusion pipelines. - -This package contains concrete implementations of diffusion pipelines for -specific models or architectures. -""" - -__all__ = [] \ No newline at end of file diff --git a/fastvideo/pipelines/implementations/hunyuan/__init__.py b/fastvideo/pipelines/implementations/hunyuan/__init__.py index 976e6fe6c..9e90d55d8 100644 --- a/fastvideo/pipelines/implementations/hunyuan/__init__.py +++ b/fastvideo/pipelines/implementations/hunyuan/__init__.py @@ -5,9 +5,9 @@ This package contains implementations of diffusion pipelines for HunYuan models. """ from fastvideo.pipelines.implementations.hunyuan.hunyuan_pipeline import ( - HunyuanVideoPipeline, + EntryClass, ) __all__ = [ - "HunyuanVideoPipeline", + "EntryClass", ] \ No newline at end of file diff --git a/fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py b/fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py index a2b71fec1..e606f1d01 100644 --- a/fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py +++ b/fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py @@ -8,14 +8,14 @@ from fastvideo.inference_args import InferenceArgs from fastvideo.pipelines.implementations.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE from fastvideo.logger import init_logger from typing import Tuple, Optional - +from fastvideo.models.dits import HunyuanVideoTransformer3DModel import os import glob import json from fastvideo.loader.fsdp_load import load_fsdp_model - from fastvideo.models.hunyuan.text_encoder import TextEncoder - +from fastvideo.models.vaes import AutoencoderKLHunyuanVideo +from safetensors.torch import load_file as safetensors_load_file from fastvideo.platforms import current_platform logger = init_logger(__name__) @@ -113,7 +113,7 @@ class HunyuanPipelineLoader(PipelineLoader): """Custom transformer loading for Hunyuan""" # Path to model files - # TODO(PY): remove this hardcode + # TODO(PY/Will): remove this hardcode path = "data/hunyuanvideo_community/transformer" path = Path(path) @@ -130,7 +130,7 @@ class HunyuanPipelineLoader(PipelineLoader): config.pop("_class_name") if "_diffusers_version" in config: config.pop("_diffusers_version") - + # Find all safetensors files safetensors_list = glob.glob(os.path.join(str(path), "*.safetensors")) if not safetensors_list: @@ -139,8 +139,9 @@ class HunyuanPipelineLoader(PipelineLoader): logger.info(f"Loading model from {len(safetensors_list)} safetensors files in {path}") # Load the model using FSDP loader + # TOOD(Will) new registry for DiT model = load_fsdp_model( - model_name="HunyuanVideoTransformer3DModel", + model_cls=HunyuanVideoTransformer3DModel, init_params=config, weight_dir_list=safetensors_list, device=self.device, @@ -181,6 +182,49 @@ class HunyuanPipelineLoader(PipelineLoader): return model def load_vae(self, inference_args: InferenceArgs): + """Custom VAE loading for Hunyuan""" + # TODO(will): replace this with abstracted model + use_v1_loader = True + if use_v1_loader: + return self.load_vae_v1(inference_args) + else: + return self.load_vae_v0(inference_args) + + def load_vae_v1(self, inference_args: InferenceArgs): + """Custom VAE loading for Hunyuan""" + path = "data/hunyuanvideo_community/vae" + path = Path(path) + # TODO(PY/Will): remove this hardcode + # Load config file + config_path = path / "config.json" + if not config_path.exists(): + raise ValueError(f"Config file not found at {config_path}") + + with open(config_path, "r") as f: + config = json.load(f) + + # Clean up config + + class_name = config.pop("_class_name") + if "_diffusers_version" in config: + config.pop("_diffusers_version") + vae = AutoencoderKLHunyuanVideo(**config).to(self.device) + + # Find all safetensors files + safetensors_list = glob.glob(os.path.join(str(path), "*.safetensors")) + # TODO(PY) + assert len(safetensors_list) == 1, f"Found {len(safetensors_list)} safetensors files in {path}" + loaded = safetensors_load_file(safetensors_list[0]) + vae.load_state_dict(loaded) + dtype = PRECISION_TO_TYPE[inference_args.vae_precision] + vae = vae.eval().to(dtype) + + return vae, {"s_ratio": config["spatial_compression_ratio"], "t_ratio": config["temporal_compression_ratio"]} + + + + + def load_vae_v0(self, inference_args: InferenceArgs): """Custom VAE loading for Hunyuan""" # TODO(will): replace this with abstracted model from fastvideo.models.hunyuan.vae import load_vae @@ -188,7 +232,7 @@ class HunyuanPipelineLoader(PipelineLoader): inference_args.vae, inference_args.vae_precision, logger=logger, - device=self.device if not inference_args.use_cpu_offload else "cpu", + device=self.device ) vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio} return vae, vae_kwargs diff --git a/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py b/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py index d634bea20..77a643fdc 100644 --- a/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py +++ b/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py @@ -5,7 +5,7 @@ This module contains an implementation of the HunYuan video diffusion pipeline using the modular pipeline architecture. """ -from typing import Any, Dict, List, Optional, Union, Tuple +from typing import Union, Any, Dict import torch from diffusers.image_processor import VaeImageProcessor @@ -20,14 +20,16 @@ from fastvideo.pipelines.stages import ( DecodingStage, PostProcessingStage, ) -from fastvideo.pipelines import register_pipeline from fastvideo.pipelines.stages.prompt_encoding import PromptEncodingStage # from fastvideo.pipelines.stages.timestep_preparation import FlowMatchingTimestepPreparationStage from fastvideo.inference_args import InferenceArgs # from fastvideo.pipelines.composed.composed_pipeline_base import DiffusionPipelineOutput from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.models.text_encoder import TextEncoder +# TODO(will): move PRECISION_TO_TYPE to better place +from .constants import PROMPT_TEMPLATE, PRECISION_TO_TYPE + from diffusers.utils import BaseOutput -from typing import Union, Optional, Tuple, List import numpy as np from dataclasses import dataclass @@ -49,11 +51,75 @@ class DiffusionPipelineOutput(BaseOutput): videos: Union[torch.Tensor, np.ndarray] -@register_pipeline("hunyuan-video") class HunyuanVideoPipeline(ComposedPipelineBase): + def initialize_encoders(self, modules: Dict[str, Any], inference_args: InferenceArgs): + """ + Initialize the encoders. Will remove the encoders/tokenizers modules from the + modules. Will add the TextEncoder or ImageEncoder to the modules. + """ + if inference_args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0) + elif inference_args.prompt_template is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template].get("crop_start", 0) + else: + crop_start = 0 + + max_length = inference_args.text_len + crop_start + + # prompt_template + prompt_template = (PROMPT_TEMPLATE[inference_args.prompt_template] if inference_args.prompt_template is not None else None) + + # prompt_template_video + prompt_template_video = (PROMPT_TEMPLATE[inference_args.prompt_template_video] + if inference_args.prompt_template_video is not None else None) + + encoder_1 = modules.pop("text_encoder") + assert encoder_1 is not None, "Text encoder is not found" + encoder_1.to(inference_args.device) + encoder_1.to(dtype=PRECISION_TO_TYPE[inference_args.text_encoder_precision]) + encoder_1.requires_grad_(False) + + tokenizer_1 = modules.pop("tokenizer") + assert tokenizer_1 is not None, "Tokenizer is not found" + + text_encoder = TextEncoder( + text_encoder=encoder_1, + tokenizer=tokenizer_1, + # text_encoder_type="text_encoder" + max_length=max_length, + # text_encoder_precision=inference_args.text_encoder_precision, + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=inference_args.hidden_state_skip_layer, + apply_final_norm=inference_args.apply_final_norm, + reproduce=inference_args.reproduce, + device=inference_args.device if not inference_args.use_cpu_offload else "cpu", + ) + + encoder_2 = modules.pop("text_encoder_2") + assert encoder_2 is not None, "Text encoder 2 is not found" + encoder_2.to(inference_args.device) + encoder_2.to(dtype=PRECISION_TO_TYPE[inference_args.text_encoder_precision]) + encoder_2.requires_grad_(False) + + tokenizer_2 = modules.pop("tokenizer_2") + assert tokenizer_2 is not None, "Tokenizer 2 is not found" + + text_encoder_2 = TextEncoder( + text_encoder=encoder_2, + tokenizer=tokenizer_2, + # text_encoder_type="text_encoder_2", + max_length=inference_args.text_len_2, + # text_encoder_precision=inference_args.text_encoder_precision, + reproduce=inference_args.reproduce, + device=inference_args.device if not inference_args.use_cpu_offload else "cpu", + ) + modules["text_encoder"] = text_encoder + modules["text_encoder_2"] = text_encoder_2 + + def setup_pipeline(self, inference_args: InferenceArgs): - self._stages = [] self.add_stage("input_validation_stage", InputValidationStage()) self.add_stage("prompt_encoding_stage_primary", @@ -76,7 +142,7 @@ class HunyuanVideoPipeline(ComposedPipelineBase): assert len(self._modules) > 0, "Pipeline modules are not set" - vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1) + vae_scale_factor = 2**(len(self.vae.block_out_channels) - 1) inference_args.vae_scale_factor = vae_scale_factor self.image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) @@ -123,3 +189,5 @@ class HunyuanVideoPipeline(ComposedPipelineBase): batch = self.decoding_stage(batch, inference_args) return DiffusionPipelineOutput(videos=batch.videos) + +EntryClass = HunyuanVideoPipeline \ No newline at end of file diff --git a/fastvideo/pipelines/loader.py b/fastvideo/pipelines/loader.py index 146b3232e..a65a5ad2a 100644 --- a/fastvideo/pipelines/loader.py +++ b/fastvideo/pipelines/loader.py @@ -246,7 +246,15 @@ class PipelineLoader(ABC): solver=inference_args.flow_solver, ) return None - + + @staticmethod + def load_pipeline_from_config(config: dict): + """Load a pipeline from a config.""" + pipeline_cls = resolve_pipeline_cls(inference_args) + assert issubclass(pipeline_cls, ComposedPipelineBase) + logger.info(f"Pipeline class: {pipeline_cls}") + + def get_pipeline_loader(inference_args: InferenceArgs) -> PipelineLoader: """Get a pipeline loader based on the inference args.""" diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index 218d77a4b..9f61a37b0 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -1,69 +1,91 @@ -""" -Registry for pipeline implementations. +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/models/registry.py +# and https://github.com/sgl-project/sglang/blob/v0.4.3/python/sglang/srt/models/registry.py -This module provides a registry for pipeline implementations, allowing them to be -looked up by name. -""" +import importlib +import pkgutil +from dataclasses import dataclass, field +from functools import lru_cache +from typing import AbstractSet, Dict, List, Optional, Tuple, Type, Union -from typing import Dict, Type, Any, Optional, Callable, TypeVar, cast +from fastvideo.logger import init_logger -# Define a type variable for the pipeline class -T = TypeVar('T') +import torch.nn as nn -class PipelineRegistry: - """Registry for pipeline implementations.""" - - _registry: Dict[str, Type[Any]] = {} - - @classmethod - def register(cls, name: str) -> Callable[[Type[T]], Type[T]]: - """ - Register a pipeline implementation. - - Args: - name: The name to register the pipeline under. - - Returns: - A decorator that registers the pipeline. - """ - def decorator(pipeline_cls: Type[T]) -> Type[T]: - cls._registry[name] = pipeline_cls - return pipeline_cls - - return decorator - - @classmethod - def get(cls, name: str) -> Optional[Type[Any]]: - """ - Get a pipeline implementation by name. - - Args: - name: The name of the pipeline. - - Returns: - The pipeline class, or None if not found. - """ - return cls._registry.get(name) - - @classmethod - def list(cls) -> Dict[str, Type[Any]]: - """ - List all registered pipelines. - - Returns: - A dictionary of pipeline names to pipeline classes. - """ - return cls._registry.copy() +logger = init_logger(__name__) -def register_pipeline(name: str) -> Callable[[Type[T]], Type[T]]: - """ - Register a pipeline implementation. - - Args: - name: The name to register the pipeline under. - - Returns: - A decorator that registers the pipeline. - """ - return PipelineRegistry.register(name) \ No newline at end of file +@dataclass +class _PipelineRegistry: + # Keyed by pipeline_arch + pipelines: Dict[str, Union[Type[nn.Module], str]] = field(default_factory=dict) + + def get_supported_archs(self) -> AbstractSet[str]: + return self.pipelines.keys() + + def _raise_for_unsupported(self, architectures: List[str]): + all_supported_archs = self.get_supported_archs() + + if any(arch in all_supported_archs for arch in architectures): + raise ValueError( + f"Pipeline architectures {architectures} failed " + "to be inspected. Please check the logs for more details." + ) + + raise ValueError( + f"Pipeline architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}" + ) + + def _try_load_pipeline_cls(self, pipeline_arch: str) -> Optional[Type[nn.Module]]: + if pipeline_arch not in self.pipelines: + return None + + return self.pipelines[pipeline_arch] + + def resolve_pipeline_cls( + self, + architecture: str, + ) -> Tuple[Type[nn.Module], str]: + if not architecture: + logger.warning("No pipeline architecture is specified") + + pipeline_cls = self._try_load_pipeline_cls(architecture) + if pipeline_cls is not None: + return (pipeline_cls, architecture) + + return self._raise_for_unsupported(architecture) + + +@lru_cache() +def import_pipeline_classes(): + pipeline_arch_name_to_cls = {} + package_name = "fastvideo.pipelines.implementations" + package = importlib.import_module(package_name) + for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."): + if ispkg: + try: + module = importlib.import_module(name) + except Exception as e: + logger.warning(f"Ignore import error when loading {name}. " f"{e}") + continue + if hasattr(module, "EntryClass"): + entry = module.EntryClass + print(entry) + print(entry.__name__) + if isinstance( + entry, list + ): # To support multiple pipeline classes in one module + for tmp in entry: + assert ( + tmp.__name__ not in pipeline_arch_name_to_cls + ), f"Duplicated pipeline implementation for {tmp.__name__}" + pipeline_arch_name_to_cls[tmp.__name__] = tmp + else: + assert ( + entry.__name__ not in pipeline_arch_name_to_cls + ), f"Duplicated pipeline implementation for {entry.__name__}" + pipeline_arch_name_to_cls[entry.__name__] = entry + return pipeline_arch_name_to_cls + + +PipelineRegistry = _PipelineRegistry(import_pipeline_classes()) \ No newline at end of file diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 61775ef11..6863cc327 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -3,17 +3,11 @@ Decoding stage for diffusion pipelines. """ import torch -from typing import Optional, Union, List, Tuple -import numpy as np - from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.inference_args import InferenceArgs from fastvideo.models.hunyuan.constants import PRECISION_TO_TYPE -from fastvideo.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D from fastvideo.logger import init_logger -from diffusers.utils import BaseOutput - logger = init_logger(__name__) @@ -50,18 +44,6 @@ class DecodingStage(PipelineStage): vae_dtype = PRECISION_TO_TYPE[inference_args.vae_precision] vae_autocast_enabled = (vae_dtype != torch.float32) and not inference_args.disable_autocast - # Handle different latent shapes - expand_temporal_dim = False - if len(latents.shape) == 4: - if isinstance(self.vae, AutoencoderKLCausal3D): - latents = latents.unsqueeze(2) - expand_temporal_dim = True - elif len(latents.shape) == 5: - pass - else: - raise ValueError( - f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}." - ) # Apply scaling/shifting if needed if (hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor): @@ -73,11 +55,9 @@ class DecodingStage(PipelineStage): with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): if inference_args.vae_tiling: self.vae.enable_tiling() - image = self.vae.decode(latents, return_dict=False, generator=batch.generator)[0] + image = self.vae.decode(latents) + - # Handle temporal dimension if needed - if expand_temporal_dim or image.shape[2] == 1: - image = image.squeeze(2) # Normalize image to [0, 1] range image = (image / 2 + 0.5).clamp(0, 1) diff --git a/fastvideo/pipelines/stages/prompt_encoding.py b/fastvideo/pipelines/stages/prompt_encoding.py index 4785ffe8f..f3205fbc9 100644 --- a/fastvideo/pipelines/stages/prompt_encoding.py +++ b/fastvideo/pipelines/stages/prompt_encoding.py @@ -10,7 +10,7 @@ import torch from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.inference_args import InferenceArgs from fastvideo.pipelines.stages import PipelineStage -from fastvideo.models.hunyuan.text_encoder import TextEncoder +from fastvideo.models.text_encoder import TextEncoder from fastvideo.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/sample/v1_fastvideo_inference.py b/fastvideo/sample/v1_fastvideo_inference.py index 7eec79cc2..a86105003 100644 --- a/fastvideo/sample/v1_fastvideo_inference.py +++ b/fastvideo/sample/v1_fastvideo_inference.py @@ -14,38 +14,31 @@ from fastvideo.inference_args import InferenceArgs from fastvideo.utils.utils import FlexibleArgumentParser from fastvideo.distributed import init_distributed_environment, initialize_model_parallel from fastvideo.logger import init_logger -import torch.distributed as dist from fastvideo.utils.parallel_states import initialize_sequence_parallel_state logger = init_logger(__name__) -def initialize_distributed(): - local_rank = int(os.getenv("RANK", 0)) - world_size = int(os.getenv("WORLD_SIZE", 1)) - print("world_size", world_size) - torch.cuda.set_device(local_rank) - dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank) - initialize_sequence_parallel_state(world_size) - def main(inference_args: InferenceArgs): # initialize_distributed() # print(nccl_info.sp_size) - # local_rank = int(os.environ.get("LOCAL_RANK", 0)) - # rank = int(os.environ.get("RANK", 0)) - # world_size = int(os.environ.get("WORLD_SIZE", 1)) - + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + torch.cuda.set_device(local_rank) # logger.info(f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}") - # init_distributed_environment( - # world_size=world_size, - # rank=rank, - # local_rank=local_rank - # ) - # print(inference_args.sp_size) + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=local_rank + ) + print(inference_args.sp_size) # Initialize tensor model parallel groups - # initialize_model_parallel( - # sequence_model_parallel_size=inference_args.sp_size - # ) - initialize_distributed() + initialize_model_parallel( + sequence_model_parallel_size=inference_args.sp_size, + tensor_model_parallel_size=inference_args.sp_size + ) + # initialize_sequence_parallel_state(world_size) + # initialize_distributed() print('Creating engine') diff --git a/fastvideo/sample/v2_fastvideo_inference.py b/fastvideo/sample/v2_fastvideo_inference.py deleted file mode 100644 index fe47e09e9..000000000 --- a/fastvideo/sample/v2_fastvideo_inference.py +++ /dev/null @@ -1,102 +0,0 @@ -import argparse -import os -from pathlib import Path - -import imageio -import numpy as np -import torch -import torchvision -from einops import rearrange - -# Fix the import path -from fastvideo.inference_engine import InferenceEngine -from fastvideo.inference_args import InferenceArgs -from fastvideo.utils.utils import FlexibleArgumentParser -from fastvideo.distributed import init_distributed_environment, initialize_model_parallel -from fastvideo.logger import init_logger -import torch.distributed as dist -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state -logger = init_logger(__name__) - -def initialize_distributed(): - local_rank = int(os.getenv("RANK", 0)) - world_size = int(os.getenv("WORLD_SIZE", 1)) - print("world_size", world_size) - torch.cuda.set_device(local_rank) - dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank) - initialize_sequence_parallel_state(world_size) - -def main(inference_args: InferenceArgs): - # initialize_distributed() - # print(nccl_info.sp_size) - local_rank = int(os.environ.get("LOCAL_RANK", 0)) - rank = int(os.environ.get("RANK", 0)) - world_size = int(os.environ.get("WORLD_SIZE", 1)) - torch.cuda.set_device(local_rank) - # logger.info(f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}") - init_distributed_environment( - world_size=world_size, - rank=rank, - local_rank=local_rank - ) - print(inference_args.sp_size) - - # Initialize tensor model parallel groups - initialize_model_parallel( - sequence_model_parallel_size=inference_args.sp_size - ) - initialize_sequence_parallel_state(world_size) - # initialize_distributed() - - - print('Creating engine') - # Create inference object using the updated API - engine = InferenceEngine.create_engine( - inference_args, - ) - print('Engine created') - # return - - # Load prompts - if inference_args.prompt.endswith('.txt'): - with open(inference_args.prompt) as f: - prompts = [line.strip() for line in f.readlines()] - else: - prompts = [inference_args.prompt] - - # Process each prompt - for prompt in prompts: - outputs = engine.run( - prompt=prompt, - inference_args=inference_args, - ) - - # Process outputs - videos = rearrange(outputs["samples"], "b c t h w -> t b c h w") - frames = [] - for x in videos: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) - - # Save video - os.makedirs(os.path.dirname(inference_args.output_path), exist_ok=True) - imageio.mimsave( - os.path.join(inference_args.output_path, f"{prompt[:100]}.mp4"), - frames, - fps=inference_args.fps - ) - - -if __name__ == "__main__": - parser = FlexibleArgumentParser() - InferenceArgs.add_cli_args(parser) - args = parser.parse_args() - inference_args = InferenceArgs.from_cli_args(args) - inference_args.check_inference_args() - - # Validate arguments - if inference_args.vae_sp and not inference_args.vae_tiling: - raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.") - - main(inference_args) \ No newline at end of file diff --git a/fastvideo/v1/test_hunyuan_vae.py b/fastvideo/v1/test_hunyuan_vae.py new file mode 100644 index 000000000..596607138 --- /dev/null +++ b/fastvideo/v1/test_hunyuan_vae.py @@ -0,0 +1,161 @@ +from fastvideo.models.vaes import AutoencoderKLHunyuanVideo as MyHunyuanVAE +from diffusers import AutoencoderKLHunyuanVideo as DiffusersHunyuanVAE + +import os +import torch +import torch.nn as nn +import argparse +import numpy as np +import json +from fastvideo.logger import init_logger +from safetensors.torch import load_file +logger = init_logger(__name__) + +def initialize_identical_weights(model1, model2, seed=42): + """Initialize both models with identical weights using a fixed seed for reproducibility.""" + # Get all parameters from both models + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Initialize each layer with identical values + with torch.no_grad(): + # Initialize weights + for name1, param1 in params1.items(): + if 'weight' in name1: + # Set seed before each weight initialization + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + + for name2, param2 in params2.items(): + if 'weight' in name2: + # Reset seed to get same initialization + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + + # Initialize biases + for name1, param1 in params1.items(): + if 'bias' in name1: + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + param1.data = param1.data.to(torch.bfloat16) + + for name2, param2 in params2.items(): + if 'bias' in name2: + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + param2.data = param2.data.to(torch.bfloat16) + + logger.info("Both models initialized with identical weights in bfloat16") + return model1, model2 + +def setup_args(): + parser = argparse.ArgumentParser(description='HunyuanVAE Test') + parser.add_argument('--in-channels', type=int, default=4, + help='Number of input channels') + parser.add_argument('--out-channels', type=int, default=4, + help='Number of output channels') + parser.add_argument('--latent-channels', type=int, default=4, + help='Number of latent channels') + return parser.parse_args() + +def test_hunyuan_vae(): + args = setup_args() + + # Set fixed random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Model parameters + in_channels = args.in_channels + out_channels = args.out_channels + latent_channels = args.latent_channels + + device = torch.device("cuda:0") + print + # Initialize the two model implementations + path = "/mbz/users/hao.zhang/peiyuan/FastVideo/data/hunyuanvideo_community/vae" + config_path = os.path.join(path, "config.json") + config = json.load(open(config_path)) + config.pop("_class_name") + config.pop("_diffusers_version") + model1 = MyHunyuanVAE( + **config + ).to(torch.bfloat16) + + model2 = DiffusersHunyuanVAE( + **config + ).to(torch.bfloat16) + + loaded = load_file(os.path.join(path, "diffusion_pytorch_model.safetensors")) + model1.load_state_dict(loaded) + model2.load_state_dict(loaded) + + + # Set both models to eval mode + model1.eval() + model2.eval() + + # Move to GPU + model1 = model1.to(device) + model2 = model2.to(device) + + + model1.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_min_num_frames=8, + tile_sample_stride_height=16, + tile_sample_stride_width=16, + tile_sample_stride_num_frames=4 + ) + model2.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_min_num_frames=8, + tile_sample_stride_height=16, + tile_sample_stride_width=16, + tile_sample_stride_num_frames=4 + ) + + # Create identical inputs for both models + batch_size = 1 + + # Video input [B, C, T, H, W] + input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16) + + # Disable gradients for inference + with torch.no_grad(): + # Test encoding + logger.info("Testing encoding...") + latent1 = model1.encode(input_tensor).mean + print("--------------------------------") + latent2 = model2.encode(input_tensor).latent_dist.mean + # Check if latents have the same shape + assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" + assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" + # Check if latents are similar + max_diff_encode = torch.max(torch.abs(latent1 - latent2)) + mean_diff_encode = torch.mean(torch.abs(latent1 - latent2)) + logger.info(f"Maximum difference between encoded latents: {max_diff_encode.item()}") + logger.info(f"Mean difference between encoded latents: {mean_diff_encode.item()}") + assert max_diff_encode < 1e-4, f"Encoded latents differ significantly: max diff = {max_diff_encode.item()}" + # Test decoding + logger.info("Testing decoding...") + output1 = model1.decode(latent1) + output2 = model2.decode(latent2).sample + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + + # Check if outputs are similar + max_diff_decode = torch.max(torch.abs(output1 - output2)) + mean_diff_decode = torch.mean(torch.abs(output1 - output2)) + logger.info(f"Maximum difference between decoded outputs: {max_diff_decode.item()}") + logger.info(f"Mean difference between decoded outputs: {mean_diff_decode.item()}") + assert max_diff_decode < 1e-4, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}" + + + logger.info("Test passed! Both VAE implementations produce similar outputs.") + logger.info("Test completed successfully") + +if __name__ == "__main__": + test_hunyuan_vae() \ No newline at end of file diff --git a/fastvideo/v1/test_hunyuanvideo.py b/fastvideo/v1/test_hunyuanvideo.py index ed68b8d89..8921a6811 100644 --- a/fastvideo/v1/test_hunyuanvideo.py +++ b/fastvideo/v1/test_hunyuanvideo.py @@ -20,7 +20,7 @@ from fastvideo.utils.parallel_states import initialize_sequence_parallel_state from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed.device_mesh import init_device_mesh from fastvideo.loader.fsdp_load import shard_model -from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT +from fastvideo.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer from fastvideo.models.hunyuan_hf.modeling_hunyuan import HunyuanVideoTransformer3DModel logger = init_logger(__name__) @@ -116,7 +116,7 @@ def test_hunyuanvideo_distributed(): patch_size = [1, 2, 2] torch.cuda.set_device(f"cuda:{local_rank}") # Initialize the two model implementations - model1 = HunyuanVideoDiT( + model1 = HunyuanVideoDit( patch_size=2, patch_size_t=1, in_channels=4, diff --git a/fastvideo/v1/test_hunyuanvideo_load.py b/fastvideo/v1/test_hunyuanvideo_load.py index 6bf8075d9..b56f38859 100644 --- a/fastvideo/v1/test_hunyuanvideo_load.py +++ b/fastvideo/v1/test_hunyuanvideo_load.py @@ -15,7 +15,7 @@ from fastvideo.distributed.parallel_state import ( cleanup_dist_env_and_memory ) import json -from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT +from fastvideo.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit from fastvideo.models.hunyuan.modules.models import HUNYUAN_VIDEO_CONFIG from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer from fastvideo.loader.fsdp_load import load_fsdp_model diff --git a/scripts/inference/v1_inference_hunyuan.sh b/scripts/inference/v1_inference_hunyuan.sh index 1b229ed8c..5d3293c64 100644 --- a/scripts/inference/v1_inference_hunyuan.sh +++ b/scripts/inference/v1_inference_hunyuan.sh @@ -1,7 +1,8 @@ #!/bin/bash num_gpus=4 -export MODEL_BASE=data/FastHunyuan +export MODEL_BASE=FastVideo/FastHunyuan-diffusers +# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ fastvideo/sample/v1_fastvideo_inference.py \ --sp_size 4 \ @@ -15,7 +16,6 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ --flow-reverse \ --prompt ./assets/prompt.txt \ --seed 1024 \ - --output_path outputs_video/hunyuan/vae_sp_v1/ \ + --output_path outputs_video/hunyuan/vae_sp_v2/ \ --model_path $MODEL_BASE \ - --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ --vae-sp diff --git a/scripts/inference/v2_inference_hunyuan.sh b/scripts/inference/v2_inference_hunyuan.sh deleted file mode 100644 index a0eda35d8..000000000 --- a/scripts/inference/v2_inference_hunyuan.sh +++ /dev/null @@ -1,21 +0,0 @@ -#!/bin/bash - -num_gpus=4 -export MODEL_BASE=data/FastHunyuan -torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ - fastvideo/sample/v2_fastvideo_inference.py \ - --sp_size 4 \ - --height 720 \ - --width 1280 \ - --num_frames 125 \ - --num_inference_steps 6 \ - --guidance_scale 1 \ - --embedded_cfg_scale 6 \ - --flow_shift 17 \ - --flow-reverse \ - --prompt ./assets/prompt.txt \ - --seed 1024 \ - --output_path outputs_video/hunyuan/vae_sp_v2/ \ - --model_path $MODEL_BASE \ - --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ - --vae-sp From b631546e183d33a70e1e169f06c223a9b50e9a1b Mon Sep 17 00:00:00 2001 From: William Lin Date: Fri, 14 Mar 2025 12:06:05 -0700 Subject: [PATCH 02/12] move refactor to fastvideo/v1 (#265) --- fastvideo/models/flash_attn_no_pad.py | 28 + fastvideo/models/hunyuan/inference.py | 22 - fastvideo/models/hunyuan/modules/__init__.py | 42 +- fastvideo/models/hunyuan/modules/attenion.py | 14 +- fastvideo/models/hunyuan/modules/models.py | 14 +- .../models/hunyuan/modules/posemb_layers.py | 1 - .../models/hunyuan/modules/token_refiner.py | 5 +- .../hunyuan/vae/autoencoder_kl_causal_3d.py | 16 +- .../models/hunyuan_hf/modeling_hunyuan.py | 836 +++++++++++++ .../models/hunyuan_hf/pipeline_hunyuan.py | 691 +++++++++++ .../mochi_hf/convert_diffusers_to_mochi.py | 361 ++++++ .../models/mochi_hf/mochi_latents_utils.py | 45 + fastvideo/models/mochi_hf/modeling_mochi.py | 663 +++++++++++ fastvideo/models/mochi_hf/norm.py | 128 ++ fastvideo/models/mochi_hf/pipeline_mochi.py | 757 ++++++++++++ fastvideo/models/stepvideo/__init__.py | 7 + fastvideo/models/stepvideo/__version__.py | 1 + fastvideo/models/stepvideo/config.py | 174 +++ .../models/stepvideo/diffusion/scheduler.py | 220 ++++ .../stepvideo/diffusion/video_pipeline.py | 325 +++++ .../models/stepvideo/modules/attentions.py | 96 ++ fastvideo/models/stepvideo/modules/blocks.py | 296 +++++ fastvideo/models/stepvideo/modules/model.py | 198 ++++ .../models/stepvideo/modules/normalization.py | 312 +++++ fastvideo/models/stepvideo/modules/rope.py | 90 ++ fastvideo/models/stepvideo/parallel.py | 21 + .../models/stepvideo/text_encoder/__init__.py | 12 + .../models/stepvideo/text_encoder/clip.py | 36 + .../stepvideo/text_encoder/flashattention.py | 45 + .../models/stepvideo/text_encoder/stepllm.py | 291 +++++ .../stepvideo/text_encoder/tokenizer.py | 209 ++++ fastvideo/models/stepvideo/utils/__init__.py | 2 + .../models/stepvideo/utils/quantization.py | 117 ++ fastvideo/models/stepvideo/utils/utils.py | 60 + .../models/stepvideo/utils/video_process.py | 51 + fastvideo/models/stepvideo/vae/vae.py | 975 +++++++++++++++ fastvideo/pipelines/stages/__init__.py | 28 - fastvideo/sample/sample_t2v_hunyuan.py | 23 +- fastvideo/{ => v1}/attention/__init__.py | 0 .../attention/backends}/__init__.py | 0 fastvideo/v1/attention/backends/abstract.py | 300 +++++ fastvideo/v1/attention/backends/flash_attn.py | 1042 +++++++++++++++++ fastvideo/v1/attention/backends/utils.py | 619 ++++++++++ fastvideo/{ => v1}/attention/flash_attn.py | 4 +- fastvideo/{ => v1}/attention/layer.py | 2 +- fastvideo/{ => v1}/distributed/__init__.py | 0 .../{ => v1}/distributed/communication_op.py | 0 .../device_communicators}/__init__.py | 0 .../base_device_communicator.py | 0 .../device_communicators/cuda_communicator.py | 2 +- .../device_communicators/pynccl.py | 8 +- .../device_communicators/pynccl_wrapper.py | 4 +- .../{ => v1}/distributed/parallel_state.py | 22 +- fastvideo/{ => v1}/distributed/utils.py | 4 +- fastvideo/{ => v1}/envs.py | 0 fastvideo/{ => v1}/inference_args.py | 4 +- fastvideo/{ => v1}/inference_engine.py | 16 +- fastvideo/{loader => v1/layers}/__init__.py | 0 fastvideo/{ => v1}/layers/activation.py | 0 fastvideo/{ => v1}/layers/conv.py | 0 fastvideo/{ => v1}/layers/layernorm.py | 0 fastvideo/{ => v1}/layers/linear.py | 8 +- fastvideo/{ => v1}/layers/mlp.py | 4 +- fastvideo/{ => v1}/layers/rotary_embedding.py | 2 +- fastvideo/{ => v1}/layers/utils.py | 0 fastvideo/{ => v1}/layers/visual_embedding.py | 6 +- .../layers/vocab_parallel_embedding.py | 8 +- .../base.py => v1/loader/__init__.py} | 0 fastvideo/{ => v1}/loader/fsdp_load.py | 2 +- fastvideo/{ => v1}/loader/loader.py | 14 +- fastvideo/{ => v1}/loader/utils.py | 2 +- fastvideo/{ => v1}/loader/weight_utils.py | 6 +- fastvideo/{ => v1}/logger.py | 10 +- fastvideo/{ => v1}/logging_utils/__init__.py | 2 +- fastvideo/{ => v1}/logging_utils/formatter.py | 0 fastvideo/{ => v1}/models/__init__.py | 8 +- fastvideo/{ => v1}/models/dits/base.py | 0 .../{ => v1}/models/dits/hunyuanvideo.py | 18 +- fastvideo/{ => v1}/models/encoders/clip.py | 8 +- fastvideo/{ => v1}/models/encoders/encoder.py | 0 fastvideo/{ => v1}/models/encoders/llama.py | 16 +- fastvideo/{ => v1}/models/encoders/utils.py | 4 +- fastvideo/{ => v1}/models/encoders/vision.py | 6 +- .../{ => v1}/models/hf_transformer_utils.py | 2 +- fastvideo/v1/models/hunyuan/__init__.py | 0 .../models}/hunyuan/constants.py | 0 .../v1/models/hunyuan/diffusion/__init__.py | 3 + .../hunyuan/diffusion/pipelines/__init__.py | 2 + .../pipelines/pipeline_hunyuan_video.py | 931 +++++++++++++++ .../hunyuan/diffusion/schedulers/__init__.py | 2 + .../scheduling_flow_match_discrete.py | 239 ++++ fastvideo/v1/models/hunyuan/idle_config.py | 380 ++++++ fastvideo/v1/models/hunyuan/inference.py | 504 ++++++++ .../v1/models/hunyuan/modules/__init__.py | 25 + .../hunyuan/modules/activation_layers.py | 23 + .../v1/models/hunyuan/modules/attenion.py | 128 ++ .../v1/models/hunyuan/modules/embed_layers.py | 150 +++ .../v1/models/hunyuan/modules/mlp_layers.py | 107 ++ fastvideo/v1/models/hunyuan/modules/models.py | 657 +++++++++++ .../models/hunyuan/modules/modulate_layers.py | 152 +++ .../v1/models/hunyuan/modules/norm_layers.py | 78 ++ .../models/hunyuan/modules/posemb_layers.py | 290 +++++ .../models/hunyuan/modules/token_refiner.py | 199 ++++ fastvideo/v1/models/hunyuan/prompt_rewrite.py | 52 + .../models/hunyuan/text_encoder/__init__.py | 323 +++++ fastvideo/v1/models/hunyuan/utils/__init__.py | 0 .../v1/models/hunyuan/utils/data_utils.py | 14 + .../v1/models/hunyuan/utils/file_utils.py | 71 ++ fastvideo/v1/models/hunyuan/utils/helpers.py | 41 + ...preprocess_text_encoder_tokenizer_utils.py | 41 + fastvideo/v1/models/hunyuan/vae/__init__.py | 64 + .../hunyuan/vae/autoencoder_kl_causal_3d.py | 776 ++++++++++++ .../hunyuan/vae/unet_causal_3d_blocks.py | 760 ++++++++++++ fastvideo/v1/models/hunyuan/vae/vae.py | 342 ++++++ fastvideo/{ => v1}/models/loader.py | 24 +- fastvideo/{ => v1}/models/parameter.py | 6 +- fastvideo/{ => v1}/models/registry.py | 8 +- fastvideo/v1/models/schedulers/base.py | 0 .../scheduling_flow_match_euler_discrete.py | 0 fastvideo/{ => v1}/models/text_encoder.py | 2 +- fastvideo/{ => v1}/models/utils.py | 4 +- fastvideo/{ => v1}/models/vaes/common.py | 2 +- fastvideo/{ => v1}/models/vaes/hunyuanvae.py | 6 +- fastvideo/{ => v1}/pipelines/__init__.py | 20 +- .../{ => v1}/pipelines/composed/__init__.py | 2 +- .../composed/composed_pipeline_base.py | 8 +- .../pipelines}/hunyuan/constants.py | 0 .../pipelines/hunyuan/hunyuan_loader.py | 16 +- .../pipelines/hunyuan/hunyuan_video.py | 14 +- .../implementations/hunyuan/__init__.py | 2 +- .../implementations/hunyuan/constants.py | 89 ++ .../implementations/hunyuan/hunyuan_loader.py | 22 +- .../hunyuan/hunyuan_pipeline.py | 18 +- fastvideo/{ => v1}/pipelines/loader.py | 18 +- fastvideo/{ => v1}/pipelines/pipeline_base.py | 20 +- .../{ => v1}/pipelines/pipeline_batch_info.py | 0 .../{ => v1}/pipelines/pipeline_registry.py | 4 +- fastvideo/v1/pipelines/stages/__init__.py | 28 + fastvideo/{ => v1}/pipelines/stages/base.py | 8 +- .../{ => v1}/pipelines/stages/conditioning.py | 8 +- .../{ => v1}/pipelines/stages/decoding.py | 10 +- .../{ => v1}/pipelines/stages/denoising.py | 14 +- .../pipelines/stages/input_validation.py | 8 +- .../pipelines/stages/latent_preparation.py | 8 +- .../pipelines/stages/post_processing.py | 8 +- .../pipelines/stages/prompt_encoding.py | 10 +- .../pipelines/stages/timestep_preparation.py | 8 +- fastvideo/{ => v1}/platforms/__init__.py | 0 fastvideo/{ => v1}/platforms/cuda.py | 8 +- fastvideo/{ => v1}/platforms/interface.py | 4 +- .../{ => v1}/sample/v1_fastvideo_inference.py | 11 +- .../v1/{ => tests}/distributed_example.py | 2 +- fastvideo/v1/{ => tests}/distributed_run.sh | 0 fastvideo/v1/tests/test_clip_encoder.py | 168 +++ fastvideo/v1/{ => tests}/test_hunyuan_vae.py | 4 +- fastvideo/v1/{ => tests}/test_hunyuanvideo.py | 14 +- .../v1/{ => tests}/test_hunyuanvideo_load.py | 14 +- fastvideo/v1/tests/test_llama_encoder.py | 0 fastvideo/v1/{ => tests}/tp_example.py | 8 +- fastvideo/v1/utils/checkpoint.py | 286 +++++ fastvideo/v1/utils/communications.py | 308 +++++ fastvideo/v1/utils/dataset_utils.py | 342 ++++++ fastvideo/v1/utils/env_utils.py | 38 + fastvideo/v1/utils/fsdp_util.py | 132 +++ fastvideo/v1/utils/load.py | 339 ++++++ fastvideo/v1/utils/logging_.py | 24 + fastvideo/v1/utils/optimizer.py | 68 ++ fastvideo/v1/utils/parallel_states.py | 63 + fastvideo/{ => v1}/utils/utils.py | 6 +- fastvideo/v1/utils/validation.py | 308 +++++ scripts/inference/inference_hunyuan.sh | 40 +- scripts/inference/v1_inference_hunyuan.sh | 2 +- 172 files changed, 17897 insertions(+), 406 deletions(-) create mode 100644 fastvideo/models/flash_attn_no_pad.py create mode 100644 fastvideo/models/hunyuan_hf/modeling_hunyuan.py create mode 100644 fastvideo/models/hunyuan_hf/pipeline_hunyuan.py create mode 100644 fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py create mode 100644 fastvideo/models/mochi_hf/mochi_latents_utils.py create mode 100644 fastvideo/models/mochi_hf/modeling_mochi.py create mode 100644 fastvideo/models/mochi_hf/norm.py create mode 100644 fastvideo/models/mochi_hf/pipeline_mochi.py create mode 100644 fastvideo/models/stepvideo/__init__.py create mode 100644 fastvideo/models/stepvideo/__version__.py create mode 100644 fastvideo/models/stepvideo/config.py create mode 100644 fastvideo/models/stepvideo/diffusion/scheduler.py create mode 100755 fastvideo/models/stepvideo/diffusion/video_pipeline.py create mode 100755 fastvideo/models/stepvideo/modules/attentions.py create mode 100755 fastvideo/models/stepvideo/modules/blocks.py create mode 100755 fastvideo/models/stepvideo/modules/model.py create mode 100755 fastvideo/models/stepvideo/modules/normalization.py create mode 100755 fastvideo/models/stepvideo/modules/rope.py create mode 100644 fastvideo/models/stepvideo/parallel.py create mode 100644 fastvideo/models/stepvideo/text_encoder/__init__.py create mode 100755 fastvideo/models/stepvideo/text_encoder/clip.py create mode 100755 fastvideo/models/stepvideo/text_encoder/flashattention.py create mode 100755 fastvideo/models/stepvideo/text_encoder/stepllm.py create mode 100755 fastvideo/models/stepvideo/text_encoder/tokenizer.py create mode 100755 fastvideo/models/stepvideo/utils/__init__.py create mode 100644 fastvideo/models/stepvideo/utils/quantization.py create mode 100755 fastvideo/models/stepvideo/utils/utils.py create mode 100644 fastvideo/models/stepvideo/utils/video_process.py create mode 100755 fastvideo/models/stepvideo/vae/vae.py delete mode 100644 fastvideo/pipelines/stages/__init__.py rename fastvideo/{ => v1}/attention/__init__.py (100%) rename fastvideo/{distributed/device_communicators => v1/attention/backends}/__init__.py (100%) create mode 100644 fastvideo/v1/attention/backends/abstract.py create mode 100755 fastvideo/v1/attention/backends/flash_attn.py create mode 100644 fastvideo/v1/attention/backends/utils.py rename fastvideo/{ => v1}/attention/flash_attn.py (97%) rename fastvideo/{ => v1}/attention/layer.py (98%) rename fastvideo/{ => v1}/distributed/__init__.py (100%) rename fastvideo/{ => v1}/distributed/communication_op.py (100%) rename fastvideo/{layers => v1/distributed/device_communicators}/__init__.py (100%) rename fastvideo/{ => v1}/distributed/device_communicators/base_device_communicator.py (100%) rename fastvideo/{ => v1}/distributed/device_communicators/cuda_communicator.py (98%) rename fastvideo/{ => v1}/distributed/device_communicators/pynccl.py (97%) rename fastvideo/{ => v1}/distributed/device_communicators/pynccl_wrapper.py (99%) rename fastvideo/{ => v1}/distributed/parallel_state.py (98%) rename fastvideo/{ => v1}/distributed/utils.py (99%) rename fastvideo/{ => v1}/envs.py (100%) rename fastvideo/{ => v1}/inference_args.py (99%) rename fastvideo/{ => v1}/inference_engine.py (94%) rename fastvideo/{loader => v1/layers}/__init__.py (100%) rename fastvideo/{ => v1}/layers/activation.py (100%) rename fastvideo/{ => v1}/layers/conv.py (100%) rename fastvideo/{ => v1}/layers/layernorm.py (100%) rename fastvideo/{ => v1}/layers/linear.py (99%) rename fastvideo/{ => v1}/layers/mlp.py (90%) rename fastvideo/{ => v1}/layers/rotary_embedding.py (99%) rename fastvideo/{ => v1}/layers/utils.py (100%) rename fastvideo/{ => v1}/layers/visual_embedding.py (96%) rename fastvideo/{ => v1}/layers/vocab_parallel_embedding.py (98%) rename fastvideo/{models/schedulers/base.py => v1/loader/__init__.py} (100%) rename fastvideo/{ => v1}/loader/fsdp_load.py (99%) rename fastvideo/{ => v1}/loader/loader.py (97%) rename fastvideo/{ => v1}/loader/utils.py (99%) rename fastvideo/{ => v1}/loader/weight_utils.py (99%) rename fastvideo/{ => v1}/logger.py (96%) rename fastvideo/{ => v1}/logging_utils/__init__.py (54%) rename fastvideo/{ => v1}/logging_utils/formatter.py (100%) rename fastvideo/{ => v1}/models/__init__.py (87%) rename fastvideo/{ => v1}/models/dits/base.py (100%) rename fastvideo/{ => v1}/models/dits/hunyuanvideo.py (97%) rename fastvideo/{ => v1}/models/encoders/clip.py (98%) rename fastvideo/{ => v1}/models/encoders/encoder.py (100%) rename fastvideo/{ => v1}/models/encoders/llama.py (97%) rename fastvideo/{ => v1}/models/encoders/utils.py (99%) rename fastvideo/{ => v1}/models/encoders/vision.py (97%) rename fastvideo/{ => v1}/models/hf_transformer_utils.py (98%) create mode 100644 fastvideo/v1/models/hunyuan/__init__.py rename fastvideo/{pipelines => v1/models}/hunyuan/constants.py (100%) create mode 100644 fastvideo/v1/models/hunyuan/diffusion/__init__.py create mode 100644 fastvideo/v1/models/hunyuan/diffusion/pipelines/__init__.py create mode 100644 fastvideo/v1/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py create mode 100644 fastvideo/v1/models/hunyuan/diffusion/schedulers/__init__.py create mode 100644 fastvideo/v1/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py create mode 100644 fastvideo/v1/models/hunyuan/idle_config.py create mode 100644 fastvideo/v1/models/hunyuan/inference.py create mode 100644 fastvideo/v1/models/hunyuan/modules/__init__.py create mode 100644 fastvideo/v1/models/hunyuan/modules/activation_layers.py create mode 100644 fastvideo/v1/models/hunyuan/modules/attenion.py create mode 100644 fastvideo/v1/models/hunyuan/modules/embed_layers.py create mode 100644 fastvideo/v1/models/hunyuan/modules/mlp_layers.py create mode 100644 fastvideo/v1/models/hunyuan/modules/models.py create mode 100644 fastvideo/v1/models/hunyuan/modules/modulate_layers.py create mode 100644 fastvideo/v1/models/hunyuan/modules/norm_layers.py create mode 100644 fastvideo/v1/models/hunyuan/modules/posemb_layers.py create mode 100644 fastvideo/v1/models/hunyuan/modules/token_refiner.py create mode 100644 fastvideo/v1/models/hunyuan/prompt_rewrite.py create mode 100644 fastvideo/v1/models/hunyuan/text_encoder/__init__.py create mode 100644 fastvideo/v1/models/hunyuan/utils/__init__.py create mode 100644 fastvideo/v1/models/hunyuan/utils/data_utils.py create mode 100644 fastvideo/v1/models/hunyuan/utils/file_utils.py create mode 100644 fastvideo/v1/models/hunyuan/utils/helpers.py create mode 100644 fastvideo/v1/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py create mode 100644 fastvideo/v1/models/hunyuan/vae/__init__.py create mode 100644 fastvideo/v1/models/hunyuan/vae/autoencoder_kl_causal_3d.py create mode 100644 fastvideo/v1/models/hunyuan/vae/unet_causal_3d_blocks.py create mode 100644 fastvideo/v1/models/hunyuan/vae/vae.py rename fastvideo/{ => v1}/models/loader.py (94%) rename fastvideo/{ => v1}/models/parameter.py (98%) rename fastvideo/{ => v1}/models/registry.py (97%) create mode 100644 fastvideo/v1/models/schedulers/base.py rename fastvideo/{ => v1}/models/schedulers/scheduling_flow_match_euler_discrete.py (100%) rename fastvideo/{ => v1}/models/text_encoder.py (99%) rename fastvideo/{ => v1}/models/utils.py (96%) rename fastvideo/{ => v1}/models/vaes/common.py (99%) rename fastvideo/{ => v1}/models/vaes/hunyuanvae.py (99%) rename fastvideo/{ => v1}/pipelines/__init__.py (93%) rename fastvideo/{ => v1}/pipelines/composed/__init__.py (84%) rename fastvideo/{ => v1}/pipelines/composed/composed_pipeline_base.py (95%) rename fastvideo/{pipelines/implementations => v1/pipelines}/hunyuan/constants.py (100%) rename fastvideo/{ => v1}/pipelines/hunyuan/hunyuan_loader.py (95%) rename fastvideo/{ => v1}/pipelines/hunyuan/hunyuan_video.py (93%) rename fastvideo/{ => v1}/pipelines/implementations/hunyuan/__init__.py (69%) create mode 100644 fastvideo/v1/pipelines/implementations/hunyuan/constants.py rename fastvideo/{ => v1}/pipelines/implementations/hunyuan/hunyuan_loader.py (94%) rename fastvideo/{ => v1}/pipelines/implementations/hunyuan/hunyuan_pipeline.py (92%) rename fastvideo/{ => v1}/pipelines/loader.py (93%) rename fastvideo/{ => v1}/pipelines/pipeline_base.py (97%) rename fastvideo/{ => v1}/pipelines/pipeline_batch_info.py (100%) rename fastvideo/{ => v1}/pipelines/pipeline_registry.py (96%) create mode 100644 fastvideo/v1/pipelines/stages/__init__.py rename fastvideo/{ => v1}/pipelines/stages/base.py (94%) rename fastvideo/{ => v1}/pipelines/stages/conditioning.py (90%) rename fastvideo/{ => v1}/pipelines/stages/decoding.py (88%) rename fastvideo/{ => v1}/pipelines/stages/denoising.py (95%) rename fastvideo/{ => v1}/pipelines/stages/input_validation.py (94%) rename fastvideo/{ => v1}/pipelines/stages/latent_preparation.py (94%) rename fastvideo/{ => v1}/pipelines/stages/post_processing.py (85%) rename fastvideo/{ => v1}/pipelines/stages/prompt_encoding.py (94%) rename fastvideo/{ => v1}/pipelines/stages/timestep_preparation.py (93%) rename fastvideo/{ => v1}/platforms/__init__.py (100%) rename fastvideo/{ => v1}/platforms/cuda.py (97%) rename fastvideo/{ => v1}/platforms/interface.py (97%) rename fastvideo/{ => v1}/sample/v1_fastvideo_inference.py (87%) rename fastvideo/v1/{ => tests}/distributed_example.py (90%) rename fastvideo/v1/{ => tests}/distributed_run.sh (100%) create mode 100644 fastvideo/v1/tests/test_clip_encoder.py rename fastvideo/v1/{ => tests}/test_hunyuan_vae.py (98%) rename fastvideo/v1/{ => tests}/test_hunyuanvideo.py (94%) rename fastvideo/v1/{ => tests}/test_hunyuanvideo_load.py (94%) create mode 100644 fastvideo/v1/tests/test_llama_encoder.py rename fastvideo/v1/{ => tests}/tp_example.py (96%) create mode 100644 fastvideo/v1/utils/checkpoint.py create mode 100644 fastvideo/v1/utils/communications.py create mode 100644 fastvideo/v1/utils/dataset_utils.py create mode 100644 fastvideo/v1/utils/env_utils.py create mode 100644 fastvideo/v1/utils/fsdp_util.py create mode 100644 fastvideo/v1/utils/load.py create mode 100644 fastvideo/v1/utils/logging_.py create mode 100644 fastvideo/v1/utils/optimizer.py create mode 100644 fastvideo/v1/utils/parallel_states.py rename fastvideo/{ => v1}/utils/utils.py (98%) create mode 100644 fastvideo/v1/utils/validation.py diff --git a/fastvideo/models/flash_attn_no_pad.py b/fastvideo/models/flash_attn_no_pad.py new file mode 100644 index 000000000..f686e70a6 --- /dev/null +++ b/fastvideo/models/flash_attn_no_pad.py @@ -0,0 +1,28 @@ +from einops import rearrange +from flash_attn import flash_attn_varlen_qkvpacked_func +from flash_attn.bert_padding import pad_input, unpad_input + + +def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None): + # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27 + batch_size = qkv.shape[0] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask) + + x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) + output_unpad = flash_attn_varlen_qkvpacked_func( + x_unpad, + cu_seqlens, + max_s, + dropout_p, + softmax_scale=softmax_scale, + causal=causal, + ) + 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/models/hunyuan/inference.py b/fastvideo/models/hunyuan/inference.py index c034f330c..62694e6d9 100644 --- a/fastvideo/models/hunyuan/inference.py +++ b/fastvideo/models/hunyuan/inference.py @@ -58,22 +58,6 @@ class Inference(object): args (argparse.Namespace): The arguments for the pipeline. device (int): The device for inference. Default is 0. """ - from fastvideo.pipelines.loader import get_pipeline_loader - loader = get_pipeline_loader(args) - vae, vae_kwargs, text_encoder, text_encoder_2, transformer, scheduler = \ - loader.load_components(args) - return cls( - args=args, - vae=vae, - vae_kwargs=vae_kwargs, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - model=transformer, - use_cpu_offload=args.use_cpu_offload, - device=device, - logger=logger, - parallel_args=None, - ) # ======================================================================== logger.info(f"Got text-to-video model root path: {pretrained_model_path}") @@ -223,7 +207,6 @@ class Inference(object): else: raise ValueError(f"Invalid model path: {dit_weight}") - print('model_path', model_path) if not model_path.exists(): raise ValueError(f"model_path not exists: {model_path}") logger.info(f"Loading torch model {model_path}...") @@ -313,11 +296,6 @@ class HunyuanVideoSampler(Inference): progress_bar_config=None, data_type="video", ): - print("Loading diffusion pipeline...") - # print('vae', vae) - print('text_encoder', text_encoder) - print('text_encoder_2', text_encoder_2) - print('model', model) """Load the denoising scheduler for inference.""" if scheduler is None: if args.denoise_type == "flow": diff --git a/fastvideo/models/hunyuan/modules/__init__.py b/fastvideo/models/hunyuan/modules/__init__.py index 3bb947c4e..c85b51a4c 100644 --- a/fastvideo/models/hunyuan/modules/__init__.py +++ b/fastvideo/models/hunyuan/modules/__init__.py @@ -1,25 +1,25 @@ -# from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer +from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer -# def load_model(args, in_channels, out_channels, factor_kwargs): -# """load hunyuan video model +def load_model(args, in_channels, out_channels, factor_kwargs): + """load hunyuan video model -# Args: -# args (dict): model args -# in_channels (int): input channels number -# out_channels (int): output channels number -# factor_kwargs (dict): factor kwargs + Args: + args (dict): model args + in_channels (int): input channels number + out_channels (int): output channels number + factor_kwargs (dict): factor kwargs -# Returns: -# model (nn.Module): The hunyuan video model -# """ -# if args.model in HUNYUAN_VIDEO_CONFIG.keys(): -# model = HYVideoDiffusionTransformer( -# in_channels=in_channels, -# out_channels=out_channels, -# **HUNYUAN_VIDEO_CONFIG[args.model], -# **factor_kwargs, -# ) -# return model -# else: -# raise NotImplementedError() + Returns: + model (nn.Module): The hunyuan video model + """ + if args.model in HUNYUAN_VIDEO_CONFIG.keys(): + model = HYVideoDiffusionTransformer( + in_channels=in_channels, + out_channels=out_channels, + **HUNYUAN_VIDEO_CONFIG[args.model], + **factor_kwargs, + ) + return model + else: + raise NotImplementedError() diff --git a/fastvideo/models/hunyuan/modules/attenion.py b/fastvideo/models/hunyuan/modules/attenion.py index a5760c1a0..699b3ae1d 100644 --- a/fastvideo/models/hunyuan/modules/attenion.py +++ b/fastvideo/models/hunyuan/modules/attenion.py @@ -9,7 +9,6 @@ except ImportError: sliding_tile_attention = None from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad -from flash_attn import flash_attn_func from fastvideo.utils.communications import all_gather, all_to_all_4D from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info @@ -95,17 +94,14 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy= hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2) else: - txt_length = text_mask.sum().int().item() - query = torch.cat([query, encoder_query[:, :txt_length]], dim=1) - key = torch.cat([key, encoder_key[:, :txt_length]], dim=1) - value = torch.cat([value, encoder_value[:, :txt_length]], dim=1) + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) # B, S, 3, H, D - # qkv = torch.stack([query, key, value], dim=2) + qkv = torch.stack([query, key, value], dim=2) attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) - # hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) - hidden_states = flash_attn_func(query, key, value, causal=False, dropout_p=0.0, softmax_scale=None) - hidden_states = torch.cat([hidden_states, encoder_query[:, txt_length:, :]], dim=1) + hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length), dim=1) diff --git a/fastvideo/models/hunyuan/modules/models.py b/fastvideo/models/hunyuan/modules/models.py index 0d1c9e084..6b1706ad0 100644 --- a/fastvideo/models/hunyuan/modules/models.py +++ b/fastvideo/models/hunyuan/modules/models.py @@ -131,12 +131,12 @@ class MMDoubleStreamBlock(nn.Module): # Prepare image for attention. img_modulated = self.img_norm1(img) img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale) - img_qkv = self.img_attn_qkv(img_modulated) img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) # 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) + # Apply RoPE if needed. if freqs_cis is not None: @@ -148,12 +148,12 @@ class MMDoubleStreamBlock(nn.Module): shrink_head(freqs_cis[0], dim=0), shrink_head(freqs_cis[1], dim=0), ) + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" img_q, img_k = img_qq, img_kk - # Prepare txt for attention. txt_modulated = self.txt_norm1(txt) txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale) @@ -162,6 +162,7 @@ class MMDoubleStreamBlock(nn.Module): # Apply QK-Norm if needed. txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + attn = parallel_attention( (img_q, txt_q), (img_k, txt_k), @@ -173,7 +174,9 @@ class MMDoubleStreamBlock(nn.Module): ) # attention computation end + img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:] + # Calculate the img blocks. img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate) img = img + apply_gate( @@ -540,6 +543,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): # text modulation vec = vec + self.vector_in(text_states_2) + # guidance modulation if self.guidance_embed: if guidance is None: @@ -547,6 +551,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): # our timestep_embedding is merged into guidance_in(TimestepEmbedder) vec = vec + self.guidance_in(guidance) + # Embed image and text. img = self.img_in(img) if self.text_projection == "linear": @@ -555,11 +560,13 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None) else: raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}") + txt_seq_len = txt.shape[1] img_seq_len = img.shape[1] freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None # --------------------- Pass through DiT blocks ------------------------ + for index, block in enumerate(self.double_blocks): double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]] img, txt = block(*double_block_args) @@ -580,9 +587,12 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): x = block(*single_block_args) if output_features and _ % output_features_stride == 0: features_list.append(x[:, :img_seq_len, ...]) + img = x[:, :img_seq_len, ...] + # ---------------------------- Final layer ------------------------------ img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + img = self.unpatchify(img, tt, th, tw) assert not return_dict, "return_dict is not supported." if output_features: diff --git a/fastvideo/models/hunyuan/modules/posemb_layers.py b/fastvideo/models/hunyuan/modules/posemb_layers.py index 314aa126d..34dcdb5d5 100644 --- a/fastvideo/models/hunyuan/modules/posemb_layers.py +++ b/fastvideo/models/hunyuan/modules/posemb_layers.py @@ -160,7 +160,6 @@ def apply_rotary_emb( cos, sin = cos.to(xq.device), sin.to(xq.device) # real * cos - imag * sin # imag * cos + real * sin - xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) else: diff --git a/fastvideo/models/hunyuan/modules/token_refiner.py b/fastvideo/models/hunyuan/modules/token_refiner.py index 1bc127fea..6aafc1dba 100644 --- a/fastvideo/models/hunyuan/modules/token_refiner.py +++ b/fastvideo/models/hunyuan/modules/token_refiner.py @@ -77,6 +77,7 @@ class IndividualTokenRefinerBlock(nn.Module): # Self-Attention attn = attention(q, k, v, attn_mask=attn_mask) + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) # FFN Layer @@ -191,9 +192,11 @@ class SingleTokenRefiner(nn.Module): else: mask_float = mask.float().unsqueeze(-1) # [b, s1, 1] context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) - context_aware_representations = self.c_embedder(context_aware_representations.to(x.dtype)) + context_aware_representations = self.c_embedder(context_aware_representations) c = timestep_aware_representations + context_aware_representations x = self.input_embedder(x) + x = self.individual_token_refiner(x, c, mask) + return x diff --git a/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py b/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py index 0a14520a1..8c2bf87fb 100644 --- a/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py +++ b/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py @@ -21,15 +21,11 @@ from math import prod from typing import Dict, Optional, Tuple, Union import torch -import torch.nn as nn import torch.distributed as dist +import torch.nn as nn from diffusers.configuration_utils import ConfigMixin, register_to_config from fastvideo.utils.parallel_states import nccl_info -from fastvideo.distributed.parallel_state import get_sp_group -from fastvideo.logger import init_logger - -logger = init_logger(__name__) try: # This diffusers is modified and packed in the mirror. @@ -311,7 +307,6 @@ class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: assert len(z.shape) == 5, "The input tensor should have 5 dimensions." - logger.info(f"use_parallel: {self.use_parallel}") if self.use_parallel: return self.parallel_tiled_decode(z, return_dict=return_dict) @@ -581,12 +576,7 @@ class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): """ Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs """ - # world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group - - sp_group = get_sp_group() - world_size = sp_group.world_size - rank = sp_group.rank_in_group - + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group B, C, T, H, W = z.shape # Calculate parameters @@ -660,8 +650,6 @@ class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): gathered_results = torch.zeros_like(padded_results).repeat( world_size, *[1] * len(padded_results.shape)).contiguous( ) # use contiguous to make sure it won't copy data in the following operations - # gathered_results = sp_group.all_gather(padded_results) - # gathered_dim_metadata = sp_group.broadcast_object_list(gathered_dim_metadata, local_dim_metadata) dist.all_gather_into_tensor(gathered_results, padded_results) dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) # Process gathered results diff --git a/fastvideo/models/hunyuan_hf/modeling_hunyuan.py b/fastvideo/models/hunyuan_hf/modeling_hunyuan.py new file mode 100644 index 000000000..39b1a815a --- /dev/null +++ b/fastvideo/models/hunyuan_hf/modeling_hunyuan.py @@ -0,0 +1,836 @@ +# Copyright 2024 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, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin +from diffusers.models.attention import FeedForward +from diffusers.models.attention_processor import Attention, AttentionProcessor +from diffusers.models.embeddings import (CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, + get_1d_rotary_pos_embed) +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle +from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers + +from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad +from fastvideo.utils.communications import all_gather, all_to_all_4D +from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) + + +class HunyuanVideoAttnProcessor2_0: + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError( + "HunyuanVideoAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + sequence_length = hidden_states.size(1) + encoder_sequence_length = encoder_hidden_states.size(1) + if attn.add_q_proj is None and encoder_hidden_states is not None: + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + # 1. QKV projections + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2) + key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2) + value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2) + + # 2. QK normalization + if attn.norm_q is not None: + query = attn.norm_q(query).to(value) + if attn.norm_k is not None: + key = attn.norm_k(key).to(value) + + image_rotary_emb = ( + shrink_head(image_rotary_emb[0], dim=0), + shrink_head(image_rotary_emb[1], dim=0), + ) + + # 3. Rotational positional embeddings applied to latent stream + if image_rotary_emb is not None: + from diffusers.models.embeddings import apply_rotary_emb + + if attn.add_q_proj is None and encoder_hidden_states is not None: + query = torch.cat( + [ + apply_rotary_emb(query[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb), + query[:, :, -encoder_hidden_states.shape[1]:], + ], + dim=2, + ) + key = torch.cat( + [ + apply_rotary_emb(key[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb), + key[:, :, -encoder_hidden_states.shape[1]:], + ], + dim=2, + ) + else: + query = apply_rotary_emb(query, image_rotary_emb) + key = apply_rotary_emb(key, image_rotary_emb) + + # 4. Encoder condition QKV projection and normalization + if attn.add_q_proj is not None and encoder_hidden_states is not None: + encoder_query = attn.add_q_proj(encoder_hidden_states) + encoder_key = attn.add_k_proj(encoder_hidden_states) + encoder_value = attn.add_v_proj(encoder_hidden_states) + + encoder_query = encoder_query.unflatten(2, (attn.heads, -1)).transpose(1, 2) + encoder_key = encoder_key.unflatten(2, (attn.heads, -1)).transpose(1, 2) + encoder_value = encoder_value.unflatten(2, (attn.heads, -1)).transpose(1, 2) + + if attn.norm_added_q is not None: + encoder_query = attn.norm_added_q(encoder_query).to(encoder_value) + if attn.norm_added_k is not None: + encoder_key = attn.norm_added_k(encoder_key).to(encoder_value) + + query = torch.cat([query, encoder_query], dim=2) + key = torch.cat([key, encoder_key], dim=2) + value = torch.cat([value, encoder_value], dim=2) + + if get_sequence_parallel_state(): + query_img, query_txt = query[:, :, :sequence_length, :], query[:, :, sequence_length:, :] + key_img, key_txt = key[:, :, :sequence_length, :], key[:, :, sequence_length:, :] + value_img, value_txt = value[:, :, :sequence_length, :], value[:, :, sequence_length:, :] + query_img = all_to_all_4D(query_img, scatter_dim=1, gather_dim=2) # + key_img = all_to_all_4D(key_img, scatter_dim=1, gather_dim=2) + value_img = all_to_all_4D(value_img, scatter_dim=1, gather_dim=2) + + query_txt = shrink_head(query_txt, dim=1) + key_txt = shrink_head(key_txt, dim=1) + value_txt = shrink_head(value_txt, dim=1) + query = torch.cat([query_img, query_txt], dim=2) + key = torch.cat([key_img, key_txt], dim=2) + value = torch.cat([value_img, value_txt], dim=2) + + query = query.unsqueeze(2) + key = key.unsqueeze(2) + value = value.unsqueeze(2) + qkv = torch.cat([query, key, value], dim=2) + qkv = qkv.transpose(1, 3) + + # 5. Attention + attention_mask = attention_mask[:, 0, :] + seq_len = qkv.shape[1] + attn_len = attention_mask.shape[1] + attention_mask = F.pad(attention_mask, (seq_len - attn_len, 0), value=True) + + hidden_states = flash_attn_no_pad(qkv, attention_mask, causal=False, dropout_p=0.0, softmax_scale=None) + + if get_sequence_parallel_state(): + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( + (sequence_length * nccl_info.sp_size, encoder_sequence_length), dim=1) + hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) + encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous() + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.flatten(2, 3) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + else: + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + # 6. Output projection + if encoder_hidden_states is not None: + hidden_states, encoder_hidden_states = ( + hidden_states[:, :-encoder_hidden_states.shape[1]], + hidden_states[:, -encoder_hidden_states.shape[1]:], + ) + + if encoder_hidden_states is not None: + if getattr(attn, "to_out", None) is not None: + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + if getattr(attn, "to_add_out", None) is not None: + encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + + return hidden_states, encoder_hidden_states + + +class HunyuanVideoPatchEmbed(nn.Module): + + def __init__( + self, + patch_size: Union[int, Tuple[int, int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + ) -> None: + super().__init__() + + patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size + self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.proj(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) # BCFHW -> BNC + return hidden_states + + +class HunyuanVideoAdaNorm(nn.Module): + + def __init__(self, in_features: int, out_features: Optional[int] = None) -> None: + super().__init__() + + out_features = out_features or 2 * in_features + self.linear = nn.Linear(in_features, out_features) + self.nonlinearity = nn.SiLU() + + def forward(self, + temb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + temb = self.linear(self.nonlinearity(temb)) + gate_msa, gate_mlp = temb.chunk(2, dim=1) + gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1) + return gate_msa, gate_mlp + + +class HunyuanVideoIndividualTokenRefinerBlock(nn.Module): + + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_width_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + heads=num_attention_heads, + dim_head=attention_head_dim, + bias=attention_bias, + ) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) + self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate) + + self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + norm_hidden_states = self.norm1(hidden_states) + + attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=None, + attention_mask=attention_mask, + ) + + gate_msa, gate_mlp = self.norm_out(temb) + hidden_states = hidden_states + attn_output * gate_msa + + ff_output = self.ff(self.norm2(hidden_states)) + hidden_states = hidden_states + ff_output * gate_mlp + + return hidden_states + + +class HunyuanVideoIndividualTokenRefiner(nn.Module): + + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + num_layers: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + self.refiner_blocks = nn.ModuleList([ + HunyuanVideoIndividualTokenRefinerBlock( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + attention_bias=attention_bias, + ) for _ in range(num_layers) + ]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> None: + self_attn_mask = None + if attention_mask is not None: + batch_size = attention_mask.shape[0] + seq_len = attention_mask.shape[1] + attention_mask = attention_mask.to(hidden_states.device).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) + self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() + self_attn_mask[:, :, :, 0] = True + + for block in self.refiner_blocks: + hidden_states = block(hidden_states, temb, self_attn_mask) + + return hidden_states + + +class HunyuanVideoTokenRefiner(nn.Module): + + def __init__( + self, + in_channels: int, + num_attention_heads: int, + attention_head_dim: int, + num_layers: int, + mlp_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.time_text_embed = CombinedTimestepTextProjEmbeddings(embedding_dim=hidden_size, + pooled_projection_dim=in_channels) + self.proj_in = nn.Linear(in_channels, hidden_size, bias=True) + self.token_refiner = HunyuanVideoIndividualTokenRefiner( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + num_layers=num_layers, + mlp_width_ratio=mlp_ratio, + mlp_drop_rate=mlp_drop_rate, + attention_bias=attention_bias, + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + attention_mask: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + if attention_mask is None: + pooled_projections = hidden_states.mean(dim=1) + else: + original_dtype = hidden_states.dtype + mask_float = attention_mask.float().unsqueeze(-1) + pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1) + pooled_projections = pooled_projections.to(original_dtype) + + temb = self.time_text_embed(timestep, pooled_projections) + hidden_states = self.proj_in(hidden_states) + hidden_states = self.token_refiner(hidden_states, temb, attention_mask) + + return hidden_states + + +class HunyuanVideoRotaryPosEmbed(nn.Module): + + def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None: + super().__init__() + + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.rope_dim = rope_dim + self.theta = theta + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + rope_sizes = [ + num_frames * nccl_info.sp_size // self.patch_size_t, height // self.patch_size, width // self.patch_size + ] + + axes_grids = [] + for i in range(3): + # Note: The following line diverges from original behaviour. We create the grid on the device, whereas + # original implementation creates it on CPU and then moves it to device. This results in numerical + # differences in layerwise debugging outputs, but visually it is the same. + grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32) + axes_grids.append(grid) + grid = torch.meshgrid(*axes_grids, indexing="ij") # [W, H, T] + grid = torch.stack(grid, dim=0) # [3, W, H, T] + + freqs = [] + for i in range(3): + freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True) + freqs.append(freq) + + freqs_cos = torch.cat([f[0] for f in freqs], dim=1) # (W * H * T, D / 2) + freqs_sin = torch.cat([f[1] for f in freqs], dim=1) # (W * H * T, D / 2) + return freqs_cos, freqs_sin + + +class HunyuanVideoSingleTransformerBlock(nn.Module): + + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 4.0, + qk_norm: str = "rms_norm", + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + mlp_dim = int(hidden_size * mlp_ratio) + + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=hidden_size, + bias=True, + processor=HunyuanVideoAttnProcessor2_0(), + qk_norm=qk_norm, + eps=1e-6, + pre_only=True, + ) + + self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm") + self.proj_mlp = nn.Linear(hidden_size, mlp_dim) + self.act_mlp = nn.GELU(approximate="tanh") + self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + text_seq_length = encoder_hidden_states.shape[1] + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + residual = hidden_states + + # 1. Input normalization + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) + + norm_hidden_states, norm_encoder_hidden_states = ( + norm_hidden_states[:, :-text_seq_length, :], + norm_hidden_states[:, -text_seq_length:, :], + ) + + # 2. Attention + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + attn_output = torch.cat([attn_output, context_attn_output], dim=1) + + # 3. Modulation and residual connection + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) + hidden_states = gate.unsqueeze(1) * self.proj_out(hidden_states) + hidden_states = hidden_states + residual + + hidden_states, encoder_hidden_states = ( + hidden_states[:, :-text_seq_length, :], + hidden_states[:, -text_seq_length:, :], + ) + return hidden_states, encoder_hidden_states + + +class HunyuanVideoTransformerBlock(nn.Module): + + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float, + qk_norm: str = "rms_norm", + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm") + self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm") + + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + added_kv_proj_dim=hidden_size, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=hidden_size, + context_pre_only=False, + bias=True, + processor=HunyuanVideoAttnProcessor2_0(), + qk_norm=qk_norm, + eps=1e-6, + ) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") + + self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # 1. Input normalization + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( + encoder_hidden_states, emb=temb) + + # 2. Joint attention + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=freqs_cis, + ) + + # 3. Modulation and residual connection + hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1) + + norm_hidden_states = self.norm2(hidden_states) + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + + # 4. Feed-forward + ff_output = self.ff(norm_hidden_states) + context_ff_output = self.ff_context(norm_encoder_hidden_states) + + hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output + encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output + + return hidden_states, encoder_hidden_states + + +class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): + r""" + A Transformer model for video-like data used in [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo). + + Args: + in_channels (`int`, defaults to `16`): + The number of channels in the input. + out_channels (`int`, defaults to `16`): + The number of channels in the output. + num_attention_heads (`int`, defaults to `24`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, defaults to `128`): + The number of channels in each head. + num_layers (`int`, defaults to `20`): + The number of layers of dual-stream blocks to use. + num_single_layers (`int`, defaults to `40`): + The number of layers of single-stream blocks to use. + num_refiner_layers (`int`, defaults to `2`): + The number of layers of refiner blocks to use. + mlp_ratio (`float`, defaults to `4.0`): + The ratio of the hidden layer size to the input size in the feedforward network. + patch_size (`int`, defaults to `2`): + The size of the spatial patches to use in the patch embedding layer. + patch_size_t (`int`, defaults to `1`): + The size of the tmeporal patches to use in the patch embedding layer. + qk_norm (`str`, defaults to `rms_norm`): + The normalization to use for the query and key projections in the attention layers. + guidance_embeds (`bool`, defaults to `True`): + Whether to use guidance embeddings in the model. + text_embed_dim (`int`, defaults to `4096`): + Input dimension of text embeddings from the text encoder. + pooled_projection_dim (`int`, defaults to `768`): + The dimension of the pooled projection of the text embeddings. + rope_theta (`float`, defaults to `256.0`): + The value of theta to use in the RoPE layer. + rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`): + The dimensions of the axes to use in the RoPE layer. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 16, + out_channels: int = 16, + num_attention_heads: int = 24, + attention_head_dim: int = 128, + num_layers: int = 20, + num_single_layers: int = 40, + num_refiner_layers: int = 2, + mlp_ratio: float = 4.0, + patch_size: int = 2, + patch_size_t: int = 1, + qk_norm: str = "rms_norm", + guidance_embeds: bool = True, + text_embed_dim: int = 4096, + pooled_projection_dim: int = 768, + rope_theta: float = 256.0, + rope_axes_dim: Tuple[int] = (16, 56, 56), + ) -> None: + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels or in_channels + + # 1. Latent and condition embedders + self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim) + self.context_embedder = HunyuanVideoTokenRefiner(text_embed_dim, + num_attention_heads, + attention_head_dim, + num_layers=num_refiner_layers) + self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim) + + # 2. RoPE + self.rope = HunyuanVideoRotaryPosEmbed(patch_size, patch_size_t, rope_axes_dim, rope_theta) + + # 3. Dual stream transformer blocks + self.transformer_blocks = nn.ModuleList([ + HunyuanVideoTransformerBlock(num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm) + for _ in range(num_layers) + ]) + + # 4. Single stream transformer blocks + self.single_transformer_blocks = nn.ModuleList([ + HunyuanVideoSingleTransformerBlock(num_attention_heads, + attention_head_dim, + mlp_ratio=mlp_ratio, + qk_norm=qk_norm) for _ in range(num_single_layers) + ]) + + # 5. Output projection + self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels) + + self.gradient_checkpointing = False + + @property + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor() + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes.") + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_attention_mask: torch.Tensor, + guidance: torch.Tensor = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + if guidance is None: + guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16) + + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + lora_scale = attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + else: + if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: + logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.") + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p, p_t = self.config.patch_size, self.config.patch_size_t + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p + post_patch_width = width // p + + pooled_projections = encoder_hidden_states[:, 0, :self.config.pooled_projection_dim] + encoder_hidden_states = encoder_hidden_states[:, 1:] + + # 1. RoPE + image_rotary_emb = self.rope(hidden_states) + + # 2. Conditional embeddings + temb = self.time_text_embed(timestep, guidance, pooled_projections) + hidden_states = self.x_embedder(hidden_states) + encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) + + # 3. Attention mask preparation + latent_sequence_length = hidden_states.shape[1] + condition_sequence_length = encoder_hidden_states.shape[1] + sequence_length = latent_sequence_length + condition_sequence_length + attention_mask = torch.zeros(batch_size, + sequence_length, + sequence_length, + device=hidden_states.device, + dtype=torch.bool) # [B, N, N] + + effective_condition_sequence_length = encoder_attention_mask.sum(dim=1, dtype=torch.int) + effective_sequence_length = latent_sequence_length + effective_condition_sequence_length + + for i in range(batch_size): + attention_mask[i, :effective_sequence_length[i], :effective_sequence_length[i]] = True + + # 4. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + + for block in self.transformer_blocks: + hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + image_rotary_emb, + **ckpt_kwargs, + ) + + for block in self.single_transformer_blocks: + hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + image_rotary_emb, + **ckpt_kwargs, + ) + + else: + for block in self.transformer_blocks: + hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask, + image_rotary_emb) + + for block in self.single_transformer_blocks: + hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask, + image_rotary_emb) + + # 5. Output projection + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width, + -1, p_t, p, p) + hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7) + hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (hidden_states, ) + + return Transformer2DModelOutput(sample=hidden_states) diff --git a/fastvideo/models/hunyuan_hf/pipeline_hunyuan.py b/fastvideo/models/hunyuan_hf/pipeline_hunyuan.py new file mode 100644 index 000000000..a30f083e1 --- /dev/null +++ b/fastvideo/models/hunyuan_hf/pipeline_hunyuan.py @@ -0,0 +1,691 @@ +# Copyright 2024 The HunyuanVideo 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. + +import inspect +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.loaders import HunyuanVideoLoraLoaderMixin +from diffusers.models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel +from diffusers.pipelines.hunyuan_video.pipeline_output import HunyuanVideoPipelineOutput +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import logging, replace_example_docstring +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from einops import rearrange +from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast + +from fastvideo.utils.communications import all_gather +from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import torch + >>> from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel + >>> from diffusers.utils import export_to_video + + >>> model_id = "tencent/HunyuanVideo" + >>> transformer = HunyuanVideoTransformer3DModel.from_pretrained( + ... model_id, subfolder="transformer", torch_dtype=torch.bfloat16 + ... ) + >>> pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.float16) + >>> pipe.vae.enable_tiling() + >>> pipe.to("cuda") + + >>> output = pipe( + ... prompt="A cat walks on the grass, realistic", + ... height=320, + ... width=512, + ... num_frames=61, + ... num_inference_steps=30, + ... ).frames[0] + >>> export_to_video(output, "output.mp4", fps=15) + ``` +""" + +DEFAULT_PROMPT_TEMPLATE = { + "template": ("<|start_header_id|>system<|end_header_id|>\n\nDescribe 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:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"), + "crop_start": + 95, +} + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler.") + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler.") + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class HunyuanVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin): + r""" + Pipeline for text-to-video generation using HunyuanVideo. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + text_encoder ([`LlamaModel`]): + [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). + tokenizer_2 (`LlamaTokenizer`): + Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). + transformer ([`HunyuanVideoTransformer3DModel`]): + Conditional Transformer to denoise the encoded image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLHunyuanVideo`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + text_encoder_2 ([`CLIPTextModel`]): + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically + the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. + tokenizer_2 (`CLIPTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + text_encoder: LlamaModel, + tokenizer: LlamaTokenizerFast, + transformer: HunyuanVideoTransformer3DModel, + vae: AutoencoderKLHunyuanVideo, + scheduler: FlowMatchEulerDiscreteScheduler, + text_encoder_2: CLIPTextModel, + tokenizer_2: CLIPTokenizer, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + text_encoder_2=text_encoder_2, + tokenizer_2=tokenizer_2, + ) + + self.vae_scale_factor_temporal = (self.vae.temporal_compression_ratio + if hasattr(self, "vae") and self.vae is not None else 4) + self.vae_scale_factor_spatial = (self.vae.spatial_compression_ratio + if hasattr(self, "vae") and self.vae is not None else 8) + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + def _get_llama_prompt_embeds( + self, + prompt: Union[str, List[str]], + prompt_template: Dict[str, Any], + num_videos_per_prompt: int = 1, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + max_sequence_length: int = 256, + num_hidden_layers_to_skip: int = 2, + ) -> Tuple[torch.Tensor, torch.Tensor]: + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + prompt = [prompt_template["template"].format(p) for p in prompt] + + crop_start = prompt_template.get("crop_start", None) + if crop_start is None: + prompt_template_input = self.tokenizer( + prompt_template["template"], + padding="max_length", + return_tensors="pt", + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=False, + ) + crop_start = prompt_template_input["input_ids"].shape[-1] + # Remove <|eot_id|> token and placeholder {} + crop_start -= 2 + + max_sequence_length += crop_start + text_inputs = self.tokenizer( + prompt, + max_length=max_sequence_length, + padding="max_length", + truncation=True, + return_tensors="pt", + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + ) + text_input_ids = text_inputs.input_ids.to(device=device) + prompt_attention_mask = text_inputs.attention_mask.to(device=device) + + prompt_embeds = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_attention_mask, + output_hidden_states=True, + ).hidden_states[-(num_hidden_layers_to_skip + 1)] + prompt_embeds = prompt_embeds.to(dtype=dtype) + + if crop_start is not None and crop_start > 0: + prompt_embeds = prompt_embeds[:, crop_start:] + prompt_attention_mask = prompt_attention_mask[:, crop_start:] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + prompt_attention_mask = prompt_attention_mask.repeat(1, num_videos_per_prompt) + prompt_attention_mask = prompt_attention_mask.view(batch_size * num_videos_per_prompt, seq_len) + + return prompt_embeds, prompt_attention_mask + + def _get_clip_prompt_embeds( + self, + prompt: Union[str, List[str]], + num_videos_per_prompt: int = 1, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + max_sequence_length: int = 77, + ) -> torch.Tensor: + device = device or self._execution_device + dtype = dtype or self.text_encoder_2.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + text_inputs = self.tokenizer_2( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1]) + logger.warning("The following part of your input was truncated because CLIP can only handle sequences up to" + f" {max_sequence_length} tokens: {removed_text}") + + prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output + + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, -1) + + return prompt_embeds + + def encode_prompt( + self, + prompt: Union[str, List[str]], + prompt_2: Union[str, List[str]] = None, + prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + num_videos_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + pooled_prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + max_sequence_length: int = 256, + ): + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds( + prompt, + prompt_template, + num_videos_per_prompt, + device=device, + dtype=dtype, + max_sequence_length=max_sequence_length, + ) + + if pooled_prompt_embeds is None: + if prompt_2 is None and pooled_prompt_embeds is None: + prompt_2 = prompt + pooled_prompt_embeds = self._get_clip_prompt_embeds( + prompt, + num_videos_per_prompt, + device=device, + dtype=dtype, + max_sequence_length=77, + ) + + return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask + + def check_inputs( + self, + prompt, + prompt_2, + height, + width, + prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + prompt_template=None, + ): + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs + for k in callback_on_step_end_tensor_inputs): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two.") + elif prompt_2 is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two.") + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.") + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): + raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") + + if prompt_template is not None: + if not isinstance(prompt_template, dict): + raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}") + if "template" not in prompt_template: + raise ValueError( + f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}") + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: 32, + height: int = 720, + width: int = 1280, + num_frames: int = 129, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + shape = ( + batch_size, + num_channels_latents, + num_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators.") + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.vae.enable_slicing() + + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_slicing() + + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + self.vae.enable_tiling() + + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_tiling() + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Union[str, List[str]] = None, + height: int = 720, + width: int = 1280, + num_frames: int = 129, + num_inference_steps: int = 50, + sigmas: List[float] = None, + guidance_scale: float = 6.0, + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + pooled_prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, + MultiPipelineCallbacks]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + max_sequence_length: int = 256, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + will be used instead. + height (`int`, defaults to `720`): + The height in pixels of the generated image. + width (`int`, defaults to `1280`): + The width in pixels of the generated image. + num_frames (`int`, defaults to `129`): + The number of frames in the generated video. + num_inference_steps (`int`, defaults to `50`): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, defaults to `6.0`): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. Note that the only available HunyuanVideo model is + CFG-distilled, which means that traditional guidance between unconditional and conditional latent is + not applied. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + + Examples: + + Returns: + [`~HunyuanVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned, otherwise a `tuple` is returned + where the first element is a list with the generated images and the second element is a list of `bool`s + indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + height, + width, + prompt_embeds, + callback_on_step_end_tensor_inputs, + prompt_template, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + + device = self._execution_device + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # 3. Encode input prompt + prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt( + prompt=prompt, + prompt_2=prompt, + prompt_template=prompt_template, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + device=device, + max_sequence_length=max_sequence_length, + ) + + transformer_dtype = self.transformer.dtype + prompt_embeds = prompt_embeds.to(transformer_dtype) + prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) + if pooled_prompt_embeds is not None: + pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype) + + # 4. Prepare timesteps + sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + ) + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_latent_frames, + torch.float32, + device, + generator, + latents, + ) + # check sequence_parallel + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + if get_sequence_parallel_state(): + latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous() + latents = latents[:, :, rank, :, :, :] + + # 6. Prepare guidance condition + guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0 + + # 7. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + latent_model_input = latents.to(transformer_dtype) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + if pooled_prompt_embeds.shape[-1] != prompt_embeds.shape[-1]: + pooled_prompt_embeds_padding = F.pad( + pooled_prompt_embeds, + (0, prompt_embeds.shape[2] - pooled_prompt_embeds.shape[1]), + value=0, + ).unsqueeze(1) + encoder_hidden_states = torch.cat([pooled_prompt_embeds_padding, prompt_embeds], dim=1) + + noise_pred = self.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=encoder_hidden_states, # [1, 257, 4096] + timestep=timestep, + encoder_attention_mask=prompt_attention_mask, + guidance=guidance, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if get_sequence_parallel_state(): + latents = all_gather(latents, dim=2) + + if not output_type == "latent": + latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor + video = self.vae.decode(latents, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + else: + video = latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video, ) + + return HunyuanVideoPipelineOutput(frames=video) diff --git a/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py b/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py new file mode 100644 index 000000000..4249ba971 --- /dev/null +++ b/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py @@ -0,0 +1,361 @@ +import argparse +import os + +import torch +from safetensors.torch import save_file + +parser = argparse.ArgumentParser() +parser.add_argument("--diffusers_path", required=True, type=str) +parser.add_argument("--transformer_path", type=str, default=None, help="Path to save transformer model") +parser.add_argument("--vae_encoder_path", type=str, default=None, help="Path to save VAE encoder model") +parser.add_argument("--vae_decoder_path", type=str, default=None, help="Path to save VAE decoder model") + +args = parser.parse_args() + + +def reverse_scale_shift(weight, dim): + scale, shift = weight.chunk(2, dim=0) + new_weight = torch.cat([shift, scale], dim=0) + return new_weight + + +def reverse_proj_gate(weight): + gate, proj = weight.chunk(2, dim=0) + new_weight = torch.cat([proj, gate], dim=0) + return new_weight + + +def convert_diffusers_transformer_to_mochi(state_dict): + original_state_dict = state_dict.copy() + new_state_dict = {} + + # Convert patch_embed + new_state_dict["x_embedder.proj.weight"] = original_state_dict.pop("patch_embed.proj.weight") + new_state_dict["x_embedder.proj.bias"] = original_state_dict.pop("patch_embed.proj.bias") + + # Convert time_embed + new_state_dict["t_embedder.mlp.0.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.weight") + new_state_dict["t_embedder.mlp.0.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.bias") + new_state_dict["t_embedder.mlp.2.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.weight") + new_state_dict["t_embedder.mlp.2.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.bias") + new_state_dict["t5_y_embedder.to_kv.weight"] = original_state_dict.pop("time_embed.pooler.to_kv.weight") + new_state_dict["t5_y_embedder.to_kv.bias"] = original_state_dict.pop("time_embed.pooler.to_kv.bias") + new_state_dict["t5_y_embedder.to_q.weight"] = original_state_dict.pop("time_embed.pooler.to_q.weight") + new_state_dict["t5_y_embedder.to_q.bias"] = original_state_dict.pop("time_embed.pooler.to_q.bias") + new_state_dict["t5_y_embedder.to_out.weight"] = original_state_dict.pop("time_embed.pooler.to_out.weight") + new_state_dict["t5_y_embedder.to_out.bias"] = original_state_dict.pop("time_embed.pooler.to_out.bias") + new_state_dict["t5_yproj.weight"] = original_state_dict.pop("time_embed.caption_proj.weight") + new_state_dict["t5_yproj.bias"] = original_state_dict.pop("time_embed.caption_proj.bias") + + # Convert transformer blocks + num_layers = 48 + for i in range(num_layers): + block_prefix = f"transformer_blocks.{i}." + new_prefix = f"blocks.{i}." + + # norm1 + new_state_dict[new_prefix + "mod_x.weight"] = original_state_dict.pop(block_prefix + "norm1.linear.weight") + new_state_dict[new_prefix + "mod_x.bias"] = original_state_dict.pop(block_prefix + "norm1.linear.bias") + + if i < num_layers - 1: + new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(block_prefix + + "norm1_context.linear.weight") + new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(block_prefix + + "norm1_context.linear.bias") + else: + new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(block_prefix + + "norm1_context.linear_1.weight") + new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(block_prefix + + "norm1_context.linear_1.bias") + + # Visual attention + q = original_state_dict.pop(block_prefix + "attn1.to_q.weight") + k = original_state_dict.pop(block_prefix + "attn1.to_k.weight") + v = original_state_dict.pop(block_prefix + "attn1.to_v.weight") + qkv_weight = torch.cat([q, k, v], dim=0) + new_state_dict[new_prefix + "attn.qkv_x.weight"] = qkv_weight + + new_state_dict[new_prefix + "attn.q_norm_x.weight"] = original_state_dict.pop(block_prefix + + "attn1.norm_q.weight") + new_state_dict[new_prefix + "attn.k_norm_x.weight"] = original_state_dict.pop(block_prefix + + "attn1.norm_k.weight") + new_state_dict[new_prefix + "attn.proj_x.weight"] = original_state_dict.pop(block_prefix + + "attn1.to_out.0.weight") + new_state_dict[new_prefix + "attn.proj_x.bias"] = original_state_dict.pop(block_prefix + "attn1.to_out.0.bias") + + # Context attention + q = original_state_dict.pop(block_prefix + "attn1.add_q_proj.weight") + k = original_state_dict.pop(block_prefix + "attn1.add_k_proj.weight") + v = original_state_dict.pop(block_prefix + "attn1.add_v_proj.weight") + qkv_weight = torch.cat([q, k, v], dim=0) + new_state_dict[new_prefix + "attn.qkv_y.weight"] = qkv_weight + + new_state_dict[new_prefix + "attn.q_norm_y.weight"] = original_state_dict.pop(block_prefix + + "attn1.norm_added_q.weight") + new_state_dict[new_prefix + "attn.k_norm_y.weight"] = original_state_dict.pop(block_prefix + + "attn1.norm_added_k.weight") + if i < num_layers - 1: + new_state_dict[new_prefix + "attn.proj_y.weight"] = original_state_dict.pop(block_prefix + + "attn1.to_add_out.weight") + new_state_dict[new_prefix + "attn.proj_y.bias"] = original_state_dict.pop(block_prefix + + "attn1.to_add_out.bias") + + # MLP + new_state_dict[new_prefix + "mlp_x.w1.weight"] = reverse_proj_gate( + original_state_dict.pop(block_prefix + "ff.net.0.proj.weight")) + new_state_dict[new_prefix + "mlp_x.w2.weight"] = original_state_dict.pop(block_prefix + "ff.net.2.weight") + if i < num_layers - 1: + new_state_dict[new_prefix + "mlp_y.w1.weight"] = reverse_proj_gate( + original_state_dict.pop(block_prefix + "ff_context.net.0.proj.weight")) + new_state_dict[new_prefix + "mlp_y.w2.weight"] = original_state_dict.pop(block_prefix + + "ff_context.net.2.weight") + + # Output layers + new_state_dict["final_layer.mod.weight"] = reverse_scale_shift(original_state_dict.pop("norm_out.linear.weight"), + dim=0) + new_state_dict["final_layer.mod.bias"] = reverse_scale_shift(original_state_dict.pop("norm_out.linear.bias"), dim=0) + new_state_dict["final_layer.linear.weight"] = original_state_dict.pop("proj_out.weight") + new_state_dict["final_layer.linear.bias"] = original_state_dict.pop("proj_out.bias") + + new_state_dict["pos_frequencies"] = original_state_dict.pop("pos_frequencies") + + print("Remaining Keys:", original_state_dict.keys()) + + return new_state_dict + + +def convert_diffusers_vae_to_mochi(state_dict): + original_state_dict = state_dict.copy() + encoder_state_dict = {} + decoder_state_dict = {} + + # Convert encoder + prefix = "encoder." + + encoder_state_dict["layers.0.weight"] = original_state_dict.pop(f"{prefix}proj_in.weight") + encoder_state_dict["layers.0.bias"] = original_state_dict.pop(f"{prefix}proj_in.bias") + + # Convert block_in + for i in range(3): + encoder_state_dict[f"layers.{i+1}.stack.0.weight"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight") + encoder_state_dict[f"layers.{i+1}.stack.0.bias"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias") + encoder_state_dict[f"layers.{i+1}.stack.2.weight"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.conv1.conv.weight") + encoder_state_dict[f"layers.{i+1}.stack.2.bias"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.conv1.conv.bias") + encoder_state_dict[f"layers.{i+1}.stack.3.weight"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight") + encoder_state_dict[f"layers.{i+1}.stack.3.bias"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias") + encoder_state_dict[f"layers.{i+1}.stack.5.weight"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.conv2.conv.weight") + encoder_state_dict[f"layers.{i+1}.stack.5.bias"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.conv2.conv.bias") + + # Convert down_blocks + down_block_layers = [3, 4, 6] + for block in range(3): + encoder_state_dict[f"layers.{block+4}.layers.0.weight"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.conv_in.conv.weight") + encoder_state_dict[f"layers.{block+4}.layers.0.bias"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.conv_in.conv.bias") + + for i in range(down_block_layers[block]): + # Convert resnets + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.weight"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.bias"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.weight"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.bias"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.weight"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.bias"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.weight"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.bias"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias") + + # Convert attentions + q = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_q.weight") + k = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_k.weight") + v = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_v.weight") + qkv_weight = torch.cat([q, k, v], dim=0) + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.qkv.weight"] = qkv_weight + + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.weight"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.weight") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.bias"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.bias") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.weight"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.weight") + encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.bias"] = original_state_dict.pop( + f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.bias") + + # Convert block_out + for i in range(3): + encoder_state_dict[f"layers.{i+7}.stack.0.weight"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight") + encoder_state_dict[f"layers.{i+7}.stack.0.bias"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias") + encoder_state_dict[f"layers.{i+7}.stack.2.weight"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.conv1.conv.weight") + encoder_state_dict[f"layers.{i+7}.stack.2.bias"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.conv1.conv.bias") + encoder_state_dict[f"layers.{i+7}.stack.3.weight"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight") + encoder_state_dict[f"layers.{i+7}.stack.3.bias"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias") + encoder_state_dict[f"layers.{i+7}.stack.5.weight"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.conv2.conv.weight") + encoder_state_dict[f"layers.{i+7}.stack.5.bias"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.conv2.conv.bias") + + q = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_q.weight") + k = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_k.weight") + v = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_v.weight") + qkv_weight = torch.cat([q, k, v], dim=0) + encoder_state_dict[f"layers.{i+7}.attn_block.attn.qkv.weight"] = qkv_weight + + encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.weight"] = original_state_dict.pop( + f"{prefix}block_out.attentions.{i}.to_out.0.weight") + encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.bias"] = original_state_dict.pop( + f"{prefix}block_out.attentions.{i}.to_out.0.bias") + encoder_state_dict[f"layers.{i+7}.attn_block.norm.weight"] = original_state_dict.pop( + f"{prefix}block_out.norms.{i}.norm_layer.weight") + encoder_state_dict[f"layers.{i+7}.attn_block.norm.bias"] = original_state_dict.pop( + f"{prefix}block_out.norms.{i}.norm_layer.bias") + + # Convert output layers + encoder_state_dict["output_norm.weight"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.weight") + encoder_state_dict["output_norm.bias"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.bias") + encoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight") + + # Convert decoder + prefix = "decoder." + + decoder_state_dict["blocks.0.0.weight"] = original_state_dict.pop(f"{prefix}conv_in.weight") + decoder_state_dict["blocks.0.0.bias"] = original_state_dict.pop(f"{prefix}conv_in.bias") + + # Convert block_in + for i in range(3): + decoder_state_dict[f"blocks.0.{i+1}.stack.0.weight"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight") + decoder_state_dict[f"blocks.0.{i+1}.stack.0.bias"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias") + decoder_state_dict[f"blocks.0.{i+1}.stack.2.weight"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.conv1.conv.weight") + decoder_state_dict[f"blocks.0.{i+1}.stack.2.bias"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.conv1.conv.bias") + decoder_state_dict[f"blocks.0.{i+1}.stack.3.weight"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight") + decoder_state_dict[f"blocks.0.{i+1}.stack.3.bias"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias") + decoder_state_dict[f"blocks.0.{i+1}.stack.5.weight"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.conv2.conv.weight") + decoder_state_dict[f"blocks.0.{i+1}.stack.5.bias"] = original_state_dict.pop( + f"{prefix}block_in.resnets.{i}.conv2.conv.bias") + + # Convert up_blocks + up_block_layers = [6, 4, 3] + for block in range(3): + for i in range(up_block_layers[block]): + decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.weight"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight") + decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.bias"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias") + decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.weight"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight") + decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.bias"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias") + decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.weight"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight") + decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.bias"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias") + decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.weight"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight") + decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.bias"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias") + decoder_state_dict[f"blocks.{block+1}.proj.weight"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.proj.weight") + decoder_state_dict[f"blocks.{block+1}.proj.bias"] = original_state_dict.pop( + f"{prefix}up_blocks.{block}.proj.bias") + + # Convert block_out + for i in range(3): + decoder_state_dict[f"blocks.4.{i}.stack.0.weight"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight") + decoder_state_dict[f"blocks.4.{i}.stack.0.bias"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias") + decoder_state_dict[f"blocks.4.{i}.stack.2.weight"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.conv1.conv.weight") + decoder_state_dict[f"blocks.4.{i}.stack.2.bias"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.conv1.conv.bias") + decoder_state_dict[f"blocks.4.{i}.stack.3.weight"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight") + decoder_state_dict[f"blocks.4.{i}.stack.3.bias"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias") + decoder_state_dict[f"blocks.4.{i}.stack.5.weight"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.conv2.conv.weight") + decoder_state_dict[f"blocks.4.{i}.stack.5.bias"] = original_state_dict.pop( + f"{prefix}block_out.resnets.{i}.conv2.conv.bias") + + # Convert output layers + decoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight") + decoder_state_dict["output_proj.bias"] = original_state_dict.pop(f"{prefix}proj_out.bias") + + return encoder_state_dict, decoder_state_dict + + +def ensure_safetensors_extension(path): + if not path.endswith(".safetensors"): + path = path + ".safetensors" + return path + + +def ensure_directory_exists(path): + directory = os.path.dirname(path) + if directory: + os.makedirs(directory, exist_ok=True) + + +def main(args): + from diffusers import MochiPipeline + + pipe = MochiPipeline.from_pretrained(args.diffusers_path) + + if args.transformer_path: + transformer_path = ensure_safetensors_extension(args.transformer_path) + ensure_directory_exists(transformer_path) + + print("Converting transformer model...") + transformer_state_dict = convert_diffusers_transformer_to_mochi(pipe.transformer.state_dict()) + save_file(transformer_state_dict, transformer_path) + print(f"Saved transformer to {transformer_path}") + + if args.vae_encoder_path and args.vae_decoder_path: + encoder_path = ensure_safetensors_extension(args.vae_encoder_path) + decoder_path = ensure_safetensors_extension(args.vae_decoder_path) + + ensure_directory_exists(encoder_path) + ensure_directory_exists(decoder_path) + + print("Converting VAE models...") + encoder_state_dict, decoder_state_dict = convert_diffusers_vae_to_mochi(pipe.vae.state_dict()) + + save_file(encoder_state_dict, encoder_path) + print(f"Saved VAE encoder to {encoder_path}") + + save_file(decoder_state_dict, decoder_path) + print(f"Saved VAE decoder to {decoder_path}") + elif args.vae_encoder_path or args.vae_decoder_path: + print("Warning: Both VAE encoder and decoder paths must be specified to convert VAE models.") + + +if __name__ == "__main__": + main(args) diff --git a/fastvideo/models/mochi_hf/mochi_latents_utils.py b/fastvideo/models/mochi_hf/mochi_latents_utils.py new file mode 100644 index 000000000..c554d56c7 --- /dev/null +++ b/fastvideo/models/mochi_hf/mochi_latents_utils.py @@ -0,0 +1,45 @@ +import torch + +mochi_latents_mean = torch.tensor([ + -0.06730895953510081, + -0.038011381506090416, + -0.07477820912866141, + -0.05565264470995561, + 0.012767231469026969, + -0.04703542746246419, + 0.043896967884726704, + -0.09346305707025976, + -0.09918314763016893, + -0.008729793427399178, + -0.011931556316503654, + -0.0321993391887285, +]).view(1, 12, 1, 1, 1) +mochi_latents_std = torch.tensor([ + 0.9263795028493863, + 0.9248894543193766, + 0.9393059390890617, + 0.959253732819592, + 0.8244560132752793, + 0.917259975397747, + 0.9294154431013696, + 1.3720942357788521, + 0.881393668867029, + 0.9168315692124348, + 0.9185249279345552, + 0.9274757570805041, +]).view(1, 12, 1, 1, 1) +mochi_scaling_factor = 1.0 + + +def normalize_dit_input(model_type, latents): + if model_type == "mochi": + latents_mean = mochi_latents_mean.to(latents.device, latents.dtype) + latents_std = mochi_latents_std.to(latents.device, latents.dtype) + latents = (latents - latents_mean) / latents_std + return latents + elif model_type == "hunyuan_hf": + return latents * 0.476986 + elif model_type == "hunyuan": + return latents * 0.476986 + else: + raise NotImplementedError(f"model_type {model_type} not supported") diff --git a/fastvideo/models/mochi_hf/modeling_mochi.py b/fastvideo/models/mochi_hf/modeling_mochi.py new file mode 100644 index 000000000..b3da2ec8c --- /dev/null +++ b/fastvideo/models/mochi_hf/modeling_mochi.py @@ -0,0 +1,663 @@ +# Copyright 2024 The Genmo 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, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import PeftAdapterMixin +from diffusers.models.attention import FeedForward as HF_FeedForward +from diffusers.models.attention_processor import Attention +from diffusers.models.embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import AdaLayerNormContinuous +from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers +from diffusers.utils.torch_utils import maybe_allow_in_graph +from liger_kernel.ops.swiglu import LigerSiLUMulFunction + +from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad +from fastvideo.models.mochi_hf.norm import (MochiLayerNormContinuous, MochiModulatedRMSNorm, MochiRMSNorm, + MochiRMSNormZero) +from fastvideo.utils.communications import all_gather, all_to_all_4D +from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class FeedForward(HF_FeedForward): + + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = "geglu", + final_dropout: bool = False, + inner_dim=None, + bias: bool = True, + ): + super().__init__(dim, dim_out, mult, dropout, activation_fn, final_dropout, inner_dim, bias) + assert activation_fn == "swiglu" + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.net[0].proj(hidden_states) + hidden_states, gate = hidden_states.chunk(2, dim=-1) + + return self.net[2](LigerSiLUMulFunction.apply(gate, hidden_states)) + + +class MochiAttention(nn.Module): + + def __init__( + self, + query_dim: int, + processor: "MochiAttnProcessor2_0", + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + added_kv_proj_dim: Optional[int] = None, + added_proj_bias: Optional[bool] = True, + out_dim: int = None, + out_context_dim: int = None, + out_bias: bool = True, + context_pre_only: bool = False, + eps: float = 1e-5, + ): + super().__init__() + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.out_dim = out_dim if out_dim is not None else query_dim + self.out_context_dim = out_context_dim if out_context_dim else query_dim + self.context_pre_only = context_pre_only + + self.heads = out_dim // dim_head if out_dim is not None else heads + + self.norm_q = MochiRMSNorm(dim_head, eps) + self.norm_k = MochiRMSNorm(dim_head, eps) + self.norm_added_q = MochiRMSNorm(dim_head, eps) + self.norm_added_k = MochiRMSNorm(dim_head, eps) + + self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias) + + self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + if self.context_pre_only is not None: + self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(nn.Dropout(dropout)) + + if not self.context_pre_only: + self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias) + + self.processor = processor + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + **kwargs, + ): + return self.processor( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + + +class MochiAttnProcessor2_0: + """Attention processor used in Mochi.""" + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # [b, s, h * d] + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + # [b, s, h=24, d=128] + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + # [b, 256, h * d] + encoder_query = attn.add_q_proj(encoder_hidden_states) + encoder_key = attn.add_k_proj(encoder_hidden_states) + encoder_value = attn.add_v_proj(encoder_hidden_states) + + # [b, 256, h=24, d=128] + encoder_query = encoder_query.unflatten(2, (attn.heads, -1)) + encoder_key = encoder_key.unflatten(2, (attn.heads, -1)) + encoder_value = encoder_value.unflatten(2, (attn.heads, -1)) + + if attn.norm_added_q is not None: + encoder_query = attn.norm_added_q(encoder_query) + if attn.norm_added_k is not None: + encoder_key = attn.norm_added_k(encoder_key) + + if image_rotary_emb is not None: + freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1] + # shard the head dimension + if get_sequence_parallel_state(): + # B, S, H, D to (S, B,) H, D + # batch_size, seq_len, attn_heads, head_dim + query = all_to_all_4D(query, scatter_dim=2, gather_dim=1) + key = all_to_all_4D(key, scatter_dim=2, gather_dim=1) + value = all_to_all_4D(value, scatter_dim=2, gather_dim=1) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) + + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + if image_rotary_emb is not None: + freqs_cos = shrink_head(freqs_cos, dim=1) + freqs_sin = shrink_head(freqs_sin, dim=1) + + if image_rotary_emb is not None: + + def apply_rotary_emb(x, freqs_cos, freqs_sin): + x_even = x[..., 0::2].float() + x_odd = x[..., 1::2].float() + cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype) + sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype) + + return torch.stack([cos, sin], dim=-1).flatten(-2) + + query = apply_rotary_emb(query, freqs_cos, freqs_sin) + key = apply_rotary_emb(key, freqs_cos, freqs_sin) + + # query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) + # encoder_query, encoder_key, encoder_value = ( + # encoder_query.transpose(1, 2), + # encoder_key.transpose(1, 2), + # encoder_value.transpose(1, 2), + # ) + # [b, s, h, d] + sequence_length = query.size(1) + encoder_sequence_length = encoder_query.size(1) + + # H + query = torch.cat([query, encoder_query], dim=1).unsqueeze(2) + key = torch.cat([key, encoder_key], dim=1).unsqueeze(2) + value = torch.cat([value, encoder_value], dim=1).unsqueeze(2) + # B, S, 3, H, D + qkv = torch.cat([query, key, value], dim=2) + + attn_mask = encoder_attention_mask[:, :].bool() + attn_mask = F.pad(attn_mask, (sequence_length, 0), value=True) + hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) + + # hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask = None, dropout_p=0.0, is_causal=False) + + # valid_lengths = encoder_attention_mask.sum(dim=1) + sequence_length + # def no_padding_mask(score, b, h, q_idx, kv_idx): + # return torch.where(kv_idx < valid_lengths[b],score, -float("inf")) + + # hidden_states = flex_attention(query, key, value, score_mod=no_padding_mask) + if get_sequence_parallel_state(): + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( + (sequence_length, encoder_sequence_length), dim=1) + # B, S, H, D + hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) + encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous() + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.flatten(2, 3) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + else: + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( + (sequence_length, encoder_sequence_length), dim=1) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + if hasattr(attn, "to_add_out"): + encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + + return hidden_states, encoder_hidden_states + + +@maybe_allow_in_graph +class MochiTransformerBlock(nn.Module): + r""" + Transformer block used in [Mochi](https://huggingface.co/genmo/mochi-1-preview). + + Args: + dim (`int`): + The number of channels in the input and output. + num_attention_heads (`int`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`): + The number of channels in each head. + qk_norm (`str`, defaults to `"rms_norm"`): + The normalization layer to use. + activation_fn (`str`, defaults to `"swiglu"`): + Activation function to use in feed-forward. + context_pre_only (`bool`, defaults to `False`): + Whether or not to process context-related conditions with additional layers. + eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + pooled_projection_dim: int, + qk_norm: str = "rms_norm", + activation_fn: str = "swiglu", + context_pre_only: bool = False, + eps: float = 1e-6, + ) -> None: + super().__init__() + + self.context_pre_only = context_pre_only + self.ff_inner_dim = (4 * dim * 2) // 3 + self.ff_context_inner_dim = (4 * pooled_projection_dim * 2) // 3 + + self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False) + + if not context_pre_only: + self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False) + else: + self.norm1_context = MochiLayerNormContinuous( + embedding_dim=pooled_projection_dim, + conditioning_embedding_dim=dim, + eps=eps, + ) + + self.attn1 = MochiAttention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + bias=False, + added_kv_proj_dim=pooled_projection_dim, + added_proj_bias=False, + out_dim=dim, + out_context_dim=pooled_projection_dim, + context_pre_only=context_pre_only, + processor=MochiAttnProcessor2_0(), + eps=1e-5, + ) + + # TODO(aryan): norm_context layers are not needed when `context_pre_only` is True + self.norm2 = MochiModulatedRMSNorm(eps=eps) + self.norm2_context = (MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None) + + self.norm3 = MochiModulatedRMSNorm(eps) + self.norm3_context = (MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None) + + self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False) + self.ff_context = None + if not context_pre_only: + self.ff_context = FeedForward( + pooled_projection_dim, + inner_dim=self.ff_context_inner_dim, + activation_fn=activation_fn, + bias=False, + ) + + self.norm4 = MochiModulatedRMSNorm(eps=eps) + self.norm4_context = MochiModulatedRMSNorm(eps=eps) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: Optional[torch.Tensor] = None, + output_attn=False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) + + if not self.context_pre_only: + ( + norm_encoder_hidden_states, + enc_gate_msa, + enc_scale_mlp, + enc_gate_mlp, + ) = self.norm1_context(encoder_hidden_states, temb) + else: + norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb) + + attn_hidden_states, context_attn_hidden_states = self.attn1( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + ) + + hidden_states = hidden_states + self.norm2(attn_hidden_states, torch.tanh(gate_msa).unsqueeze(1)) + norm_hidden_states = self.norm3(hidden_states, (1 + scale_mlp.unsqueeze(1).to(torch.float32))) + ff_output = self.ff(norm_hidden_states) + hidden_states = hidden_states + self.norm4(ff_output, torch.tanh(gate_mlp).unsqueeze(1)) + + if not self.context_pre_only: + encoder_hidden_states = encoder_hidden_states + self.norm2_context(context_attn_hidden_states, + torch.tanh(enc_gate_msa).unsqueeze(1)) + norm_encoder_hidden_states = self.norm3_context( + encoder_hidden_states, + (1 + enc_scale_mlp.unsqueeze(1).to(torch.float32)), + ) + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states + self.norm4_context(context_ff_output, + torch.tanh(enc_gate_mlp).unsqueeze(1)) + + if not output_attn: + attn_hidden_states = None + return hidden_states, encoder_hidden_states, attn_hidden_states + + +class MochiRoPE(nn.Module): + r""" + RoPE implementation used in [Mochi](https://huggingface.co/genmo/mochi-1-preview). + + Args: + base_height (`int`, defaults to `192`): + Base height used to compute interpolation scale for rotary positional embeddings. + base_width (`int`, defaults to `192`): + Base width used to compute interpolation scale for rotary positional embeddings. + """ + + def __init__(self, base_height: int = 192, base_width: int = 192) -> None: + super().__init__() + + self.target_area = base_height * base_width + + def _centers(self, start, stop, num, device, dtype) -> torch.Tensor: + edges = torch.linspace(start, stop, num + 1, device=device, dtype=dtype) + return (edges[:-1] + edges[1:]) / 2 + + def _get_positions( + self, + num_frames: int, + height: int, + width: int, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + scale = (self.target_area / (height * width))**0.5 + t = torch.arange(num_frames * nccl_info.sp_size, device=device, dtype=dtype) + h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype) + w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype) + + grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij") + + positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3) + return positions + + def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: + with torch.autocast(freqs.device.type, enabled=False): + # Always run ROPE freqs computation in FP32 + freqs = torch.einsum( + "nd,dhf->nhf", # codespell:ignore + pos.to(torch.float32), # codespell:ignore + freqs.to(torch.float32)) + freqs_cos = torch.cos(freqs) + freqs_sin = torch.sin(freqs) + return freqs_cos, freqs_sin + + def forward( + self, + pos_frequencies: torch.Tensor, + num_frames: int, + height: int, + width: int, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + pos = self._get_positions(num_frames, height, width, device, dtype) + rope_cos, rope_sin = self._create_rope(pos_frequencies, pos) + return rope_cos, rope_sin + + +@maybe_allow_in_graph +class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): + r""" + A Transformer model for video-like data introduced in [Mochi](https://huggingface.co/genmo/mochi-1-preview). + + Args: + patch_size (`int`, defaults to `2`): + The size of the patches to use in the patch embedding layer. + num_attention_heads (`int`, defaults to `24`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, defaults to `128`): + The number of channels in each head. + num_layers (`int`, defaults to `48`): + The number of layers of Transformer blocks to use. + in_channels (`int`, defaults to `12`): + The number of channels in the input. + out_channels (`int`, *optional*, defaults to `None`): + The number of channels in the output. + qk_norm (`str`, defaults to `"rms_norm"`): + The normalization layer to use. + text_embed_dim (`int`, defaults to `4096`): + Input dimension of text embeddings from the text encoder. + time_embed_dim (`int`, defaults to `256`): + Output dimension of timestep embeddings. + activation_fn (`str`, defaults to `"swiglu"`): + Activation function to use in feed-forward. + max_sequence_length (`int`, defaults to `256`): + The maximum sequence length of text embeddings supported. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + patch_size: int = 2, + num_attention_heads: int = 24, + attention_head_dim: int = 128, + num_layers: int = 48, + pooled_projection_dim: int = 1536, + in_channels: int = 12, + out_channels: Optional[int] = None, + qk_norm: str = "rms_norm", + text_embed_dim: int = 4096, + time_embed_dim: int = 256, + activation_fn: str = "swiglu", + max_sequence_length: int = 256, + ) -> None: + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels or in_channels + + self.patch_embed = PatchEmbed( + patch_size=patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + pos_embed_type=None, + ) + + self.time_embed = MochiCombinedTimestepCaptionEmbedding( + embedding_dim=inner_dim, + pooled_projection_dim=pooled_projection_dim, + text_embed_dim=text_embed_dim, + time_embed_dim=time_embed_dim, + num_attention_heads=8, + ) + + self.pos_frequencies = nn.Parameter(torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0)) + self.rope = MochiRoPE() + + self.transformer_blocks = nn.ModuleList([ + MochiTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + pooled_projection_dim=pooled_projection_dim, + qk_norm=qk_norm, + activation_fn=activation_fn, + context_pre_only=i == num_layers - 1, + ) for i in range(num_layers) + ]) + + self.norm_out = AdaLayerNormContinuous( + inner_dim, + inner_dim, + elementwise_affine=False, + eps=1e-6, + norm_type="layer_norm", + ) + self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels) + + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_attention_mask: torch.Tensor, + output_features=False, + output_features_stride=8, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = False, + ) -> torch.Tensor: + assert (return_dict is False), "return_dict is not supported in MochiTransformer3DModel" + + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + lora_scale = attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + else: + if (attention_kwargs is not None and attention_kwargs.get("scale", None) is not None): + logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.") + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p = self.config.patch_size + + post_patch_height = height // p + post_patch_width = width // p + # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux + timestep = 1000 - timestep + temb, encoder_hidden_states = self.time_embed( + timestep, + encoder_hidden_states, + encoder_attention_mask, + hidden_dtype=hidden_states.dtype, + ) + + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) + hidden_states = self.patch_embed(hidden_states) + hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2) + + image_rotary_emb = self.rope( + self.pos_frequencies, + num_frames, + post_patch_height, + post_patch_width, + device=hidden_states.device, + dtype=torch.float32, + ) + attn_outputs_list = [] + for i, block in enumerate(self.transformer_blocks): + if self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = ({"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}) + ( + hidden_states, + encoder_hidden_states, + attn_outputs, + ) = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + encoder_attention_mask, + temb, + image_rotary_emb, + output_features, + **ckpt_kwargs, + ) + else: + hidden_states, encoder_hidden_states, attn_outputs = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + temb=temb, + image_rotary_emb=image_rotary_emb, + output_attn=output_features, + ) + if i % output_features_stride == 0: + attn_outputs_list.append(attn_outputs) + + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1) + hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5) + output = hidden_states.reshape(batch_size, -1, num_frames, height, width) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not output_features: + attn_outputs_list = None + else: + attn_outputs_list = torch.stack(attn_outputs_list, dim=0) + # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux + return (-output, attn_outputs_list) diff --git a/fastvideo/models/mochi_hf/norm.py b/fastvideo/models/mochi_hf/norm.py new file mode 100644 index 000000000..d91bb95f1 --- /dev/null +++ b/fastvideo/models/mochi_hf/norm.py @@ -0,0 +1,128 @@ +# Copyright 2024 The Genmo 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 Tuple + +import torch +import torch.nn as nn + + +class MochiModulatedRMSNorm(nn.Module): + + def __init__(self, eps: float): + super().__init__() + + self.eps = eps + + def forward(self, hidden_states, scale=None): + hidden_states_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + if scale is not None: + hidden_states = hidden_states * scale + + hidden_states = hidden_states.to(hidden_states_dtype) + + return hidden_states + + +class MochiRMSNorm(nn.Module): + + def __init__(self, dim, eps: float, elementwise_affine=True): + super().__init__() + + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim)) + else: + self.weight = None + + def forward(self, hidden_states): + hidden_states_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + if self.weight is not None: + # convert into half-precision if necessary + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + hidden_states = hidden_states * self.weight + hidden_states = hidden_states.to(hidden_states_dtype) + + return hidden_states + + +class MochiLayerNormContinuous(nn.Module): + + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + eps=1e-5, + bias=True, + ): + super().__init__() + + # AdaLN + self.silu = nn.SiLU() + self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) + self.norm = MochiModulatedRMSNorm(eps=eps) + + def forward( + self, + x: torch.Tensor, + conditioning_embedding: torch.Tensor, + ) -> torch.Tensor: + input_dtype = x.dtype + + # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) + scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) + x = self.norm(x, (1 + scale.unsqueeze(1).to(torch.float32))) + + return x.to(input_dtype) + + +class MochiRMSNormZero(nn.Module): + r""" + Adaptive RMS Norm used in Mochi. + Parameters: + embedding_dim (`int`): The size of each embedding vector. + """ + + def __init__( + self, + embedding_dim: int, + hidden_dim: int, + eps: float = 1e-5, + elementwise_affine: bool = False, + ) -> None: + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, hidden_dim) + self.norm = MochiModulatedRMSNorm(eps=eps) + + def forward(self, hidden_states: torch.Tensor, + emb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hidden_states_dtype = hidden_states.dtype + + emb = self.linear(self.silu(emb)) + scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) + + hidden_states = self.norm(hidden_states, (1 + scale_msa[:, None].to(torch.float32))) + hidden_states = hidden_states.to(hidden_states_dtype) + + return hidden_states, gate_msa, scale_mlp, gate_mlp diff --git a/fastvideo/models/mochi_hf/pipeline_mochi.py b/fastvideo/models/mochi_hf/pipeline_mochi.py new file mode 100644 index 000000000..f42192407 --- /dev/null +++ b/fastvideo/models/mochi_hf/pipeline_mochi.py @@ -0,0 +1,757 @@ +# Copyright 2024 Black Forest Labs 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. + +import copy +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.loaders import Mochi1LoraLoaderMixin +from diffusers.models.autoencoders import AutoencoderKL +from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from einops import rearrange +from transformers import T5EncoderModel, T5TokenizerFast + +from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel +from fastvideo.utils.communications import all_gather +from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import MochiPipeline + >>> from diffusers.utils import export_to_video + + >>> pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + >>> prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k." + >>> frames = pipe(prompt, num_inference_steps=28, guidance_scale=3.5).frames[0] + >>> export_to_video(frames, "mochi.mp4") + ``` +""" + + +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.16, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# from: https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77 +def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None): + if linear_steps is None: + linear_steps = num_steps // 2 + linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)] + threshold_noise_step_diff = linear_steps - threshold_noise * num_steps + quadratic_steps = num_steps - linear_steps + quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2) + linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2) + const = quadratic_coef * (linear_steps**2) + quadratic_sigma_schedule = [ + quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps) + ] + sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + sigma_schedule = [1.0 - x for x in sigma_schedule] + return sigma_schedule + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler.") + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler.") + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): + r""" + The mochi pipeline for text-to-video generation. + + Reference: https://github.com/genmoai/models + + Args: + transformer ([`MochiTransformer3DModel`]): + Conditional Transformer architecture to denoise the encoded video latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`T5EncoderModel`]): + [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically + the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant. + tokenizer (`CLIPTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). + tokenizer (`T5TokenizerFast`): + Second Tokenizer of class + [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKL, + text_encoder: T5EncoderModel, + tokenizer: T5TokenizerFast, + transformer: MochiTransformer3DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + ) + self.vae_spatial_scale_factor = 8 + self.vae_temporal_scale_factor = 6 + self.patch_size = 2 + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_scale_factor) + self.tokenizer_max_length = (self.tokenizer.model_max_length + if hasattr(self, "tokenizer") and self.tokenizer is not None else 77) + self.default_height = 480 + self.default_width = 848 + + # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds + def _get_t5_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 256, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.bool().to(device) + + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1]) + logger.warning("The following part of your input was truncated because `max_sequence_length` is set to " + f" {max_sequence_length} tokens: {removed_text}") + + prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + return prompt_embeds, prompt_attention_mask + + # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Optional[Union[str, List[str]]] = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + max_sequence_length: int = 256, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = (batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt) + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}.") + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`.") + + ( + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) + + def check_inputs( + self, + prompt, + height, + width, + callback_on_step_end_tensor_inputs=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs + for k in callback_on_step_end_tensor_inputs): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two.") + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.") + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if (negative_prompt_embeds is not None and negative_prompt_attention_mask is None): + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}.") + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}.") + + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.vae.enable_slicing() + + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_slicing() + + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + self.vae.enable_tiling() + + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_tiling() + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + num_frames, + dtype, + device, + generator, + latents=None, + ): + height = height // self.vae_spatial_scale_factor + width = width // self.vae_spatial_scale_factor + num_frames = (num_frames - 1) // self.vae_temporal_scale_factor + 1 + + shape = (batch_size, num_channels_latents, num_frames, height, width) + + if latents is not None: + return latents.to(device=device, dtype=dtype) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators.") + + latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32) + latents = latents.to(dtype) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_frames: int = 19, + num_inference_steps: int = 64, + timesteps: List[int] = None, + guidance_scale: float = 4.5, + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 256, + return_all_states=False, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_frames (`int`, defaults to 16): + The number of video frames to generate + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + guidance_scale (`float`, defaults to `4.5`): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.mochi.MochiPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to `256`): + Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.mochi.MochiPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.mochi.MochiPipelineOutput`] is returned, otherwise a `tuple` + is returned where the first element is a list with the generated images. + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + height = height or self.default_height + width = width or self.default_width + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + height=height, + width=width, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 3. Prepare text embeddings + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + max_sequence_length=max_sequence_length, + device=device, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + prompt_embeds.dtype, + device, + generator, + latents, + ) + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + if get_sequence_parallel_state(): + latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous() + latents = latents[:, :, rank, :, :, :] + + original_noise = copy.deepcopy(latents) + # 5. Prepare timestep + # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77 + threshold_noise = 0.025 + sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise) + sigmas = np.array(sigmas) + # check if of type FlowMatchEulerDiscreteScheduler + if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + ) + else: + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 6. Denoising loop + self._progress_bar_config = {"disable": nccl_info.rank_within_group != 0} + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + latent_model_input = (torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype) + + noise_pred = self.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + timestep=timestep, + encoder_attention_mask=prompt_attention_mask, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + + # Mochi CFG + Sampling runs in FP32 + noise_pred = noise_pred.to(torch.float32) + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0] + latents = latents.to(latents_dtype) + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if get_sequence_parallel_state(): + latents = all_gather(latents, dim=2) + # latents_shape = list(latents.shape) + # full_shape = [latents_shape[0] * world_size] + latents_shape[1:] + # all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device) + # torch.distributed.all_gather_into_tensor(all_latents, latents) + # latents_list = list(all_latents.chunk(world_size, dim=0)) + # latents = torch.cat(latents_list, dim=2) + + if output_type == "latent": + video = latents + else: + # unscale/denormalize the latents + # denormalize with the mean and std if available and not None + has_latents_mean = (hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None) + has_latents_std = (hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None) + if has_latents_mean and has_latents_std: + latents_mean = (torch.tensor(self.vae.config.latents_mean).view(1, 12, 1, 1, + 1).to(latents.device, latents.dtype)) + latents_std = (torch.tensor(self.vae.config.latents_std).view(1, 12, 1, 1, + 1).to(latents.device, latents.dtype)) + latents = (latents * latents_std / self.vae.config.scaling_factor + latents_mean) + else: + latents = latents / self.vae.config.scaling_factor + + video = self.vae.decode(latents, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + if return_all_states: + # Pay extra attention here: + # prompt_embeds with shape torch.Size([2, 256]), where prompt_embeds[1] is the prompt_embeds for the actual prompt + # prompt_embeds[0] is for negative prompt + return original_noise, video, latents, prompt_embeds, prompt_attention_mask + + if not return_dict: + return (video, ) + + return MochiPipelineOutput(frames=video) diff --git a/fastvideo/models/stepvideo/__init__.py b/fastvideo/models/stepvideo/__init__.py new file mode 100644 index 000000000..e919871df --- /dev/null +++ b/fastvideo/models/stepvideo/__init__.py @@ -0,0 +1,7 @@ +import os + +os.environ["NCCL_DEBUG"] = "ERROR" + +from .diffusion.scheduler import * +from .diffusion.video_pipeline import * +from .modules.model import * \ No newline at end of file diff --git a/fastvideo/models/stepvideo/__version__.py b/fastvideo/models/stepvideo/__version__.py new file mode 100644 index 000000000..3f5c4a7d6 --- /dev/null +++ b/fastvideo/models/stepvideo/__version__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/fastvideo/models/stepvideo/config.py b/fastvideo/models/stepvideo/config.py new file mode 100644 index 000000000..d7adb50bd --- /dev/null +++ b/fastvideo/models/stepvideo/config.py @@ -0,0 +1,174 @@ +import argparse + + +def parse_args(namespace=None): + parser = argparse.ArgumentParser(description="StepVideo inference script") + + parser = add_extra_models_args(parser) + parser = add_denoise_schedule_args(parser) + parser = add_inference_args(parser) + parser = add_parallel_args(parser) + + args = parser.parse_args(namespace=namespace) + + return args + + +def add_extra_models_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Extra models args, including vae, text encoders and tokenizers)") + + group.add_argument( + "--vae_url", + type=str, + default='127.0.0.1', + help="vae url.", + ) + group.add_argument( + "--caption_url", + type=str, + default='127.0.0.1', + help="caption url.", + ) + + return parser + + +def add_denoise_schedule_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Denoise schedule args") + + # Flow Matching + group.add_argument( + "--time_shift", + type=float, + default=7.0, + help="Shift factor for flow matching schedulers.", + ) + group.add_argument( + "--flow_reverse", + action="store_true", + help="If reverse, learning/sampling from t=1 -> t=0.", + ) + group.add_argument( + "--flow_solver", + type=str, + default="euler", + help="Solver for flow matching.", + ) + + return parser + + +def add_inference_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Inference args") + + # ======================== Model loads ======================== + group.add_argument( + "--model_dir", + type=str, + default="./ckpts", + help="Root path of all the models, including t2v models and extra models.", + ) + group.add_argument( + "--model_resolution", + type=str, + default="540p", + choices=["540p"], + help="Root path of all the models, including t2v models and extra models.", + ) + group.add_argument( + "--use-cpu-offload", + action="store_true", + help="Use CPU offload for the model load.", + ) + + # ======================== Inference general setting ======================== + group.add_argument( + "--batch_size", + type=int, + default=1, + help="Batch size for inference and evaluation.", + ) + group.add_argument( + "--infer_steps", + type=int, + default=50, + help="Number of denoising steps for inference.", + ) + group.add_argument( + "--save_path", + type=str, + default="./results", + help="Path to save the generated samples.", + ) + group.add_argument( + "--name_suffix", + type=str, + default="", + help="Suffix for the names of saved samples.", + ) + group.add_argument( + "--num_videos", + type=int, + default=1, + help="Number of videos to generate for each prompt.", + ) + # ---sample size--- + group.add_argument( + "--num_frames", + type=int, + default=204, + help="How many frames to sample from a video. ", + ) + group.add_argument( + "--height", + type=int, + default=544, + help="The height of video sample", + ) + group.add_argument( + "--width", + type=int, + default=992, + help="The width of video sample", + ) + # --- prompt --- + group.add_argument( + "--prompt", + type=str, + default=None, + help="Prompt for sampling during evaluation.", + ) + group.add_argument("--seed", type=int, default=1234, help="Seed for evaluation.") + + # Classifier-Free Guidance + group.add_argument("--pos_magic", + type=str, + default="超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。", + help="Positive magic prompt for sampling.") + group.add_argument("--neg_magic", + type=str, + default="画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。", + help="Negative magic prompt for sampling.") + group.add_argument("--cfg_scale", type=float, default=9.0, help="Classifier free guidance scale.") + + return parser + + +def add_parallel_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Parallel args") + + # ======================== Model loads ======================== + group.add_argument( + "--ulysses_degree", + type=int, + default=8, + help="Ulysses degree.", + ) + group.add_argument( + "--ring_degree", + type=int, + default=1, + help="Ulysses degree.", + ) + + return parser diff --git a/fastvideo/models/stepvideo/diffusion/scheduler.py b/fastvideo/models/stepvideo/diffusion/scheduler.py new file mode 100644 index 000000000..1246bbf97 --- /dev/null +++ b/fastvideo/models/stepvideo/diffusion/scheduler.py @@ -0,0 +1,220 @@ +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput, logging + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + reverse (`bool`, defaults to `True`): + Whether to reverse the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + reverse: bool = False, + solver: str = "euler", + device: Union[str, torch.device] = None, + ): + sigmas = torch.linspace(1, 0, num_train_timesteps + 1) + + if not reverse: + sigmas = sigmas.flip(0) + + self.sigmas = sigmas + # the value fed to model + self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32) + + self._step_index = None + self._begin_index = None + + self.device = device + + self.supported_solver = ["euler"] + if solver not in self.supported_solver: + raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}") + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + time_shift: float = 13.0, + device: Union[str, torch.device] = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + """ + device = device or self.device + self.num_inference_steps = num_inference_steps + + sigmas = torch.linspace(1, 0, num_inference_steps + 1, device=device) + sigmas = self.sd3_time_shift(sigmas, time_shift) + + if not self.config.reverse: + sigmas = 1 - sigmas + + self.sigmas = sigmas + self.timesteps = sigmas[:-1] + + # Reset step index + self._step_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + return sample + + def sd3_time_shift(self, t: torch.Tensor, time_shift: float = 13.0): + return (time_shift * t) / (1 + (time_shift - 1) * t) + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + return_dict: bool = False, + ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor)): + raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep."), ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] + + if self.config.solver == "euler": + prev_sample = sample + model_output.to(torch.float32) * dt + else: + raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}") + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return prev_sample + + return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps diff --git a/fastvideo/models/stepvideo/diffusion/video_pipeline.py b/fastvideo/models/stepvideo/diffusion/video_pipeline.py new file mode 100755 index 000000000..264c2c356 --- /dev/null +++ b/fastvideo/models/stepvideo/diffusion/video_pipeline.py @@ -0,0 +1,325 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. + +import asyncio +import pickle +from dataclasses import dataclass +from typing import Dict, List, Optional, Union + +import numpy as np +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.utils import BaseOutput + +from fastvideo.models.stepvideo.diffusion.scheduler import FlowMatchDiscreteScheduler +from fastvideo.models.stepvideo.modules.model import StepVideoModel +from fastvideo.models.stepvideo.utils import VideoProcessor + + +def call_api_gen(url, api, port=8080): + url = f"http://{url}:{port}/{api}-api" + import aiohttp + + async def _fn(samples, *args, **kwargs): + if api == 'vae': + data = { + "samples": samples, + } + elif api == 'caption': + data = { + "prompts": samples, + } + else: + raise Exception(f"Not supported api: {api}...") + + async with aiohttp.ClientSession() as sess: + data_bytes = pickle.dumps(data) + async with sess.get(url, data=data_bytes, timeout=12000) as response: + result = bytearray() + while not response.content.at_eof(): + chunk = await response.content.read(1024) + result += chunk + response_data = pickle.loads(result) + return response_data + + return _fn + + +@dataclass +class StepVideoPipelineOutput(BaseOutput): + video: Union[torch.Tensor, np.ndarray] + + +class StepVideoPipeline(DiffusionPipeline): + r""" + Pipeline for text-to-video generation using StepVideo. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + transformer ([`StepVideoModel`]): + Conditional Transformer to denoise the encoded image latents. + scheduler ([`FlowMatchDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae_url: + remote vae server's url. + caption_url: + remote caption (stepllm and clip) server's url. + """ + + def __init__( + self, + transformer: StepVideoModel, + scheduler: FlowMatchDiscreteScheduler, + vae_url: str = '127.0.0.1', + caption_url: str = '127.0.0.1', + save_path: str = './results', + name_suffix: str = '', + ): + super().__init__() + + self.register_modules( + transformer=transformer, + scheduler=scheduler, + ) + + self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 8 + self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 16 + self.video_processor = VideoProcessor(save_path, name_suffix) + + self.vae_url = vae_url + self.caption_url = caption_url + self.setup_api(self.vae_url, self.caption_url) + + def setup_api(self, vae_url, caption_url): + self.vae_url = vae_url + self.caption_url = caption_url + self.caption = call_api_gen(caption_url, 'caption') + self.vae = call_api_gen(vae_url, 'vae') + return self + + def encode_prompt( + self, + prompt: str, + neg_magic: str = '', + pos_magic: str = '', + ): + device = self._execution_device + prompts = [prompt + pos_magic] + bs = len(prompts) + prompts += [neg_magic] * bs + + data = asyncio.run(self.caption(prompts)) + prompt_embeds, prompt_attention_mask, clip_embedding = data['y'].to(device), data['y_mask'].to( + device), data['clip_embedding'].to(device) + + return prompt_embeds, clip_embedding, prompt_attention_mask + + def decode_vae(self, samples): + samples = asyncio.run(self.vae(samples.cpu())) + return samples + + def check_inputs(self, num_frames, width, height): + num_frames = max(num_frames // 17 * 17, 1) + width = max(width // 16 * 16, 16) + height = max(height // 16 * 16, 16) + return num_frames, width, height + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: 64, + height: int = 544, + width: int = 992, + num_frames: int = 204, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_frames, width, height = self.check_inputs(num_frames, width, height) + shape = ( + batch_size, + max(num_frames // 17 * 3, 1), + num_channels_latents, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) # b,f,c,h,w + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators.") + + if generator is None: + generator = torch.Generator(device=self._execution_device) + + latents = torch.randn(shape, generator=generator, device=device, dtype=dtype) + return latents + + @torch.inference_mode() + def __call__( + self, + prompt: Union[str, List[str]] = None, + height: int = 544, + width: int = 992, + num_frames: int = 204, + num_inference_steps: int = 50, + guidance_scale: float = 9.0, + time_shift: float = 13.0, + neg_magic: str = "", + pos_magic: str = "", + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: Optional[str] = "mp4", + output_file_name: Optional[str] = "", + return_dict: bool = True, + mask_strategy: Optional[Dict[str, list]] = None, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, defaults to `544`): + The height in pixels of the generated image. + width (`int`, defaults to `992`): + The width in pixels of the generated image. + num_frames (`int`, defaults to `204`): + The number of frames in the generated video. + num_inference_steps (`int`, defaults to `50`): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, defaults to `9.0`): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + output_file_name(`str`, *optional*`): + The output mp4 file name. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`StepVideoPipelineOutput`] instead of a plain tuple. + + Examples: + + Returns: + [`~StepVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`StepVideoPipelineOutput`] is returned, otherwise a `tuple` is returned + where the first element is a list with the generated images and the second element is a list of `bool`s + indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. + """ + + # 1. Check inputs. Raise error if not correct + device = self._execution_device + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + prompt_embeds, prompt_embeds_2, prompt_attention_mask = self.encode_prompt( + prompt=prompt, + neg_magic=neg_magic, + pos_magic=pos_magic, + ) + + transformer_dtype = self.transformer.dtype + prompt_embeds = prompt_embeds.to(transformer_dtype) + prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) + prompt_embeds_2 = prompt_embeds_2.to(transformer_dtype) + + # 4. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps=num_inference_steps, time_shift=time_shift, device=device) + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + torch.bfloat16, + device, + generator, + latents, + ) + + def dict_to_3d_list(best_masks, t_max=50, l_max=48, h_max=48): + result = [[[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max)] + if best_masks is None: + return result + for key, value in best_masks.items(): + timestep, layer, head = map(int, key.split('_')) + result[timestep][layer][head] = value + return result + + mask_strategy = dict_to_3d_list(mask_strategy) + + #best_mask_selections = None + # 7. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(self.scheduler.timesteps): + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(transformer_dtype) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype) + + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + encoder_hidden_states_2=prompt_embeds_2, + return_dict=False, + mask_strategy=mask_strategy[i], + ) + # perform guidance + if do_classifier_free_guidance: + noise_pred_text, noise_pred_uncond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(model_output=noise_pred, timestep=t, sample=latents) + + progress_bar.update() + + if not torch.distributed.is_initialized() or int(torch.distributed.get_rank()) == 0: + if not output_type == "latent": + video = self.decode_vae(latents) + video = self.video_processor.postprocess_video(video, + output_file_name=output_file_name, + output_type=output_type) + else: + video = latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video, ) + + return StepVideoPipelineOutput(video=video) diff --git a/fastvideo/models/stepvideo/modules/attentions.py b/fastvideo/models/stepvideo/modules/attentions.py new file mode 100755 index 000000000..1011fae5c --- /dev/null +++ b/fastvideo/models/stepvideo/modules/attentions.py @@ -0,0 +1,96 @@ +import torch +import torch.nn as nn +from einops import rearrange +from flash_attn import flash_attn_func + +try: + from st_attn import sliding_tile_attention +except ImportError: + print("Could not load Sliding Tile Attention.") + sliding_tile_attention = None + +from fastvideo.utils.communications import all_to_all_4D +from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info + + +class Attention(nn.Module): + + def __init__(self): + super().__init__() + + def attn_processor(self, attn_type): + if attn_type == 'torch': + return self.torch_attn_func + elif attn_type == 'parallel': + return self.parallel_attn_func + else: + raise Exception('Not supported attention type...') + + def tile(self, x, sp_size): + x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=36 // sp_size, h=48, w=48) + return rearrange(x, + "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", + n_t=6, + n_h=6, + n_w=6, + ts_t=6, + ts_h=8, + ts_w=8) + + def untile(self, x, sp_size): + x = rearrange(x, + "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d", + n_t=6, + n_h=6, + n_w=6, + ts_t=6, + ts_h=8, + ts_w=8) + return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=36 // sp_size, h=48, w=48) + + def torch_attn_func(self, q, k, v, attn_mask=None, causal=False, drop_rate=0.0, **kwargs): + + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(q.dtype) + + if attn_mask is not None and attn_mask.ndim == 3: ## no head + n_heads = q.shape[2] + attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) + + q, k, v = map(lambda x: rearrange(x, 'b s h d -> b h s d'), (q, k, v)) + x = torch.nn.functional.scaled_dot_product_attention(q, + k, + v, + attn_mask=attn_mask, + dropout_p=drop_rate, + is_causal=causal) + x = rearrange(x, 'b h s d -> b s h d') + return x + + def parallel_attn_func(self, q, k, v, causal=False, mask_strategy=None, **kwargs): + if get_sequence_parallel_state(): + q = all_to_all_4D(q, scatter_dim=2, gather_dim=1) + k = all_to_all_4D(k, scatter_dim=2, gather_dim=1) + v = all_to_all_4D(v, scatter_dim=2, gather_dim=1) + + if mask_strategy[0] is not None: + q = self.tile(q, nccl_info.sp_size).transpose(1, 2).contiguous() + k = self.tile(k, nccl_info.sp_size).transpose(1, 2).contiguous() + v = self.tile(v, nccl_info.sp_size).transpose(1, 2).contiguous() + + head_num = q.size(1) # 48 // sp_size + current_rank = nccl_info.rank_within_group + + start_head = current_rank * head_num + windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)] + + x = sliding_tile_attention(q, k, v, windows, 0, False).transpose(1, 2).contiguous() + x = self.untile(x, nccl_info.sp_size) + else: + x = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False) + + if get_sequence_parallel_state(): + x = all_to_all_4D(x, scatter_dim=1, gather_dim=2) + + x = x.to(q.dtype) + return x diff --git a/fastvideo/models/stepvideo/modules/blocks.py b/fastvideo/models/stepvideo/modules/blocks.py new file mode 100755 index 000000000..befcfec7e --- /dev/null +++ b/fastvideo/models/stepvideo/modules/blocks.py @@ -0,0 +1,296 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + +from fastvideo.models.stepvideo.modules.attentions import Attention +from fastvideo.models.stepvideo.modules.normalization import RMSNorm +from fastvideo.models.stepvideo.modules.rope import RoPE3D + + +class SelfAttention(Attention): + + def __init__(self, hidden_dim, head_dim, bias=False, with_rope=True, with_qk_norm=True, attn_type='torch'): + super().__init__() + self.head_dim = head_dim + self.n_heads = hidden_dim // head_dim + + self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=bias) + self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias) + + self.with_rope = with_rope + self.with_qk_norm = with_qk_norm + if self.with_qk_norm: + self.q_norm = RMSNorm(head_dim, elementwise_affine=True) + self.k_norm = RMSNorm(head_dim, elementwise_affine=True) + + if self.with_rope: + self.rope_3d = RoPE3D(freq=1e4, F0=1.0, scaling_factor=1.0) + self.rope_ch_split = [64, 32, 32] + + self.core_attention = self.attn_processor(attn_type=attn_type) + self.parallel = attn_type == 'parallel' + + def apply_rope3d(self, x, fhw_positions, rope_ch_split, parallel=True): + x = self.rope_3d(x, fhw_positions, rope_ch_split, parallel) + return x + + def forward(self, x, cu_seqlens=None, max_seqlen=None, rope_positions=None, attn_mask=None, mask_strategy=None): + xqkv = self.wqkv(x) + xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim) + + xq, xk, xv = torch.split(xqkv, [self.head_dim] * 3, dim=-1) ## seq_len, n, dim + + if self.with_qk_norm: + xq = self.q_norm(xq) + xk = self.k_norm(xk) + + if self.with_rope: + xq = self.apply_rope3d(xq, rope_positions, self.rope_ch_split, parallel=self.parallel) + xk = self.apply_rope3d(xk, rope_positions, self.rope_ch_split, parallel=self.parallel) + + output = self.core_attention(xq, + xk, + xv, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + attn_mask=attn_mask, + mask_strategy=mask_strategy) + output = rearrange(output, 'b s h d -> b s (h d)') + output = self.wo(output) + + return output + + +class CrossAttention(Attention): + + def __init__(self, hidden_dim, head_dim, bias=False, with_qk_norm=True, attn_type='torch'): + super().__init__() + self.head_dim = head_dim + self.n_heads = hidden_dim // head_dim + + self.wq = nn.Linear(hidden_dim, hidden_dim, bias=bias) + self.wkv = nn.Linear(hidden_dim, hidden_dim * 2, bias=bias) + self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias) + + self.with_qk_norm = with_qk_norm + if self.with_qk_norm: + self.q_norm = RMSNorm(head_dim, elementwise_affine=True) + self.k_norm = RMSNorm(head_dim, elementwise_affine=True) + + self.core_attention = self.attn_processor(attn_type=attn_type) + + def forward(self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, attn_mask=None): + xq = self.wq(x) + xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim) + + xkv = self.wkv(encoder_hidden_states) + xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim) + + xk, xv = torch.split(xkv, [self.head_dim] * 2, dim=-1) ## seq_len, n, dim + + if self.with_qk_norm: + xq = self.q_norm(xq) + xk = self.k_norm(xk) + + output = self.core_attention(xq, xk, xv, attn_mask=attn_mask) + + output = rearrange(output, 'b s h d -> b s (h d)') + output = self.wo(output) + + return output + + +class GELU(nn.Module): + r""" + GELU activation function with tanh approximation support with `approximate="tanh"`. + + Parameters: + dim_in (`int`): The number of channels in the input. + dim_out (`int`): The number of channels in the output. + approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation. + bias (`bool`, defaults to True): Whether to use a bias in the linear layer. + """ + + def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out, bias=bias) + self.approximate = approximate + + def gelu(self, gate: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(gate, approximate=self.approximate) + + def forward(self, hidden_states): + hidden_states = self.proj(hidden_states) + hidden_states = self.gelu(hidden_states) + return hidden_states + + +class FeedForward(nn.Module): + + def __init__( + self, + dim: int, + inner_dim: Optional[int] = None, + dim_out: Optional[int] = None, + mult: int = 4, + bias: bool = False, + ): + super().__init__() + inner_dim = dim * mult if inner_dim is None else inner_dim + dim_out = dim if dim_out is None else dim_out + self.net = nn.ModuleList([ + GELU(dim, inner_dim, approximate="tanh", bias=bias), + nn.Identity(), + nn.Linear(inner_dim, dim_out, bias=bias) + ]) + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +def modulate(x, scale, shift): + x = x * (1 + scale) + shift + return x + + +def gate(x, gate): + x = gate * x + return x + + +class StepVideoTransformerBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + upcast_attention (`bool`, *optional*): + Whether to upcast the attention computation to float32. This is useful for mixed precision training. + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_type (`str`, *optional*, defaults to `"layer_norm"`): + The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + attention_type (`str`, *optional*, defaults to `"default"`): + The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. + positional_embeddings (`str`, *optional*, defaults to `None`): + The type of positional embeddings to apply to. + num_positional_embeddings (`int`, *optional*, defaults to `None`): + The maximum number of positional embeddings to apply. + """ + + def __init__(self, + dim: int, + attention_head_dim: int, + norm_eps: float = 1e-5, + ff_inner_dim: Optional[int] = None, + ff_bias: bool = False, + attention_type: str = 'parallel'): + super().__init__() + self.dim = dim + self.norm1 = nn.LayerNorm(dim, eps=norm_eps) + self.attn1 = SelfAttention(dim, + attention_head_dim, + bias=False, + with_rope=True, + with_qk_norm=True, + attn_type=attention_type) + + self.norm2 = nn.LayerNorm(dim, eps=norm_eps) + self.attn2 = CrossAttention(dim, attention_head_dim, bias=False, with_qk_norm=True, attn_type='torch') + + self.ff = FeedForward(dim=dim, inner_dim=ff_inner_dim, dim_out=dim, bias=ff_bias) + + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + @torch.no_grad() + def forward(self, + q: torch.Tensor, + kv: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + attn_mask=None, + rope_positions: list = None, + mask_strategy=None) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (torch.clone(chunk) for chunk in ( + self.scale_shift_table[None] + timestep.reshape(-1, 6, self.dim)).chunk(6, dim=1)) + + scale_shift_q = modulate(self.norm1(q), scale_msa, shift_msa) + + attn_q = self.attn1(scale_shift_q, rope_positions=rope_positions, mask_strategy=mask_strategy) + + q = gate(attn_q, gate_msa) + q + + attn_q = self.attn2(q, kv, attn_mask) + + q = attn_q + q + + scale_shift_q = modulate(self.norm2(q), scale_mlp, shift_mlp) + + ff_output = self.ff(scale_shift_q) + + q = gate(ff_output, gate_mlp) + q + + return q + + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding""" + + def __init__( + self, + patch_size=64, + in_channels=3, + embed_dim=768, + layer_norm=False, + flatten=True, + bias=True, + ): + super().__init__() + + self.flatten = flatten + self.layer_norm = layer_norm + + self.proj = nn.Conv2d(in_channels, + embed_dim, + kernel_size=(patch_size, patch_size), + stride=patch_size, + bias=bias) + + def forward(self, latent): + latent = self.proj(latent).to(latent.dtype) + if self.flatten: + latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC + if self.layer_norm: + latent = self.norm(latent) + + return latent diff --git a/fastvideo/models/stepvideo/modules/model.py b/fastvideo/models/stepvideo/modules/model.py new file mode 100755 index 000000000..2864403ad --- /dev/null +++ b/fastvideo/models/stepvideo/modules/model.py @@ -0,0 +1,198 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +from typing import Dict, Optional + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from einops import rearrange, repeat +from torch import nn + +from fastvideo.models.stepvideo.modules.blocks import PatchEmbed, StepVideoTransformerBlock +from fastvideo.models.stepvideo.modules.normalization import AdaLayerNormSingle, PixArtAlphaTextProjection +from fastvideo.models.stepvideo.parallel import parallel_forward +from fastvideo.models.stepvideo.utils import with_empty_init + + +class StepVideoModel(ModelMixin, ConfigMixin): + _no_split_modules = ["StepVideoTransformerBlock", "PatchEmbed"] + + @with_empty_init + @register_to_config + def __init__( + self, + num_attention_heads: int = 48, + attention_head_dim: int = 128, + in_channels: int = 64, + out_channels: Optional[int] = 64, + num_layers: int = 48, + dropout: float = 0.0, + patch_size: int = 1, + norm_type: str = "ada_norm_single", + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + use_additional_conditions: Optional[bool] = False, + caption_channels: Optional[int] | list | tuple = [6144, 1024], + attention_type: Optional[str] = "parallel", + ): + super().__init__() + + # Set some common variables used across the board. + self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim + self.out_channels = in_channels if out_channels is None else out_channels + + self.use_additional_conditions = use_additional_conditions + + self.pos_embed = PatchEmbed( + patch_size=patch_size, + in_channels=self.config.in_channels, + embed_dim=self.inner_dim, + ) + + self.transformer_blocks = nn.ModuleList([ + StepVideoTransformerBlock(dim=self.inner_dim, + attention_head_dim=self.config.attention_head_dim, + attention_type=attention_type) for _ in range(self.config.num_layers) + ]) + + # 3. Output blocks. + self.norm_out = nn.LayerNorm(self.inner_dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) + self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels) + self.patch_size = patch_size + + self.adaln_single = AdaLayerNormSingle(self.inner_dim, use_additional_conditions=self.use_additional_conditions) + + if isinstance(self.config.caption_channels, int): + caption_channel = self.config.caption_channels + else: + caption_channel, clip_channel = self.config.caption_channels + self.clip_projection = nn.Linear(clip_channel, self.inner_dim) + + self.caption_norm = nn.LayerNorm(caption_channel, eps=norm_eps, elementwise_affine=norm_elementwise_affine) + + self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channel, hidden_size=self.inner_dim) + + self.parallel = attention_type == 'parallel' + + def patchfy(self, hidden_states): + hidden_states = rearrange(hidden_states, 'b f c h w -> (b f) c h w') + hidden_states = self.pos_embed(hidden_states) + return hidden_states + + def prepare_attn_mask(self, encoder_attention_mask, encoder_hidden_states, q_seqlen): + kv_seqlens = encoder_attention_mask.sum(dim=1).int() + mask = torch.zeros([len(kv_seqlens), q_seqlen, max(kv_seqlens)], + dtype=torch.bool, + device=encoder_attention_mask.device) + encoder_hidden_states = encoder_hidden_states[:, :max(kv_seqlens)] + for i, kv_len in enumerate(kv_seqlens): + mask[i, :, :kv_len] = 1 + return encoder_hidden_states, mask + + @parallel_forward + def block_forward(self, + hidden_states, + encoder_hidden_states=None, + timestep=None, + rope_positions=None, + attn_mask=None, + parallel=True, + mask_strategy=None): + + for i, block in enumerate(self.transformer_blocks): + hidden_states = block(hidden_states, + encoder_hidden_states, + timestep=timestep, + attn_mask=attn_mask, + rope_positions=rope_positions, + mask_strategy=mask_strategy[i]) + + return hidden_states + + @torch.inference_mode() + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_hidden_states_2: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + added_cond_kwargs: Dict[str, torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + fps: torch.Tensor = None, + return_dict: bool = True, + mask_strategy=None, + ): + assert hidden_states.ndim == 5 + "hidden_states's shape should be (bsz, f, ch, h ,w)" + + bsz, frame, _, height, width = hidden_states.shape + height, width = height // self.patch_size, width // self.patch_size + + hidden_states = self.patchfy(hidden_states) + len_frame = hidden_states.shape[1] + + if self.use_additional_conditions: + added_cond_kwargs = { + "resolution": torch.tensor([(height, width)] * bsz, + device=hidden_states.device, + dtype=hidden_states.dtype), + "nframe": torch.tensor([frame] * bsz, device=hidden_states.device, dtype=hidden_states.dtype), + "fps": fps + } + else: + added_cond_kwargs = {} + + timestep, embedded_timestep = self.adaln_single(timestep, added_cond_kwargs=added_cond_kwargs) + + encoder_hidden_states = self.caption_projection(self.caption_norm(encoder_hidden_states)) + + if encoder_hidden_states_2 is not None and hasattr(self, 'clip_projection'): + clip_embedding = self.clip_projection(encoder_hidden_states_2) + encoder_hidden_states = torch.cat([clip_embedding, encoder_hidden_states], dim=1) + + hidden_states = rearrange(hidden_states, '(b f) l d-> b (f l) d', b=bsz, f=frame, l=len_frame).contiguous() + encoder_hidden_states, attn_mask = self.prepare_attn_mask(encoder_attention_mask, + encoder_hidden_states, + q_seqlen=frame * len_frame) + + hidden_states = self.block_forward(hidden_states, + encoder_hidden_states, + timestep=timestep, + rope_positions=[frame, height, width], + attn_mask=attn_mask, + parallel=self.parallel, + mask_strategy=mask_strategy) + + hidden_states = rearrange(hidden_states, 'b (f l) d -> (b f) l d', b=bsz, f=frame, l=len_frame) + + embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d', f=frame).contiguous() + + shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) + # Modulation + hidden_states = hidden_states * (1 + scale) + shift + hidden_states = self.proj_out(hidden_states) + + # unpatchify + hidden_states = hidden_states.reshape(shape=(-1, height, width, self.patch_size, self.patch_size, + self.out_channels)) + + hidden_states = rearrange(hidden_states, 'n h w p q c -> n c h p w q') + output = hidden_states.reshape(shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)) + + output = rearrange(output, '(b f) c h w -> b f c h w', f=frame) + + if return_dict: + return {'x': output} + return output diff --git a/fastvideo/models/stepvideo/modules/normalization.py b/fastvideo/models/stepvideo/modules/normalization.py new file mode 100755 index 000000000..18950be82 --- /dev/null +++ b/fastvideo/models/stepvideo/modules/normalization.py @@ -0,0 +1,312 @@ +import math +from typing import Dict, Optional, Tuple + +import torch +import torch.nn as nn + + +class RMSNorm(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): + """ + 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 + + +ACTIVATION_FUNCTIONS = { + "swish": nn.SiLU(), + "silu": nn.SiLU(), + "mish": nn.Mish(), + "gelu": nn.GELU(), + "relu": nn.ReLU(), +} + + +def get_activation(act_fn: str) -> nn.Module: + """Helper function to get activation function from string. + + Args: + act_fn (str): Name of activation function. + + Returns: + nn.Module: Activation function. + """ + + act_fn = act_fn.lower() + if act_fn in ACTIVATION_FUNCTIONS: + return ACTIVATION_FUNCTIONS[act_fn] + else: + raise ValueError(f"Unsupported activation function: {act_fn}") + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +): + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. + + :param timesteps: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the + embeddings. :return: an [N x dim] Tensor of positional embeddings. + """ + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class Timesteps(nn.Module): + + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + + def forward(self, timesteps): + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + ) + return t_emb + + +class TimestepEmbedding(nn.Module): + + def __init__(self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + post_act_fn: Optional[str] = None, + cond_proj_dim=None, + sample_proj_bias=True): + super().__init__() + linear_cls = nn.Linear + + self.linear_1 = linear_cls( + in_channels, + time_embed_dim, + bias=sample_proj_bias, + ) + + if cond_proj_dim is not None: + self.cond_proj = linear_cls( + cond_proj_dim, + in_channels, + bias=False, + ) + else: + self.cond_proj = None + + self.act = get_activation(act_fn) + + if out_dim is not None: + time_embed_dim_out = out_dim + else: + time_embed_dim_out = time_embed_dim + + self.linear_2 = linear_cls( + time_embed_dim, + time_embed_dim_out, + bias=sample_proj_bias, + ) + + if post_act_fn is None: + self.post_act = None + else: + self.post_act = get_activation(post_act_fn) + + def forward(self, sample, condition=None): + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): + + def __init__(self, embedding_dim, size_emb_dim, use_additional_conditions: bool = False): + super().__init__() + + self.outdim = size_emb_dim + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.use_additional_conditions = use_additional_conditions + if self.use_additional_conditions: + self.additional_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.resolution_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim) + self.nframe_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + self.fps_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timestep, resolution=None, nframe=None, fps=None): + hidden_dtype = next(self.timestep_embedder.parameters()).dtype + + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + + if self.use_additional_conditions: + batch_size = timestep.shape[0] + resolution_emb = self.additional_condition_proj(resolution.flatten()).to(hidden_dtype) + resolution_emb = self.resolution_embedder(resolution_emb).reshape(batch_size, -1) + nframe_emb = self.additional_condition_proj(nframe.flatten()).to(hidden_dtype) + nframe_emb = self.nframe_embedder(nframe_emb).reshape(batch_size, -1) + conditioning = timesteps_emb + resolution_emb + nframe_emb + + if fps is not None: + fps_emb = self.additional_condition_proj(fps.flatten()).to(hidden_dtype) + fps_emb = self.fps_embedder(fps_emb).reshape(batch_size, -1) + conditioning = conditioning + fps_emb + else: + conditioning = timesteps_emb + + return conditioning + + +class AdaLayerNormSingle(nn.Module): + r""" + Norm layer adaptive layer norm single (adaLN-single). + + As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + use_additional_conditions (`bool`): To use additional conditions for normalization or not. + """ + + def __init__(self, embedding_dim: int, use_additional_conditions: bool = False, time_step_rescale=1000): + super().__init__() + + self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim, + size_emb_dim=embedding_dim // 2, + use_additional_conditions=use_additional_conditions) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + self.time_step_rescale = time_step_rescale ## timestep usually in [0, 1], we rescale it to [0,1000] for stability + + def forward( + self, + timestep: torch.Tensor, + added_cond_kwargs: Dict[str, torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + embedded_timestep = self.emb(timestep * self.time_step_rescale, **added_cond_kwargs) + + out = self.linear(self.silu(embedded_timestep)) + + return out, embedded_timestep + + +class PixArtAlphaTextProjection(nn.Module): + """ + Projects caption embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_features, hidden_size): + super().__init__() + self.linear_1 = nn.Linear( + in_features, + hidden_size, + bias=True, + ) + self.act_1 = nn.GELU(approximate="tanh") + self.linear_2 = nn.Linear( + hidden_size, + hidden_size, + bias=True, + ) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states diff --git a/fastvideo/models/stepvideo/modules/rope.py b/fastvideo/models/stepvideo/modules/rope.py new file mode 100755 index 000000000..c9b094d66 --- /dev/null +++ b/fastvideo/models/stepvideo/modules/rope.py @@ -0,0 +1,90 @@ +import torch + +from fastvideo.utils.parallel_states import nccl_info + + +class RoPE1D: + + def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0): + self.base = freq + self.F0 = F0 + self.scaling_factor = scaling_factor + self.cache = {} + + def get_cos_sin(self, D, seq_len, device, dtype): + if (D, seq_len, device, dtype) not in self.cache: + inv_freq = 1.0 / (self.base**(torch.arange(0, D, 2).float().to(device) / D)) + t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) + freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype) + freqs = torch.cat((freqs, freqs), dim=-1) + cos = freqs.cos() # (Seq, Dim) + sin = freqs.sin() + self.cache[D, seq_len, device, dtype] = (cos, sin) + return self.cache[D, seq_len, device, dtype] + + @staticmethod + def rotate_half(x): + x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:] + return torch.cat((-x2, x1), dim=-1) + + def apply_rope1d(self, tokens, pos1d, cos, sin): + assert pos1d.ndim == 2 + cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :] + sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :] + return (tokens * cos) + (self.rotate_half(tokens) * sin) + + def __call__(self, tokens, positions): + """ + input: + * tokens: batch_size x ntokens x nheads x dim + * positions: batch_size x ntokens (t position of each token) + output: + * tokens after applying RoPE2D (batch_size x ntokens x nheads x dim) + """ + D = tokens.size(3) + assert positions.ndim == 2 # Batch, Seq + cos, sin = self.get_cos_sin(D, int(positions.max()) + 1, tokens.device, tokens.dtype) + tokens = self.apply_rope1d(tokens, positions, cos, sin) + return tokens + + +class RoPE3D(RoPE1D): + + def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0): + super(RoPE3D, self).__init__(freq, F0, scaling_factor) + self.position_cache = {} + + def get_mesh_3d(self, rope_positions, bsz): + f, h, w = rope_positions + + if f"{f}-{h}-{w}" not in self.position_cache: + x = torch.arange(f, device='cpu') + y = torch.arange(h, device='cpu') + z = torch.arange(w, device='cpu') + self.position_cache[f"{f}-{h}-{w}"] = torch.cartesian_prod(x, y, z).view(1, f * h * w, 3).expand(bsz, -1, 3) + return self.position_cache[f"{f}-{h}-{w}"] + + def __call__(self, tokens, rope_positions, ch_split, parallel=False): + """ + input: + * tokens: batch_size x ntokens x nheads x dim + * rope_positions: list of (f, h, w) + output: + * tokens after applying RoPE2D (batch_size x ntokens x nheads x dim) + """ + assert sum(ch_split) == tokens.size(-1) + + mesh_grid = self.get_mesh_3d(rope_positions, bsz=tokens.shape[0]) + out = [] + for i, (D, x) in enumerate(zip(ch_split, torch.split(tokens, ch_split, dim=-1))): + cos, sin = self.get_cos_sin(D, int(mesh_grid.max()) + 1, tokens.device, tokens.dtype) + + if parallel: + mesh = torch.chunk(mesh_grid[:, :, i], nccl_info.sp_size, dim=1)[nccl_info.rank_within_group].clone() + else: + mesh = mesh_grid[:, :, i].clone() + x = self.apply_rope1d(x, mesh.to(tokens.device), cos, sin) + out.append(x) + + tokens = torch.cat(out, dim=-1) + return tokens diff --git a/fastvideo/models/stepvideo/parallel.py b/fastvideo/models/stepvideo/parallel.py new file mode 100644 index 000000000..9ce6ee3ce --- /dev/null +++ b/fastvideo/models/stepvideo/parallel.py @@ -0,0 +1,21 @@ +import torch + +from fastvideo.utils.communications import all_gather +from fastvideo.utils.parallel_states import nccl_info + + +def parallel_forward(fn_): + + def wrapTheFunction(_, hidden_states, *args, **kwargs): + if kwargs['parallel']: + hidden_states = torch.chunk(hidden_states, nccl_info.sp_size, dim=-2)[nccl_info.rank_within_group] + kwargs['attn_mask'] = torch.chunk(kwargs['attn_mask'], nccl_info.sp_size, + dim=-2)[nccl_info.rank_within_group] + output = fn_(_, hidden_states, *args, **kwargs) + + if kwargs['parallel']: + output = all_gather(output.contiguous(), dim=-2) + + return output + + return wrapTheFunction diff --git a/fastvideo/models/stepvideo/text_encoder/__init__.py b/fastvideo/models/stepvideo/text_encoder/__init__.py new file mode 100644 index 000000000..bb56aad46 --- /dev/null +++ b/fastvideo/models/stepvideo/text_encoder/__init__.py @@ -0,0 +1,12 @@ +import os + +import torch + +from fastvideo.models.stepvideo.config import parse_args + +try: + args = parse_args() + torch.ops.load_library( + os.path.join(args.model_dir, 'lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so')) +except Exception as err: + print(err) diff --git a/fastvideo/models/stepvideo/text_encoder/clip.py b/fastvideo/models/stepvideo/text_encoder/clip.py new file mode 100755 index 000000000..8e0386eef --- /dev/null +++ b/fastvideo/models/stepvideo/text_encoder/clip.py @@ -0,0 +1,36 @@ +import os + +import torch +import torch.nn as nn +from transformers import BertModel, BertTokenizer + + +class HunyuanClip(nn.Module): + """ + Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py + hunyuan's clip used BertModel and BertTokenizer, so we copy it. + """ + + def __init__(self, model_dir, max_length=77): + super(HunyuanClip, self).__init__() + + self.max_length = max_length + self.tokenizer = BertTokenizer.from_pretrained(os.path.join(model_dir, 'tokenizer')) + self.text_encoder = BertModel.from_pretrained(os.path.join(model_dir, 'clip_text_encoder')) + + @torch.no_grad + def forward(self, prompts, with_mask=True): + self.device = next(self.text_encoder.parameters()).device + text_inputs = self.tokenizer( + prompts, + padding="max_length", + max_length=self.max_length, + truncation=True, + return_attention_mask=True, + return_tensors="pt", + ) + prompt_embeds = self.text_encoder( + text_inputs.input_ids.to(self.device), + attention_mask=text_inputs.attention_mask.to(self.device) if with_mask else None, + ) + return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output diff --git a/fastvideo/models/stepvideo/text_encoder/flashattention.py b/fastvideo/models/stepvideo/text_encoder/flashattention.py new file mode 100755 index 000000000..647698f4b --- /dev/null +++ b/fastvideo/models/stepvideo/text_encoder/flashattention.py @@ -0,0 +1,45 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import torch + + +def flash_attn_func(q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=True, + return_attn_probs=False, + tp_group_rank=0, + tp_group_size=1): + softmax_scale = q.size(-1)**(-0.5) if softmax_scale is None else softmax_scale + return torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale, causal, return_attn_probs, None, + tp_group_rank, tp_group_size)[0] + + +class FlashSelfAttention(torch.nn.Module): + + def __init__( + self, + attention_dropout=0.0, + ): + super().__init__() + self.dropout_p = attention_dropout + + def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None): + if cu_seqlens is None: + output = flash_attn_func(q, k, v, dropout_p=self.dropout_p) + else: + raise ValueError('cu_seqlens is not supported!') + + return output diff --git a/fastvideo/models/stepvideo/text_encoder/stepllm.py b/fastvideo/models/stepvideo/text_encoder/stepllm.py new file mode 100755 index 000000000..7e971477a --- /dev/null +++ b/fastvideo/models/stepvideo/text_encoder/stepllm.py @@ -0,0 +1,291 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import os +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers.modeling_utils import PretrainedConfig, PreTrainedModel + +from fastvideo.models.stepvideo.modules.normalization import RMSNorm +from fastvideo.models.stepvideo.text_encoder.flashattention import FlashSelfAttention +from fastvideo.models.stepvideo.text_encoder.tokenizer import LLaMaEmbedding, Wrapped_StepChatTokenizer +from fastvideo.models.stepvideo.utils import with_empty_init + + +def safediv(n, d): + q, r = divmod(n, d) + assert r == 0 + return q + + +class MultiQueryAttention(nn.Module): + + def __init__(self, cfg, layer_id=None): + super().__init__() + + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.max_seq_len = cfg.seq_length + self.use_flash_attention = cfg.use_flash_attn + assert self.use_flash_attention, 'FlashAttention is required!' + + self.n_groups = cfg.num_attention_groups + self.tp_size = 1 + self.n_local_heads = cfg.num_attention_heads + self.n_local_groups = self.n_groups + + self.wqkv = nn.Linear( + cfg.hidden_size, + cfg.hidden_size + self.head_dim * 2 * self.n_groups, + bias=False, + ) + self.wo = nn.Linear( + cfg.hidden_size, + cfg.hidden_size, + bias=False, + ) + + assert self.use_flash_attention, 'non-Flash attention not supported yet.' + self.core_attention = FlashSelfAttention(attention_dropout=cfg.attention_dropout) + + self.layer_id = layer_id + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor], + cu_seqlens: Optional[torch.Tensor], + max_seq_len: Optional[torch.Tensor], + ): + seqlen, bsz, dim = x.shape + xqkv = self.wqkv(x) + + xq, xkv = torch.split( + xqkv, + (dim // self.tp_size, self.head_dim * 2 * self.n_groups // self.tp_size), + dim=-1, + ) + + # gather on 1st dimension + xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim) + xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim) + xk, xv = xkv.chunk(2, -1) + + # rotary embedding + flash attn + xq = rearrange(xq, "s b h d -> b s h d") + xk = rearrange(xk, "s b h d -> b s h d") + xv = rearrange(xv, "s b h d -> b s h d") + + q_per_kv = self.n_local_heads // self.n_local_groups + if q_per_kv > 1: + b, s, h, d = xk.size() + if h == 1: + xk = xk.expand(b, s, q_per_kv, d) + xv = xv.expand(b, s, q_per_kv, d) + else: + ''' To cover the cases where h > 1, we have + the following implementation, which is equivalent to: + xk = xk.repeat_interleave(q_per_kv, dim=-2) + xv = xv.repeat_interleave(q_per_kv, dim=-2) + but can avoid calling aten::item() that involves cpu. + ''' + idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten() + xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + + if self.use_flash_attention: + output = self.core_attention(xq, xk, xv, cu_seqlens=cu_seqlens, max_seq_len=max_seq_len) + # reduce-scatter only support first dimension now + output = rearrange(output, "b s h d -> s b (h d)").contiguous() + else: + xq, xk, xv = [rearrange(x, "b s ... -> s b ...").contiguous() for x in (xq, xk, xv)] + output = self.core_attention(xq, xk, xv, mask) + output = self.wo(output) + return output + + +class FeedForward(nn.Module): + + def __init__( + self, + cfg, + dim: int, + hidden_dim: int, + layer_id: int, + multiple_of: int = 256, + ): + super().__init__() + + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + + def swiglu(x): + x = torch.chunk(x, 2, dim=-1) + return F.silu(x[0]) * x[1] + + self.swiglu = swiglu + + self.w1 = nn.Linear( + dim, + 2 * hidden_dim, + bias=False, + ) + self.w2 = nn.Linear( + hidden_dim, + dim, + bias=False, + ) + + def forward(self, x): + x = self.swiglu(self.w1(x)) + output = self.w2(x) + return output + + +class TransformerBlock(nn.Module): + + def __init__(self, cfg, layer_id: int): + super().__init__() + + self.n_heads = cfg.num_attention_heads + self.dim = cfg.hidden_size + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.attention = MultiQueryAttention( + cfg, + layer_id=layer_id, + ) + + self.feed_forward = FeedForward( + cfg, + dim=cfg.hidden_size, + hidden_dim=cfg.ffn_hidden_size, + layer_id=layer_id, + ) + self.layer_id = layer_id + self.attention_norm = RMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + self.ffn_norm = RMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor], + cu_seqlens: Optional[torch.Tensor], + max_seq_len: Optional[torch.Tensor], + ): + residual = self.attention.forward(self.attention_norm(x), mask, cu_seqlens, max_seq_len) + h = x + residual + ffn_res = self.feed_forward.forward(self.ffn_norm(h)) + out = h + ffn_res + return out + + +class Transformer(nn.Module): + + def __init__( + self, + config, + max_seq_size=8192, + ): + super().__init__() + self.num_layers = config.num_layers + self.layers = self._build_layers(config) + + def _build_layers(self, config): + layers = torch.nn.ModuleList() + for layer_id in range(self.num_layers): + layers.append(TransformerBlock( + config, + layer_id=layer_id + 1, + )) + return layers + + def forward( + self, + hidden_states, + attention_mask, + cu_seqlens=None, + max_seq_len=None, + ): + + if max_seq_len is not None and not isinstance(max_seq_len, torch.Tensor): + max_seq_len = torch.tensor(max_seq_len, dtype=torch.int32, device="cpu") + + for lid, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + attention_mask, + cu_seqlens, + max_seq_len, + ) + return hidden_states + + +class Step1Model(PreTrainedModel): + config_class = PretrainedConfig + + @with_empty_init + def __init__( + self, + config, + ): + super().__init__(config) + self.tok_embeddings = LLaMaEmbedding(config) + self.transformer = Transformer(config) + + def forward( + self, + input_ids=None, + attention_mask=None, + ): + + hidden_states = self.tok_embeddings(input_ids) + + hidden_states = self.transformer( + hidden_states, + attention_mask, + ) + return hidden_states + + +class STEP1TextEncoder(torch.nn.Module): + + def __init__(self, model_dir, max_length=320): + super(STEP1TextEncoder, self).__init__() + self.max_length = max_length + self.text_tokenizer = Wrapped_StepChatTokenizer(os.path.join(model_dir, 'step1_chat_tokenizer.model')) + text_encoder = Step1Model.from_pretrained(model_dir) + self.text_encoder = text_encoder.eval().to(torch.bfloat16) + + @torch.no_grad + def forward(self, prompts, with_mask=True, max_length=None): + self.device = next(self.text_encoder.parameters()).device + with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16): + if type(prompts) is str: + prompts = [prompts] + + txt_tokens = self.text_tokenizer(prompts, + max_length=max_length or self.max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + y = self.text_encoder(txt_tokens.input_ids.to(self.device), + attention_mask=txt_tokens.attention_mask.to(self.device) if with_mask else None) + y_mask = txt_tokens.attention_mask + return y.transpose(0, 1), y_mask diff --git a/fastvideo/models/stepvideo/text_encoder/tokenizer.py b/fastvideo/models/stepvideo/text_encoder/tokenizer.py new file mode 100755 index 000000000..62ca0d608 --- /dev/null +++ b/fastvideo/models/stepvideo/text_encoder/tokenizer.py @@ -0,0 +1,209 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +from typing import List + +import torch +import torch.nn as nn + + +class LLaMaEmbedding(nn.Module): + """Language model embeddings. + + Arguments: + hidden_size: hidden size + vocab_size: vocabulary size + max_sequence_length: maximum size of sequence. This + is used for positional embedding + embedding_dropout_prob: dropout probability for embeddings + init_method: weight initialization method + num_tokentypes: size of the token-type embeddings. 0 value + will ignore this embedding + """ + + def __init__( + self, + cfg, + ): + super().__init__() + self.hidden_size = cfg.hidden_size + self.params_dtype = cfg.params_dtype + self.fp32_residual_connection = cfg.fp32_residual_connection + self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32 + self.word_embeddings = torch.nn.Embedding( + cfg.padded_vocab_size, + self.hidden_size, + ) + self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) + + def forward(self, input_ids): + # Embeddings. + if self.embedding_weights_in_fp32: + self.word_embeddings = self.word_embeddings.to(torch.float32) + embeddings = self.word_embeddings(input_ids) + if self.embedding_weights_in_fp32: + embeddings = embeddings.to(self.params_dtype) + self.word_embeddings = self.word_embeddings.to(self.params_dtype) + + # Data format change to avoid explicit transposes : [b s h] --> [s b h]. + embeddings = embeddings.transpose(0, 1).contiguous() + + # If the input flag for fp32 residual connection is set, convert for float. + if self.fp32_residual_connection: + embeddings = embeddings.float() + + # Dropout. + embeddings = self.embedding_dropout(embeddings) + + return embeddings + + +class StepChatTokenizer: + """Step Chat Tokenizer""" + + def __init__( + self, + model_file, + name="StepChatTokenizer", + bot_token="<|BOT|>", # Begin of Turn + eot_token="<|EOT|>", # End of Turn + call_start_token="<|CALL_START|>", # Call Start + call_end_token="<|CALL_END|>", # Call End + think_start_token="<|THINK_START|>", # Think Start + think_end_token="<|THINK_END|>", # Think End + mask_start_token="<|MASK_1e69f|>", # Mask start + mask_end_token="<|UNMASK_1e69f|>", # Mask end + ): + import sentencepiece + + self._tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file) + + self._vocab = {} + self._inv_vocab = {} + + self._special_tokens = {} + self._inv_special_tokens = {} + + self._t5_tokens = [] + + for idx in range(self._tokenizer.get_piece_size()): + text = self._tokenizer.id_to_piece(idx) + self._inv_vocab[idx] = text + self._vocab[text] = idx + + if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown(idx): + self._special_tokens[text] = idx + self._inv_special_tokens[idx] = text + + self._unk_id = self._tokenizer.unk_id() + self._bos_id = self._tokenizer.bos_id() + self._eos_id = self._tokenizer.eos_id() + + for token in [bot_token, eot_token, call_start_token, call_end_token, think_start_token, think_end_token]: + assert token in self._vocab, f"Token '{token}' not found in tokenizer" + assert token in self._special_tokens, f"Token '{token}' is not a special token" + + for token in [mask_start_token, mask_end_token]: + assert token in self._vocab, f"Token '{token}' not found in tokenizer" + + self._bot_id = self._tokenizer.piece_to_id(bot_token) + self._eot_id = self._tokenizer.piece_to_id(eot_token) + self._call_start_id = self._tokenizer.piece_to_id(call_start_token) + self._call_end_id = self._tokenizer.piece_to_id(call_end_token) + self._think_start_id = self._tokenizer.piece_to_id(think_start_token) + self._think_end_id = self._tokenizer.piece_to_id(think_end_token) + self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token) + self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token) + + self._underline_id = self._tokenizer.piece_to_id("\u2581") + + @property + def vocab(self): + return self._vocab + + @property + def inv_vocab(self): + return self._inv_vocab + + @property + def vocab_size(self): + return self._tokenizer.vocab_size() + + def tokenize(self, text: str) -> List[int]: + return self._tokenizer.encode_as_ids(text) + + def detokenize(self, token_ids: List[int]) -> str: + return self._tokenizer.decode_ids(token_ids) + + +class Tokens: + + def __init__(self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len) -> None: + self.input_ids = input_ids + self.attention_mask = attention_mask + self.cu_input_ids = cu_input_ids + self.cu_seqlens = cu_seqlens + self.max_seq_len = max_seq_len + + def to(self, device): + self.input_ids = self.input_ids.to(device) + self.attention_mask = self.attention_mask.to(device) + self.cu_input_ids = self.cu_input_ids.to(device) + self.cu_seqlens = self.cu_seqlens.to(device) + return self + + +class Wrapped_StepChatTokenizer(StepChatTokenizer): + + def __call__(self, text, max_length=320, padding="max_length", truncation=True, return_tensors="pt"): + # [bos, ..., eos, pad, pad, ..., pad] + self.BOS = 1 + self.EOS = 2 + self.PAD = 2 + out_tokens = [] + attn_mask = [] + if len(text) == 0: + part_tokens = [self.BOS] + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) + else: + for part in text: + part_tokens = self.tokenize(part) + part_tokens = part_tokens[:(max_length - 2)] # leave 2 space for bos and eos + part_tokens = [self.BOS] + part_tokens + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) + + out_tokens = torch.tensor(out_tokens, dtype=torch.long) + attn_mask = torch.tensor(attn_mask, dtype=torch.long) + + # padding y based on tp size + padded_len = 0 + padded_flag = True if padded_len > 0 else False + if padded_flag: + pad_tokens = torch.tensor([[self.PAD] * max_length], device=out_tokens.device) + pad_attn_mask = torch.tensor([[1] * padded_len + [0] * (max_length - padded_len)], device=attn_mask.device) + out_tokens = torch.cat([out_tokens, pad_tokens], dim=0) + attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0) + + # cu_seqlens + cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0) + seqlen = attn_mask.sum(dim=1).tolist() + cu_seqlens = torch.cumsum(torch.tensor([0] + seqlen), 0).to(device=out_tokens.device, dtype=torch.int32) + max_seq_len = max(seqlen) + return Tokens(out_tokens, cu_out_tokens, attn_mask, cu_seqlens, max_seq_len) diff --git a/fastvideo/models/stepvideo/utils/__init__.py b/fastvideo/models/stepvideo/utils/__init__.py new file mode 100755 index 000000000..a77226b13 --- /dev/null +++ b/fastvideo/models/stepvideo/utils/__init__.py @@ -0,0 +1,2 @@ +from .utils import * +from .video_process import * \ No newline at end of file diff --git a/fastvideo/models/stepvideo/utils/quantization.py b/fastvideo/models/stepvideo/utils/quantization.py new file mode 100644 index 000000000..e87e38670 --- /dev/null +++ b/fastvideo/models/stepvideo/utils/quantization.py @@ -0,0 +1,117 @@ +# from stepvideo.diffusion.video_pipeline import StepVideoPipeline +import torch +import torch.nn as nn +from torch.nn import functional as F + + +def get_fp_maxval(bits=8, mantissa_bit=3, sign_bits=1): + _bits = torch.tensor(bits) + _mantissa_bit = torch.tensor(mantissa_bit) + _sign_bits = torch.tensor(sign_bits) + M = torch.clamp(torch.round(_mantissa_bit), 1, _bits - _sign_bits) + E = _bits - _sign_bits - M + bias = 2**(E - 1) - 1 + mantissa = 1 + for i in range(mantissa_bit - 1): + mantissa += 1 / (2**(i + 1)) + maxval = mantissa * 2**(2**E - 1 - bias) + return maxval + + +def quantize_to_fp8(x, bits=8, mantissa_bit=3, sign_bits=1): + """ + Default is E4M3. + """ + bits = torch.tensor(bits) + mantissa_bit = torch.tensor(mantissa_bit) + sign_bits = torch.tensor(sign_bits) + M = torch.clamp(torch.round(mantissa_bit), 1, bits - sign_bits) + E = bits - sign_bits - M + bias = 2**(E - 1) - 1 + mantissa = 1 + for i in range(mantissa_bit - 1): + mantissa += 1 / (2**(i + 1)) + maxval = mantissa * 2**(2**E - 1 - bias) + minval = -maxval + minval = -maxval if sign_bits == 1 else torch.zeros_like(maxval) + input_clamp = torch.min(torch.max(x, minval), maxval) + log_scales = torch.clamp((torch.floor(torch.log2(torch.abs(input_clamp)) + bias)).detach(), 1.0) + log_scales = 2.0**(log_scales - M - bias.type(x.dtype)) + # dequant + qdq_out = torch.round(input_clamp / log_scales) * log_scales + return qdq_out, log_scales + + +def fp8_tensor_quant(x, scale, bits=8, mantissa_bit=3, sign_bits=1): + for i in range(len(x.shape) - 1): + scale = scale.unsqueeze(-1) + new_x = x / scale + quant_dequant_x, log_scales = quantize_to_fp8(new_x, bits=bits, mantissa_bit=mantissa_bit, sign_bits=sign_bits) + return quant_dequant_x, scale, log_scales + + +def fp8_activation_dequant(qdq_out, scale, dtype): + qdq_out = qdq_out.type(dtype) + quant_dequant_x = qdq_out * scale.to(dtype) + return quant_dequant_x + + +def fp8_linear_forward(cls, original_dtype, input): + weight_dtype = cls.weight.dtype + ##### + if cls.weight.dtype != torch.float8_e4m3fn: + assert False + maxval = get_fp_maxval() + scale = torch.max(torch.abs(cls.weight.flatten())) / maxval + linear_weight, scale, log_scales = fp8_tensor_quant(cls.weight, scale) + linear_weight = linear_weight.to(torch.float8_e4m3fn) + weight_dtype = linear_weight.dtype + else: + scale = cls.fp8_scale.to(cls.weight.device) + linear_weight = cls.weight + ##### + + if weight_dtype == torch.float8_e4m3fn: + if True or len(input.shape) == 3: + cls_dequant = fp8_activation_dequant(linear_weight, scale, original_dtype) + if cls.bias is not None: + print(f"input dtype: {input.dtype}") + print(f"cls_dequant dtype: {cls_dequant.dtype}") + print(f"cls.bias dtype: {cls.bias.dtype}") + + output = F.linear(input, cls_dequant, cls.bias) + else: + output = F.linear(input, cls_dequant) + return output + else: + return cls.original_forward(input.to(original_dtype)) + else: + return cls.original_forward(input) + + +def convert_fp8_linear(module, original_dtype, params_to_keep={}): + setattr(module, "fp8_matmul_enabled", True) + fp8_layers = [] + scale_dict = {} + counter = 0 + for key, layer in module.named_modules(): + if isinstance(layer, nn.Linear) and 'transformer_blocks' in key: + print(f"Converting {key} to FP8") + fp8_layers.append(key) + original_forward = layer.forward + maxval = get_fp_maxval() + scale = torch.max(torch.abs(layer.weight.flatten())) / maxval + + original_weight = layer.weight.data # Store a reference to the original weights + quantized_weight, scale, _ = fp8_tensor_quant(original_weight, scale) + scale_dict[key] = scale + layer.weight = torch.nn.Parameter(quantized_weight.to(torch.float8_e4m3fn)) + del original_weight # Delete the reference to the original weights + torch.cuda.empty_cache() + + # print(f"layer weight dtype: {layer.weight.dtype} for layer {key}") + setattr(layer, "fp8_scale", scale.to(dtype=original_dtype)) + setattr(layer, "original_forward", original_forward) + setattr(layer, "forward", lambda input, m=layer: fp8_linear_forward(m, original_dtype, input)) + counter += 1 + return scale_dict diff --git a/fastvideo/models/stepvideo/utils/utils.py b/fastvideo/models/stepvideo/utils/utils.py new file mode 100755 index 000000000..835393277 --- /dev/null +++ b/fastvideo/models/stepvideo/utils/utils.py @@ -0,0 +1,60 @@ +import random +from functools import wraps + +import numpy as np +import torch +import torch.utils._device + + +def setup_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +class EmptyInitOnDevice(torch.overrides.TorchFunctionMode): + + def __init__(self, device=None): + self.device = device + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + if getattr(func, '__module__', None) == 'torch.nn.init': + if 'tensor' in kwargs: + return kwargs['tensor'] + else: + return args[0] + if self.device is not None and func in torch.utils._device._device_constructors( + ) and kwargs.get('device') is None: + kwargs['device'] = self.device + return func(*args, **kwargs) + + +def with_empty_init(func): + + @wraps(func) + def wrapper(*args, **kwargs): + with EmptyInitOnDevice('cpu'): + return func(*args, **kwargs) + + return wrapper + + +def culens2mask(cu_seqlens=None, cu_seqlens_kv=None, max_seqlen=None, max_seqlen_kv=None, is_causal=False): + assert len(cu_seqlens) == len(cu_seqlens_kv) + "q k v should have same bsz..." + bsz = len(cu_seqlens) - 1 + seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + + attn_mask = torch.zeros(bsz, max_seqlen, max_seqlen_kv, dtype=torch.bool) + for i, (seq_len, seq_len_kv) in enumerate(zip(seqlens, seqlens_kv)): + if is_causal: + attn_mask[i, :seq_len, :seq_len_kv] = torch.triu(torch.ones(seq_len, seq_len_kv), diagonal=1).bool() + else: + attn_mask[i, :seq_len, :seq_len_kv] = torch.ones([seq_len, seq_len_kv], dtype=torch.bool) + + return attn_mask diff --git a/fastvideo/models/stepvideo/utils/video_process.py b/fastvideo/models/stepvideo/utils/video_process.py new file mode 100644 index 000000000..93e6af546 --- /dev/null +++ b/fastvideo/models/stepvideo/utils/video_process.py @@ -0,0 +1,51 @@ +import os + +import imageio +import numpy as np +import torch + + +class VideoProcessor: + + def __init__(self, save_path: str = './results', name_suffix: str = ''): + self.save_path = save_path + os.makedirs(self.save_path, exist_ok=True) + self.name_suffix = name_suffix + + def crop2standard540p(self, vid_array): + _, height, width, _ = vid_array.shape + height_center = height // 2 + width_center = width // 2 + if width_center > height_center: ## horizon mode + return vid_array[:, height_center - 270:height_center + 270, width_center - 480:width_center + 480] + elif width_center < height_center: ## portrait mode + return vid_array[:, height_center - 480:height_center + 480, width_center - 270:width_center + 270] + else: + return vid_array + + def save_imageio_video(self, video_array: np.array, output_filename: str, fps=25, codec='libx264'): + + ffmpeg_params = [ + "-vf", + "atadenoise=0a=0.1:0b=0.1:1a=0.1:1b=0.1", # denoise + ] + + with imageio.get_writer(output_filename, fps=fps, codec=codec, ffmpeg_params=ffmpeg_params) as vid_writer: + for img_array in video_array: + vid_writer.append_data(img_array) + + def postprocess_video(self, video_tensor, output_file_name='', output_type="mp4", crop2standard540p=True): + if len(self.name_suffix) == 0: + video_path = os.path.join(self.save_path, f"{output_file_name}.{output_type}") + else: + video_path = os.path.join(self.save_path, f"{output_file_name}-{self.name_suffix}.{output_type}") + + video_tensor = torch.cat([t for t in video_tensor], dim=-2) + video_tensor = (video_tensor.cpu().clamp(-1, 1) + 1) * 127.5 + video_array = video_tensor.clamp(0, 255).to(torch.uint8).numpy().transpose(0, 2, 3, 1) + + if crop2standard540p: + video_array = self.crop2standard540p(video_array) + + self.save_imageio_video(video_array, video_path) + print(f"Saved the generated video in {video_path}") diff --git a/fastvideo/models/stepvideo/vae/vae.py b/fastvideo/models/stepvideo/vae/vae.py new file mode 100755 index 000000000..1dbf532c0 --- /dev/null +++ b/fastvideo/models/stepvideo/vae/vae.py @@ -0,0 +1,975 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import torch +from einops import rearrange +from torch import nn +from torch.nn import functional as F + +from fastvideo.models.stepvideo.utils import with_empty_init + + +def base_group_norm(x, norm_layer, act_silu=False, channel_last=False): + if hasattr(base_group_norm, 'spatial') and base_group_norm.spatial: + assert channel_last + x_shape = x.shape + x = x.flatten(0, 1) + if channel_last: + # Permute to NCHW format + x = x.permute(0, 3, 1, 2) + + out = F.group_norm(x.contiguous(), norm_layer.num_groups, norm_layer.weight, norm_layer.bias, norm_layer.eps) + if act_silu: + out = F.silu(out) + + if channel_last: + # Permute back to NHWC format + out = out.permute(0, 2, 3, 1) + + out = out.view(x_shape) + else: + if channel_last: + # Permute to NCHW format + x = x.permute(0, 3, 1, 2) + out = F.group_norm(x.contiguous(), norm_layer.num_groups, norm_layer.weight, norm_layer.bias, norm_layer.eps) + if act_silu: + out = F.silu(out) + if channel_last: + # Permute back to NHWC format + out = out.permute(0, 2, 3, 1) + return out + + +def base_conv2d(x, conv_layer, channel_last=False, residual=None): + if channel_last: + x = x.permute(0, 3, 1, 2) # NHWC to NCHW + out = F.conv2d(x, conv_layer.weight, conv_layer.bias, stride=conv_layer.stride, padding=conv_layer.padding) + if residual is not None: + if channel_last: + residual = residual.permute(0, 3, 1, 2) # NHWC to NCHW + out += residual + if channel_last: + out = out.permute(0, 2, 3, 1) # NCHW to NHWC + return out + + +def base_conv3d(x, conv_layer, channel_last=False, residual=None, only_return_output=False): + if only_return_output: + size = cal_outsize(x.shape, conv_layer.weight.shape, conv_layer.stride, conv_layer.padding) + return torch.empty(size, device=x.device, dtype=x.dtype) + if channel_last: + x = x.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW + out = F.conv3d(x, conv_layer.weight, conv_layer.bias, stride=conv_layer.stride, padding=conv_layer.padding) + if residual is not None: + if channel_last: + residual = residual.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW + out += residual + if channel_last: + out = out.permute(0, 2, 3, 4, 1) # NCDHW to NDHWC + return out + + +def cal_outsize(input_sizes, kernel_sizes, stride, padding): + stride_d, stride_h, stride_w = stride + padding_d, padding_h, padding_w = padding + dilation_d, dilation_h, dilation_w = 1, 1, 1 + + in_d = input_sizes[1] + in_h = input_sizes[2] + in_w = input_sizes[3] + + kernel_d = kernel_sizes[2] + kernel_h = kernel_sizes[3] + kernel_w = kernel_sizes[4] + out_channels = kernel_sizes[0] + + out_d = calc_out_(in_d, padding_d, dilation_d, kernel_d, stride_d) + out_h = calc_out_(in_h, padding_h, dilation_h, kernel_h, stride_h) + out_w = calc_out_(in_w, padding_w, dilation_w, kernel_w, stride_w) + size = [input_sizes[0], out_d, out_h, out_w, out_channels] + return size + + +def calc_out_(in_size, padding, dilation, kernel, stride): + return (in_size + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1 + + +def base_conv3d_channel_last(x, conv_layer, residual=None): + in_numel = x.numel() + out_numel = int(x.numel() * conv_layer.out_channels / conv_layer.in_channels) + if (in_numel >= 2**30) or (out_numel >= 2**30): + assert conv_layer.stride[0] == 1, "time split asks time stride = 1" + + B, T, H, W, C = x.shape + K = conv_layer.kernel_size[0] + + chunks = 4 + chunk_size = T // chunks + + if residual is None: + out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual, only_return_output=True) + else: + out_nhwc = residual + + assert B == 1 + for i in range(chunks): + if i == chunks - 1: + xi = x[:1, chunk_size * i:] + out_nhwci = out_nhwc[:1, chunk_size * i:] + else: + xi = x[:1, chunk_size * i:chunk_size * (i + 1) + K - 1] + out_nhwci = out_nhwc[:1, chunk_size * i:chunk_size * (i + 1)] + if residual is not None: + if i == chunks - 1: + ri = residual[:1, chunk_size * i:] + else: + ri = residual[:1, chunk_size * i:chunk_size * (i + 1)] + else: + ri = None + out_nhwci.copy_(base_conv3d(xi, conv_layer, channel_last=True, residual=ri)) + else: + out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual) + return out_nhwc + + +class Upsample2D(nn.Module): + + def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + + if use_conv: + self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1) + else: + assert "Not Supported" + self.conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1) + + def forward(self, x, output_size=None): + assert x.shape[-1] == self.channels + + if self.use_conv_transpose: + return self.conv(x) + + if output_size is None: + x = F.interpolate(x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last), + scale_factor=2.0, + mode='nearest').permute(0, 2, 3, 1).contiguous() + else: + x = F.interpolate(x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last), + size=output_size, + mode='nearest').permute(0, 2, 3, 1).contiguous() + + # x = self.conv(x) + x = base_conv2d(x, self.conv, channel_last=True) + return x + + +class Downsample2D(nn.Module): + + def __init__(self, channels, use_conv=False, out_channels=None, padding=1): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = 2 + + if use_conv: + self.conv = nn.Conv2d(self.channels, self.out_channels, 3, stride=stride, padding=padding) + else: + assert self.channels == self.out_channels + self.conv = nn.AvgPool2d(kernel_size=stride, stride=stride) + + def forward(self, x): + assert x.shape[-1] == self.channels + if self.use_conv and self.padding == 0: + pad = (0, 0, 0, 1, 0, 1) + x = F.pad(x, pad, mode="constant", value=0) + + assert x.shape[-1] == self.channels + # x = self.conv(x) + x = base_conv2d(x, self.conv, channel_last=True) + return x + + +class CausalConv(nn.Module): + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs): + super().__init__() + + if isinstance(kernel_size, int): + kernel_size = kernel_size if isinstance(kernel_size, tuple) else ((kernel_size, ) * 3) + time_kernel_size, height_kernel_size, width_kernel_size = kernel_size + + self.dilation = kwargs.pop('dilation', 1) + self.stride = kwargs.pop('stride', 1) + if isinstance(self.stride, int): + self.stride = (self.stride, 1, 1) + time_pad = self.dilation * (time_kernel_size - 1) + max((1 - self.stride[0]), 0) + height_pad = height_kernel_size // 2 + width_pad = width_kernel_size // 2 + self.time_causal_padding = (width_pad, width_pad, height_pad, height_pad, time_pad, 0) + self.time_uncausal_padding = (width_pad, width_pad, height_pad, height_pad, 0, 0) + + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs) + self.is_first_run = True + + def forward(self, x, is_init=True, residual=None): + x = nn.functional.pad(x, self.time_causal_padding if is_init else self.time_uncausal_padding) + + x = self.conv(x) + if residual is not None: + x.add_(residual) + return x + + +class ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert out_channels * factor**3 % in_channels == 0 + self.repeats = out_channels * factor**3 // in_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view(x.size(0), self.out_channels, self.factor, self.factor, self.factor, x.size(2), x.size(3), x.size(4)) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view(x.size(0), self.out_channels, + x.size(2) * self.factor, + x.size(4) * self.factor, + x.size(6) * self.factor) + x = x[:, :, self.factor - 1:, :, :] + return x + + +class ConvPixelShuffleUpSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + factor: int, + ): + super().__init__() + self.factor = factor + out_ratio = factor**3 + self.conv = CausalConv(in_channels, out_channels * out_ratio, kernel_size=kernel_size) + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = self.conv(x, is_init) + x = self.pixel_shuffle_3d(x, self.factor) + return x + + @staticmethod + def pixel_shuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: + batch_size, channels, depth, height, width = x.size() + new_channels = channels // (factor**3) + new_depth = depth * factor + new_height = height * factor + new_width = width * factor + + x = x.view(batch_size, new_channels, factor, factor, factor, depth, height, width) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view(batch_size, new_channels, new_depth, new_height, new_width) + x = x[:, :, factor - 1:, :, :] + return x + + +class ConvPixelUnshuffleDownSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + factor: int, + ): + super().__init__() + self.factor = factor + out_ratio = factor**3 + assert out_channels % out_ratio == 0 + self.conv = CausalConv(in_channels, out_channels // out_ratio, kernel_size=kernel_size) + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = self.conv(x, is_init) + x = self.pixel_unshuffle_3d(x, self.factor) + return x + + @staticmethod + def pixel_unshuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: + pad = (0, 0, 0, 0, factor - 1, 0) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view(B, C, D // factor, factor, H // factor, factor, W // factor, factor) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * factor**3, D // factor, H // factor, W // factor) + return x + + +class PixelUnshuffleChannelAveragingDownSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert in_channels * factor**3 % out_channels == 0 + self.group_size = in_channels * factor**3 // out_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + pad = (0, 0, 0, 0, self.factor - 1, 0) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view(B, C, D // self.factor, self.factor, H // self.factor, self.factor, W // self.factor, self.factor) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * self.factor**3, D // self.factor, H // self.factor, W // self.factor) + x = x.view(B, self.out_channels, self.group_size, D // self.factor, H // self.factor, W // self.factor) + x = x.mean(dim=2) + return x + + +def base_group_norm_with_zero_pad(x, norm_layer, act_silu=True, pad_size=2): + out_shape = list(x.shape) + out_shape[1] += pad_size + out = torch.empty(out_shape, dtype=x.dtype, device=x.device) + out[:, pad_size:] = base_group_norm(x, norm_layer, act_silu=act_silu, channel_last=True) + out[:, :pad_size] = 0 + return out + + +class CausalConvChannelLast(CausalConv): + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs): + super().__init__(chan_in, chan_out, kernel_size, **kwargs) + + self.time_causal_padding = (0, 0) + self.time_causal_padding + self.time_uncausal_padding = (0, 0) + self.time_uncausal_padding + + def forward(self, x, is_init=True, residual=None): + if self.is_first_run: + self.is_first_run = False + # self.conv.weight = nn.Parameter(self.conv.weight.permute(0,2,3,4,1).contiguous()) + + x = nn.functional.pad(x, self.time_causal_padding if is_init else self.time_uncausal_padding) + + x = base_conv3d_channel_last(x, self.conv, residual=residual) + return x + + +class CausalConvAfterNorm(CausalConv): + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs): + super().__init__(chan_in, chan_out, kernel_size, **kwargs) + + if self.time_causal_padding == (1, 1, 1, 1, 2, 0): + self.conv = nn.Conv3d(chan_in, + chan_out, + kernel_size, + stride=self.stride, + dilation=self.dilation, + padding=(0, 1, 1), + **kwargs) + else: + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs) + self.is_first_run = True + + def forward(self, x, is_init=True, residual=None): + if self.is_first_run: + self.is_first_run = False + + if self.time_causal_padding == (1, 1, 1, 1, 2, 0): + pass + else: + x = nn.functional.pad(x, self.time_causal_padding).contiguous() + + x = base_conv3d_channel_last(x, self.conv, residual=residual) + return x + + +class AttnBlock(nn.Module): + + def __init__(self, in_channels): + super().__init__() + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels) + self.q = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.k = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.v = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.proj_out = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + + def attention(self, x, is_init=True): + x = base_group_norm(x, self.norm, act_silu=False, channel_last=True) + q = self.q(x, is_init) + k = self.k(x, is_init) + v = self.v(x, is_init) + + b, t, h, w, c = q.shape + q, k, v = map(lambda x: rearrange(x, "b t h w c -> b 1 (t h w) c"), (q, k, v)) + x = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) + x = rearrange(x, "b 1 (t h w) c -> b t h w c", t=t, h=h, w=w) + + return x + + def forward(self, x): + x = x.permute(0, 2, 3, 4, 1).contiguous() + h = self.attention(x) + x = self.proj_out(h, residual=x) + x = x.permute(0, 4, 1, 2, 3) + return x + + +class Resnet3DBlock(nn.Module): + + def __init__( + self, + in_channels, + out_channels=None, + temb_channels=512, + conv_shortcut=False, + ): + super().__init__() + + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels) + self.conv1 = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) + if temb_channels > 0: + self.temb_proj = nn.Linear(temb_channels, out_channels) + + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels) + self.conv2 = CausalConvAfterNorm(out_channels, out_channels, kernel_size=3) + + assert conv_shortcut is False + self.use_conv_shortcut = conv_shortcut + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) + else: + self.nin_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=1) + + def forward(self, x, temb=None, is_init=True): + x = x.permute(0, 2, 3, 4, 1).contiguous() + + h = base_group_norm_with_zero_pad(x, self.norm1, act_silu=True, pad_size=2) + h = self.conv1(h) + if temb is not None: + h = h + self.temb_proj(nn.functional.silu(temb))[:, :, None, None] + + x = self.nin_shortcut(x) if self.in_channels != self.out_channels else x + + h = base_group_norm_with_zero_pad(h, self.norm2, act_silu=True, pad_size=2) + x = self.conv2(h, residual=x) + + x = x.permute(0, 4, 1, 2, 3) + return x + + +class Downsample3D(nn.Module): + + def __init__(self, in_channels, with_conv, stride): + super().__init__() + + self.with_conv = with_conv + if with_conv: + self.conv = CausalConv(in_channels, in_channels, kernel_size=3, stride=stride) + + def forward(self, x, is_init=True): + if self.with_conv: + x = self.conv(x, is_init) + else: + x = nn.functional.avg_pool3d(x, kernel_size=2, stride=2) + return x + + +class VideoEncoder(nn.Module): + + def __init__( + self, + ch=32, + ch_mult=(4, 8, 16, 16), + num_res_blocks=2, + in_channels=3, + z_channels=16, + double_z=True, + down_sampling_layer=[1, 2], + resamp_with_conv=True, + version=1, + ): + super().__init__() + + temb_ch = 0 + + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = CausalConv(in_channels, ch, kernel_size=3) + self.down_sampling_layer = down_sampling_layer + + in_ch_mult = (1, ) + tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + if i_level in self.down_sampling_layer: + down.downsample = Downsample3D(block_in, resamp_with_conv, stride=(2, 2, 2)) + else: + down.downsample = Downsample2D(block_in, resamp_with_conv, padding=0) #DIFF + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) + self.version = version + if version == 2: + channels = 4 * z_channels * 2**3 + self.conv_patchify = ConvPixelUnshuffleDownSampleLayer3D(block_in, channels, kernel_size=3, factor=2) + self.shortcut_pathify = PixelUnshuffleChannelAveragingDownSampleLayer3D(block_in, channels, 2) + self.shortcut_out = PixelUnshuffleChannelAveragingDownSampleLayer3D( + channels, 2 * z_channels if double_z else z_channels, 1) + self.conv_out = CausalConvChannelLast(channels, 2 * z_channels if double_z else z_channels, kernel_size=3) + else: + self.conv_out = CausalConvAfterNorm(block_in, 2 * z_channels if double_z else z_channels, kernel_size=3) + + @torch.inference_mode() + def forward(self, x, video_frame_num, is_init=True): + # timestep embedding + temb = None + + t = video_frame_num + + # downsampling + h = self.conv_in(x, is_init) + + # make it real channel last, but behave like normal layout + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](h, temb, is_init) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + + if i_level != self.num_resolutions - 1: + if isinstance(self.down[i_level].downsample, Downsample2D): + _, _, t, _, _ = h.shape + h = rearrange(h, "b c t h w -> (b t) h w c", t=t) + h = self.down[i_level].downsample(h) + h = rearrange(h, "(b t) h w c -> b c t h w", t=t) + else: + h = self.down[i_level].downsample(h, is_init) + + h = self.mid.block_1(h, temb, is_init) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb, is_init) + + h = h.permute(0, 2, 3, 4, 1).contiguous() # b c l h w -> b l h w c + if self.version == 2: + h = base_group_norm(h, self.norm_out, act_silu=True, channel_last=True) + h = h.permute(0, 4, 1, 2, 3).contiguous() + shortcut = self.shortcut_pathify(h, is_init) + h = self.conv_patchify(h, is_init) + h = h.add_(shortcut) + shortcut = self.shortcut_out(h, is_init).permute(0, 2, 3, 4, 1) + h = self.conv_out(h.permute(0, 2, 3, 4, 1).contiguous(), is_init) + h = h.add_(shortcut) + else: + h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2) + h = self.conv_out(h, is_init) + h = h.permute(0, 4, 1, 2, 3) # b l h w c -> b c l h w + + h = rearrange(h, "b c t h w -> b t c h w") + return h + + +class Res3DBlockUpsample(nn.Module): + + def __init__(self, input_filters, num_filters, down_sampling_stride, down_sampling=False): + super().__init__() + + self.input_filters = input_filters + self.num_filters = num_filters + + self.act_ = nn.SiLU(inplace=True) + + self.conv1 = CausalConvChannelLast(num_filters, num_filters, kernel_size=[3, 3, 3]) + self.norm1 = nn.GroupNorm(32, num_filters) + + self.conv2 = CausalConvChannelLast(num_filters, num_filters, kernel_size=[3, 3, 3]) + self.norm2 = nn.GroupNorm(32, num_filters) + + self.down_sampling = down_sampling + if down_sampling: + self.down_sampling_stride = down_sampling_stride + else: + self.down_sampling_stride = [1, 1, 1] + + if num_filters != input_filters or down_sampling: + self.conv3 = CausalConvChannelLast(input_filters, + num_filters, + kernel_size=[1, 1, 1], + stride=self.down_sampling_stride) + self.norm3 = nn.GroupNorm(32, num_filters) + + def forward(self, x, is_init=False): + x = x.permute(0, 2, 3, 4, 1).contiguous() + + residual = x + + h = self.conv1(x, is_init) + h = base_group_norm(h, self.norm1, act_silu=True, channel_last=True) + + h = self.conv2(h, is_init) + h = base_group_norm(h, self.norm2, act_silu=False, channel_last=True) + + if self.down_sampling or self.num_filters != self.input_filters: + x = self.conv3(x, is_init) + x = base_group_norm(x, self.norm3, act_silu=False, channel_last=True) + + h.add_(x) + h = self.act_(h) + if residual is not None: + h.add_(residual) + + h = h.permute(0, 4, 1, 2, 3) + return h + + +class Upsample3D(nn.Module): + + def __init__(self, in_channels, scale_factor=2): + super().__init__() + + self.scale_factor = scale_factor + self.conv3d = Res3DBlockUpsample(input_filters=in_channels, + num_filters=in_channels, + down_sampling_stride=(1, 1, 1), + down_sampling=False) + + def forward(self, x, is_init=True, is_split=True): + b, c, t, h, w = x.shape + + # x = x.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3).to(memory_format=torch.channels_last_3d) + if is_split: + split_size = c // 8 + x_slices = torch.split(x, split_size, dim=1) + x = [nn.functional.interpolate(x, scale_factor=self.scale_factor) for x in x_slices] + x = torch.cat(x, dim=1) + else: + x = nn.functional.interpolate(x, scale_factor=self.scale_factor) + + x = self.conv3d(x, is_init) + return x + + +class VideoDecoder(nn.Module): + + def __init__( + self, + ch=128, + z_channels=16, + out_channels=3, + ch_mult=(1, 2, 4, 4), + num_res_blocks=2, + temporal_up_layers=[2, 3], + temporal_downsample=4, + resamp_with_conv=True, + version=1, + ): + super().__init__() + + temb_ch = 0 + + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.temporal_downsample = temporal_downsample + + block_in = ch * ch_mult[self.num_resolutions - 1] + self.version = version + if version == 2: + channels = 4 * z_channels * 2**3 + self.conv_in = CausalConv(z_channels, channels, kernel_size=3) + self.shortcut_in = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(z_channels, channels, 1) + self.conv_unpatchify = ConvPixelShuffleUpSampleLayer3D(channels, block_in, kernel_size=3, factor=2) + self.shortcut_unpathify = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(channels, block_in, 2) + else: + self.conv_in = CausalConv(z_channels, block_in, kernel_size=3) + + # middle + self.mid = nn.Module() + self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) + + # upsampling + self.up_id = len(temporal_up_layers) + self.video_frame_num = 1 + self.cur_video_frame_num = self.video_frame_num // 2**self.up_id + 1 + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append(Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + if i_level in temporal_up_layers: + up.upsample = Upsample3D(block_in) + self.cur_video_frame_num = self.cur_video_frame_num * 2 + else: + up.upsample = Upsample2D(block_in, resamp_with_conv) + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) + self.conv_out = CausalConvAfterNorm(block_in, out_channels, kernel_size=3) + + @torch.inference_mode() + def forward(self, z, is_init=True): + z = rearrange(z, "b t c h w -> b c t h w") + + h = self.conv_in(z, is_init=is_init) + if self.version == 2: + shortcut = self.shortcut_in(z, is_init=is_init) + h = h.add_(shortcut) + shortcut = self.shortcut_unpathify(h, is_init=is_init) + h = self.conv_unpatchify(h, is_init=is_init) + h = h.add_(shortcut) + + temb = None + + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.mid.block_1(h, temb, is_init=is_init) + h = self.mid.attn_1(h) + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.mid.block_2(h, temb, is_init=is_init) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.up[i_level].block[i_block](h, temb, is_init=is_init) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + if isinstance(self.up[i_level].upsample, Upsample2D): + B = h.size(0) + h = h.permute(0, 2, 3, 4, 1).flatten(0, 1) + h = self.up[i_level].upsample(h) + h = h.unflatten(0, (B, -1)).permute(0, 4, 1, 2, 3) + else: + h = self.up[i_level].upsample(h, is_init=is_init) + + # end + h = h.permute(0, 2, 3, 4, 1) # b c l h w -> b l h w c + h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2) + h = self.conv_out(h) + h = h.permute(0, 4, 1, 2, 3) + + if is_init: + h = h[:, :, (self.temporal_downsample - 1):] + return h + + +def rms_norm(input, normalized_shape, eps=1e-6): + dtype = input.dtype + input = input.to(torch.float32) + variance = input.pow(2).flatten(-len(normalized_shape)).mean(-1)[(..., ) + (None, ) * len(normalized_shape)] + input = input * torch.rsqrt(variance + eps) + return input.to(dtype) + + +class DiagonalGaussianDistribution(object): + + def __init__(self, parameters, deterministic=False, rms_norm_mean=False, only_return_mean=False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=-3) #N,[X],C,H,W + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + self.deterministic = deterministic + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean, + device=self.parameters.device, + dtype=self.parameters.dtype) + if rms_norm_mean: + self.mean = rms_norm(self.mean, self.mean.size()[1:]) + self.only_return_mean = only_return_mean + + def sample(self, generator=None): + # make sure sample is on the same device + # as the parameters and has same dtype + sample = torch.randn(self.mean.shape, generator=generator, device=self.parameters.device) + sample = sample.to(dtype=self.parameters.dtype) + x = self.mean + self.std * sample + if self.only_return_mean: + return self.mean + else: + return x + + +class AutoencoderKL(nn.Module): + + @with_empty_init + def __init__( + self, + in_channels=3, + out_channels=3, + z_channels=16, + num_res_blocks=2, + model_path=None, + weight_dict={}, + world_size=1, + version=1, + ): + super().__init__() + + self.frame_len = 17 + self.latent_len = 3 if version == 2 else 5 + + base_group_norm.spatial = True if version == 2 else False + + self.encoder = VideoEncoder( + in_channels=in_channels, + z_channels=z_channels, + num_res_blocks=num_res_blocks, + version=version, + ) + + self.decoder = VideoDecoder( + z_channels=z_channels, + out_channels=out_channels, + num_res_blocks=num_res_blocks, + version=version, + ) + + if model_path is not None: + weight_dict = self.init_from_ckpt(model_path) + if len(weight_dict) != 0: + self.load_from_dict(weight_dict) + self.convert_channel_last() + + self.world_size = world_size + + def init_from_ckpt(self, model_path): + from safetensors import safe_open + p = {} + with safe_open(model_path, framework="pt", device="cpu") as f: + for k in f.keys(): + tensor = f.get_tensor(k) + if k.startswith("decoder.conv_out."): + k = k.replace("decoder.conv_out.", "decoder.conv_out.conv.") + p[k] = tensor + return p + + def load_from_dict(self, p): + self.load_state_dict(p) + + def convert_channel_last(self): + #Conv2d NCHW->NHWC + pass + + def naive_encode(self, x, is_init_image=True): + b, len, c, h, w = x.size() + x = rearrange(x, 'b l c h w -> b c l h w').contiguous() + z = self.encoder(x, len, True) # 下采样[1, 4, 8, 16, 16] + return z + + @torch.inference_mode() + def encode(self, x): + # b (nc cf) c h w -> (b nc) cf c h w -> encode -> (b nc) cf c h w -> b (nc cf) c h w + chunks = list(x.split(self.frame_len, dim=1)) + for i in range(len(chunks)): + chunks[i] = self.naive_encode(chunks[i], True) + z = torch.cat(chunks, dim=1) + + posterior = DiagonalGaussianDistribution(z) + return posterior.sample() + + def decode_naive(self, z, is_init=True): + z = z.to(next(self.decoder.parameters()).dtype) + dec = self.decoder(z, is_init) + return dec + + @torch.inference_mode() + def decode(self, z): + # b (nc cf) c h w -> (b nc) cf c h w -> decode -> (b nc) c cf h w -> b (nc cf) c h w + chunks = list(z.split(self.latent_len, dim=1)) + + if self.world_size > 1: + chunks_total_num = len(chunks) + max_num_per_rank = (chunks_total_num + self.world_size - 1) // self.world_size + rank = torch.distributed.get_rank() + chunks_ = chunks[max_num_per_rank * rank:max_num_per_rank * (rank + 1)] + if len(chunks_) < max_num_per_rank: + chunks_.extend(chunks[:max_num_per_rank - len(chunks_)]) + chunks = chunks_ + + for i in range(len(chunks)): + chunks[i] = self.decode_naive(chunks[i], True).permute(0, 2, 1, 3, 4) + x = torch.cat(chunks, dim=1) + + if self.world_size > 1: + x_ = torch.empty([x.size(0), (self.world_size * max_num_per_rank) * self.frame_len, *x.shape[2:]], + dtype=x.dtype, + device=x.device) + torch.distributed.all_gather_into_tensor(x_, x) + x = x_[:, :chunks_total_num * self.frame_len] + + x = self.mix(x) + return x + + def mix(self, x): + remain_scale = 0.6 + mix_scale = 1. - remain_scale + front = slice(self.frame_len - 1, x.size(1) - 1, self.frame_len) + back = slice(self.frame_len, x.size(1), self.frame_len) + x[:, back] = x[:, back] * remain_scale + x[:, front] * mix_scale + x[:, front] = x[:, front] * remain_scale + x[:, back] * mix_scale + return x diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py deleted file mode 100644 index db8612a98..000000000 --- a/fastvideo/pipelines/stages/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -""" -Pipeline stages for diffusion models. - -This package contains the various stages that can be composed to create -complete diffusion pipelines. -""" - -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.input_validation import InputValidationStage -from fastvideo.pipelines.stages.prompt_encoding import PromptEncodingStage -from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage -from fastvideo.pipelines.stages.latent_preparation import LatentPreparationStage -from fastvideo.pipelines.stages.conditioning import ConditioningStage -from fastvideo.pipelines.stages.denoising import DenoisingStage -from fastvideo.pipelines.stages.decoding import DecodingStage -from fastvideo.pipelines.stages.post_processing import PostProcessingStage - -__all__ = [ - "PipelineStage", - "InputValidationStage", - "PromptEncodingStage", - "TimestepPreparationStage", - "LatentPreparationStage", - "ConditioningStage", - "DenoisingStage", - "DecodingStage", - "PostProcessingStage", -] \ No newline at end of file diff --git a/fastvideo/sample/sample_t2v_hunyuan.py b/fastvideo/sample/sample_t2v_hunyuan.py index 40cce827d..0e595dc3c 100644 --- a/fastvideo/sample/sample_t2v_hunyuan.py +++ b/fastvideo/sample/sample_t2v_hunyuan.py @@ -11,8 +11,7 @@ from einops import rearrange from fastvideo.models.hunyuan.inference import HunyuanVideoSampler from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info -from fastvideo.inference_args import InferenceArgs -from fastvideo.utils.utils import FlexibleArgumentParser + def initialize_distributed(): local_rank = int(os.getenv("RANK", 0)) @@ -23,21 +22,21 @@ def initialize_distributed(): initialize_sequence_parallel_state(world_size) -def main(inference_args: InferenceArgs): +def main(args): initialize_distributed() print(nccl_info.sp_size) - print(inference_args) - models_root_path = Path(inference_args.model_path) + print(args) + models_root_path = Path(args.model_path) if not models_root_path.exists(): raise ValueError(f"`models_root` not exists: {models_root_path}") # Create save folder to save the samples - save_path = inference_args.output_path + save_path = args.output_path os.makedirs(os.path.dirname(save_path), exist_ok=True) # Load models - hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=inference_args) + hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args) # Get the updated args args = hunyuan_video_sampler.args @@ -74,15 +73,7 @@ def main(inference_args: InferenceArgs): if __name__ == "__main__": - parser = FlexibleArgumentParser() - InferenceArgs.add_cli_args(parser) - args = parser.parse_args() - inference_args = InferenceArgs.from_cli_args(args) - # process for vae sequence parallel - if inference_args.vae_sp and not inference_args.vae_tiling: - raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.") - main(inference_args) - exit() + parser = argparse.ArgumentParser() # Basic parameters parser.add_argument("--prompt", type=str, help="prompt file for inference") diff --git a/fastvideo/attention/__init__.py b/fastvideo/v1/attention/__init__.py similarity index 100% rename from fastvideo/attention/__init__.py rename to fastvideo/v1/attention/__init__.py diff --git a/fastvideo/distributed/device_communicators/__init__.py b/fastvideo/v1/attention/backends/__init__.py similarity index 100% rename from fastvideo/distributed/device_communicators/__init__.py rename to fastvideo/v1/attention/backends/__init__.py diff --git a/fastvideo/v1/attention/backends/abstract.py b/fastvideo/v1/attention/backends/abstract.py new file mode 100644 index 000000000..0cd95e074 --- /dev/null +++ b/fastvideo/v1/attention/backends/abstract.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from contextlib import contextmanager +from dataclasses import dataclass, fields +from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, + Protocol, Set, Tuple, Type, TypeVar) + +import torch + +from vllm.multimodal import MultiModalPlaceholderMap + +if TYPE_CHECKING: + from vllm.worker.model_runner_base import (ModelRunnerBase, + ModelRunnerInputBase, + ModelRunnerInputBuilderBase) + + +class AttentionType: + """ + Attention type. + Use string to be compatible with `torch.compile`. + """ + # Decoder attention between previous layer Q/K/V + DECODER = "decoder" + # Encoder attention between previous layer Q/K/V for encoder-decoder + ENCODER = "encoder" + # Encoder attention between previous layer Q/K/V + ENCODER_ONLY = "encoder_only" + # Attention between dec. Q and enc. K/V for encoder-decoder + ENCODER_DECODER = "encoder_decoder" + + +class AttentionBackend(ABC): + """Abstract class for attention backends.""" + # For some attention backends, we allocate an output tensor before + # calling the custom op. When piecewise cudagraph is enabled, this + # makes sure the output tensor is allocated inside the cudagraph. + accept_output_buffer: bool = False + + @staticmethod + @abstractmethod + def get_name() -> str: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_impl_cls() -> Type["AttentionImpl"]: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_metadata_cls() -> Type["AttentionMetadata"]: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_state_cls() -> Type["AttentionState"]: + raise NotImplementedError + + @classmethod + def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata": + return cls.get_metadata_cls()(*args, **kwargs) + + @staticmethod + @abstractmethod + def get_builder_cls() -> Type["AttentionMetadataBuilder"]: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + ) -> Tuple[int, ...]: + raise NotImplementedError + + @staticmethod + @abstractmethod + def swap_blocks( + src_kv_cache: torch.Tensor, + dst_kv_cache: torch.Tensor, + src_to_dst: torch.Tensor, + ) -> None: + raise NotImplementedError + + @staticmethod + @abstractmethod + def copy_blocks( + kv_caches: List[torch.Tensor], + src_to_dists: torch.Tensor, + ) -> None: + raise NotImplementedError + + def advance_step(self, model_input: "ModelRunnerInputBase", + sampled_token_ids: Optional[torch.Tensor], + block_size: int, num_seqs: int, num_queries: int) -> None: + raise NotImplementedError + + +@dataclass +class AttentionMetadata: + """Attention metadata for prefill and decode batched together.""" + # Total number of prefill requests. + num_prefills: int + # Number of prefill tokens. + num_prefill_tokens: int + # Number of decode tokens. Note that it is equivalent to the number of + # decode requests. + num_decode_tokens: int + # (num_tokens,). The indices of the token slots that input tokens will be + # stored into. E.g., if `slot_mapping` is [35, 2, 17] and the block size + # is 16, the three tokens are stored in the 3rd slot in block 2, 2nd slot + # in block 0, and 1st slot in block 1, respectively. + slot_mapping: torch.Tensor + + # The index maps that relate multi-modal embeddings to the corresponding + # placeholders. + # + # N.B. These aren't really related to attention and don't belong on this + # type -- this is just a temporary solution to make them available to + # `model_executable`. + multi_modal_placeholder_index_maps: Optional[Dict[ + str, MultiModalPlaceholderMap.IndexMap]] + + # Enable/disable KV scales calculation. This is so that we can disable the + # calculation until after prefill and cuda graph capture. + enable_kv_scales_calculation: bool + + @property + @abstractmethod + def prefill_metadata(self) -> Optional["AttentionMetadata"]: + """Return the attention metadata that's required to run prefill + attention.""" + pass + + @property + @abstractmethod + def decode_metadata(self) -> Optional["AttentionMetadata"]: + """Return the attention metadata that's required to run decode + attention.""" + pass + + def asdict_zerocopy(self, + skip_fields: Optional[Set[str]] = None + ) -> Dict[str, Any]: + """Similar to dataclasses.asdict, but avoids deepcopying.""" + if skip_fields is None: + skip_fields = set() + # Note that if we add dataclasses as fields, they will need + # similar handling. + return { + field.name: getattr(self, field.name) + for field in fields(self) if field.name not in skip_fields + } + + +T = TypeVar("T", bound=AttentionMetadata) + + +class AttentionState(ABC, Generic[T]): + """Holds attention backend-specific objects reused during the + lifetime of the model runner.""" + + @abstractmethod + def __init__(self, runner: "ModelRunnerBase"): + ... + + @abstractmethod + @contextmanager + def graph_capture(self, max_batch_size: int): + """Context manager used when capturing CUDA graphs.""" + yield + + @abstractmethod + def graph_clone(self, batch_size: int) -> "AttentionState[T]": + """Clone attention state to save in CUDA graph metadata.""" + ... + + @abstractmethod + def graph_capture_get_metadata_for_batch( + self, + batch_size: int, + is_encoder_decoder_model: bool = False) -> T: + """Get attention metadata for CUDA graph capture of batch_size.""" + ... + + @abstractmethod + def get_graph_input_buffers( + self, + attn_metadata: T, + is_encoder_decoder_model: bool = False) -> Dict[str, Any]: + """Get attention-specific input buffers for CUDA graph capture.""" + ... + + @abstractmethod + def prepare_graph_input_buffers( + self, + input_buffers: Dict[str, Any], + attn_metadata: T, + is_encoder_decoder_model: bool = False) -> None: + """In-place modify input buffers dict for CUDA graph replay.""" + ... + + @abstractmethod + def begin_forward(self, model_input: "ModelRunnerInputBase") -> None: + """Prepare state for forward pass.""" + ... + + +class AttentionMetadataBuilder(ABC, Generic[T]): + """Abstract class for attention metadata builders.""" + + @abstractmethod + def __init__(self, input_builder: "ModelRunnerInputBuilderBase") -> None: + """Create the builder, remember some configuration and parameters.""" + raise NotImplementedError + + @abstractmethod + def prepare(self) -> None: + """Prepare for one batch.""" + raise NotImplementedError + + @abstractmethod + def build(self, seq_lens: List[int], query_lens: List[int], + cuda_graph_pad_size: int, batch_size: int) -> T: + """Build attention metadata with on-device tensors.""" + raise NotImplementedError + + +class AttentionLayer(Protocol): + + _k_scale: torch.Tensor + _v_scale: torch.Tensor + _k_scale_float: float + _v_scale_float: float + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + ... + + +class AttentionImpl(ABC, Generic[T]): + + @abstractmethod + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: Optional[int] = None, + alibi_slopes: Optional[List[float]] = None, + sliding_window: Optional[int] = None, + kv_cache_dtype: str = "auto", + blocksparse_params: Optional[Dict[str, Any]] = None, + logits_soft_cap: Optional[float] = None, + attn_type: str = AttentionType.DECODER, + ) -> None: + raise NotImplementedError + + @abstractmethod + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: T, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + raise NotImplementedError + + +class MLAAttentionImpl(AttentionImpl[T], Generic[T]): + + @abstractmethod + def forward( + self, + layer: AttentionLayer, + hidden_states_or_cq: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: T, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + raise NotImplementedError + + +def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: + return kv_cache_dtype != "auto" diff --git a/fastvideo/v1/attention/backends/flash_attn.py b/fastvideo/v1/attention/backends/flash_attn.py new file mode 100755 index 000000000..c9c1e0c35 --- /dev/null +++ b/fastvideo/v1/attention/backends/flash_attn.py @@ -0,0 +1,1042 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Attention layer with FlashAttention.""" +from collections import defaultdict +from dataclasses import dataclass +from itertools import accumulate +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type + +import torch +import math + +from vllm import _custom_ops as ops +# yapf conflicts with isort for this block +# yapf: disable +from fastvideo.v1.attention.backends.abstract import (AttentionBackend, AttentionImpl, + AttentionLayer, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, + is_quantized_kv_cache) +# yapf: enable +from fastvideo.v1.attention.backends.utils import ( + PAD_SLOT_ID, CommonAttentionState, compute_slot_mapping, + compute_slot_mapping_start_idx, get_flash_attn_version, + get_num_prefill_decode_query_kv_tokens, get_seq_len_block_table_args, + is_all_cross_attn_metadata_set, is_all_encoder_attn_metadata_set, + is_block_tables_empty) +from fastvideo.v1.logger import init_logger +# from vllm.multimodal import MultiModalPlaceholderMap +from vllm.utils import async_tensor_h2d, make_tensor_with_pad +from vllm.vllm_flash_attn import (flash_attn_varlen_func, + flash_attn_with_kvcache) + +if TYPE_CHECKING: + from vllm.worker.model_runner import (ModelInputForGPUBuilder, + ModelInputForGPUWithSamplingMetadata) + +logger = init_logger(__name__) + + +class FlashAttentionBackend(AttentionBackend): + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> List[int]: + return [32, 64, 96, 128, 160, 192, 224, 256] + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN" + + @staticmethod + def get_impl_cls() -> Type["FlashAttentionImpl"]: + return FlashAttentionImpl + + @staticmethod + def get_metadata_cls() -> Type["AttentionMetadata"]: + return FlashAttentionMetadata + + @staticmethod + def get_builder_cls() -> Type["FlashAttentionMetadataBuilder"]: + return FlashAttentionMetadataBuilder + + @staticmethod + def get_state_cls() -> Type["CommonAttentionState"]: + return CommonAttentionState + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + ) -> Tuple[int, ...]: + if block_size % 16 != 0: + raise ValueError("Block size must be a multiple of 16.") + return (2, num_blocks, block_size, num_kv_heads, head_size) + + @staticmethod + def swap_blocks( + src_kv_cache: torch.Tensor, + dst_kv_cache: torch.Tensor, + src_to_dst: torch.Tensor, + ) -> None: + src_key_cache = src_kv_cache[0] + dst_key_cache = dst_kv_cache[0] + ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst) + src_value_cache = src_kv_cache[1] + dst_value_cache = dst_kv_cache[1] + ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst) + + @staticmethod + def copy_blocks( + kv_caches: List[torch.Tensor], + src_to_dists: torch.Tensor, + ) -> None: + key_caches = [kv_cache[0] for kv_cache in kv_caches] + value_caches = [kv_cache[1] for kv_cache in kv_caches] + + ops.copy_blocks(key_caches, value_caches, src_to_dists) + + +@dataclass +class FlashAttentionMetadata(AttentionMetadata): + """Metadata for FlashAttentionBackend. + + NOTE: Any python object stored here is not updated when it is + cuda-graph replayed. If you have values that need to be changed + dynamically, it should be stored in tensor. The tensor has to be + updated from `CUDAGraphRunner.forward` API. + """ + # (batch_size,). The sequence length per sequence. Sequence length means + # the computed tokens + new tokens None if it is a decoding. + seq_lens: Optional[List[int]] + # seq_lens stored as a tensor. + seq_lens_tensor: Optional[torch.Tensor] + + # NOTE(sang): Definition of context_len, query_len, and seq_len. + # |---------- N-1 iteration --------| + # |---------------- N iteration ---------------------| + # |- tokenA -|......................|-- newTokens ---| + # |---------- context_len ----------| + # |-------------------- seq_len ---------------------| + # |-- query_len ---| + + # Maximum sequence length among prefill batch. 0 if there are decoding + # requests only. + max_prefill_seq_len: int + # Maximum sequence length among decode batch. 0 if there are prefill + # requests only. + max_decode_seq_len: int + # (batch_size,) A tensor of context lengths (tokens that are computed + # so far). + context_lens_tensor: Optional[torch.Tensor] + + # (batch_size, max_blocks_per_seq). + # Block addresses per sequence. (Seq id -> list of physical block) + # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks + # in the kv cache. Each block can contain up to block_size tokens. + # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph + # captured. + block_tables: Optional[torch.Tensor] + + # Whether or not if cuda graph is enabled. + # Cuda-graph is currently enabled for decoding only. + # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention. + + use_cuda_graph: bool + + # Maximum query length in the batch. + max_query_len: Optional[int] = None + + # Max number of query tokens among request in the batch. + max_decode_query_len: Optional[int] = None + + # (batch_size + 1,). The cumulative subquery lengths of the sequences in + # the batch, used to index into subquery. E.g., if the subquery length + # is [4, 6], it is [0, 4, 10]. + query_start_loc: Optional[torch.Tensor] = None + # (batch_size + 1,). The cumulative sequence lengths of the sequences in + # the batch, used to index into sequence. E.g., if the sequence length is + # [4, 6], it is [0, 4, 10]. + seq_start_loc: Optional[torch.Tensor] = None + + _cached_prefill_metadata: Optional["FlashAttentionMetadata"] = None + _cached_decode_metadata: Optional["FlashAttentionMetadata"] = None + + # Begin encoder attn & enc/dec cross-attn fields... + + # Encoder sequence lengths representation + encoder_seq_lens: Optional[List[int]] = None + encoder_seq_lens_tensor: Optional[torch.Tensor] = None + # (batch_size + 1,). The cumulative sequence lengths of the sequences in + # the batch, used to index into sequence. E.g., if the sequence length is + # [4, 6], it is [0, 4, 10]. + encoder_seq_start_loc: Optional[torch.Tensor] = None + # Maximum sequence length among encoder sequences + max_encoder_seq_len: Optional[int] = None + # Number of tokens input to encoder + num_encoder_tokens: Optional[int] = None + + # Cross-attention memory-mapping data structures: slot mapping + # and block tables + cross_slot_mapping: Optional[torch.Tensor] = None + cross_block_tables: Optional[torch.Tensor] = None + + @property + def is_all_encoder_attn_metadata_set(self): + ''' + All attention metadata required for encoder attention is set. + ''' + return is_all_encoder_attn_metadata_set(self) + + @property + def is_all_cross_attn_metadata_set(self): + ''' + All attention metadata required for enc/dec cross-attention is set. + + Superset of encoder attention required metadata. + ''' + return is_all_cross_attn_metadata_set(self) + + @property + def prefill_metadata(self) -> Optional["FlashAttentionMetadata"]: + if self.num_prefills == 0: + return None + + if self._cached_prefill_metadata is not None: + return self._cached_prefill_metadata + + assert ((self.seq_lens is not None) + or (self.encoder_seq_lens is not None)) + assert ((self.seq_lens_tensor is not None) + or (self.encoder_seq_lens_tensor is not None)) + + # Compute some attn_metadata fields which default to None + query_start_loc = (None if self.query_start_loc is None else + self.query_start_loc[:self.num_prefills + 1]) + slot_mapping = (None if self.slot_mapping is None else + self.slot_mapping[:self.num_prefill_tokens]) + seq_lens = (None if self.seq_lens is None else + self.seq_lens[:self.num_prefills]) + seq_lens_tensor = (None if self.seq_lens_tensor is None else + self.seq_lens_tensor[:self.num_prefills]) + seq_start_loc = (None if self.seq_start_loc is None else + self.seq_start_loc[:self.num_prefills + 1]) + context_lens_tensor = (None if self.context_lens_tensor is None else + self.context_lens_tensor[:self.num_prefills]) + block_tables = (None if self.block_tables is None else + self.block_tables[:self.num_prefills]) + + self._cached_prefill_metadata = FlashAttentionMetadata( + num_prefills=self.num_prefills, + num_prefill_tokens=self.num_prefill_tokens, + num_decode_tokens=0, + slot_mapping=slot_mapping, + multi_modal_placeholder_index_maps=self. + multi_modal_placeholder_index_maps, + enable_kv_scales_calculation=self.enable_kv_scales_calculation, + seq_lens=seq_lens, + seq_lens_tensor=seq_lens_tensor, + max_query_len=self.max_query_len, + max_prefill_seq_len=self.max_prefill_seq_len, + max_decode_query_len=0, + max_decode_seq_len=0, + query_start_loc=query_start_loc, + seq_start_loc=seq_start_loc, + context_lens_tensor=context_lens_tensor, + block_tables=block_tables, + use_cuda_graph=False, + # Begin encoder & cross attn fields below... + encoder_seq_lens=self.encoder_seq_lens, + encoder_seq_lens_tensor=self.encoder_seq_lens_tensor, + encoder_seq_start_loc=self.encoder_seq_start_loc, + max_encoder_seq_len=self.max_encoder_seq_len, + cross_slot_mapping=self.cross_slot_mapping, + cross_block_tables=self.cross_block_tables) + return self._cached_prefill_metadata + + @property + def decode_metadata(self) -> Optional["FlashAttentionMetadata"]: + if self.num_decode_tokens == 0: + return None + + if self._cached_decode_metadata is not None: + return self._cached_decode_metadata + assert ((self.seq_lens_tensor is not None) + or (self.encoder_seq_lens_tensor is not None)) + + # Compute some attn_metadata fields which default to None + slot_mapping = (None if self.slot_mapping is None else + self.slot_mapping[self.num_prefill_tokens:]) + seq_lens_tensor = (None if self.seq_lens_tensor is None else + self.seq_lens_tensor[self.num_prefills:]) + block_tables = (None if self.block_tables is None else + self.block_tables[self.num_prefills:]) + + self._cached_decode_metadata = FlashAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decode_tokens=self.num_decode_tokens, + slot_mapping=slot_mapping, + multi_modal_placeholder_index_maps=None, + enable_kv_scales_calculation=True, + seq_lens=None, + seq_lens_tensor=seq_lens_tensor, + max_decode_query_len=self.max_decode_query_len, + max_query_len=self.max_query_len, + max_prefill_seq_len=0, + max_decode_seq_len=self.max_decode_seq_len, + # Batch may be composed of prefill|decodes, adjust query start + # indices to refer to the start of decodes. E.g. + # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6]. + query_start_loc=(self.query_start_loc[self.num_prefills:] - + self.query_start_loc[self.num_prefills]) + if self.query_start_loc is not None else None, + seq_start_loc=self.seq_start_loc[self.num_prefills:] + if self.seq_start_loc is not None else None, + context_lens_tensor=None, + block_tables=block_tables, + use_cuda_graph=self.use_cuda_graph, + # Begin encoder & cross attn fields below... + encoder_seq_lens=self.encoder_seq_lens, + encoder_seq_lens_tensor=self.encoder_seq_lens_tensor, + encoder_seq_start_loc=self.encoder_seq_start_loc, + max_encoder_seq_len=self.max_encoder_seq_len, + cross_slot_mapping=self.cross_slot_mapping, + cross_block_tables=self.cross_block_tables) + return self._cached_decode_metadata + + def advance_step(self, + model_input: "ModelInputForGPUWithSamplingMetadata", + sampled_token_ids: Optional[torch.Tensor], + block_size: int, + num_seqs: int, + num_queries: int, + turn_prefills_into_decodes: bool = False): + """ + Update metadata in-place to advance one decode step. + """ + # When using cudagraph, the num_seqs is padded to the next captured + # batch sized, but num_queries tracks the actual number of requests in + # the batch. For --enforce-eager mode, num_seqs == num_queries + if num_seqs != num_queries: + assert num_seqs > num_queries + assert self.use_cuda_graph + + if turn_prefills_into_decodes: + # When Mutli-Step is enabled with Chunked-Prefill, prefills and + # decodes are scheduled together. In the first step, all the + # prefills turn into decodes. This update reflects that + # conversion. + assert self.num_decode_tokens + self.num_prefills == num_seqs + self.num_decode_tokens += self.num_prefills + self.num_prefills = 0 + self.num_prefill_tokens = 0 + self.max_prefill_seq_len = 0 + self.max_query_len = 1 + + self.slot_mapping = self.slot_mapping[:num_seqs] + else: + assert self.seq_lens is not None + assert self.max_decode_seq_len == max(self.seq_lens) + + assert self.num_prefills == 0 + assert self.num_prefill_tokens == 0 + assert self.num_decode_tokens == num_seqs + assert self.slot_mapping.shape == (num_seqs, ) + + assert self.seq_lens is not None + assert len(self.seq_lens) == num_seqs + assert self.seq_lens_tensor is not None + assert self.seq_lens_tensor.shape == (num_seqs, ) + assert self.max_query_len == 1 + assert self.max_prefill_seq_len == 0 + + assert self.query_start_loc is not None + assert self.query_start_loc.shape == (num_queries + 1, ) + assert self.seq_start_loc is not None + assert self.seq_start_loc.shape == (num_seqs + 1, ) + + assert self.context_lens_tensor is not None + assert self.context_lens_tensor.shape == (num_queries, ) + + assert self.block_tables is not None + assert self.block_tables.shape[0] == num_seqs + + # Update query lengths. Note that we update only queries and not seqs, + # since tensors may be padded due to captured cuda graph batch size + for i in range(num_queries): + self.seq_lens[i] += 1 + self.max_decode_seq_len = max(self.seq_lens) + + ops.advance_step_flashattn(num_seqs=num_seqs, + num_queries=num_queries, + block_size=block_size, + input_tokens=model_input.input_tokens, + sampled_token_ids=sampled_token_ids, + input_positions=model_input.input_positions, + seq_lens=self.seq_lens_tensor, + slot_mapping=self.slot_mapping, + block_tables=self.block_tables) + + +class FlashAttentionMetadataBuilder( + AttentionMetadataBuilder[FlashAttentionMetadata]): + + def __init__(self, input_builder: "ModelInputForGPUBuilder"): + self.input_builder = input_builder + self.runner = input_builder.runner + self.sliding_window = input_builder.sliding_window + self.block_size = input_builder.block_size + + def prepare(self): + self.slot_mapping: List[int] = [] + self.prefill_seq_lens: List[int] = [] + self.context_lens: List[int] = [] + self.block_tables: List[List[int]] = [] + self.curr_seq_lens: List[int] = [] + self.multimodal_placeholder_maps: Dict[ + str, + MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap) + self.num_prefills = 0 + self.num_prefill_tokens = 0 + self.num_decode_tokens = 0 + self.has_prefix_cache_hit = False + + def _add_seq_group( + self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup", + chunked_prefill_enabled: bool, prefix_cache_hit: bool): + """Add a sequence group to the metadata. Specifically update/append + 1. context length. + 2. block table. + 3. slot mapping. + """ + is_prompt = inter_data.is_prompt + block_tables = inter_data.block_tables + + for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len, + curr_sliding_window_block) in zip( + inter_data.seq_ids, [len(t) for t in inter_data.input_tokens], + inter_data.orig_seq_lens, inter_data.seq_lens, + inter_data.query_lens, inter_data.context_lens, + inter_data.curr_sliding_window_blocks): + self.context_lens.append(context_len) + + if is_prompt: + mm_maps = inter_data.multi_modal_placeholder_maps + if mm_maps: + for modality, placeholders in mm_maps.items(): + self.multimodal_placeholder_maps[modality].extend( + placeholders) + + self.num_prefills += 1 + self.num_prefill_tokens += token_len + self.prefill_seq_lens.append(seq_len) + else: + self.num_decode_tokens += query_len + self.curr_seq_lens.append(curr_seq_len) + + # Compute block table. + # TODO(sang): Combine chunked prefill and prefix caching by + # only allowing multiple of block_size chunk size. + # NOTE: This only works for oooooooxxx style attention. + block_table = [] + if prefix_cache_hit: + # NOTE(woosuk): For flash-attn, the block table should + # include the entries for the incoming prefill tokens. + block_table = block_tables[seq_id] + elif ((chunked_prefill_enabled or not is_prompt) + and block_tables is not None): + if curr_sliding_window_block == 0: + block_table = block_tables[seq_id] + else: + block_table = block_tables[seq_id][ + -curr_sliding_window_block:] + self.block_tables.append(block_table) + + # Compute slot mapping. + is_profile_run = is_block_tables_empty(block_tables) + start_idx = compute_slot_mapping_start_idx(is_prompt, query_len, + context_len, + self.sliding_window) + compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id, + seq_len, context_len, start_idx, + self.block_size, inter_data.block_tables) + + def _get_graph_runner_block_tables( + self, num_seqs: int, + block_tables: List[List[int]]) -> torch.Tensor: + # The shape of graph_block_tables is + # [max batch size, max context len // block size]. + max_batch_size, max_blocks = self.runner.graph_block_tables.shape + assert max_batch_size >= num_seqs + + graph_block_tables = self.runner.graph_block_tables[:num_seqs] + for i, block_table in enumerate(block_tables): + if block_table: + num_blocks = len(block_table) + if num_blocks <= max_blocks: + graph_block_tables[i, :num_blocks] = block_table + else: + # It may be possible to have more blocks allocated due + # to lookahead slots of multi-step, however, they are + # not used anyway, so can be safely ignored. + graph_block_tables[ + i, :max_blocks] = block_table[:max_blocks] + + return torch.from_numpy(graph_block_tables).to( + device=self.runner.device, non_blocking=True) + + def build(self, seq_lens: List[int], query_lens: List[int], + cuda_graph_pad_size: int, batch_size: int): + """Build attention metadata with on-device tensors. + + Args: + seq_lens: The maybe padded sequence lengths of the input sequences. + query_lens: The query lengths of the input sequences. + cuda_graph_pad_size: The padding size for cuda graph. + -1 if cuda graph is not used. + batch_size: The maybe padded batch size. + """ + prefix_cache_hit = any([ + inter_data.prefix_cache_hit + for inter_data in self.input_builder.inter_data_list + ]) + for inter_data in self.input_builder.inter_data_list: + self._add_seq_group(inter_data, + self.input_builder.chunked_prefill_enabled, + prefix_cache_hit) + + device = self.runner.device + use_captured_graph = cuda_graph_pad_size != -1 + + max_query_len = max(query_lens) + decode_query_lens = query_lens[self.num_prefills:] + if len(decode_query_lens) > 0: + max_decode_query_len = max(decode_query_lens) + else: + max_decode_query_len = 1 + max_prefill_seq_len = max(self.prefill_seq_lens, default=0) + max_decode_seq_len = max(self.curr_seq_lens, default=0) + num_decode_tokens = self.num_decode_tokens + query_start_loc = list(accumulate(query_lens, initial=0)) + seq_start_loc = list(accumulate(seq_lens, initial=0)) + + num_seqs = len(seq_lens) + if use_captured_graph: + self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size) + self.block_tables.extend([] * cuda_graph_pad_size) + num_decode_tokens = batch_size - self.num_prefill_tokens + block_tables = self._get_graph_runner_block_tables( + num_seqs, self.block_tables) + else: + block_tables = make_tensor_with_pad( + self.block_tables, + pad=0, + dtype=torch.int, + device=device, + ) + assert max_query_len > 0, ("query_lens: {}".format(query_lens)) + + assert device is not None + context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int, + device, self.runner.pin_memory) + seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device, + self.runner.pin_memory) + slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long, + device, self.runner.pin_memory) + query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32, + device, + self.runner.pin_memory) + seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32, + device, self.runner.pin_memory) + placeholder_index_maps = { + modality: placeholder_map.index_map() + for modality, placeholder_map in + self.multimodal_placeholder_maps.items() + } + + return FlashAttentionMetadata( + num_prefills=self.num_prefills, + slot_mapping=slot_mapping_tensor, + num_prefill_tokens=self.num_prefill_tokens, + num_decode_tokens=num_decode_tokens, + seq_lens=seq_lens, + multi_modal_placeholder_index_maps=placeholder_index_maps, + enable_kv_scales_calculation=True, + seq_lens_tensor=seq_lens_tensor, + max_query_len=max_query_len, + max_decode_query_len=max_decode_query_len, + max_prefill_seq_len=max_prefill_seq_len, + max_decode_seq_len=max_decode_seq_len, + query_start_loc=query_start_loc_tensor, + seq_start_loc=seq_start_loc_tensor, + context_lens_tensor=context_lens_tensor, + block_tables=block_tables, + use_cuda_graph=use_captured_graph, + ) + + +class FlashAttentionImpl(AttentionImpl): + """ + If the input tensors contain prompt tokens, the layout is as follows: + |<--------------- num_prefill_tokens ----------------->| + |<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->| + + Otherwise, the layout is as follows: + |<----------------- num_decode_tokens ------------------>| + |<--decode_0-->|..........|<--decode_M-1-->|<--padding-->| + + Generation tokens can contain padding when cuda-graph is used. + Currently, prompt tokens don't contain any padding. + + The prompts might have different lengths, while the generation tokens + always have length 1. + + If chunked prefill is enabled, prefill tokens and decode tokens can be + batched together in a flattened 1D query. + + |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->| + |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->| + + Currently, cuda graph is disabled for chunked prefill, meaning there's no + padding between prefill and decode tokens. + """ + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: Optional[List[float]], + sliding_window: Optional[int], + kv_cache_dtype: str, + blocksparse_params: Optional[Dict[str, Any]] = None, + logits_soft_cap: Optional[float] = None, + attn_type: str = AttentionType.DECODER, + ) -> None: + if blocksparse_params is not None: + raise ValueError( + "FlashAttention does not support block-sparse attention.") + self.num_heads = num_heads + self.head_size = head_size + self.scale = float(scale) + self.num_kv_heads = num_kv_heads + if alibi_slopes is not None: + alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32) + self.alibi_slopes = alibi_slopes + self.sliding_window = ((sliding_window - 1, + 0) if sliding_window is not None else (-1, -1)) + self.kv_cache_dtype = kv_cache_dtype + if is_quantized_kv_cache(self.kv_cache_dtype): + raise NotImplementedError( + "FlashAttention with FP8 KV cache not yet supported") + if logits_soft_cap is None: + # In flash-attn, setting logits_soft_cap as 0 means no soft cap. + logits_soft_cap = 0 + self.logits_soft_cap = logits_soft_cap + + assert self.num_heads % self.num_kv_heads == 0 + self.num_queries_per_kv = self.num_heads // self.num_kv_heads + + support_head_sizes = FlashAttentionBackend.get_supported_head_sizes() + if head_size not in support_head_sizes: + raise ValueError( + f"Head size {head_size} is not supported by FlashAttention. " + f"Supported head sizes are: {support_head_sizes}.") + self.attn_type = attn_type + self.vllm_flash_attn_version = get_flash_attn_version() + + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: FlashAttentionMetadata, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass with FlashAttention. + + Args: + query: shape = [num_tokens, num_heads, head_size] + key: shape = [num_tokens, num_kv_heads, head_size] + value: shape = [num_tokens, num_kv_heads, head_size] + output: shape = [num_tokens, num_heads, head_size] + kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size] + NOTE: kv_cache will be an empty tensor with shape [0] + for profiling run. + attn_metadata: Metadata for attention. + NOTE: It in-place updates the output tensor. + """ + # NOTE(woosuk): FlashAttention does not support FP8 KV cache. + assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0, ( + "key/v_scale is not supported in FlashAttention.") + + assert output is not None, "Output tensor must be provided." + + attn_type = self.attn_type + if (attn_type == AttentionType.ENCODER + and (not attn_metadata.is_all_encoder_attn_metadata_set)): + raise AttributeError("Encoder attention requires setting " + "encoder metadata attributes.") + elif (attn_type == AttentionType.ENCODER_DECODER + and (not attn_metadata.is_all_cross_attn_metadata_set)): + raise AttributeError("Encoder/decoder cross-attention " + "requires setting cross-attention " + "metadata attributes.") + + kv_cache_dtype: str = self.kv_cache_dtype + softmax_scale: float = self.scale + window_size = self.sliding_window + alibi_slopes: Optional[torch.Tensor] = self.alibi_slopes + logits_soft_cap: Optional[float] = self.logits_soft_cap + + if kv_cache.numel() > 0: + key_cache = kv_cache[0] + value_cache = kv_cache[1] + # We skip updating the KV cache under two conditions: + # a. When the Attention Type is ENCODER. In this phase, we compute + # only the encoder attention without updating the cache. + # b. When both Key and Value are None. This occurs during + # cross-attention computation in the decoding phase, where the + # KV cache is already populated with the cross-attention + # tensor. Thus, we skip cache updates during this time. + if (attn_type != AttentionType.ENCODER) and (key is not None) and ( + value is not None): + if attn_type == AttentionType.ENCODER_DECODER: + # Update cross-attention KV cache (prefill-only) + updated_slot_mapping = attn_metadata.cross_slot_mapping + else: + # Update self-attention KV cache (prefill/decode) + updated_slot_mapping = attn_metadata.slot_mapping + + # Reshape the input keys and values and store them in the cache. + # If kv_cache is not provided, the new key and value tensors are + # not cached. This happens during the initial memory + # profiling run. + torch.ops._C_cache_ops.reshape_and_cache_flash( + key, + value, + kv_cache[0], + kv_cache[1], + updated_slot_mapping.flatten(), # type: ignore[union-attr] + kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + + (num_prefill_query_tokens, num_prefill_kv_tokens, + num_decode_query_tokens) = \ + get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type) + decode_query = query[num_prefill_query_tokens:] + decode_output = output[num_prefill_query_tokens:] + # QKV for prefill. + query = query[:num_prefill_query_tokens] + prefill_output = output[:num_prefill_query_tokens] + assert query.shape[0] == num_prefill_query_tokens + assert decode_query.shape[0] == num_decode_query_tokens + + if prefill_meta := attn_metadata.prefill_metadata: + # Prompt run. + if (kv_cache.numel() == 0 or prefill_meta.block_tables is None + or prefill_meta.block_tables.numel() == 0): + # normal attention + # When block_tables are not filled, it means q and k are the + # prompt, and they have the same length. + q_seq_start_loc, q_seq_len, k_seq_start_loc, k_seq_len = \ + _get_query_key_seq_metadata(prefill_meta, True, attn_type) + + key = key[:num_prefill_kv_tokens] + value = value[:num_prefill_kv_tokens] + + flash_attn_varlen_func( + q=query, + k=key, + v=value, + cu_seqlens_q=q_seq_start_loc, + cu_seqlens_k=k_seq_start_loc, + max_seqlen_q=q_seq_len, + max_seqlen_k=k_seq_len, + softmax_scale=softmax_scale, + causal=_get_causal_option(attn_type), + window_size=window_size, + alibi_slopes=alibi_slopes, + softcap=logits_soft_cap, + out=prefill_output, + fa_version=self.vllm_flash_attn_version, + ) + else: + # prefix-enabled attention + assert attn_type == AttentionType.DECODER, ( + "Only decoder-only models support prefix caching") + assert prefill_meta.seq_lens is not None + max_seq_len = max(prefill_meta.seq_lens) + flash_attn_varlen_func( # noqa + q=query, + k=key_cache, + v=value_cache, + cu_seqlens_q=prefill_meta.query_start_loc, + max_seqlen_q=prefill_meta.max_query_len, + seqused_k=prefill_meta.seq_lens_tensor, + max_seqlen_k=max_seq_len, + softmax_scale=softmax_scale, + causal=True, + window_size=window_size, + alibi_slopes=alibi_slopes, + block_table=prefill_meta.block_tables, + softcap=logits_soft_cap, + out=prefill_output, + fa_version=self.vllm_flash_attn_version, + ) + + if decode_meta := attn_metadata.decode_metadata: + # Decoding run. + # Use flash_attn_varlen_func kernel for speculative decoding + # because different queries might have different lengths. + + assert decode_meta.max_decode_query_len is not None + # use only for actual varlen decoding + if decode_meta.max_decode_query_len > 1: + assert attn_type == AttentionType.DECODER, ( + "Only decoder-only models support max_decode_query_len > 1" + ) + flash_attn_varlen_func( + q=decode_query, + k=key_cache, + v=value_cache, + cu_seqlens_q=decode_meta.query_start_loc, + max_seqlen_q=decode_meta.max_decode_query_len, + seqused_k=decode_meta.seq_lens_tensor, + max_seqlen_k=decode_meta.max_decode_seq_len, + softmax_scale=softmax_scale, + causal=True, + window_size=window_size, + alibi_slopes=alibi_slopes, + softcap=logits_soft_cap, + block_table=decode_meta.block_tables, + out=decode_output, + fa_version=self.vllm_flash_attn_version, + ) + else: + # Use flash_attn_with_kvcache for normal decoding. + ( + seq_lens_arg, + _, + block_tables_arg, + ) = get_seq_len_block_table_args(decode_meta, False, attn_type) + flash_attn_with_kvcache( + q=decode_query.unsqueeze(1), + k_cache=key_cache, + v_cache=value_cache, + block_table=block_tables_arg, + cache_seqlens=seq_lens_arg, + softmax_scale=softmax_scale, + causal=True, + window_size=window_size, + alibi_slopes=alibi_slopes, + softcap=logits_soft_cap, + out=decode_output.unsqueeze(1), + fa_version=self.vllm_flash_attn_version, + ) + return output + + +def _get_query_key_seq_metadata( + attn_metadata, + is_prompt: bool, + attn_type: str, +) -> tuple: + """ + Returns sequence metadata for key and query based on the specified + attention type and whether input is a prompt. + + This function computes the starting locations and maximum sequence lengths + for key and query sequences for different attention types. + + Args: + attn_metadata: The attention metadata object + is_prompt (bool): A flag indicating if the input is a prompt + attn_type (AttentionType): The type of attention being used. + + Returns: + tuple: A tuple containing four integers: + - Starting location for the query sequence. + - Maximum sequence length for the query sequence. + - Starting location for the key sequence. + - Maximum sequence length for the key sequence. + + Raises: + AttributeError: If an invalid attention type is provided. + """ + if attn_type == AttentionType.DECODER: + # Decoder self-attention + # Choose max_seq_len based on whether we are in prompt_run + if is_prompt: + max_seq_len = attn_metadata.max_prefill_seq_len + else: + max_seq_len = attn_metadata.max_decode_seq_len + return (attn_metadata.seq_start_loc, max_seq_len, + attn_metadata.seq_start_loc, max_seq_len) + + elif attn_type == AttentionType.ENCODER_DECODER: + # This is cross attention between the where the key + # is the precomputed encoder attention and query + # is the input sequence. + # Choose query max length based on whether it is prompt + # or not. + if is_prompt: + max_seq_len = attn_metadata.max_prefill_seq_len + else: + max_seq_len = attn_metadata.max_decode_seq_len + return (attn_metadata.seq_start_loc, max_seq_len, + attn_metadata.encoder_seq_start_loc, + attn_metadata.max_encoder_seq_len) + elif attn_type == AttentionType.ENCODER: + # For encoder attention both the query and the key are same i.e the + # encoder sequence. + return (attn_metadata.encoder_seq_start_loc, + attn_metadata.max_encoder_seq_len, + attn_metadata.encoder_seq_start_loc, + attn_metadata.max_encoder_seq_len) + elif attn_type == AttentionType.ENCODER_ONLY: + assert is_prompt, "Should not have decode for encoder only model." + return (attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len, + attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len) + else: + raise AttributeError(f"Invalid attention type {str(attn_type)}") + + +def _get_causal_option(attn_type: str) -> bool: + """ + Determine whether the given attention type is suitable for causal + attention mechanisms. + + Args: + attn_type (AttentionType): The type of attention being evaluated + + Returns: + bool: Returns `True` if the attention type is suitable for causal + attention (i.e., not encoder, encoder-only, or encoder-decoder), + otherwise returns `False`. + """ + return not (attn_type == AttentionType.ENCODER + or attn_type == AttentionType.ENCODER_ONLY + or attn_type == AttentionType.ENCODER_DECODER) + + +class TextEncoderFlashAttentionImpl: + """Simplified Flash Attention implementation for text encoders in diffusion models. + + This implementation removes KV caching and other generation-specific features, + focusing only on the encoder attention pattern needed for text encoders in + diffusion models like DiT. + """ + + def __init__( + self, + num_heads: int, + head_size: int, + scale: Optional[float] = None, + ) -> None: + """Initialize the text encoder flash attention implementation. + + Args: + num_heads: Number of attention heads + head_size: Size of each attention head + scale: Attention scale factor. If None, defaults to 1/sqrt(head_size) + """ + self.num_heads = num_heads + self.head_size = head_size + self.num_kv_heads = num_heads # For encoder, typically num_kv_heads = num_heads + + # Set default scale if not provided + if scale is None: + scale = 1.0 / math.sqrt(head_size) + self.scale = float(scale) + + # Set defaults for optional parameters + self.alibi_slopes = None + self.sliding_window = (-1, -1) # No sliding window by default + self.logits_soft_cap = 0 # No soft cap by default + + # Always use encoder-only attention type + self.attn_type = AttentionType.ENCODER_ONLY + self.vllm_flash_attn_version = get_flash_attn_version() + + # Verify head size is supported + support_head_sizes = FlashAttentionBackend.get_supported_head_sizes() + if head_size not in support_head_sizes: + raise ValueError( + f"Head size {head_size} is not supported by FlashAttention. " + f"Supported head sizes are: {support_head_sizes}.") + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + seq_lens: Optional[List[int]] = None, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass with FlashAttention for text encoder. + + Args: + query: shape = [num_tokens, num_heads, head_size] + key: shape = [num_tokens, num_kv_heads, head_size] + value: shape = [num_tokens, num_kv_heads, head_size] + seq_lens: List of sequence lengths for batched input + output: shape = [num_tokens, num_heads, head_size] + Pre-allocated output tensor (optional) + + Returns: + torch.Tensor: Output tensor with shape [num_tokens, num_heads, head_size] + """ + device = query.device + batch_size = len(seq_lens) if seq_lens is not None else 1 + + if output is None: + output = torch.empty_like(query) + + if seq_lens is None: + # Single sequence case - use standard flash attention + return flash_attn_varlen_func( + q=query, + k=key, + v=value, + softmax_scale=self.scale, + causal=False, # Non-causal for encoder + window_size=self.sliding_window, + alibi_slopes=self.alibi_slopes, + softcap=self.logits_soft_cap, + out=output, + fa_version=self.vllm_flash_attn_version, + ) + else: + # Multiple sequences case - use varlen flash attention + max_seq_len = max(seq_lens) + seq_start_loc = list(accumulate(seq_lens, initial=0)) + seq_start_loc_tensor = torch.tensor(seq_start_loc, + dtype=torch.int32, + device=device) + + return flash_attn_varlen_func( + q=query, + k=key, + v=value, + cu_seqlens_q=seq_start_loc_tensor, + cu_seqlens_k=seq_start_loc_tensor, + max_seqlen_q=max_seq_len, + max_seqlen_k=max_seq_len, + softmax_scale=self.scale, + causal=False, # Non-causal for encoder + window_size=self.sliding_window, + alibi_slopes=self.alibi_slopes, + softcap=self.logits_soft_cap, + out=output, + fa_version=self.vllm_flash_attn_version, + ) diff --git a/fastvideo/v1/attention/backends/utils.py b/fastvideo/v1/attention/backends/utils.py new file mode 100644 index 000000000..4374b5422 --- /dev/null +++ b/fastvideo/v1/attention/backends/utils.py @@ -0,0 +1,619 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Attention backend utils""" +from collections import defaultdict +from contextlib import contextmanager +from itertools import accumulate +from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Type, TypeVar, Union + +import numpy as np +import torch + +from vllm import envs +from vllm.attention import (AttentionMetadata, AttentionMetadataBuilder, + AttentionState) +from vllm.attention.backends.abstract import AttentionType +from vllm.logger import init_logger +from vllm.multimodal import MultiModalPlaceholderMap +from vllm.platforms import current_platform +from vllm.utils import async_tensor_h2d, make_tensor_with_pad + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm.worker.model_runner_base import ModelRunnerBase + +# Error string(s) for encoder/decoder +# unsupported attention scenarios +STR_NOT_IMPL_ENC_DEC_ROCM_HIP = ("ROCm/HIP is not currently supported " + "with encoder/decoder models.") + +PAD_SLOT_ID = -1 + +# Switch to numpy implementation of compute_slot_mapping +# if we have at least this many elements. Could be tuned further. +_COMPUTE_SLOT_MAPPING_NUMPY_NUMEL = 256 + +if TYPE_CHECKING: + from vllm.worker.model_runner import ModelInputForGPUBuilder + + +def is_block_tables_empty(block_tables: Union[None, Dict]): + """ + Check if block_tables is None or a dictionary with all None values. + """ + if block_tables is None: + return True + return (isinstance(block_tables, dict) + and all(value is None for value in block_tables.values())) + + +def compute_slot_mapping_start_idx(is_prompt: bool, query_len: int, + context_len: int, sliding_window: int): + """ + Compute the start index of slot mapping. + """ + start_idx = 0 + if is_prompt and sliding_window is not None: + start_idx = max(0, query_len - sliding_window) + return start_idx + + +def _compute_slot_mapping_python(slot_mapping: List[int], + block_table: List[int], range_start: int, + range_end: int, block_size: int): + for i in range(range_start, range_end): + block_number = block_table[i // block_size] + block_offset = i % block_size + slot = block_number * block_size + block_offset + slot_mapping.append(slot) + + +def _compute_slot_mapping_numpy(slot_mapping: List[int], + block_table: List[int], range_start: int, + range_end: int, block_size: int): + block_table_array = np.array(block_table) + idx = np.arange(range_start, range_end) + block_offset = idx % block_size + idx //= block_size + seq_slot_mapping_array = block_table_array[idx] + seq_slot_mapping_array *= block_size + seq_slot_mapping_array += block_offset + slot_mapping.extend(seq_slot_mapping_array) + + +def compute_slot_mapping(is_profile_run: bool, slot_mapping: List[int], + seq_id: int, seq_len: int, context_len: int, + start_idx: int, block_size: int, + block_tables: Dict[int, List[int]]): + """ + Compute slot mapping. + """ + if is_profile_run: + # During memory profiling, the block tables are not + # initialized yet. In this case, we just use a dummy + # slot mapping. + # In embeddings, the block tables are {seq_id: None}. + slot_mapping.extend([PAD_SLOT_ID] * seq_len) + return + + # Mask the [0, start_idx) tokens of the prompt with + # PAD_SLOT_ID, where start_idx is max(0, seq_len - + # sliding_window). For example, if the prompt len is 10, + # sliding window is 8, and block size is 4, the first two + # tokens are masked and the slot mapping will be + # [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1]. + padding_mask_len = max(0, start_idx - context_len) + slot_mapping.extend([PAD_SLOT_ID] * padding_mask_len) + + range_start = max(start_idx, context_len) + range_end = seq_len + numel = range_end - range_start + block_table = block_tables[seq_id] + + # numpy implementation will be faster than python if we have + # many elements, otherwise it will be slower. + if numel < _COMPUTE_SLOT_MAPPING_NUMPY_NUMEL: + _compute_slot_mapping_python(slot_mapping, block_table, range_start, + range_end, block_size) + else: + _compute_slot_mapping_numpy(slot_mapping, block_table, range_start, + range_end, block_size) + + +TAttentionMetadata = TypeVar("TAttentionMetadata", bound='AttentionMetadata') + + +class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]): + + _metadata_cls: Type[TAttentionMetadata] + + def __init__(self, input_builder: "ModelInputForGPUBuilder"): + self.input_builder = input_builder + self.runner = input_builder.runner + + self.sliding_window = input_builder.sliding_window + self.block_size = input_builder.block_size + + def prepare(self): + self.slot_mapping: List[int] = [] + self.prefill_seq_lens: List[int] = [] + self.context_lens: List[int] = [] + self.block_tables: List[List[int]] = [] + self.curr_seq_lens: List[int] = [] + self.multimodal_placeholder_maps: Dict[ + str, + MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap) + self.num_prefills = 0 + self.num_prefill_tokens = 0 + self.num_decode_tokens = 0 + + def _add_seq_group( + self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup", + chunked_prefill_enabled: bool): + is_prompt = inter_data.is_prompt + block_tables = inter_data.block_tables + + for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len, + curr_sliding_window_block) in zip( + inter_data.seq_ids, [len(t) for t in inter_data.input_tokens], + inter_data.orig_seq_lens, inter_data.seq_lens, + inter_data.query_lens, inter_data.context_lens, + inter_data.curr_sliding_window_blocks): + self.context_lens.append(context_len) + if is_prompt: + mm_maps = inter_data.multi_modal_placeholder_maps + if mm_maps: + for modality, placeholders in mm_maps.items(): + self.multimodal_placeholder_maps[modality].extend( + placeholders) + + self.num_prefills += 1 + self.num_prefill_tokens += token_len + self.prefill_seq_lens.append(seq_len) + else: + assert query_len == 1, ( + "seq_len: {}, context_len: {}, query_len: {}".format( + seq_len, context_len, query_len)) + self.num_decode_tokens += query_len + self.curr_seq_lens.append(curr_seq_len) + + # Compute block table. + # TODO(sang): Combine chunked prefill and prefix caching by + # only allowing multiple of block_size chunk size. + # NOTE: This only works for oooooooxxx style attention. + block_table = [] + if inter_data.prefix_cache_hit: + block_table = block_tables[seq_id] + elif ((chunked_prefill_enabled or not is_prompt) + and block_tables is not None): + if curr_sliding_window_block == 0: + block_table = block_tables[seq_id] + else: + block_table = block_tables[seq_id][ + -curr_sliding_window_block:] + self.block_tables.append(block_table) + + # Compute slot mapping. + is_profile_run = is_block_tables_empty(block_tables) + start_idx = compute_slot_mapping_start_idx(is_prompt, query_len, + context_len, + self.sliding_window) + compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id, + seq_len, context_len, start_idx, + self.block_size, inter_data.block_tables) + + def build(self, seq_lens: List[int], query_lens: List[int], + cuda_graph_pad_size: int, batch_size: int): + """Build attention metadata with on-device tensors. + + Args: + seq_lens: The maybe padded sequence lengths of the input sequences. + query_lens: The query lengths of the input sequences. + cuda_graph_pad_size: The padding size for cuda graph. + -1 if cuda graph is not used. + batch_size: The maybe padded batch size. + """ + for inter_data in self.input_builder.inter_data_list: + self._add_seq_group(inter_data, + self.input_builder.chunked_prefill_enabled) + + device = self.runner.device + use_captured_graph = cuda_graph_pad_size != -1 + + max_query_len = max(query_lens) + max_prefill_seq_len = max(self.prefill_seq_lens, default=0) + max_decode_seq_len = max(self.curr_seq_lens, default=0) + num_decode_tokens = self.num_decode_tokens + query_start_loc = list(accumulate(query_lens, initial=0)) + seq_start_loc = list(accumulate(seq_lens, initial=0)) + + if use_captured_graph: + self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size) + self.block_tables.extend([] * cuda_graph_pad_size) + num_decode_tokens = batch_size + + # The shape of graph_block_tables is + # [max batch size, max context len // block size]. + input_block_tables = self.runner.graph_block_tables[:batch_size] + for i, block_table in enumerate(self.block_tables): + if block_table: + input_block_tables[i, :len(block_table)] = block_table + block_tables = torch.from_numpy(input_block_tables).to( + device, non_blocking=True) + else: + block_tables = make_tensor_with_pad( + self.block_tables, + pad=0, + dtype=torch.int, + device=device, + ) + assert max_query_len > 0, "query_lens: {}".format(query_lens) + + assert device is not None + context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int, + device, self.runner.pin_memory) + seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device, + self.runner.pin_memory) + slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long, + device, self.runner.pin_memory) + query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32, + device, + self.runner.pin_memory) + seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32, + device, self.runner.pin_memory) + placeholder_index_maps = { + modality: placeholder_map.index_map() + for modality, placeholder_map in + self.multimodal_placeholder_maps.items() + } + + return self._metadata_cls( # type: ignore + num_prefills=self.num_prefills, + slot_mapping=slot_mapping_tensor, + multi_modal_placeholder_index_maps=placeholder_index_maps, + enable_kv_scales_calculation=True, + num_prefill_tokens=self.num_prefill_tokens, + num_decode_tokens=num_decode_tokens, + seq_lens=seq_lens, + seq_lens_tensor=seq_lens_tensor, + max_query_len=max_query_len, + max_prefill_seq_len=max_prefill_seq_len, + max_decode_seq_len=max_decode_seq_len, + query_start_loc=query_start_loc_tensor, + seq_start_loc=seq_start_loc_tensor, + context_lens_tensor=context_lens_tensor, + block_tables=block_tables, + use_cuda_graph=use_captured_graph, + ) + + +class CommonAttentionState(AttentionState): + + def __init__(self, runner: "ModelRunnerBase"): + self.runner = runner + self._is_graph_capturing = False + + @contextmanager + def graph_capture(self, max_batch_size: int): + + self._is_graph_capturing = True + + self._graph_slot_mapping = torch.full((max_batch_size, ), + PAD_SLOT_ID, + dtype=torch.long, + device=self.runner.device) + self._graph_seq_lens = torch.ones(max_batch_size, + dtype=torch.int32, + device=self.runner.device) + self._graph_block_tables = torch.from_numpy( + self.runner.graph_block_tables).to(device=self.runner.device) + + yield + + self._is_graph_capturing = False + del self._graph_slot_mapping + del self._graph_seq_lens + del self._graph_block_tables + + def graph_clone(self, batch_size: int) -> "CommonAttentionState": + assert self._is_graph_capturing + return self.__class__(self.runner) + + def graph_capture_get_metadata_for_batch( + self, batch_size: int, is_encoder_decoder_model: bool = False): + assert self._is_graph_capturing + attn_metadata = self.runner.attn_backend.make_metadata( + num_prefills=0, + num_prefill_tokens=0, + num_decode_tokens=batch_size, + slot_mapping=self._graph_slot_mapping[:batch_size], + multi_modal_placeholder_index_maps=None, + enable_kv_scales_calculation=True, + seq_lens=None, + seq_lens_tensor=self._graph_seq_lens[:batch_size], + max_query_len=1, + max_decode_query_len=1, + max_prefill_seq_len=0, + max_decode_seq_len=self.runner.max_seq_len_to_capture, + query_start_loc=None, + seq_start_loc=None, + context_lens_tensor=None, + block_tables=self._graph_block_tables[:batch_size], + use_cuda_graph=True, + ) + if is_encoder_decoder_model: + # The encoder decoder model works only with XFormers and + # Flash Attention backend. Assert the same. + assert self.runner.attn_backend.get_name() in\ + ["XFORMERS", "FLASH_ATTN"], \ + f"Expected attn_backend name to be either 'XFORMERS' or " \ + f"'FLASH_ATTN', but "\ + f"got '{self.runner.attn_backend.get_name()}'" + self._update_captured_metadata_for_enc_dec_model( + batch_size=batch_size, attn_metadata=attn_metadata) + + return attn_metadata + + def get_graph_input_buffers( + self, + attn_metadata, + is_encoder_decoder_model: bool = False) -> Dict[str, Any]: + input_buffers = { + "slot_mapping": attn_metadata.slot_mapping, + "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor, + "block_tables": attn_metadata.decode_metadata.block_tables, + } + if is_encoder_decoder_model: + # The encoder decoder model works only with XFormers and + # Flash Attention backend. Assert the same. + assert self.runner.attn_backend.get_name() in\ + ["XFORMERS", "FLASH_ATTN"], \ + f"Expected attn_backend name to be either 'XFORMERS' or "\ + f"'FLASH_ATTN', but "\ + f"got '{self.runner.attn_backend.get_name()}'" + self._add_additonal_input_buffers_for_enc_dec_model( + attn_metadata=attn_metadata, input_buffers=input_buffers) + return input_buffers + + def prepare_graph_input_buffers( + self, + input_buffers, + attn_metadata, + is_encoder_decoder_model: bool = False) -> None: + input_buffers["seq_lens_tensor"].copy_( + attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True) + input_buffers["block_tables"].copy_( + attn_metadata.decode_metadata.block_tables, non_blocking=True) + if is_encoder_decoder_model: + # The encoder decoder model works only with XFormers and + # Flash Attention backend. Assert the same. + assert self.runner.attn_backend.get_name() in\ + ["XFORMERS", "FLASH_ATTN"], \ + f"Expected attn_backend name to be either 'XFORMERS' or "\ + f"'FLASH_ATTN', but "\ + f"got '{self.runner.attn_backend.get_name()}'" + self._prepare_input_buffers_for_enc_dec_model( + attn_metadata, input_buffers) + + def begin_forward(self, model_input) -> None: + return + + def _update_captured_metadata_for_enc_dec_model(self, batch_size: int, + attn_metadata): + """ + Updates the attention metadata parameters for CUDA graph capture in an + encoder-decoder model. + + This method modifies attention-related tensors and metadata required + for CUDA graph capture in encoder-decoder models. Specifically, it + updates the cross-attention and encoder sequence tensors in the + AttentionMetadata object. + """ + # During decode phase the cross_slot_mapping will be empty. Hence set + # an empty tensor for CUDA Graph capture. + attn_metadata.cross_slot_mapping = torch.tensor( + [], dtype=torch.int).cuda() + attn_metadata.cross_block_tables = torch.full( + (batch_size, self.runner.get_max_block_per_batch()), + 1, + dtype=torch.int).cuda() + attn_metadata.encoder_seq_lens = torch.full((batch_size, ), + 1, + dtype=torch.int).cuda() + attn_metadata.encoder_seq_lens_tensor = torch.full( + (batch_size, ), 1, dtype=torch.int).cuda() + attn_metadata.max_encoder_seq_len = self.runner.max_seq_len_to_capture + attn_metadata.num_encoder_tokens = 0 + + def _add_additonal_input_buffers_for_enc_dec_model( + self, attn_metadata, input_buffers: Dict[str, Any]): + """ + Saves additional input buffers specific to the encoder-decoder model + from the attention metadata. + + This method extracts and stores encoder-decoder related input buffers + from the `attn_metadata` into the `input_buffers` dictionary. The + buffers include encoder sequence lengths, cross-slot mappings, and + cross-block tables, which are essential for the encoder-decoder model + during CUDA graph replay. + """ + input_buffers["encoder_seq_lens_tensor"] = ( + attn_metadata.decode_metadata.encoder_seq_lens_tensor) + input_buffers["cross_slot_mapping"] = ( + attn_metadata.decode_metadata.cross_slot_mapping) + input_buffers["cross_block_tables"] = ( + attn_metadata.decode_metadata.cross_block_tables) + + def _prepare_input_buffers_for_enc_dec_model(self, attn_metadata, + input_buffers: Dict[str, + Any]): + """ + Populates input buffers with data from the encoder-decoder model's + attention metadata. + + This method fills the input buffers with encoder-decoder specific + tensors. It copies data from the `attn_metadata` and keyword arguments + (`kwargs`) into corresponding buffers in the `input_buffers` dictionary. + The copied data includes attention-related metadata as well as input + IDs and positional information for the encoder. + """ + input_buffers["encoder_seq_lens_tensor"].copy_( + attn_metadata.decode_metadata.encoder_seq_lens_tensor, + non_blocking=True) + input_buffers["cross_slot_mapping"].copy_( + attn_metadata.decode_metadata.cross_slot_mapping, + non_blocking=True) + input_buffers["cross_block_tables"].copy_( + attn_metadata.decode_metadata.cross_block_tables, + non_blocking=True) + + +def is_all_encoder_attn_metadata_set(attn_metadata): + ''' + All attention metadata required for encoder attention is set. + ''' + return ((attn_metadata.encoder_seq_lens is not None) + and (attn_metadata.encoder_seq_lens_tensor is not None) + and (attn_metadata.max_encoder_seq_len is not None)) + + +def is_all_cross_attn_metadata_set(attn_metadata): + ''' + All attention metadata required for enc/dec cross-attention is set. + + Superset of encoder attention required metadata. + ''' + return (attn_metadata.is_all_encoder_attn_metadata_set + and (attn_metadata.cross_slot_mapping is not None) + and (attn_metadata.cross_block_tables is not None)) + + +def get_seq_len_block_table_args( + attn_metadata, + is_prompt: bool, + attn_type: str, +) -> tuple: + ''' + The particular choice of sequence-length- and block-table-related + attributes which should be extracted from attn_metadata is dependent + on the type of attention operation. + + Decoder attn -> select entirely decoder self-attention-related fields + Encoder/decoder cross-attn -> select encoder sequence lengths & + cross-attn block-tables fields + Encoder attn -> select encoder sequence lengths fields & no block tables + + Arguments: + + * attn_metadata: Attention metadata structure associated with attention op + * is_prompt: True if prefill, False otherwise + * attn_type: encoder attention, decoder self-attention, + encoder/decoder cross-attention + + Returns: + + * Appropriate sequence-lengths tensor + * Appropriate max sequence-length scalar + * Appropriate block tables (or None) + ''' + + if attn_type == AttentionType.DECODER: + # Decoder self-attention + # Choose max_seq_len based on whether we are in prompt_run + if is_prompt: + max_seq_len = attn_metadata.max_prefill_seq_len + else: + max_seq_len = attn_metadata.max_decode_seq_len + return (attn_metadata.seq_lens_tensor, max_seq_len, + attn_metadata.block_tables) + elif attn_type == AttentionType.ENCODER_DECODER: + # Enc/dec cross-attention KVs match encoder sequence length; + # cross-attention utilizes special "cross" block tables + return (attn_metadata.encoder_seq_lens_tensor, + attn_metadata.max_encoder_seq_len, + attn_metadata.cross_block_tables) + elif attn_type == AttentionType.ENCODER: + # No block tables associated with encoder attention + return (attn_metadata.encoder_seq_lens_tensor, + attn_metadata.max_encoder_seq_len, None) + else: + raise AttributeError(f"Invalid attention type {str(attn_type)}") + + +def get_num_prefill_decode_query_kv_tokens( + attn_metadata, + attn_type: str, +) -> Tuple[int, int, int]: + """ + Calculate the number of prefill and decode tokens for query, key/value + based on the attention metadata and the specified attention type. + + Args: + attn_metadata (FlashAttentionMetadata): Attention Metadata object. + attn_type (AttentionType): The type of attention being used. + Returns: + Tuple[int, int, int]: A tuple containing three integers: + - The number of prefill query tokens. + - The number of prefill key/value tokens. + - The number of decode query tokens. + + Raises: + AssertionError: If the number of encoder tokens in `attn_metadata` + is `None` when required for the calculations. + """ + num_prefill_query_tokens = 0 + num_decode_query_tokens = 0 + num_prefill_kv_tokens = 0 + if attn_type == AttentionType.ENCODER: + # Encoder attention is only invoked during prefill phase. + # The same input servers a both query and key. + assert attn_metadata.num_encoder_tokens is not None + num_prefill_query_tokens = attn_metadata.num_encoder_tokens + num_prefill_kv_tokens = attn_metadata.num_encoder_tokens + num_decode_query_tokens = 0 + elif attn_type == AttentionType.ENCODER_DECODER: + assert attn_metadata.num_encoder_tokens is not None + num_prefill_query_tokens = attn_metadata.num_prefill_tokens + # The key is the encoder/cross-attention. + num_prefill_kv_tokens = attn_metadata.num_encoder_tokens + num_decode_query_tokens = attn_metadata.num_decode_tokens + else: # attn_type == AttentionType.DECODER or + # attn_type == AttentionType.ENCODER_ONLY + num_prefill_query_tokens = attn_metadata.num_prefill_tokens + num_prefill_kv_tokens = attn_metadata.num_prefill_tokens + num_decode_query_tokens = attn_metadata.num_decode_tokens + + return (num_prefill_query_tokens, num_prefill_kv_tokens, + num_decode_query_tokens) + + +def get_flash_attn_version(): + try: + from vllm.vllm_flash_attn.flash_attn_interface import ( + fa_version_unsupported_reason, is_fa_version_supported) + + # if hopper default to FA3, otherwise stick to FA2 for now + # TODO(lucas): profile FA3 on ampere to see if it makes sense to + # use FA3 as default for both + if current_platform.get_device_capability()[0] == 9: + fa_version = 3 if is_fa_version_supported(3) else 2 + else: + fa_version = 2 + + if envs.VLLM_FLASH_ATTN_VERSION is not None: + assert envs.VLLM_FLASH_ATTN_VERSION in [2, 3] + fa_version = envs.VLLM_FLASH_ATTN_VERSION + if (current_platform.get_device_capability()[0] == 10 + and envs.VLLM_FLASH_ATTN_VERSION == 3): + logger.warning("Cannot use FA version 3 on Blackwell platform", + "defaulting to FA version 2.") + fa_version = 2 + + if not is_fa_version_supported(fa_version): + logger.error("Cannot use FA version %d is not supported due to %s", + fa_version, fa_version_unsupported_reason(fa_version)) + + assert is_fa_version_supported(fa_version) + return fa_version + except (ImportError, AssertionError): + return None diff --git a/fastvideo/attention/flash_attn.py b/fastvideo/v1/attention/flash_attn.py similarity index 97% rename from fastvideo/attention/flash_attn.py rename to fastvideo/v1/attention/flash_attn.py index 7d570d522..d8c7711b1 100644 --- a/fastvideo/attention/flash_attn.py +++ b/fastvideo/v1/attention/flash_attn.py @@ -3,8 +3,8 @@ from typing import List, Optional import torch import torch.nn as nn -from fastvideo.distributed.communication_op import sequence_model_parallel_all_to_all_4D, sequence_model_parallel_all_gather -from fastvideo.distributed.parallel_state import get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size +from fastvideo.v1.distributed.communication_op import sequence_model_parallel_all_to_all_4D, sequence_model_parallel_all_gather +from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size from flash_attn import flash_attn_func, flash_attn_varlen_func diff --git a/fastvideo/attention/layer.py b/fastvideo/v1/attention/layer.py similarity index 98% rename from fastvideo/attention/layer.py rename to fastvideo/v1/attention/layer.py index 3c1920327..2fa49c7b9 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/v1/attention/layer.py @@ -2,7 +2,7 @@ import torch import torch.nn as nn from typing import Optional, Dict, Any, Type -from fastvideo.attention.flash_attn import LocalAttention, DistributedAttention +from fastvideo.v1.attention.flash_attn import LocalAttention, DistributedAttention class AttentionBackend: diff --git a/fastvideo/distributed/__init__.py b/fastvideo/v1/distributed/__init__.py similarity index 100% rename from fastvideo/distributed/__init__.py rename to fastvideo/v1/distributed/__init__.py diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/v1/distributed/communication_op.py similarity index 100% rename from fastvideo/distributed/communication_op.py rename to fastvideo/v1/distributed/communication_op.py diff --git a/fastvideo/layers/__init__.py b/fastvideo/v1/distributed/device_communicators/__init__.py similarity index 100% rename from fastvideo/layers/__init__.py rename to fastvideo/v1/distributed/device_communicators/__init__.py diff --git a/fastvideo/distributed/device_communicators/base_device_communicator.py b/fastvideo/v1/distributed/device_communicators/base_device_communicator.py similarity index 100% rename from fastvideo/distributed/device_communicators/base_device_communicator.py rename to fastvideo/v1/distributed/device_communicators/base_device_communicator.py diff --git a/fastvideo/distributed/device_communicators/cuda_communicator.py b/fastvideo/v1/distributed/device_communicators/cuda_communicator.py similarity index 98% rename from fastvideo/distributed/device_communicators/cuda_communicator.py rename to fastvideo/v1/distributed/device_communicators/cuda_communicator.py index 264c0e3bb..50494566c 100644 --- a/fastvideo/distributed/device_communicators/cuda_communicator.py +++ b/fastvideo/v1/distributed/device_communicators/cuda_communicator.py @@ -33,7 +33,7 @@ class CudaCommunicator(DeviceCommunicatorBase): # lazy import to avoid documentation build error # from vllm.distributed.device_communicators.custom_all_reduce import ( # CustomAllreduce) - from fastvideo.distributed.device_communicators.pynccl import ( + from fastvideo.v1.distributed.device_communicators.pynccl import ( PyNcclCommunicator) self.pynccl_comm: Optional[PyNcclCommunicator] = None diff --git a/fastvideo/distributed/device_communicators/pynccl.py b/fastvideo/v1/distributed/device_communicators/pynccl.py similarity index 97% rename from fastvideo/distributed/device_communicators/pynccl.py rename to fastvideo/v1/distributed/device_communicators/pynccl.py index 2ab53e02e..8d1bbd7c9 100644 --- a/fastvideo/distributed/device_communicators/pynccl.py +++ b/fastvideo/v1/distributed/device_communicators/pynccl.py @@ -8,12 +8,12 @@ import torch import torch.distributed as dist from torch.distributed import ProcessGroup, ReduceOp -from fastvideo.distributed.device_communicators.pynccl_wrapper import ( +from fastvideo.v1.distributed.device_communicators.pynccl_wrapper import ( NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum, ncclRedOpTypeEnum, ncclUniqueId) -from fastvideo.distributed.utils import StatelessProcessGroup -from fastvideo.logger import init_logger -from fastvideo.utils.utils import current_stream +from fastvideo.v1.distributed.utils import StatelessProcessGroup +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils.utils import current_stream logger = init_logger(__name__) diff --git a/fastvideo/distributed/device_communicators/pynccl_wrapper.py b/fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py similarity index 99% rename from fastvideo/distributed/device_communicators/pynccl_wrapper.py rename to fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py index 97c3a7a34..6a0d8179a 100644 --- a/fastvideo/distributed/device_communicators/pynccl_wrapper.py +++ b/fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py @@ -30,8 +30,8 @@ from typing import Any, Dict, List, Optional import torch from torch.distributed import ReduceOp -from fastvideo.logger import init_logger -from fastvideo.utils.utils import find_nccl_library +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils.utils import find_nccl_library logger = init_logger(__name__) diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/v1/distributed/parallel_state.py similarity index 98% rename from fastvideo/distributed/parallel_state.py rename to fastvideo/v1/distributed/parallel_state.py index 56ef861f5..0d00bbe33 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/v1/distributed/parallel_state.py @@ -40,15 +40,15 @@ import torch import torch.distributed from torch.distributed import Backend, ProcessGroup -import fastvideo.envs as envs -from fastvideo.distributed.device_communicators.base_device_communicator import ( +import fastvideo.v1.envs as envs +from fastvideo.v1.distributed.device_communicators.base_device_communicator import ( DeviceCommunicatorBase) -from fastvideo.distributed.utils import StatelessProcessGroup -from fastvideo.logger import init_logger -# from fastvideo.utils import (direct_register_custom_op, resolve_obj_by_qualname, +from fastvideo.v1.distributed.utils import StatelessProcessGroup +from fastvideo.v1.logger import init_logger +# from fastvideo.v1.utils import (direct_register_custom_op, resolve_obj_by_qualname, # supports_custom_op) -from fastvideo.distributed.device_communicators.cuda_communicator import ( +from fastvideo.v1.distributed.device_communicators.cuda_communicator import ( CudaCommunicator) @@ -192,7 +192,7 @@ class GroupCoordinator: assert self.cpu_group is not None assert self.device_group is not None - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform # TODO: fix it for other platforms if current_platform.is_cuda_alike(): @@ -220,7 +220,7 @@ class GroupCoordinator: # self.mq_broadcaster = MessageQueue.create_from_process_group( # self.cpu_group, 1 << 22, 6) - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform # self.use_custom_op_call = current_platform.is_cuda_alike() self.use_custom_op_call = False @@ -270,7 +270,7 @@ class GroupCoordinator: # only cuda uses this function, # so we don't abstract it into the base class maybe_ca_context = nullcontext() - from fastvideo.distributed.device_communicators.cuda_communicator import ( + from fastvideo.v1.distributed.device_communicators.cuda_communicator import ( CudaCommunicator) if self.device_communicator is not None: assert isinstance(self.device_communicator, CudaCommunicator) @@ -1063,7 +1063,7 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False): import ray # Lazy import Ray ray.shutdown() gc.collect() - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform if not current_platform.is_cpu(): torch.cuda.empty_cache() try: @@ -1345,7 +1345,7 @@ for different models: ```python import torch -from fastvideo.distributed.parallel_state import ( +from fastvideo.v1.distributed.parallel_state import ( init_distributed_environment, initialize_tensor_parallel_group, initialize_sequence_parallel_group, diff --git a/fastvideo/distributed/utils.py b/fastvideo/v1/distributed/utils.py similarity index 99% rename from fastvideo/distributed/utils.py rename to fastvideo/v1/distributed/utils.py index 3b311a7e6..4286ab819 100644 --- a/fastvideo/distributed/utils.py +++ b/fastvideo/v1/distributed/utils.py @@ -18,8 +18,8 @@ from torch.distributed.distributed_c10d import (Backend, PrefixStore, is_nccl_available) from torch.distributed.rendezvous import rendezvous -import fastvideo.envs as envs -from fastvideo.logger import init_logger +import fastvideo.v1.envs as envs +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/envs.py b/fastvideo/v1/envs.py similarity index 100% rename from fastvideo/envs.py rename to fastvideo/v1/envs.py diff --git a/fastvideo/inference_args.py b/fastvideo/v1/inference_args.py similarity index 99% rename from fastvideo/inference_args.py rename to fastvideo/v1/inference_args.py index df8678294..9cd9923f8 100644 --- a/fastvideo/inference_args.py +++ b/fastvideo/v1/inference_args.py @@ -17,8 +17,8 @@ import argparse import dataclasses -from fastvideo.logger import init_logger -from fastvideo.utils.utils import FlexibleArgumentParser +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils.utils import FlexibleArgumentParser from typing import List, Optional diff --git a/fastvideo/inference_engine.py b/fastvideo/v1/inference_engine.py similarity index 94% rename from fastvideo/inference_engine.py rename to fastvideo/v1/inference_engine.py index 2d526d7ec..cd5b2da86 100644 --- a/fastvideo/inference_engine.py +++ b/fastvideo/v1/inference_engine.py @@ -13,14 +13,14 @@ from PIL import Image import uuid import random -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines import ComposedPipelineBase, build_pipeline -from fastvideo.pipelines.pipeline_registry import PipelineRegistry -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.logger import init_logger +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines import ComposedPipelineBase, build_pipeline +from fastvideo.v1.pipelines.pipeline_registry import PipelineRegistry +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.logger import init_logger # TODO(will): remove -from fastvideo.models.hunyuan.utils.data_utils import align_to -from fastvideo.models.hunyuan.constants import NEGATIVE_PROMPT +from fastvideo.v1.models.hunyuan.utils.data_utils import align_to +from fastvideo.v1.models.hunyuan.constants import NEGATIVE_PROMPT logger = init_logger(__name__) @@ -158,7 +158,7 @@ class InferenceEngine: negative_prompt = [negative_prompt.strip()] - # from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler + # from fastvideo.v1.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler # scheduler = FlowMatchDiscreteScheduler( # shift=flow_shift, # reverse=self.inference_args.flow_reverse, diff --git a/fastvideo/loader/__init__.py b/fastvideo/v1/layers/__init__.py similarity index 100% rename from fastvideo/loader/__init__.py rename to fastvideo/v1/layers/__init__.py diff --git a/fastvideo/layers/activation.py b/fastvideo/v1/layers/activation.py similarity index 100% rename from fastvideo/layers/activation.py rename to fastvideo/v1/layers/activation.py diff --git a/fastvideo/layers/conv.py b/fastvideo/v1/layers/conv.py similarity index 100% rename from fastvideo/layers/conv.py rename to fastvideo/v1/layers/conv.py diff --git a/fastvideo/layers/layernorm.py b/fastvideo/v1/layers/layernorm.py similarity index 100% rename from fastvideo/layers/layernorm.py rename to fastvideo/v1/layers/layernorm.py diff --git a/fastvideo/layers/linear.py b/fastvideo/v1/layers/linear.py similarity index 99% rename from fastvideo/layers/linear.py rename to fastvideo/v1/layers/linear.py index c8550d09e..365261101 100644 --- a/fastvideo/layers/linear.py +++ b/fastvideo/v1/layers/linear.py @@ -8,25 +8,25 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter, UninitializedParameter -from fastvideo.distributed import (divide, get_tensor_model_parallel_rank, +from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, split_tensor_along_last_dim, tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce) -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger # TODO(will): remove this import by copying the definition from vLLM then # manually import each quantization method we want to use. Refer to SGLang from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase) # yapf: disable -from fastvideo.models.parameter import (BasevLLMParameter, +from fastvideo.v1.models.parameter import (BasevLLMParameter, BlockQuantScaleParameter, PackedColumnParameter, PackedvLLMParameter, PerTensorScaleParameter, RowvLLMParameter) # yapf: enable -from fastvideo.models.utils import set_weight_attrs +from fastvideo.v1.models.utils import set_weight_attrs logger = init_logger(__name__) diff --git a/fastvideo/layers/mlp.py b/fastvideo/v1/layers/mlp.py similarity index 90% rename from fastvideo/layers/mlp.py rename to fastvideo/v1/layers/mlp.py index db2d2ce01..68529c33a 100644 --- a/fastvideo/layers/mlp.py +++ b/fastvideo/v1/layers/mlp.py @@ -1,5 +1,5 @@ -from fastvideo.layers.linear import ReplicatedLinear -from fastvideo.layers.activation import get_act_fn +from fastvideo.v1.layers.linear import ReplicatedLinear +from fastvideo.v1.layers.activation import get_act_fn import torch import torch.nn as nn from typing import Optional diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/v1/layers/rotary_embedding.py similarity index 99% rename from fastvideo/layers/rotary_embedding.py rename to fastvideo/v1/layers/rotary_embedding.py index 974cd23cf..f557ba1ae 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/v1/layers/rotary_embedding.py @@ -30,7 +30,7 @@ import torch.nn as nn from transformers import PretrainedConfig from vllm.model_executor.custom_op import CustomOp -from fastvideo.distributed.parallel_state import get_sp_group +from fastvideo.v1.distributed.parallel_state import get_sp_group def _rotate_neox(x: torch.Tensor) -> torch.Tensor: x1 = x[..., :x.shape[-1] // 2] diff --git a/fastvideo/layers/utils.py b/fastvideo/v1/layers/utils.py similarity index 100% rename from fastvideo/layers/utils.py rename to fastvideo/v1/layers/utils.py diff --git a/fastvideo/layers/visual_embedding.py b/fastvideo/v1/layers/visual_embedding.py similarity index 96% rename from fastvideo/layers/visual_embedding.py rename to fastvideo/v1/layers/visual_embedding.py index 1f5f8a19f..13a61aef1 100644 --- a/fastvideo/layers/visual_embedding.py +++ b/fastvideo/v1/layers/visual_embedding.py @@ -1,10 +1,10 @@ import torch import torch.nn as nn import math -from fastvideo.layers.activation import get_act_fn -from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.v1.layers.activation import get_act_fn +from fastvideo.v1.layers.linear import ReplicatedLinear from typing import Optional -from fastvideo.layers.mlp import MLP +from fastvideo.v1.layers.mlp import MLP class PatchEmbed(nn.Module): """2D Image to Patch Embedding diff --git a/fastvideo/layers/vocab_parallel_embedding.py b/fastvideo/v1/layers/vocab_parallel_embedding.py similarity index 98% rename from fastvideo/layers/vocab_parallel_embedding.py rename to fastvideo/v1/layers/vocab_parallel_embedding.py index 44d7793ab..338fb2ac9 100644 --- a/fastvideo/layers/vocab_parallel_embedding.py +++ b/fastvideo/v1/layers/vocab_parallel_embedding.py @@ -7,14 +7,14 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter, UninitializedParameter -from fastvideo.distributed import (divide, get_tensor_model_parallel_rank, +from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce) from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding) -from fastvideo.models.parameter import BasevLLMParameter -from fastvideo.models.utils import set_weight_attrs -from fastvideo.platforms import current_platform +from fastvideo.v1.models.parameter import BasevLLMParameter +from fastvideo.v1.models.utils import set_weight_attrs +from fastvideo.v1.platforms import current_platform DEFAULT_VOCAB_PADDING_SIZE = 64 diff --git a/fastvideo/models/schedulers/base.py b/fastvideo/v1/loader/__init__.py similarity index 100% rename from fastvideo/models/schedulers/base.py rename to fastvideo/v1/loader/__init__.py diff --git a/fastvideo/loader/fsdp_load.py b/fastvideo/v1/loader/fsdp_load.py similarity index 99% rename from fastvideo/loader/fsdp_load.py rename to fastvideo/v1/loader/fsdp_load.py index 5e2a0b3ac..e2096b0ac 100644 --- a/fastvideo/loader/fsdp_load.py +++ b/fastvideo/v1/loader/fsdp_load.py @@ -8,7 +8,7 @@ from itertools import chain import torch from torch import nn from torch.distributed import DeviceMesh, init_device_mesh -from fastvideo.distributed.parallel_state import get_sequence_model_parallel_world_size +from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_world_size from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed._tensor import distribute_tensor from torch.nn.modules.module import _IncompatibleKeys diff --git a/fastvideo/loader/loader.py b/fastvideo/v1/loader/loader.py similarity index 97% rename from fastvideo/loader/loader.py rename to fastvideo/v1/loader/loader.py index de28e5d63..09a8c670e 100644 --- a/fastvideo/loader/loader.py +++ b/fastvideo/v1/loader/loader.py @@ -26,17 +26,17 @@ import torch from torch import nn from transformers.utils import SAFE_WEIGHTS_INDEX_NAME -from fastvideo.logger import init_logger -from fastvideo.loader.utils import set_default_torch_dtype -from fastvideo.loader.weight_utils import ( +from fastvideo.v1.logger import init_logger +from fastvideo.v1.loader.utils import set_default_torch_dtype +from fastvideo.v1.loader.weight_utils import ( download_safetensors_index_file_from_hf, download_weights_from_hf, filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, pt_weights_iterator, safetensors_weights_iterator) -from fastvideo.platforms import current_platform -from fastvideo.models.registry import ModelRegistry +from fastvideo.v1.platforms import current_platform +from fastvideo.v1.models.registry import ModelRegistry -from fastvideo.inference_args import InferenceArgs +from fastvideo.v1.inference_args import InferenceArgs @contextmanager def device_loading_context(module: torch.nn.Module, @@ -126,7 +126,7 @@ def _initialize_model( # Initialize the model based on the model_name if model_name == "CLIPVisionModel": - from fastvideo.models.encoders.clip import CLIPVisionModel + from fastvideo.v1.models.encoders.clip import CLIPVisionModel from transformers import CLIPVisionConfig # If the config is not a CLIPVisionConfig, convert it diff --git a/fastvideo/loader/utils.py b/fastvideo/v1/loader/utils.py similarity index 99% rename from fastvideo/loader/utils.py rename to fastvideo/v1/loader/utils.py index 2917fb403..a105ceb06 100644 --- a/fastvideo/loader/utils.py +++ b/fastvideo/v1/loader/utils.py @@ -10,7 +10,7 @@ from torch import nn from transformers.dynamic_module_utils import get_class_from_dynamic_module from vllm.config import ModelConfig, ModelImpl -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger from vllm.model_executor.models import ModelRegistry from vllm.model_executor.models.adapters import (as_classification_model, as_embedding_model, diff --git a/fastvideo/loader/weight_utils.py b/fastvideo/v1/loader/weight_utils.py similarity index 99% rename from fastvideo/loader/weight_utils.py rename to fastvideo/v1/loader/weight_utils.py index 7e514bd3c..e03c82fa7 100644 --- a/fastvideo/loader/weight_utils.py +++ b/fastvideo/v1/loader/weight_utils.py @@ -26,11 +26,11 @@ from safetensors.torch import load_file, safe_open, save_file from tqdm.auto import tqdm # from vllm.config import LoadConfig, ModelConfig -from fastvideo.distributed import get_tensor_model_parallel_rank -from fastvideo.logger import init_logger +from fastvideo.v1.distributed import get_tensor_model_parallel_rank +from fastvideo.v1.logger import init_logger # from vllm.model_executor.layers.quantization import (QuantizationConfig, # get_quantization_config) -from fastvideo.platforms import current_platform +from fastvideo.v1.platforms import current_platform # from vllm.utils import PlaceholderModule # try: diff --git a/fastvideo/logger.py b/fastvideo/v1/logger.py similarity index 96% rename from fastvideo/logger.py rename to fastvideo/v1/logger.py index 8855baa85..180c19bfb 100644 --- a/fastvideo/logger.py +++ b/fastvideo/v1/logger.py @@ -4,7 +4,7 @@ # Copyright 2023 The vLLM Authors. # Copyright 2023 The FastVideo Authors. -"""Logging configuration for fastvideo.""" +"""Logging configuration for fastvideo.v1.""" import datetime import json import logging @@ -17,7 +17,7 @@ from os import path from types import MethodType from typing import Any, Optional, cast -import fastvideo.envs as envs +import fastvideo.v1.envs as envs FASTVIDEO_CONFIGURE_LOGGING = envs.FASTVIDEO_CONFIGURE_LOGGING FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH @@ -31,7 +31,7 @@ _DATE_FORMAT = "%m-%d %H:%M:%S" DEFAULT_LOGGING_CONFIG = { "formatters": { "fastvideo": { - "class": "fastvideo.logging_utils.NewLineFormatter", + "class": "fastvideo.v1.logging_utils.NewLineFormatter", "datefmt": _DATE_FORMAT, "format": _FORMAT, }, @@ -124,8 +124,8 @@ def _configure_fastvideo_root_logger() -> None: for formatter in logging_config.get("formatters", {}).values(): # This provides backwards compatibility after #10134. - if formatter.get("class") == "fastvideo.logging.NewLineFormatter": - formatter["class"] = "fastvideo.logging_utils.NewLineFormatter" + if formatter.get("class") == "fastvideo.v1.logging.NewLineFormatter": + formatter["class"] = "fastvideo.v1.logging_utils.NewLineFormatter" if logging_config: dictConfig(logging_config) diff --git a/fastvideo/logging_utils/__init__.py b/fastvideo/v1/logging_utils/__init__.py similarity index 54% rename from fastvideo/logging_utils/__init__.py rename to fastvideo/v1/logging_utils/__init__.py index 784da68c4..4a2511681 100644 --- a/fastvideo/logging_utils/__init__.py +++ b/fastvideo/v1/logging_utils/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -from fastvideo.logging_utils.formatter import NewLineFormatter +from fastvideo.v1.logging_utils.formatter import NewLineFormatter __all__ = [ "NewLineFormatter", diff --git a/fastvideo/logging_utils/formatter.py b/fastvideo/v1/logging_utils/formatter.py similarity index 100% rename from fastvideo/logging_utils/formatter.py rename to fastvideo/v1/logging_utils/formatter.py diff --git a/fastvideo/models/__init__.py b/fastvideo/v1/models/__init__.py similarity index 87% rename from fastvideo/models/__init__.py rename to fastvideo/v1/models/__init__.py index a6c336936..e6667e618 100644 --- a/fastvideo/models/__init__.py +++ b/fastvideo/v1/models/__init__.py @@ -2,9 +2,9 @@ import torch.nn as nn from typing import Dict, Any -from fastvideo.loader.loader import get_model_loader -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.loader.loader import get_model_loader +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger import os from transformers import PretrainedConfig @@ -39,7 +39,7 @@ def get_scheduler(module_path: str, architecture: str, inference_args: Inference # TODO(will): add schedulers to register or create a new scheduler registry # TODO(will): default to config file but allow override through # inference args. Currently only uses inference args. - from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchDiscreteScheduler + from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchDiscreteScheduler return FlowMatchDiscreteScheduler( shift=inference_args.flow_shift, reverse=inference_args.flow_reverse, diff --git a/fastvideo/models/dits/base.py b/fastvideo/v1/models/dits/base.py similarity index 100% rename from fastvideo/models/dits/base.py rename to fastvideo/v1/models/dits/base.py diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py similarity index 97% rename from fastvideo/models/dits/hunyuanvideo.py rename to fastvideo/v1/models/dits/hunyuanvideo.py index 49bd3e163..5f5e71f28 100644 --- a/fastvideo/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/models/dits/hunyuanvideo.py @@ -1,16 +1,16 @@ import torch import torch.nn as nn from typing import Optional, Tuple, List -from fastvideo.attention.flash_attn import DistributedAttention, LocalAttention -from fastvideo.layers.linear import ReplicatedLinear -from fastvideo.layers.layernorm import LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift -from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection, unpatchify -from fastvideo.layers.rotary_embedding import _apply_rotary_emb, get_rotary_pos_embed -from fastvideo.distributed.parallel_state import get_sequence_model_parallel_world_size, get_sequence_model_parallel_rank +from fastvideo.v1.attention.flash_attn import DistributedAttention, LocalAttention +from fastvideo.v1.layers.linear import ReplicatedLinear +from fastvideo.v1.layers.layernorm import LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift +from fastvideo.v1.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection, unpatchify +from fastvideo.v1.layers.rotary_embedding import _apply_rotary_emb, get_rotary_pos_embed +from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_world_size, get_sequence_model_parallel_rank # TODO: RMSNorm .... -from fastvideo.models.hunyuan.modules.norm_layers import RMSNorm -from fastvideo.layers.mlp import MLP -from fastvideo.models.dits.base import BaseDiT +from fastvideo.v1.models.hunyuan.modules.norm_layers import RMSNorm +from fastvideo.v1.layers.mlp import MLP +from fastvideo.v1.models.dits.base import BaseDiT diff --git a/fastvideo/models/encoders/clip.py b/fastvideo/v1/models/encoders/clip.py similarity index 98% rename from fastvideo/models/encoders/clip.py rename to fastvideo/v1/models/encoders/clip.py index 67c76163e..0960c3821 100644 --- a/fastvideo/models/encoders/clip.py +++ b/fastvideo/v1/models/encoders/clip.py @@ -9,14 +9,14 @@ from transformers import CLIPVisionConfig, CLIPTextConfig from transformers.modeling_outputs import BaseModelOutputWithPooling from vllm.attention.layer import MultiHeadAttention -from fastvideo.distributed import divide, get_tensor_model_parallel_world_size -from fastvideo.layers.activation import get_act_fn -from fastvideo.layers.linear import (ColumnParallelLinear, +from fastvideo.v1.distributed import divide, get_tensor_model_parallel_world_size +from fastvideo.v1.layers.activation import get_act_fn +from fastvideo.v1.layers.linear import (ColumnParallelLinear, QKVParallelLinear, RowParallelLinear) # TODO: support quantization # from vllm.model_executor.layers.quantization import QuantizationConfig -from fastvideo.loader.weight_utils import default_weight_loader +from fastvideo.v1.loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import SupportsQuant from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs diff --git a/fastvideo/models/encoders/encoder.py b/fastvideo/v1/models/encoders/encoder.py similarity index 100% rename from fastvideo/models/encoders/encoder.py rename to fastvideo/v1/models/encoders/encoder.py diff --git a/fastvideo/models/encoders/llama.py b/fastvideo/v1/models/encoders/llama.py similarity index 97% rename from fastvideo/models/encoders/llama.py rename to fastvideo/v1/models/encoders/llama.py index f1ca25b21..e02226e47 100644 --- a/fastvideo/models/encoders/llama.py +++ b/fastvideo/v1/models/encoders/llama.py @@ -29,19 +29,19 @@ from torch import nn from transformers import LlamaConfig from transformers.modeling_outputs import BaseModelOutputWithPast -from fastvideo.attention import TextEncoderFlashAttention +from fastvideo.v1.attention import TextEncoderFlashAttention -from fastvideo.distributed import get_tensor_model_parallel_world_size -from fastvideo.layers.activation import SiluAndMul -from fastvideo.layers.layernorm import RMSNorm -from fastvideo.layers.linear import (MergedColumnParallelLinear, +from fastvideo.v1.distributed import get_tensor_model_parallel_world_size +from fastvideo.v1.layers.activation import SiluAndMul +from fastvideo.v1.layers.layernorm import RMSNorm +from fastvideo.v1.layers.linear import (MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear) # from vllm.model_executor.layers.quantization import QuantizationConfig -from fastvideo.layers.rotary_embedding import get_rope -from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding -from fastvideo.loader.weight_utils import ( +from fastvideo.v1.layers.rotary_embedding import get_rope +from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.v1.loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name) from .utils import (extract_layer_index) diff --git a/fastvideo/models/encoders/utils.py b/fastvideo/v1/models/encoders/utils.py similarity index 99% rename from fastvideo/models/encoders/utils.py rename to fastvideo/v1/models/encoders/utils.py index f26828c23..db1d24fe5 100644 --- a/fastvideo/models/encoders/utils.py +++ b/fastvideo/v1/models/encoders/utils.py @@ -10,8 +10,8 @@ import torch.nn as nn from torch.func import functional_call from transformers import PretrainedConfig -from fastvideo.logger import init_logger -from fastvideo.loader.weight_utils import default_weight_loader +from fastvideo.v1.logger import init_logger +from fastvideo.v1.loader.weight_utils import default_weight_loader logger = init_logger(__name__) diff --git a/fastvideo/models/encoders/vision.py b/fastvideo/v1/models/encoders/vision.py similarity index 97% rename from fastvideo/models/encoders/vision.py rename to fastvideo/v1/models/encoders/vision.py index c9ff3e3a6..b0921f2a4 100644 --- a/fastvideo/models/encoders/vision.py +++ b/fastvideo/v1/models/encoders/vision.py @@ -6,11 +6,11 @@ from typing import Final, Generic, Optional, Protocol, TypeVar, Union import torch from transformers import PretrainedConfig -import fastvideo.envs as envs +import fastvideo.v1.envs as envs from vllm.attention.selector import (backend_name_to_enum, get_global_forced_attn_backend) -from fastvideo.logger import init_logger -from fastvideo.platforms import _Backend, current_platform +from fastvideo.v1.logger import init_logger +from fastvideo.v1.platforms import _Backend, current_platform logger = init_logger(__name__) diff --git a/fastvideo/models/hf_transformer_utils.py b/fastvideo/v1/models/hf_transformer_utils.py similarity index 98% rename from fastvideo/models/hf_transformer_utils.py rename to fastvideo/v1/models/hf_transformer_utils.py index b4a7ed042..7c22b971c 100644 --- a/fastvideo/models/hf_transformer_utils.py +++ b/fastvideo/v1/models/hf_transformer_utils.py @@ -31,7 +31,7 @@ from transformers import ( ) from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES -# from fastvideo.models.configs import ChatGLMConfig, DbrxConfig, ExaoneConfig, Qwen2_5_VLConfig +# from fastvideo.v1.models.configs import ChatGLMConfig, DbrxConfig, ExaoneConfig, Qwen2_5_VLConfig _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { # ChatGLMConfig.model_type: ChatGLMConfig, diff --git a/fastvideo/v1/models/hunyuan/__init__.py b/fastvideo/v1/models/hunyuan/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/pipelines/hunyuan/constants.py b/fastvideo/v1/models/hunyuan/constants.py similarity index 100% rename from fastvideo/pipelines/hunyuan/constants.py rename to fastvideo/v1/models/hunyuan/constants.py diff --git a/fastvideo/v1/models/hunyuan/diffusion/__init__.py b/fastvideo/v1/models/hunyuan/diffusion/__init__.py new file mode 100644 index 000000000..754bc48ec --- /dev/null +++ b/fastvideo/v1/models/hunyuan/diffusion/__init__.py @@ -0,0 +1,3 @@ +# ruff: noqa: F401 +from .pipelines import HunyuanVideoPipeline +from .schedulers import FlowMatchDiscreteScheduler diff --git a/fastvideo/v1/models/hunyuan/diffusion/pipelines/__init__.py b/fastvideo/v1/models/hunyuan/diffusion/pipelines/__init__.py new file mode 100644 index 000000000..fa109e14d --- /dev/null +++ b/fastvideo/v1/models/hunyuan/diffusion/pipelines/__init__.py @@ -0,0 +1,2 @@ +# ruff: noqa: F401 +from .pipeline_hunyuan_video import HunyuanVideoPipeline diff --git a/fastvideo/v1/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py b/fastvideo/v1/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py new file mode 100644 index 000000000..2efc9fecf --- /dev/null +++ b/fastvideo/v1/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py @@ -0,0 +1,931 @@ +# Copyright 2024 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. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== +import inspect +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.configuration_utils import FrozenDict +from diffusers.image_processor import VaeImageProcessor +from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin +from diffusers.models import AutoencoderKL +from diffusers.models.lora import adjust_lora_scale_text_encoder +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, deprecate, logging, replace_example_docstring, + scale_lora_layers) +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange + +from fastvideo.v1.utils.communications import all_gather +from fastvideo.v1.utils.parallel_states import get_sequence_parallel_state, nccl_info + +from ...constants import PRECISION_TO_TYPE +from ...modules import HYVideoDiffusionTransformer +from ...text_encoder import TextEncoder +from ...vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """""" + + +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + """ + Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = (guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg) + return noise_cfg + + +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler.") + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler.") + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +@dataclass +class HunyuanVideoPipelineOutput(BaseOutput): + videos: Union[torch.Tensor, np.ndarray] + + +class HunyuanVideoPipeline(DiffusionPipeline): + r""" + Pipeline for text-to-video generation using HunyuanVideo. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. + text_encoder ([`TextEncoder`]): + Frozen text-encoder. + text_encoder_2 ([`TextEncoder`]): + Frozen text-encoder_2. + transformer ([`HYVideoDiffusionTransformer`]): + A `HYVideoDiffusionTransformer` to denoise the encoded video latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" + _optional_components = ["text_encoder_2"] + _exclude_from_cpu_offload = ["transformer"] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: TextEncoder, + transformer: HYVideoDiffusionTransformer, + scheduler: KarrasDiffusionSchedulers, + text_encoder_2: Optional[TextEncoder] = None, + progress_bar_config: Dict[str, Any] = None, + args=None, + ): + super().__init__() + + # ========================================================================================== + if progress_bar_config is None: + progress_bar_config = {} + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + self._progress_bar_config.update(progress_bar_config) + + self.args = args + # ========================================================================================== + + if (hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1): + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" + f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " + "to update the config accordingly as leaving `steps_offset` might led to incorrect results" + " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," + " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" + " file") + deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(scheduler.config) + new_config["steps_offset"] = 1 + scheduler._internal_dict = FrozenDict(new_config) + + if (hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True): + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." + " `clip_sample` should be set to False in the configuration file. Please make sure to update the" + " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" + " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" + " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file") + deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(scheduler.config) + new_config["clip_sample"] = False + scheduler._internal_dict = FrozenDict(new_config) + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + transformer=transformer, + scheduler=scheduler, + text_encoder_2=text_encoder_2, + ) + self.vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + def encode_prompt( + self, + prompt, + device, + num_videos_per_prompt, + do_classifier_free_guidance, + negative_prompt=None, + prompt_embeds: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_attention_mask: Optional[torch.Tensor] = None, + lora_scale: Optional[float] = None, + clip_skip: Optional[int] = None, + text_encoder: Optional[TextEncoder] = None, + data_type: Optional[str] = "image", + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + num_videos_per_prompt (`int`): + number of videos that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the video generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + attention_mask (`torch.Tensor`, *optional*): + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + negative_attention_mask (`torch.Tensor`, *optional*): + lora_scale (`float`, *optional*): + A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + text_encoder (TextEncoder, *optional*): + data_type (`str`, *optional*): + """ + if text_encoder is None: + text_encoder = self.text_encoder + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, LoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if not USE_PEFT_BACKEND: + adjust_lora_scale_text_encoder(text_encoder.model, lora_scale) + else: + scale_lora_layers(text_encoder.model, lora_scale) + + if prompt_embeds is None: + # textual inversion: process multi-vector tokens if necessary + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, text_encoder.tokenizer) + + text_inputs = text_encoder.text2tokens(prompt, data_type=data_type) + if clip_skip is None: + prompt_outputs = text_encoder.encode(text_inputs, data_type=data_type, device=device) + prompt_embeds = prompt_outputs.hidden_state + else: + prompt_outputs = text_encoder.encode( + text_inputs, + output_hidden_states=True, + data_type=data_type, + device=device, + ) + # Access the `hidden_states` first, that contains a tuple of + # all the hidden states from the encoder layers. Then index into + # the tuple to access the hidden states from the desired layer. + prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)] + # We also need to apply the final LayerNorm here to not mess with the + # representations. The `last_hidden_states` that we typically use for + # obtaining the final prompt representations passes through the LayerNorm + # layer. + prompt_embeds = text_encoder.model.text_model.final_layer_norm(prompt_embeds) + + attention_mask = prompt_outputs.attention_mask + if attention_mask is not None: + attention_mask = attention_mask.to(device) + bs_embed, seq_len = attention_mask.shape + attention_mask = attention_mask.repeat(1, num_videos_per_prompt) + attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len) + + if text_encoder is not None: + prompt_embeds_dtype = text_encoder.dtype + elif self.transformer is not None: + prompt_embeds_dtype = self.transformer.dtype + else: + prompt_embeds_dtype = prompt_embeds.dtype + + prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) + + if prompt_embeds.ndim == 2: + bs_embed, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, -1) + else: + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) + + return ( + prompt_embeds, + negative_prompt_embeds, + attention_mask, + negative_attention_mask, + ) + + def decode_latents(self, latents, enable_tiling=True): + deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead" + deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False) + + latents = 1 / self.vae.config.scaling_factor * latents + if enable_tiling: + self.vae.enable_tiling() + image = self.vae.decode(latents, return_dict=False)[0] + image = (image / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + if image.ndim == 4: + image = image.cpu().permute(0, 2, 3, 1).float() + else: + image = image.cpu().float() + return image + + def prepare_extra_func_kwargs(self, func, kwargs): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + extra_step_kwargs = {} + + for k, v in kwargs.items(): + accepts = k in set(inspect.signature(func).parameters.keys()) + if accepts: + extra_step_kwargs[k] = v + return extra_step_kwargs + + def check_inputs( + self, + prompt, + height, + width, + video_length, + callback_steps, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + vae_ver="88-4c-sd", + ): + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if video_length is not None: + if "884" in vae_ver: + if video_length != 1 and (video_length - 1) % 4 != 0: + raise ValueError(f"`video_length` has to be 1 or a multiple of 4 but is {video_length}.") + elif "888" in vae_ver: + if video_length != 1 and (video_length - 1) % 8 != 0: + raise ValueError(f"`video_length` has to be 1 or a multiple of 8 but is {video_length}.") + + if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0): + raise ValueError(f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}.") + if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs + for k in callback_on_step_end_tensor_inputs): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two.") + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.") + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError(f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}.") + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + video_length, + dtype, + device, + generator, + latents=None, + ): + shape = ( + batch_size, + num_channels_latents, + video_length, + int(height) // self.vae_scale_factor, + int(width) // self.vae_scale_factor, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators.") + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler + if hasattr(self.scheduler, "init_noise_sigma"): + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + return latents + + # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding + def get_guidance_scale_embedding( + self, + w: torch.Tensor, + embedding_dim: int = 512, + dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + """ + See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 + + Args: + w (`torch.Tensor`): + Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. + embedding_dim (`int`, *optional*, defaults to 512): + Dimension of the embeddings to generate. + dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): + Data type of the generated embeddings. + + Returns: + `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. + """ + assert len(w.shape) == 1 + w = w * 1000.0 + + half_dim = embedding_dim // 2 + emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) + emb = w.to(dtype)[:, None] * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) + if embedding_dim % 2 == 1: # zero pad + emb = torch.nn.functional.pad(emb, (0, 1)) + assert emb.shape == (w.shape[0], embedding_dim) + return emb + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + @property + def clip_skip(self): + return self._clip_skip + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + @property + def do_classifier_free_guidance(self): + # return self._guidance_scale > 1 and self.transformer.config.time_cond_proj_dim is None + return self._guidance_scale > 1 + + @property + def cross_attention_kwargs(self): + return self._cross_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]], + height: int, + width: int, + video_length: int, + data_type: str = "video", + num_inference_steps: int = 50, + timesteps: List[int] = None, + sigmas: List[float] = None, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_videos_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, + MultiPipelineCallbacks, ]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + vae_ver: str = "88-4c-sd", + enable_tiling: bool = False, + enable_vae_sp: bool = False, + n_tokens: Optional[int] = None, + embedded_guidance_scale: Optional[float] = None, + mask_strategy: Optional[Dict[str, list]] = None, + **kwargs, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`): + The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + video_length (`int`): + The number of frames in the generated video. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 7.5): + A higher guidance scale value encourages the model to generate images closely linked to the text + `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide what to not include in image generation. If not defined, you need to + pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies + to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If + not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. + + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a + plain tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in + [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + guidance_rescale (`float`, *optional*, defaults to 0.0): + Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + + Examples: + + Returns: + [`~HunyuanVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list with the generated images and the + second element is a list of `bool`s indicating whether the corresponding generated image contains + "not-safe-for-work" (nsfw) content. + """ + callback = kwargs.pop("callback", None) + callback_steps = kwargs.pop("callback_steps", None) + + if callback is not None: + deprecate( + "callback", + "1.0.0", + "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", + ) + if callback_steps is not None: + deprecate( + "callback_steps", + "1.0.0", + "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", + ) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 0. Default height and width to unet + # height = height or self.transformer.config.sample_size * self.vae_scale_factor + # width = width or self.transformer.config.sample_size * self.vae_scale_factor + # to deal with lora scaling and other possible forward hooks + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + video_length, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + callback_on_step_end_tensor_inputs, + vae_ver=vae_ver, + ) + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + self._clip_skip = clip_skip + self._cross_attention_kwargs = cross_attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = (torch.device(f"cuda:{dist.get_rank()}") if dist.is_initialized() else self._execution_device) + + # 3. Encode input prompt + lora_scale = (self.cross_attention_kwargs.get("scale", None) + if self.cross_attention_kwargs is not None else None) + + ( + prompt_embeds, + negative_prompt_embeds, + prompt_mask, + negative_prompt_mask, + ) = self.encode_prompt( + prompt, + device, + num_videos_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + attention_mask=attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_attention_mask=negative_attention_mask, + lora_scale=lora_scale, + clip_skip=self.clip_skip, + data_type=data_type, + ) + if self.text_encoder_2 is not None: + ( + prompt_embeds_2, + negative_prompt_embeds_2, + prompt_mask_2, + negative_prompt_mask_2, + ) = self.encode_prompt( + prompt, + device, + num_videos_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=None, + attention_mask=None, + negative_prompt_embeds=None, + negative_attention_mask=None, + lora_scale=lora_scale, + clip_skip=self.clip_skip, + text_encoder=self.text_encoder_2, + data_type=data_type, + ) + else: + prompt_embeds_2 = None + negative_prompt_embeds_2 = None + prompt_mask_2 = None + negative_prompt_mask_2 = None + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + if prompt_mask is not None: + prompt_mask = torch.cat([negative_prompt_mask, prompt_mask]) + if prompt_embeds_2 is not None: + prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2]) + if prompt_mask_2 is not None: + prompt_mask_2 = torch.cat([negative_prompt_mask_2, prompt_mask_2]) + + # 4. Prepare timesteps + extra_set_timesteps_kwargs = self.prepare_extra_func_kwargs(self.scheduler.set_timesteps, + {"n_tokens": n_tokens}) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + **extra_set_timesteps_kwargs, + ) + if "884" in vae_ver: + video_length = (video_length - 1) // 4 + 1 + elif "888" in vae_ver: + video_length = (video_length - 1) // 8 + 1 + else: + video_length = video_length + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + video_length, + prompt_embeds.dtype, + device, + generator, + latents, + ) + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + if get_sequence_parallel_state(): + latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous() + latents = latents[:, :, rank, :, :, :] + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.step, + { + "generator": generator, + "eta": eta + }, + ) + + target_dtype = PRECISION_TO_TYPE[self.args.precision] + autocast_enabled = (target_dtype != torch.float32) and not self.args.disable_autocast + vae_dtype = PRECISION_TO_TYPE[self.args.vae_precision] + vae_autocast_enabled = (vae_dtype != torch.float32) and not self.args.disable_autocast + + # 7. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + def dict_to_3d_list(mask_strategy, t_max=50, l_max=60, h_max=24): + result = [[[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max)] + if mask_strategy is None: + return result + for key, value in mask_strategy.items(): + t, l, h = map(int, key.split('_')) + result[t][l][h] = value + return result + + mask_strategy = dict_to_3d_list(mask_strategy) + # if is_progress_bar: + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = (torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents) + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + t_expand = t.repeat(latent_model_input.shape[0]) + guidance_expand = (torch.tensor( + [embedded_guidance_scale] * latent_model_input.shape[0], + dtype=torch.float32, + device=device, + ).to(target_dtype) * 1000.0 if embedded_guidance_scale is not None else None) + # predict the noise residual + with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled): + # concat prompt_embeds_2 and prompt_embeds. Mismatch fill with zeros + if prompt_embeds_2.shape[-1] != prompt_embeds.shape[-1]: + prompt_embeds_2 = F.pad( + prompt_embeds_2, + (0, prompt_embeds.shape[2] - prompt_embeds_2.shape[1]), + value=0, + ).unsqueeze(1) + encoder_hidden_states = torch.cat([prompt_embeds_2, prompt_embeds], dim=1) + noise_pred = self.transformer( # For an input image (129, 192, 336) (1, 256, 256) + latent_model_input, + encoder_hidden_states, + t_expand, + prompt_mask, + mask_strategy=mask_strategy[i], + guidance=guidance_expand, + return_dict=False, + )[0] + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = rescale_noise_cfg( + noise_pred, + noise_pred_text, + guidance_rescale=self.guidance_rescale, + ) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if progress_bar is not None: + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + + if get_sequence_parallel_state(): + latents = all_gather(latents, dim=2) + + if not output_type == "latent": + expand_temporal_dim = False + if len(latents.shape) == 4: + if isinstance(self.vae, AutoencoderKLCausal3D): + latents = latents.unsqueeze(2) + expand_temporal_dim = True + elif len(latents.shape) == 5: + pass + else: + raise ValueError( + f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}.") + + if (hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor): + latents = (latents / self.vae.config.scaling_factor + self.vae.config.shift_factor) + else: + latents = latents / self.vae.config.scaling_factor + + with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): + if enable_tiling: + self.vae.enable_tiling() + if enable_vae_sp: + self.vae.enable_parallel() + image = self.vae.decode(latents, return_dict=False, generator=generator)[0] + + if expand_temporal_dim or image.shape[2] == 1: + image = image.squeeze(2) + + else: + image = latents + + image = (image / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 + image = image.cpu().float() + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return image + + return HunyuanVideoPipelineOutput(videos=image) diff --git a/fastvideo/v1/models/hunyuan/diffusion/schedulers/__init__.py b/fastvideo/v1/models/hunyuan/diffusion/schedulers/__init__.py new file mode 100644 index 000000000..2238803fe --- /dev/null +++ b/fastvideo/v1/models/hunyuan/diffusion/schedulers/__init__.py @@ -0,0 +1,2 @@ +# ruff: noqa: F401 +from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler diff --git a/fastvideo/v1/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py b/fastvideo/v1/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py new file mode 100644 index 000000000..bf032ab0e --- /dev/null +++ b/fastvideo/v1/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py @@ -0,0 +1,239 @@ +# Copyright 2024 Stability AI, Katherine Crowson 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. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== + +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput, logging + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + reverse (`bool`, defaults to `True`): + Whether to reverse the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + reverse: bool = True, + solver: str = "euler", + n_tokens: Optional[int] = None, + ): + sigmas = torch.linspace(1, 0, num_train_timesteps + 1) + + if not reverse: + sigmas = sigmas.flip(0) + + self.sigmas = sigmas + # the value fed to model + self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32) + + self._step_index = None + self._begin_index = None + + self.supported_solver = ["euler"] + if solver not in self.supported_solver: + raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}") + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device] = None, + n_tokens: int = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + """ + self.num_inference_steps = num_inference_steps + + sigmas = torch.linspace(1, 0, num_inference_steps + 1) + sigmas = self.sd3_time_shift(sigmas) + + if not self.config.reverse: + sigmas = 1 - sigmas + + self.sigmas = sigmas + self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device) + + # Reset step index + self._step_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + return sample + + def sd3_time_shift(self, t: torch.Tensor): + return (self.config.shift * t) / (1 + (self.config.shift - 1) * t) + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor)): + raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep."), ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] + + if self.config.solver == "euler": + prev_sample = sample + model_output.to(torch.float32) * dt + else: + raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}") + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return (prev_sample, ) + + return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps diff --git a/fastvideo/v1/models/hunyuan/idle_config.py b/fastvideo/v1/models/hunyuan/idle_config.py new file mode 100644 index 000000000..7d13c9d0a --- /dev/null +++ b/fastvideo/v1/models/hunyuan/idle_config.py @@ -0,0 +1,380 @@ +# ruff: noqa: F405, F403 +import argparse +import re + +from .constants import * +from .modules.models import HUNYUAN_VIDEO_CONFIG + + +def parse_args(namespace=None): + parser = argparse.ArgumentParser(description="HunyuanVideo inference script") + + parser = add_network_args(parser) + parser = add_extra_models_args(parser) + parser = add_denoise_schedule_args(parser) + parser = add_inference_args(parser) + parser = add_parallel_args(parser) + + args = parser.parse_args(namespace=namespace) + args = sanity_check_args(args) + + return args + + +def add_network_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="HunyuanVideo network args") + + # Main model + group.add_argument( + "--model", + type=str, + choices=list(HUNYUAN_VIDEO_CONFIG.keys()), + default="HYVideo-T/2-cfgdistill", + ) + group.add_argument( + "--latent-channels", + type=str, + default=16, + help="Number of latent channels of DiT. If None, it will be determined by `vae`. If provided, " + "it still needs to match the latent channels of the VAE model.", + ) + group.add_argument( + "--precision", + type=str, + default="bf16", + choices=PRECISIONS, + help="Precision mode. Options: fp32, fp16, bf16. Applied to the backbone model and optimizer.", + ) + + # RoPE + group.add_argument("--rope-theta", type=int, default=256, help="Theta used in RoPE.") + return parser + + +def add_extra_models_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Extra models args, including vae, text encoders and tokenizers)") + + # - VAE + group.add_argument( + "--vae", + type=str, + default="884-16c-hy", + choices=list(VAE_PATH), + help="Name of the VAE model.", + ) + group.add_argument( + "--vae-precision", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the VAE model.", + ) + group.add_argument( + "--vae-tiling", + action="store_true", + help="Enable tiling for the VAE model to save GPU memory.", + ) + group.set_defaults(vae_tiling=True) + + group.add_argument( + "--text-encoder", + type=str, + default="llm", + choices=list(TEXT_ENCODER_PATH), + help="Name of the text encoder model.", + ) + group.add_argument( + "--text-encoder-precision", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the text encoder model.", + ) + group.add_argument( + "--text-states-dim", + type=int, + default=4096, + help="Dimension of the text encoder hidden states.", + ) + group.add_argument("--text-len", type=int, default=256, help="Maximum length of the text input.") + group.add_argument( + "--tokenizer", + type=str, + default="llm", + choices=list(TOKENIZER_PATH), + help="Name of the tokenizer model.", + ) + group.add_argument( + "--prompt-template", + type=str, + default="dit-llm-encode", + choices=PROMPT_TEMPLATE, + help="Image prompt template for the decoder-only text encoder model.", + ) + group.add_argument( + "--prompt-template-video", + type=str, + default="dit-llm-encode-video", + choices=PROMPT_TEMPLATE, + help="Video prompt template for the decoder-only text encoder model.", + ) + group.add_argument( + "--hidden-state-skip-layer", + type=int, + default=2, + help="Skip layer for hidden states.", + ) + group.add_argument( + "--apply-final-norm", + action="store_true", + help="Apply final normalization to the used text encoder hidden states.", + ) + + # - CLIP + group.add_argument( + "--text-encoder-2", + type=str, + default="clipL", + choices=list(TEXT_ENCODER_PATH), + help="Name of the second text encoder model.", + ) + group.add_argument( + "--text-encoder-precision-2", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the second text encoder model.", + ) + group.add_argument( + "--text-states-dim-2", + type=int, + default=768, + help="Dimension of the second text encoder hidden states.", + ) + group.add_argument( + "--tokenizer-2", + type=str, + default="clipL", + choices=list(TOKENIZER_PATH), + help="Name of the second tokenizer model.", + ) + group.add_argument( + "--text-len-2", + type=int, + default=77, + help="Maximum length of the second text input.", + ) + + return parser + + +def add_denoise_schedule_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Denoise schedule args") + + group.add_argument( + "--denoise-type", + type=str, + default="flow", + help="Denoise type for noised inputs.", + ) + + # Flow Matching + group.add_argument( + "--flow-shift", + type=float, + default=7.0, + help="Shift factor for flow matching schedulers.", + ) + group.add_argument( + "--flow-reverse", + action="store_true", + help="If reverse, learning/sampling from t=1 -> t=0.", + ) + group.add_argument( + "--flow-solver", + type=str, + default="euler", + help="Solver for flow matching.", + ) + group.add_argument( + "--use-linear-quadratic-schedule", + action="store_true", + help="Use linear quadratic schedule for flow matching." + "Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)", + ) + group.add_argument( + "--linear-schedule-end", + type=int, + default=25, + help="End step for linear quadratic schedule for flow matching.", + ) + + return parser + + +def add_inference_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Inference args") + + # ======================== Model loads ======================== + group.add_argument( + "--model-base", + type=str, + default="ckpts", + help="Root path of all the models, including t2v models and extra models.", + ) + group.add_argument( + "--dit-weight", + type=str, + default="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt", + help="Path to the HunyuanVideo model. If None, search the model in the args.model_root." + "1. If it is a file, load the model directly." + "2. If it is a directory, search the model in the directory. Support two types of models: " + "1) named `pytorch_model_*.pt`" + "2) named `*_model_states.pt`, where * can be `mp_rank_00`.", + ) + group.add_argument( + "--model-resolution", + type=str, + default="540p", + choices=["540p", "720p"], + help="Root path of all the models, including t2v models and extra models.", + ) + group.add_argument( + "--load-key", + type=str, + default="module", + help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.", + ) + group.add_argument( + "--use-cpu-offload", + action="store_true", + help="Use CPU offload for the model load.", + ) + + # ======================== Inference general setting ======================== + group.add_argument( + "--batch-size", + type=int, + default=1, + help="Batch size for inference and evaluation.", + ) + group.add_argument( + "--infer-steps", + type=int, + default=50, + help="Number of denoising steps for inference.", + ) + group.add_argument( + "--disable-autocast", + action="store_true", + help="Disable autocast for denoising loop and vae decoding in pipeline sampling.", + ) + group.add_argument( + "--save-path", + type=str, + default="./results", + help="Path to save the generated samples.", + ) + group.add_argument( + "--save-path-suffix", + type=str, + default="", + help="Suffix for the directory of saved samples.", + ) + group.add_argument( + "--name-suffix", + type=str, + default="", + help="Suffix for the names of saved samples.", + ) + group.add_argument( + "--num-videos", + type=int, + default=1, + help="Number of videos to generate for each prompt.", + ) + # ---sample size--- + group.add_argument( + "--video-size", + type=int, + nargs="+", + default=(720, 1280), + help="Video size for training. If a single value is provided, it will be used for both height " + "and width. If two values are provided, they will be used for height and width " + "respectively.", + ) + group.add_argument( + "--video-length", + type=int, + default=129, + help="How many frames to sample from a video. if using 3d vae, the number should be 4n+1", + ) + # --- prompt --- + group.add_argument( + "--prompt", + type=str, + default=None, + help="Prompt for sampling during evaluation.", + ) + group.add_argument( + "--seed-type", + type=str, + default="auto", + choices=["file", "random", "fixed", "auto"], + help="Seed type for evaluation. If file, use the seed from the CSV file. If random, generate a " + "random seed. If fixed, use the fixed seed given by `--seed`. If auto, `csv` will use the " + "seed column if available, otherwise use the fixed `seed` value. `prompt` will use the " + "fixed `seed` value.", + ) + group.add_argument("--seed", type=int, default=None, help="Seed for evaluation.") + + # Classifier-Free Guidance + group.add_argument("--neg-prompt", type=str, default=None, help="Negative prompt for sampling.") + group.add_argument("--cfg-scale", type=float, default=1.0, help="Classifier free guidance scale.") + group.add_argument( + "--embedded-cfg-scale", + type=float, + default=6.0, + help="Embedded classifier free guidance scale.", + ) + + group.add_argument( + "--reproduce", + action="store_true", + help="Enable reproducibility by setting random seeds and deterministic algorithms.", + ) + + return parser + + +def add_parallel_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Parallel args") + + # ======================== Model loads ======================== + group.add_argument( + "--ulysses-degree", + type=int, + default=1, + help="Ulysses degree.", + ) + group.add_argument( + "--ring-degree", + type=int, + default=1, + help="Ulysses degree.", + ) + + return parser + + +def sanity_check_args(args): + # VAE channels + vae_pattern = r"\d{2,3}-\d{1,2}c-\w+" + if not re.match(vae_pattern, args.vae): + raise ValueError(f"Invalid VAE model: {args.vae}. Must be in the format of '{vae_pattern}'.") + vae_channels = int(args.vae.split("-")[1][:-1]) + if args.latent_channels is None: + args.latent_channels = vae_channels + if vae_channels != args.latent_channels: + raise ValueError(f"Latent channels ({args.latent_channels}) must match the VAE channels ({vae_channels}).") + return args diff --git a/fastvideo/v1/models/hunyuan/inference.py b/fastvideo/v1/models/hunyuan/inference.py new file mode 100644 index 000000000..6dcadc690 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/inference.py @@ -0,0 +1,504 @@ +import os +import random +import time +from pathlib import Path + +import torch +from loguru import logger +from safetensors.torch import load_file as safetensors_load_file + +from fastvideo.v1.models.hunyuan.constants import NEGATIVE_PROMPT, PRECISION_TO_TYPE, PROMPT_TEMPLATE +from fastvideo.v1.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline +from fastvideo.v1.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler +from fastvideo.v1.models.hunyuan.modules import load_model +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.hunyuan.utils.data_utils import align_to +from fastvideo.v1.models.hunyuan.vae import load_vae +from fastvideo.v1.utils.parallel_states import nccl_info + + +class Inference(object): + + def __init__( + self, + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=None, + pipeline=None, + use_cpu_offload=False, + device=None, + logger=None, + parallel_args=None, + ): + self.vae = vae + self.vae_kwargs = vae_kwargs + + self.text_encoder = text_encoder + self.text_encoder_2 = text_encoder_2 + + self.model = model + self.pipeline = pipeline + self.use_cpu_offload = use_cpu_offload + + self.args = args + self.device = (device if device is not None else "cuda" if torch.cuda.is_available() else "cpu") + self.logger = logger + self.parallel_args = parallel_args + + @classmethod + def from_pretrained(cls, pretrained_model_path, args, device=None, **kwargs): + """ + Initialize the Inference pipeline. + + Args: + pretrained_model_path (str or pathlib.Path): The model path, including t2v, text encoder and vae checkpoints. + args (argparse.Namespace): The arguments for the pipeline. + device (int): The device for inference. Default is 0. + """ + from fastvideo.v1.pipelines.loader import get_pipeline_loader + loader = get_pipeline_loader(args) + vae, vae_kwargs, text_encoder, text_encoder_2, transformer, scheduler = \ + loader.load_components(args) + return cls( + args=args, + vae=vae, + vae_kwargs=vae_kwargs, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + model=transformer, + use_cpu_offload=args.use_cpu_offload, + device=device, + logger=logger, + parallel_args=None, + ) + # ======================================================================== + logger.info(f"Got text-to-video model root path: {pretrained_model_path}") + + # ==================== Initialize Distributed Environment ================ + if nccl_info.sp_size > 1: + device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}") + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + + parallel_args = None # {"ulysses_degree": args.ulysses_degree, "ring_degree": args.ring_degree} + + # ======================== Get the args path ============================= + + # Disable gradient + torch.set_grad_enabled(False) + + # =========================== Build main model =========================== + logger.info("Building model...") + factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]} + in_channels = args.latent_channels + out_channels = args.latent_channels + + model = load_model( + args, + in_channels=in_channels, + out_channels=out_channels, + factor_kwargs=factor_kwargs, + ) + model = model.to(device) + model = Inference.load_state_dict(args, model, pretrained_model_path) + if args.enable_torch_compile: + model = torch.compile(model) + model.eval() + + # ============================= Build extra models ======================== + # VAE + vae, _, s_ratio, t_ratio = load_vae( + args.vae, + args.vae_precision, + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio} + + # Text encoder + if args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[args.prompt_template_video].get("crop_start", 0) + elif args.prompt_template is not None: + crop_start = PROMPT_TEMPLATE[args.prompt_template].get("crop_start", 0) + else: + crop_start = 0 + max_length = args.text_len + crop_start + + # prompt_template + prompt_template = (PROMPT_TEMPLATE[args.prompt_template] if args.prompt_template is not None else None) + + # prompt_template_video + prompt_template_video = (PROMPT_TEMPLATE[args.prompt_template_video] + if args.prompt_template_video is not None else None) + + text_encoder = TextEncoder( + text_encoder_type=args.text_encoder, + max_length=max_length, + text_encoder_precision=args.text_encoder_precision, + tokenizer_type=args.tokenizer, + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=args.hidden_state_skip_layer, + apply_final_norm=args.apply_final_norm, + reproduce=args.reproduce, + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + text_encoder_2 = None + if args.text_encoder_2 is not None: + text_encoder_2 = TextEncoder( + text_encoder_type=args.text_encoder_2, + max_length=args.text_len_2, + text_encoder_precision=args.text_encoder_precision_2, + tokenizer_type=args.tokenizer_2, + reproduce=args.reproduce, + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + + return cls( + args=args, + vae=vae, + vae_kwargs=vae_kwargs, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + model=model, + use_cpu_offload=args.use_cpu_offload, + device=device, + logger=logger, + parallel_args=parallel_args, + ) + + @staticmethod + def load_state_dict(args, model, pretrained_model_path): + load_key = args.load_key + dit_weight = Path(args.dit_weight) + + if dit_weight is None: + model_dir = pretrained_model_path / f"t2v_{args.model_resolution}" + files = list(model_dir.glob("*.pt")) + if len(files) == 0: + raise ValueError(f"No model weights found in {model_dir}") + if str(files[0]).startswith("pytorch_model_"): + model_path = dit_weight / f"pytorch_model_{load_key}.pt" + bare_model = True + elif any(str(f).endswith("_model_states.pt") for f in files): + files = [f for f in files if str(f).endswith("_model_states.pt")] + model_path = files[0] + if len(files) > 1: + logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}") + bare_model = False + else: + raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: " + f"{list(map(str, files))}. When given a directory as --dit-weight, only " + f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and " + f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a " + f"specific weight file, please provide the full path to the file.") + else: + if dit_weight.is_dir(): + files = list(dit_weight.glob("*.pt")) + if len(files) == 0: + raise ValueError(f"No model weights found in {dit_weight}") + if str(files[0]).startswith("pytorch_model_"): + model_path = dit_weight / f"pytorch_model_{load_key}.pt" + bare_model = True + elif any(str(f).endswith("_model_states.pt") for f in files): + files = [f for f in files if str(f).endswith("_model_states.pt")] + model_path = files[0] + if len(files) > 1: + logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}") + bare_model = False + else: + raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: " + f"{list(map(str, files))}. When given a directory as --dit-weight, only " + f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and " + f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a " + f"specific weight file, please provide the full path to the file.") + elif dit_weight.is_file(): + model_path = dit_weight + bare_model = "unknown" + else: + raise ValueError(f"Invalid model path: {dit_weight}") + + print('model_path', model_path) + if not model_path.exists(): + raise ValueError(f"model_path not exists: {model_path}") + logger.info(f"Loading torch model {model_path}...") + if model_path.suffix == ".safetensors": + # Use safetensors library for .safetensors files + state_dict = safetensors_load_file(model_path) + elif model_path.suffix == ".pt": + # Use torch for .pt files + state_dict = torch.load(model_path, map_location=lambda storage, loc: storage) + else: + raise ValueError(f"Unsupported file format: {model_path}") + + if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict): + bare_model = False + if bare_model is False: + if load_key in state_dict: + state_dict = state_dict[load_key] + else: + raise KeyError(f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint " + f"are: {list(state_dict.keys())}.") + model.load_state_dict(state_dict, strict=True) + return model + + @staticmethod + def parse_size(size): + if isinstance(size, int): + size = [size] + if not isinstance(size, (list, tuple)): + raise ValueError(f"Size must be an integer or (height, width), got {size}.") + if len(size) == 1: + size = [size[0], size[0]] + if len(size) != 2: + raise ValueError(f"Size must be an integer or (height, width), got {size}.") + return size + + +class HunyuanVideoSampler(Inference): + + def __init__( + self, + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=None, + pipeline=None, + use_cpu_offload=False, + device=0, + logger=None, + parallel_args=None, + ): + super().__init__( + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=text_encoder_2, + pipeline=pipeline, + use_cpu_offload=use_cpu_offload, + device=device, + logger=logger, + parallel_args=parallel_args, + ) + + self.pipeline = self.load_diffusion_pipeline( + args=args, + vae=self.vae, + text_encoder=self.text_encoder, + text_encoder_2=self.text_encoder_2, + model=self.model, + device=self.device, + ) + + self.default_negative_prompt = NEGATIVE_PROMPT + + def load_diffusion_pipeline( + self, + args, + vae, + text_encoder, + text_encoder_2, + model, + scheduler=None, + device=None, + progress_bar_config=None, + data_type="video", + ): + print("Loading diffusion pipeline...") + # print('vae', vae) + print('text_encoder', text_encoder) + print('text_encoder_2', text_encoder_2) + print('model', model) + """Load the denoising scheduler for inference.""" + if scheduler is None: + if args.denoise_type == "flow": + scheduler = FlowMatchDiscreteScheduler( + shift=args.flow_shift, + reverse=args.flow_reverse, + solver=args.flow_solver, + ) + else: + raise ValueError(f"Invalid denoise type {args.denoise_type}") + + pipeline = HunyuanVideoPipeline( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + transformer=model, + scheduler=scheduler, + progress_bar_config=progress_bar_config, + args=args, + ) + if self.use_cpu_offload: + pipeline.enable_sequential_cpu_offload() + else: + pipeline = pipeline.to(device) + + return pipeline + + @torch.no_grad() + def predict( + self, + prompt, + height=192, + width=336, + video_length=129, + seed=None, + negative_prompt=None, + infer_steps=50, + guidance_scale=6, + flow_shift=5.0, + embedded_guidance_scale=None, + batch_size=1, + num_videos_per_prompt=1, + mask_strategy=None, + **kwargs, + ): + """ + Predict the image/video from the given text. + + Args: + prompt (str or List[str]): The input text. + kwargs: + height (int): The height of the output video. Default is 192. + width (int): The width of the output video. Default is 336. + video_length (int): The frame number of the output video. Default is 129. + seed (int or List[str]): The random seed for the generation. Default is a random integer. + negative_prompt (str or List[str]): The negative text prompt. Default is an empty string. + guidance_scale (float): The guidance scale for the generation. Default is 6.0. + num_images_per_prompt (int): The number of images per prompt. Default is 1. + infer_steps (int): The number of inference steps. Default is 100. + """ + + out_dict = dict() + + # ======================================================================== + # Arguments: seed + # ======================================================================== + if isinstance(seed, torch.Tensor): + seed = seed.tolist() + if seed is None: + seeds = [random.randint(0, 1_000_000) for _ in range(batch_size * num_videos_per_prompt)] + elif isinstance(seed, int): + seeds = [seed + i for _ in range(batch_size) for i in range(num_videos_per_prompt)] + elif isinstance(seed, (list, tuple)): + if len(seed) == batch_size: + seeds = [int(seed[i]) + j for i in range(batch_size) for j in range(num_videos_per_prompt)] + elif len(seed) == batch_size * num_videos_per_prompt: + seeds = [int(s) for s in seed] + else: + raise ValueError( + f"Length of seed must be equal to number of prompt(batch_size) or " + f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}.") + else: + raise ValueError(f"Seed must be an integer, a list of integers, or None, got {seed}.") + # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... + generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds] + out_dict["seeds"] = seeds + + # ======================================================================== + # Arguments: target_width, target_height, target_video_length + # ======================================================================== + if width <= 0 or height <= 0 or video_length <= 0: + raise ValueError( + f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}" + ) + if (video_length - 1) % 4 != 0: + raise ValueError(f"`video_length-1` must be a multiple of 4, got {video_length}") + + logger.info(f"Input (height, width, video_length) = ({height}, {width}, {video_length})") + + target_height = align_to(height, 16) + target_width = align_to(width, 16) + target_video_length = video_length + + out_dict["size"] = (target_height, target_width, target_video_length) + + # ======================================================================== + # Arguments: prompt, new_prompt, negative_prompt + # ======================================================================== + if not isinstance(prompt, str): + raise TypeError(f"`prompt` must be a string, but got {type(prompt)}") + prompt = [prompt.strip()] + + # negative prompt + if negative_prompt is None or negative_prompt == "": + negative_prompt = self.default_negative_prompt + if not isinstance(negative_prompt, str): + raise TypeError(f"`negative_prompt` must be a string, but got {type(negative_prompt)}") + negative_prompt = [negative_prompt.strip()] + + # ======================================================================== + # Scheduler + # ======================================================================== + scheduler = FlowMatchDiscreteScheduler( + shift=flow_shift, + reverse=self.args.flow_reverse, + solver=self.args.flow_solver, + ) + self.pipeline.scheduler = scheduler + + if "884" in self.args.vae: + latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8] + elif "888" in self.args.vae: + latents_size = [(video_length - 1) // 8 + 1, height // 8, width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + + # ======================================================================== + # Print infer args + # ======================================================================== + debug_str = f""" + height: {target_height} + width: {target_width} + video_length: {target_video_length} + prompt: {prompt} + neg_prompt: {negative_prompt} + seed: {seed} + infer_steps: {infer_steps} + num_videos_per_prompt: {num_videos_per_prompt} + guidance_scale: {guidance_scale} + n_tokens: {n_tokens} + flow_shift: {flow_shift} + embedded_guidance_scale: {embedded_guidance_scale}""" + logger.debug(debug_str) + + # ======================================================================== + # Pipeline inference + # ======================================================================== + start_time = time.time() + samples = self.pipeline( + prompt=prompt, + height=target_height, + width=target_width, + video_length=target_video_length, + num_inference_steps=infer_steps, + guidance_scale=guidance_scale, + negative_prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + generator=generator, + output_type="pil", + n_tokens=n_tokens, + embedded_guidance_scale=embedded_guidance_scale, + data_type="video" if target_video_length > 1 else "image", + is_progress_bar=True, + vae_ver=self.args.vae, + enable_tiling=self.args.vae_tiling, + enable_vae_sp=self.args.vae_sp, + mask_strategy=mask_strategy, + )[0] + out_dict["samples"] = samples + out_dict["prompts"] = prompt + + gen_time = time.time() - start_time + logger.info(f"Success, time: {gen_time}") + + return out_dict diff --git a/fastvideo/v1/models/hunyuan/modules/__init__.py b/fastvideo/v1/models/hunyuan/modules/__init__.py new file mode 100644 index 000000000..3bb947c4e --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/__init__.py @@ -0,0 +1,25 @@ +# from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer + + +# def load_model(args, in_channels, out_channels, factor_kwargs): +# """load hunyuan video model + +# Args: +# args (dict): model args +# in_channels (int): input channels number +# out_channels (int): output channels number +# factor_kwargs (dict): factor kwargs + +# Returns: +# model (nn.Module): The hunyuan video model +# """ +# if args.model in HUNYUAN_VIDEO_CONFIG.keys(): +# model = HYVideoDiffusionTransformer( +# in_channels=in_channels, +# out_channels=out_channels, +# **HUNYUAN_VIDEO_CONFIG[args.model], +# **factor_kwargs, +# ) +# return model +# else: +# raise NotImplementedError() diff --git a/fastvideo/v1/models/hunyuan/modules/activation_layers.py b/fastvideo/v1/models/hunyuan/modules/activation_layers.py new file mode 100644 index 000000000..f8774c26c --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/activation_layers.py @@ -0,0 +1,23 @@ +import torch.nn as nn + + +def get_activation_layer(act_type): + """get activation layer + + Args: + act_type (str): the activation type + + Returns: + torch.nn.functional: the activation layer + """ + if act_type == "gelu": + return lambda: nn.GELU() + elif act_type == "gelu_tanh": + # Approximate `tanh` requires torch >= 1.13 + return lambda: nn.GELU(approximate="tanh") + elif act_type == "relu": + return nn.ReLU + elif act_type == "silu": + return nn.SiLU + else: + raise ValueError(f"Unknown activation type: {act_type}") diff --git a/fastvideo/v1/models/hunyuan/modules/attenion.py b/fastvideo/v1/models/hunyuan/modules/attenion.py new file mode 100644 index 000000000..b275523a9 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/attenion.py @@ -0,0 +1,128 @@ +import torch +import torch.nn.functional as F +from einops import rearrange + +try: + from st_attn import sliding_tile_attention +except ImportError: + print("Could not load Sliding Tile Attention.") + sliding_tile_attention = None + +from fastvideo.v1.models.flash_attn_no_pad import flash_attn_no_pad +from flash_attn import flash_attn_func +from fastvideo.v1.utils.communications import all_gather, all_to_all_4D +from fastvideo.v1.utils.parallel_states import get_sequence_parallel_state, nccl_info + + +def attention( + q, + k, + v, + drop_rate=0, + attn_mask=None, + causal=False, +): + + qkv = torch.stack([q, k, v], dim=2) + + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.bool() + + x = flash_attn_no_pad(qkv, attn_mask, causal=causal, dropout_p=drop_rate, softmax_scale=None) + + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out + + +def tile(x, sp_size): + x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80) + return rearrange(x, + "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", + n_t=5, + n_h=6, + n_w=10, + ts_t=6, + ts_h=8, + ts_w=8) + + +def untile(x, sp_size): + x = rearrange(x, + "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d", + n_t=5, + n_h=6, + n_w=10, + ts_t=6, + ts_h=8, + ts_w=8) + return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80) + + +def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=None): + query, encoder_query = q + key, encoder_key = k + value, encoder_value = v + text_length = text_mask.sum() + + if get_sequence_parallel_state(): + # batch_size, seq_len, attn_heads, head_dim + query = all_to_all_4D(query, scatter_dim=2, gather_dim=1) + key = all_to_all_4D(key, scatter_dim=2, gather_dim=1) + value = all_to_all_4D(value, scatter_dim=2, gather_dim=1) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) + + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + # [b, s, h, d] + + sequence_length = query.size(1) + encoder_sequence_length = encoder_query.size(1) + + if mask_strategy[0] is not None: + query = torch.cat([tile(query, nccl_info.sp_size), encoder_query], dim=1).transpose(1, 2) + key = torch.cat([tile(key, nccl_info.sp_size), encoder_key], dim=1).transpose(1, 2) + value = torch.cat([tile(value, nccl_info.sp_size), encoder_value], dim=1).transpose(1, 2) + + head_num = query.size(1) + current_rank = nccl_info.rank_within_group + start_head = current_rank * head_num + windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)] + + hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2) + else: + txt_length = text_mask.sum().int().item() + query = torch.cat([query, encoder_query[:, :txt_length]], dim=1) + key = torch.cat([key, encoder_key[:, :txt_length]], dim=1) + value = torch.cat([value, encoder_value[:, :txt_length]], dim=1) + # B, S, 3, H, D + # qkv = torch.stack([query, key, value], dim=2) + + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + # hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) + hidden_states = flash_attn_func(query, key, value, causal=False, dropout_p=0.0, softmax_scale=None) + hidden_states = torch.cat([hidden_states, encoder_query[:, txt_length:, :]], dim=1) + + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length), + dim=1) + + if mask_strategy[0] is not None: + hidden_states = untile(hidden_states, nccl_info.sp_size) + + if get_sequence_parallel_state(): + hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) + encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous() + + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + + attn = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + b, s, a, d = attn.shape + attn = attn.reshape(b, s, -1) + + return attn diff --git a/fastvideo/v1/models/hunyuan/modules/embed_layers.py b/fastvideo/v1/models/hunyuan/modules/embed_layers.py new file mode 100644 index 000000000..38aad89cb --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/embed_layers.py @@ -0,0 +1,150 @@ +import math + +import torch +import torch.nn as nn + +from ..utils.helpers import to_2tuple + + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + + def __init__( + self, + patch_size=16, + in_chans=3, + embed_dim=768, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + self.flatten = flatten + + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + **factory_kwargs, + ) + nn.init.xavier_uniform_(self.proj.weight.view(self.proj.weight.size(0), -1)) + if bias: + nn.init.zeros_(self.proj.bias) + + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class TextProjection(nn.Module): + """ + Projects text embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.linear_1 = nn.Linear( + in_features=in_channels, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + self.act_1 = act_layer() + self.linear_2 = nn.Linear( + in_features=hidden_size, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / + half).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb diff --git a/fastvideo/v1/models/hunyuan/modules/mlp_layers.py b/fastvideo/v1/models/hunyuan/modules/mlp_layers.py new file mode 100644 index 000000000..921b85252 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/mlp_layers.py @@ -0,0 +1,107 @@ +# Modified from timm library: +# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13 + +from functools import partial + +import torch +import torch.nn as nn + +from ..utils.helpers import to_2tuple +from .modulate_layers import modulate + + +class MLP(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__( + self, + in_channels, + hidden_channels=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_channels + hidden_channels = hidden_channels or in_channels + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer(in_channels, hidden_channels, bias=bias[0], **factory_kwargs) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = (norm_layer(hidden_channels, **factory_kwargs) if norm_layer is not None else nn.Identity()) + self.fc2 = linear_layer(hidden_channels, out_features, bias=bias[1], **factory_kwargs) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +# +class MLPEmbedder(nn.Module): + """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py""" + + def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__(self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + # Just use LayerNorm for the final layer + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + if isinstance(patch_size, int): + self.linear = nn.Linear( + hidden_size, + patch_size * patch_size * out_channels, + bias=True, + **factory_kwargs, + ) + else: + self.linear = nn.Linear( + hidden_size, + patch_size[0] * patch_size[1] * patch_size[2] * out_channels, + bias=True, + ) + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + # Here we don't distinguish between the modulate types. Just use the simple one. + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift=shift, scale=scale) + x = self.linear(x) + return x diff --git a/fastvideo/v1/models/hunyuan/modules/models.py b/fastvideo/v1/models/hunyuan/modules/models.py new file mode 100644 index 000000000..42b2b9c85 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/models.py @@ -0,0 +1,657 @@ +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models import ModelMixin +from einops import rearrange + +from fastvideo.v1.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_embed +from fastvideo.v1.utils.parallel_states import nccl_info + +from .activation_layers import get_activation_layer +from .attenion import parallel_attention +from .embed_layers import PatchEmbed, TextProjection, TimestepEmbedder +from .mlp_layers import MLP, FinalLayer, MLPEmbedder +from .modulate_layers import ModulateDiT, apply_gate, modulate +from .norm_layers import get_norm_layer +from .posemb_layers import apply_rotary_emb +from .token_refiner import SingleTokenRefiner + + +class MMDoubleStreamBlock(nn.Module): + """ + A multimodal dit block with separate modulation for + text and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206 + (Flux.1): https://github.com/black-forest-labs/flux + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float, + mlp_act_type: str = "gelu_tanh", + qk_norm: bool = True, + qk_norm_type: str = "rms", + qkv_bias: bool = False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.img_mod = ModulateDiT( + hidden_size, + factor=6, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.img_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.img_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.img_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.img_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.img_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs, + ) + + self.txt_mod = ModulateDiT( + hidden_size, + factor=6, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.txt_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + self.txt_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.txt_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.txt_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.txt_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs, + ) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + vec: torch.Tensor, + freqs_cis: tuple = None, + text_mask: torch.Tensor = None, + mask_strategy=None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(vec).chunk(6, dim=-1) + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(vec).chunk(6, dim=-1) + + # Prepare image for attention. + img_modulated = self.img_norm1(img) + img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale) + + img_qkv = self.img_attn_qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + # 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) + # Apply RoPE if needed. + if freqs_cis is not None: + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) + + freqs_cis = ( + shrink_head(freqs_cis[0], dim=0), + shrink_head(freqs_cis[1], dim=0), + ) + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + + # Prepare txt for attention. + txt_modulated = self.txt_norm1(txt) + txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale) + txt_qkv = self.txt_attn_qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + # Apply QK-Norm if needed. + txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + mask_strategy=mask_strategy, + ) + + # attention computation end + img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:] + # Calculate the img blocks. + img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate) + img = img + apply_gate( + self.img_mlp(modulate(self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale)), + gate=img_mod2_gate, + ) + + # Calculate the txt blocks. + txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate) + txt = txt + apply_gate( + self.txt_mlp(modulate(self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale)), + gate=txt_mod2_gate, + ) + return img, txt + + +class MMSingleStreamBlock(nn.Module): + """ + A DiT block with parallel linear layers as described in + https://arxiv.org/abs/2302.05442 and adapted modulation interface. + Also refer to (SD3): https://arxiv.org/abs/2403.03206 + (Flux.1): https://github.com/black-forest-labs/flux + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + qk_norm: bool = True, + qk_norm_type: str = "rms", + qk_scale: float = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.hidden_size = hidden_size + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.mlp_hidden_dim = mlp_hidden_dim + self.scale = qk_scale or head_dim**-0.5 + + # qkv and mlp_in + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + mlp_hidden_dim, **factory_kwargs) + # proj and mlp_out + self.linear2 = nn.Linear(hidden_size + mlp_hidden_dim, hidden_size, **factory_kwargs) + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.mlp_act = get_activation_layer(mlp_act_type)() + self.modulation = ModulateDiT( + hidden_size, + factor=3, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + x: torch.Tensor, + vec: torch.Tensor, + txt_len: int, + freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None, + text_mask: torch.Tensor = None, + mask_strategy=None, + ) -> torch.Tensor: + mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1) + x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale) + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + + # Apply QK-Norm if needed. + q = self.q_norm(q).to(v) + k = self.k_norm(k).to(v) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) + + freqs_cis = ( + shrink_head(freqs_cis[0], dim=0), + shrink_head(freqs_cis[1], dim=0), + ) + + img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :] + img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :] + img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :] + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + mask_strategy=mask_strategy, + ) + + # attention computation end + + # Compute activation in mlp stream, cat again and run second linear layer. + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + apply_gate(output, gate=mod_gate) + + +class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): + """ + HunyuanVideo Transformer backbone + + Inherited from ModelMixin and ConfigMixin for compatibility with diffusers' sampler StableDiffusionPipeline. + + Reference: + [1] Flux.1: https://github.com/black-forest-labs/flux + [2] MMDiT: http://arxiv.org/abs/2403.03206 + + Parameters + ---------- + args: argparse.Namespace + The arguments parsed by argparse. + patch_size: list + The size of the patch. + in_channels: int + The number of input channels. + out_channels: int + The number of output channels. + hidden_size: int + The hidden size of the transformer backbone. + heads_num: int + The number of attention heads. + mlp_width_ratio: float + The ratio of the hidden size of the MLP in the transformer block. + mlp_act_type: str + The activation function of the MLP in the transformer block. + depth_double_blocks: int + The number of transformer blocks in the double blocks. + depth_single_blocks: int + The number of transformer blocks in the single blocks. + rope_dim_list: list + The dimension of the rotary embedding for t, h, w. + qkv_bias: bool + Whether to use bias in the qkv linear layer. + qk_norm: bool + Whether to use qk norm. + qk_norm_type: str + The type of qk norm. + guidance_embed: bool + Whether to use guidance embedding for distillation. + text_projection: str + The type of the text projection, default is single_refiner. + use_attention_mask: bool + Whether to use attention mask for text encoder. + dtype: torch.dtype + The dtype of the model. + device: torch.device + The device of the model. + """ + + @register_to_config + def __init__( + self, + patch_size: list = [1, 2, 2], + in_channels: int = 4, # Should be VAE.config.latent_channels. + out_channels: int = None, + hidden_size: int = 3072, + heads_num: int = 24, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + mm_double_blocks_depth: int = 20, + mm_single_blocks_depth: int = 40, + rope_dim_list: List[int] = [16, 56, 56], + qkv_bias: bool = True, + qk_norm: bool = True, + qk_norm_type: str = "rms", + guidance_embed: bool = False, # For modulation. + text_projection: str = "single_refiner", + use_attention_mask: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + text_states_dim: int = 4096, + text_states_dim_2: int = 768, + rope_theta: int = 256, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.unpatchify_channels = self.out_channels + self.guidance_embed = guidance_embed + self.rope_dim_list = rope_dim_list + self.rope_theta = rope_theta + # Text projection. Default to linear projection. + # Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831 + self.use_attention_mask = use_attention_mask + self.text_projection = text_projection + + if hidden_size % heads_num != 0: + raise ValueError(f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}") + pe_dim = hidden_size // heads_num + if sum(rope_dim_list) != pe_dim: + raise ValueError(f"Got {rope_dim_list} but expected positional dim {pe_dim}") + self.hidden_size = hidden_size + self.heads_num = heads_num + + # image projection + self.img_in = PatchEmbed(self.patch_size, self.in_channels, self.hidden_size, **factory_kwargs) + + # text projection + if self.text_projection == "linear": + self.txt_in = TextProjection( + self.config.text_states_dim, + self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs, + ) + elif self.text_projection == "single_refiner": + self.txt_in = SingleTokenRefiner( + self.config.text_states_dim, + hidden_size, + heads_num, + depth=2, + **factory_kwargs, + ) + else: + raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}") + + # time modulation + self.time_in = TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs) + + # text modulation + self.vector_in = MLPEmbedder(self.config.text_states_dim_2, self.hidden_size, **factory_kwargs) + + # guidance modulation + self.guidance_in = (TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs) + if guidance_embed else None) + + # double blocks + self.double_blocks = nn.ModuleList([ + MMDoubleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) for _ in range(mm_double_blocks_depth) + ]) + + # single blocks + self.single_blocks = nn.ModuleList([ + MMSingleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + **factory_kwargs, + ) for _ in range(mm_single_blocks_depth) + ]) + + self.final_layer = FinalLayer( + self.hidden_size, + self.patch_size, + self.out_channels, + get_activation_layer("silu"), + **factory_kwargs, + ) + + def enable_deterministic(self): + for block in self.double_blocks: + block.enable_deterministic() + for block in self.single_blocks: + block.enable_deterministic() + + def disable_deterministic(self): + for block in self.double_blocks: + block.disable_deterministic() + for block in self.single_blocks: + block.disable_deterministic() + + def get_rotary_pos_embed(self, rope_sizes): + target_ndim = 3 + + head_dim = self.hidden_size // self.heads_num + rope_dim_list = self.rope_dim_list + if rope_dim_list is None: + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + assert (sum(rope_dim_list) == head_dim), "sum(rope_dim_list) should equal to head_dim of attention layer" + freqs_cos, freqs_sin = get_nd_rotary_pos_embed( + rope_dim_list, + rope_sizes, + theta=self.rope_theta, + use_real=True, + theta_rescale_factor=1, + ) + return freqs_cos, freqs_sin + # x: torch.Tensor, + # t: torch.Tensor, # Should be in range(0, 1000). + # text_states: torch.Tensor = None, + # text_mask: torch.Tensor = None, # Now we don't use it. + # text_states_2: Optional[torch.Tensor] = None, # Text embedding for modulation. + # guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000. + # return_dict: bool = True, + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_attention_mask: torch.Tensor, + mask_strategy=None, + output_features=False, + output_features_stride=8, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = False, + guidance=None, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + if guidance is None: + guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16) + if mask_strategy is None: + mask_strategy = [[None] * self.heads_num + for _ in range(len(self.double_blocks) + len(self.single_blocks))] + img = x = hidden_states + text_mask = encoder_attention_mask + t = timestep + txt = encoder_hidden_states[:, 1:] + text_states_2 = encoder_hidden_states[:, 0, :self.config.text_states_dim_2] + _, _, ot, oh, ow = x.shape # codespell:ignore + tt, th, tw = ( + ot // self.patch_size[0], # codespell:ignore + oh // self.patch_size[1], # codespell:ignore + ow // self.patch_size[2], # codespell:ignore + ) + original_tt = nccl_info.sp_size * tt + freqs_cos, freqs_sin = self.get_rotary_pos_embed((original_tt, th, tw)) + # Prepare modulation vectors. + vec = self.time_in(t) + + # text modulation + vec = vec + self.vector_in(text_states_2) + # guidance modulation + if self.guidance_embed: + if guidance is None: + raise ValueError("Didn't get guidance strength for guidance distilled model.") + + # our timestep_embedding is merged into guidance_in(TimestepEmbedder) + vec = vec + self.guidance_in(guidance) + # Embed image and text. + img = self.img_in(img) + if self.text_projection == "linear": + txt = self.txt_in(txt) + elif self.text_projection == "single_refiner": + txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None) + else: + raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}") + txt_seq_len = txt.shape[1] + img_seq_len = img.shape[1] + + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + # --------------------- Pass through DiT blocks ------------------------ + for index, block in enumerate(self.double_blocks): + double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]] + img, txt = block(*double_block_args) + # Merge txt and img to pass through single stream blocks. + x = torch.cat((img, txt), 1) + if output_features: + features_list = [] + if len(self.single_blocks) > 0: + for index, block in enumerate(self.single_blocks): + single_block_args = [ + x, + vec, + txt_seq_len, + (freqs_cos, freqs_sin), + text_mask, + mask_strategy[index + len(self.double_blocks)], + ] + x = block(*single_block_args) + if output_features and _ % output_features_stride == 0: + features_list.append(x[:, :img_seq_len, ...]) + img = x[:, :img_seq_len, ...] + # ---------------------------- Final layer ------------------------------ + img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + img = self.unpatchify(img, tt, th, tw) + assert not return_dict, "return_dict is not supported." + if output_features: + features_list = torch.stack(features_list, dim=0) + else: + features_list = None + return (img, features_list) + + def unpatchify(self, x, t, h, w): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = self.unpatchify_channels + pt, ph, pw = self.patch_size + assert t * h * w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum("nthwcopq->nctohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + + return imgs + + def params_count(self): + counts = { + "double": + sum([ + sum(p.numel() + for p in block.img_attn_qkv.parameters()) + sum(p.numel() + for p in block.img_attn_proj.parameters()) + + sum(p.numel() for p in block.img_mlp.parameters()) + sum(p.numel() + for p in block.txt_attn_qkv.parameters()) + + sum(p.numel() for p in block.txt_attn_proj.parameters()) + sum(p.numel() + for p in block.txt_mlp.parameters()) + for block in self.double_blocks + ]), + "single": + sum([ + sum(p.numel() for p in block.linear1.parameters()) + sum(p.numel() for p in block.linear2.parameters()) + for block in self.single_blocks + ]), + "total": + sum(p.numel() for p in self.parameters()), + } + counts["attn+mlp"] = counts["double"] + counts["single"] + return counts + + +################################################################################# +# HunyuanVideo Configs # +################################################################################# + +HUNYUAN_VIDEO_CONFIG = { + "HYVideo-T/2": { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + }, + "HYVideo-T/2-cfgdistill": { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + "guidance_embed": True, + }, +} diff --git a/fastvideo/v1/models/hunyuan/modules/modulate_layers.py b/fastvideo/v1/models/hunyuan/modules/modulate_layers.py new file mode 100644 index 000000000..c02ef3187 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/modulate_layers.py @@ -0,0 +1,152 @@ +from typing import Callable + +import torch +import torch.nn as nn + + +class ModulateDiT(nn.Module): + """Modulation layer for DiT.""" + + def __init__( + self, + hidden_size: int, + factor: int, + act_layer: Callable, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.act = act_layer() + self.linear = nn.Linear(hidden_size, factor * hidden_size, bias=True, **factory_kwargs) + # Zero-initialize the modulation + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.act(x)) + + +def modulate(x, shift=None, scale=None): + """modulate by shift and scale + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + return x * (1 + scale.unsqueeze(1)) + elif scale is None: + return x + shift.unsqueeze(1) + else: + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +def apply_gate(x, gate=None, tanh=False): + """AI is creating summary for apply_gate + + Args: + x (torch.Tensor): input tensor. + gate (torch.Tensor, optional): gate tensor. Defaults to None. + tanh (bool, optional): whether to use tanh function. Defaults to False. + + Returns: + torch.Tensor: the output tensor after apply gate. + """ + if gate is None: + return x + if tanh: + return x * gate.unsqueeze(1).tanh() + else: + return x * gate.unsqueeze(1) + + +def ckpt_wrapper(module): + + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward + + +class RMSNorm(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): + """ + 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 + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") diff --git a/fastvideo/v1/models/hunyuan/modules/norm_layers.py b/fastvideo/v1/models/hunyuan/modules/norm_layers.py new file mode 100644 index 000000000..5bc997059 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/norm_layers.py @@ -0,0 +1,78 @@ +import torch +import torch.nn as nn + + +class RMSNorm(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): + """ + 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 + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") diff --git a/fastvideo/v1/models/hunyuan/modules/posemb_layers.py b/fastvideo/v1/models/hunyuan/modules/posemb_layers.py new file mode 100644 index 000000000..314aa126d --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/posemb_layers.py @@ -0,0 +1,290 @@ +from typing import List, Tuple, Union + +import torch + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x, ) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): + """ + Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0, ) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +################################################################################# +# Rotary Positional Embedding Functions # +################################################################################# +# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80 + + +def reshape_for_broadcast( + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + x: torch.Tensor, + head_first=False, +): + """ + Reshape frequency tensor for broadcasting it with another tensor. + + This function reshapes the frequency tensor to have the same shape as the target tensor 'x' + for the purpose of broadcasting the frequency tensor during element-wise operations. + + Notes: + When using FlashMHAModified, head_first should be False. + When using Attention, head_first should be True. + + Args: + freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped. + x (torch.Tensor): Target tensor for broadcasting compatibility. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + torch.Tensor: Reshaped frequency tensor. + + Raises: + AssertionError: If the frequency tensor doesn't match the expected shape. + AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions. + """ + ndim = x.ndim + assert 0 <= 1 < ndim + + if isinstance(freqs_cis, tuple): + # freqs_cis: (cos, sin) in real space + if head_first: + assert freqs_cis[0].shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + else: + assert freqs_cis[0].shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape) + else: + # freqs_cis: values in complex space + if head_first: + assert freqs_cis.shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + else: + assert freqs_cis.shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def rotate_half(x): + x_real, x_imag = (x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)) # [B, S, H, D//2] + return torch.stack([-x_imag, x_real], dim=-1).flatten(3) + + +def apply_rotary_emb( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + head_first: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. + + This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided + frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor + is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are + returned as real tensors. + + Args: + xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] + xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] + freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + + """ + xk_out = None + if isinstance(freqs_cis, tuple): + cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D] + cos, sin = cos.to(xq.device), sin.to(xq.device) + # real * cos - imag * sin + # imag * cos + real * sin + + xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) + xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) + else: + # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex) + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # [B, S, H, D//2] + freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device) # [S, D//2] --> [1, S, 1, D//2] + # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin) + # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq) + xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) # [B, S, H, D//2] + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk) + + return xq_out, xk_out + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + use_real=False, + theta_rescale_factor: Union[float, List[float]] = 1.0, + interpolation_factor: Union[float, List[float]] = 1.0, +): + """ + This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers. + Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real + part and an imaginary part separately. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + + Returns: + pos_embed (torch.Tensor): [HW, D/2] + """ + + grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H] + + if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len( + rope_dim_list), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len( + rope_dim_list), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed( + rope_dim_list[i], + grid[i].reshape(-1), + theta, + use_real=use_real, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + ) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + if use_real: + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + else: + emb = torch.cat(embs, dim=1) # (WHD, D/2) + return emb + + +def get_1d_rotary_pos_embed( + dim: int, + pos: Union[torch.FloatTensor, int], + theta: float = 10000.0, + use_real: bool = False, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Precompute the frequency tensor for complex exponential (cis) with given dimensions. + (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + The returned tensor contains complex values in complex64 data type. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool, optional): If True, return real part and imaginary part separately. + Otherwise, return complex numbers. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + + Returns: + freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2] + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor**(dim / (dim - 2)) + + freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) # [D/2] + # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}" + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + if use_real: + freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [S, D] + return freqs_cos, freqs_sin + else: + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2] + return freqs_cis diff --git a/fastvideo/v1/models/hunyuan/modules/token_refiner.py b/fastvideo/v1/models/hunyuan/modules/token_refiner.py new file mode 100644 index 000000000..1bc127fea --- /dev/null +++ b/fastvideo/v1/models/hunyuan/modules/token_refiner.py @@ -0,0 +1,199 @@ +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + +from .activation_layers import get_activation_layer +from .attenion import attention +from .embed_layers import TextProjection, TimestepEmbedder +from .mlp_layers import MLP +from .modulate_layers import apply_gate +from .norm_layers import get_norm_layer + + +class IndividualTokenRefinerBlock(nn.Module): + + def __init__( + self, + hidden_size, + heads_num, + mlp_width_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + self.self_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.self_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.self_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + act_layer = get_activation_layer(act_type) + self.mlp = MLP( + in_channels=hidden_size, + hidden_channels=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop_rate, + **factory_kwargs, + ) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: torch.Tensor = None, + ): + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + + norm_x = self.norm1(x) + qkv = self.self_attn_qkv(norm_x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + # Apply QK-Norm if needed + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + + # Self-Attention + attn = attention(q, k, v, attn_mask=attn_mask) + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + + # FFN Layer + x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp) + + return x + + +class IndividualTokenRefiner(nn.Module): + + def __init__( + self, + hidden_size, + heads_num, + depth, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.blocks = nn.ModuleList([ + IndividualTokenRefinerBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) for _ in range(depth) + ]) + + def forward( + self, + x: torch.Tensor, + c: torch.LongTensor, + mask: Optional[torch.Tensor] = None, + ): + mask = mask.clone().bool() + # avoid attention weight become NaN + mask[:, 0] = True + for block in self.blocks: + x = block(x, c, mask) + return x + + +class SingleTokenRefiner(nn.Module): + """ + A single token refiner block for llm text embedding refine. + """ + + def __init__( + self, + in_channels, + hidden_size, + heads_num, + depth, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + attn_mode: str = "torch", + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.attn_mode = attn_mode + assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner." + + self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True, **factory_kwargs) + + act_layer = get_activation_layer(act_type) + # Build timestep embedding layer + self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs) + # Build context embedding layer + self.c_embedder = TextProjection(in_channels, hidden_size, act_layer, **factory_kwargs) + + self.individual_token_refiner = IndividualTokenRefiner( + hidden_size=hidden_size, + heads_num=heads_num, + depth=depth, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + + def forward( + self, + x: torch.Tensor, + t: torch.LongTensor, + mask: Optional[torch.LongTensor] = None, + ): + timestep_aware_representations = self.t_embedder(t) + + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.float().unsqueeze(-1) # [b, s1, 1] + context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) + context_aware_representations = self.c_embedder(context_aware_representations.to(x.dtype)) + c = timestep_aware_representations + context_aware_representations + + x = self.input_embedder(x) + x = self.individual_token_refiner(x, c, mask) + return x diff --git a/fastvideo/v1/models/hunyuan/prompt_rewrite.py b/fastvideo/v1/models/hunyuan/prompt_rewrite.py new file mode 100644 index 000000000..556b8aaf2 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/prompt_rewrite.py @@ -0,0 +1,52 @@ +normal_mode_prompt = """Normal mode - Video Recaption Task: + +You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description. + +0. Preserve ALL information, including style words and technical terms. + +1. If the input is in Chinese, translate the entire description to English. + +2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences. + +3. If the input does not include style, lighting, atmosphere, you can make reasonable associations. + +4. Output ALL must be in English. + +Given Input: +input: "{input}" +""" + +master_mode_prompt = """Master mode - Video Recaption Task: + +You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description. + +0. Preserve ALL information, including style words and technical terms. + +1. If the input is in Chinese, translate the entire description to English. + +2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences. + +3. If the input does not include style, lighting, atmosphere, you can make reasonable associations. + +4. Output ALL must be in English. + +Given Input: +input: "{input}" +""" + + +def get_rewrite_prompt(ori_prompt, mode="Normal"): + if mode == "Normal": + prompt = normal_mode_prompt.format(input=ori_prompt) + elif mode == "Master": + prompt = master_mode_prompt.format(input=ori_prompt) + else: + raise Exception("Only supports Normal and Normal", mode) + return prompt + + +ori_prompt = "一只小狗在草地上奔跑。" +normal_prompt = get_rewrite_prompt(ori_prompt, mode="Normal") +master_prompt = get_rewrite_prompt(ori_prompt, mode="Master") + +# Then you can use the normal_prompt or master_prompt to access the hunyuan-large rewrite model to get the final prompt. diff --git a/fastvideo/v1/models/hunyuan/text_encoder/__init__.py b/fastvideo/v1/models/hunyuan/text_encoder/__init__.py new file mode 100644 index 000000000..aadfaa87f --- /dev/null +++ b/fastvideo/v1/models/hunyuan/text_encoder/__init__.py @@ -0,0 +1,323 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn +from transformers import AutoModel, AutoTokenizer, CLIPTextModel, CLIPTokenizer +from transformers.utils import ModelOutput + +from ..constants import PRECISION_TO_TYPE, TEXT_ENCODER_PATH, TOKENIZER_PATH + + +def use_default(value, default): + return value if value is not None else default + + +def load_text_encoder( + text_encoder_type, + text_encoder_precision=None, + text_encoder_path=None, + logger=None, + device=None, +): + if text_encoder_path is None: + text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type] + if logger is not None: + logger.info(f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}") + + if text_encoder_type == "clipL": + text_encoder = CLIPTextModel.from_pretrained(text_encoder_path) + text_encoder.final_layer_norm = text_encoder.text_model.final_layer_norm + elif text_encoder_type == "llm": + text_encoder = AutoModel.from_pretrained(text_encoder_path, low_cpu_mem_usage=True) + text_encoder.final_layer_norm = text_encoder.norm + else: + raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") + # from_pretrained will ensure that the model is in eval mode. + + if text_encoder_precision is not None: + text_encoder = text_encoder.to(dtype=PRECISION_TO_TYPE[text_encoder_precision]) + + text_encoder.requires_grad_(False) + + if logger is not None: + logger.info(f"Text encoder to dtype: {text_encoder.dtype}") + + if device is not None: + text_encoder = text_encoder.to(device) + + return text_encoder, text_encoder_path + + +def load_tokenizer(tokenizer_type, tokenizer_path=None, padding_side="right", logger=None): + if tokenizer_path is None: + tokenizer_path = TOKENIZER_PATH[tokenizer_type] + if logger is not None: + logger.info(f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}") + + if tokenizer_type == "clipL": + tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path, max_length=77) + elif tokenizer_type == "llm": + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, padding_side=padding_side) + else: + raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}") + + return tokenizer, tokenizer_path + + +@dataclass +class TextEncoderModelOutput(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: + hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + text_outputs (`list`, *optional*, returned when `return_texts=True` is passed): + List of decoded texts. + """ + + hidden_state: torch.FloatTensor = None + attention_mask: Optional[torch.LongTensor] = None + hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None + text_outputs: Optional[list] = None + + +class TextEncoder(nn.Module): + + def __init__( + self, + text_encoder_type: str, + max_length: int, + text_encoder_precision: Optional[str] = None, + text_encoder_path: Optional[str] = None, + tokenizer_type: Optional[str] = None, + tokenizer_path: Optional[str] = None, + output_key: Optional[str] = None, + use_attention_mask: bool = True, + input_max_length: Optional[int] = None, + prompt_template: Optional[dict] = None, + prompt_template_video: Optional[dict] = None, + hidden_state_skip_layer: Optional[int] = None, + apply_final_norm: bool = False, + reproduce: bool = False, + logger=None, + device=None, + ): + super().__init__() + self.text_encoder_type = text_encoder_type + self.max_length = max_length + self.precision = text_encoder_precision + self.model_path = text_encoder_path + self.tokenizer_type = (tokenizer_type if tokenizer_type is not None else text_encoder_type) + self.tokenizer_path = (tokenizer_path if tokenizer_path is not None else text_encoder_path) + self.use_attention_mask = use_attention_mask + if prompt_template_video is not None: + assert (use_attention_mask is True), "Attention mask is True required when training videos." + self.input_max_length = (input_max_length if input_max_length is not None else max_length) + self.prompt_template = prompt_template + self.prompt_template_video = prompt_template_video + self.hidden_state_skip_layer = hidden_state_skip_layer + self.apply_final_norm = apply_final_norm + self.reproduce = reproduce + self.logger = logger + + self.use_template = self.prompt_template is not None + if self.use_template: + assert (isinstance(self.prompt_template, dict) and "template" in self.prompt_template + ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}" + assert "{}" in str(self.prompt_template["template"]), ( + "`prompt_template['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template['template']}") + + self.use_video_template = self.prompt_template_video is not None + if self.use_video_template: + if self.prompt_template_video is not None: + assert ( + isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video + ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}" + assert "{}" in str(self.prompt_template_video["template"]), ( + "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template_video['template']}") + + if "t5" in text_encoder_type: + self.output_key = output_key or "last_hidden_state" + elif "clip" in text_encoder_type: + self.output_key = output_key or "pooler_output" + elif "llm" in text_encoder_type or "glm" in text_encoder_type: + self.output_key = output_key or "last_hidden_state" + else: + raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") + + self.model, self.model_path = load_text_encoder( + text_encoder_type=self.text_encoder_type, + text_encoder_precision=self.precision, + text_encoder_path=self.model_path, + logger=self.logger, + device=device, + ) + self.dtype = self.model.dtype + self.device = self.model.device + + self.tokenizer, self.tokenizer_path = load_tokenizer( + tokenizer_type=self.tokenizer_type, + tokenizer_path=self.tokenizer_path, + padding_side="right", + logger=self.logger, + ) + + def __repr__(self): + return f"{self.text_encoder_type} ({self.precision} - {self.model_path})" + + @staticmethod + def apply_text_to_template(text, template, prevent_empty_text=True): + """ + Apply text to template. + + Args: + text (str): Input text. + template (str or list): Template string or list of chat conversation. + prevent_empty_text (bool): If True, we will prevent the user text from being empty + by adding a space. Defaults to True. + """ + if isinstance(template, str): + # Will send string to tokenizer. Used for llm + return template.format(text) + else: + raise TypeError(f"Unsupported template type: {type(template)}") + + def text2tokens(self, text, data_type="image"): + """ + Tokenize the input text. + + Args: + text (str or list): Input text. + """ + tokenize_input_type = "str" + if self.use_template: + if data_type == "image": + prompt_template = self.prompt_template["template"] + elif data_type == "video": + prompt_template = self.prompt_template_video["template"] + else: + raise ValueError(f"Unsupported data type: {data_type}") + if isinstance(text, (list, tuple)): + text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text] + if isinstance(text[0], list): + tokenize_input_type = "list" + elif isinstance(text, str): + text = self.apply_text_to_template(text, prompt_template) + if isinstance(text, list): + tokenize_input_type = "list" + else: + raise TypeError(f"Unsupported text type: {type(text)}") + + kwargs = dict( + truncation=True, + max_length=self.max_length, + padding="max_length", + return_tensors="pt", + ) + if tokenize_input_type == "str": + return self.tokenizer( + text, + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + **kwargs, + ) + elif tokenize_input_type == "list": + return self.tokenizer.apply_chat_template( + text, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **kwargs, + ) + else: + raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}") + + def encode( + self, + batch_encoding, + use_attention_mask=None, + output_hidden_states=False, + do_sample=None, + hidden_state_skip_layer=None, + return_texts=False, + data_type="image", + device=None, + ): + """ + Args: + batch_encoding (dict): Batch encoding from tokenizer. + use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask. + Defaults to None. + output_hidden_states (bool): Whether to output hidden states. If False, return the value of + self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer, + output_hidden_states will be set True. Defaults to False. + do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None. + When self.produce is False, do_sample is set to True by default. + hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer. + If None, self.output_key will be used. Defaults to None. + return_texts (bool): Whether to return the decoded texts. Defaults to False. + """ + device = self.model.device if device is None else device + use_attention_mask = use_default(use_attention_mask, self.use_attention_mask) + hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer) + do_sample = use_default(do_sample, not self.reproduce) + attention_mask = (batch_encoding["attention_mask"].to(device) if use_attention_mask else None) + outputs = self.model( + input_ids=batch_encoding["input_ids"].to(device), + attention_mask=attention_mask, + output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None, + ) + if hidden_state_skip_layer is not None: + last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)] + # Real last hidden state already has layer norm applied. So here we only apply it + # for intermediate layers. + if hidden_state_skip_layer > 0 and self.apply_final_norm: + last_hidden_state = self.model.final_layer_norm(last_hidden_state) + else: + last_hidden_state = outputs[self.output_key] + + # Remove hidden states of instruction tokens, only keep prompt tokens. + if self.use_template: + if data_type == "image": + crop_start = self.prompt_template.get("crop_start", -1) + elif data_type == "video": + crop_start = self.prompt_template_video.get("crop_start", -1) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if crop_start > 0: + last_hidden_state = last_hidden_state[:, crop_start:] + attention_mask = (attention_mask[:, crop_start:] if use_attention_mask else None) + + if output_hidden_states: + return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states) + return TextEncoderModelOutput(last_hidden_state, attention_mask) + + def forward( + self, + text, + use_attention_mask=None, + output_hidden_states=False, + do_sample=False, + hidden_state_skip_layer=None, + return_texts=False, + ): + batch_encoding = self.text2tokens(text) + return self.encode( + batch_encoding, + use_attention_mask=use_attention_mask, + output_hidden_states=output_hidden_states, + do_sample=do_sample, + hidden_state_skip_layer=hidden_state_skip_layer, + return_texts=return_texts, + ) diff --git a/fastvideo/v1/models/hunyuan/utils/__init__.py b/fastvideo/v1/models/hunyuan/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/v1/models/hunyuan/utils/data_utils.py b/fastvideo/v1/models/hunyuan/utils/data_utils.py new file mode 100644 index 000000000..524118179 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/utils/data_utils.py @@ -0,0 +1,14 @@ +import math + + +def align_to(value, alignment): + """align height, width according to alignment + + Args: + value (int): height or width + alignment (int): target alignment factor + + Returns: + int: the aligned value + """ + return int(math.ceil(value / alignment) * alignment) diff --git a/fastvideo/v1/models/hunyuan/utils/file_utils.py b/fastvideo/v1/models/hunyuan/utils/file_utils.py new file mode 100644 index 000000000..82e244f03 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/utils/file_utils.py @@ -0,0 +1,71 @@ +import os +from pathlib import Path + +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange + +CODE_SUFFIXES = { + ".py", # Python codes + ".sh", # Shell scripts + ".yaml", + ".yml", # Configuration files +} + + +def safe_dir(path): + """ + Create a directory (or the parent directory of a file) if it does not exist. + + Args: + path (str or Path): Path to the directory. + + Returns: + path (Path): Path object of the directory. + """ + path = Path(path) + path.mkdir(exist_ok=True, parents=True) + return path + + +def safe_file(path): + """ + Create the parent directory of a file if it does not exist. + + Args: + path (str or Path): Path to the file. + + Returns: + path (Path): Path object of the file. + """ + path = Path(path) + path.parent.mkdir(exist_ok=True, parents=True) + return path + + +def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=1, fps=24): + """save videos by video tensor + copy from https://github.com/guoyww/AnimateDiff/blob/e92bd5671ba62c0d774a32951453e328018b7c5b/animatediff/utils/util.py#L61 + + Args: + videos (torch.Tensor): video tensor predicted by the model + path (str): path to save video + rescale (bool, optional): rescale the video tensor from [-1, 1] to . Defaults to False. + n_rows (int, optional): Defaults to 1. + fps (int, optional): video save fps. Defaults to 8. + """ + videos = rearrange(videos, "b c t h w -> t b c h w") + outputs = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=n_rows) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = torch.clamp(x, 0, 1) + x = (x * 255).numpy().astype(np.uint8) + outputs.append(x) + + os.makedirs(os.path.dirname(path), exist_ok=True) + imageio.mimsave(path, outputs, fps=fps) diff --git a/fastvideo/v1/models/hunyuan/utils/helpers.py b/fastvideo/v1/models/hunyuan/utils/helpers.py new file mode 100644 index 000000000..7b5b8731f --- /dev/null +++ b/fastvideo/v1/models/hunyuan/utils/helpers.py @@ -0,0 +1,41 @@ +import collections.abc +from itertools import repeat + + +def _ntuple(n): + + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + x = tuple(x) + if len(x) == 1: + x = tuple(repeat(x[0], n)) + return x + return tuple(repeat(x, n)) + + return parse + + +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) + + +def as_tuple(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + if x is None or isinstance(x, (int, float, str)): + return (x, ) + else: + raise ValueError(f"Unknown type {type(x)}") + + +def as_list_of_2tuple(x): + x = as_tuple(x) + if len(x) == 1: + x = (x[0], x[0]) + assert len(x) % 2 == 0, f"Expect even length, got {len(x)}." + lst = [] + for i in range(0, len(x), 2): + lst.append((x[i], x[i + 1])) + return lst diff --git a/fastvideo/v1/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py b/fastvideo/v1/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py new file mode 100644 index 000000000..1a6f46c15 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py @@ -0,0 +1,41 @@ +import argparse + +import torch +from transformers import AutoProcessor, LlavaForConditionalGeneration + + +def preprocess_text_encoder_tokenizer(args): + + processor = AutoProcessor.from_pretrained(args.input_dir) + model = LlavaForConditionalGeneration.from_pretrained( + args.input_dir, + torch_dtype=torch.float16, + low_cpu_mem_usage=True, + ).to(0) + + model.language_model.save_pretrained(f"{args.output_dir}") + processor.tokenizer.save_pretrained(f"{args.output_dir}") + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + type=str, + required=True, + help="The path to the llava-llama-3-8b-v1_1-transformers.", + ) + parser.add_argument( + "--output_dir", + type=str, + default="", + help="The output path of the llava-llama-3-8b-text-encoder-tokenizer." + "if '', the parent dir of output will be the same as input dir.", + ) + args = parser.parse_args() + + if len(args.output_dir) == 0: + args.output_dir = "/".join(args.input_dir.split("/")[:-1]) + + preprocess_text_encoder_tokenizer(args) diff --git a/fastvideo/v1/models/hunyuan/vae/__init__.py b/fastvideo/v1/models/hunyuan/vae/__init__.py new file mode 100644 index 000000000..761c7cc1f --- /dev/null +++ b/fastvideo/v1/models/hunyuan/vae/__init__.py @@ -0,0 +1,64 @@ +from pathlib import Path + +import torch + +from ..constants import PRECISION_TO_TYPE, VAE_PATH +from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D + + +def load_vae( + vae_type: str = "884-16c-hy", + vae_precision: str = None, + sample_size: tuple = None, + vae_path: str = None, + logger=None, + device=None, +): + """the function to load the 3D VAE model + + Args: + vae_type (str): the type of the 3D VAE model. Defaults to "884-16c-hy". + vae_precision (str, optional): the precision to load vae. Defaults to None. + sample_size (tuple, optional): the tiling size. Defaults to None. + vae_path (str, optional): the path to vae. Defaults to None. + logger (_type_, optional): logger. Defaults to None. + device (_type_, optional): device to load vae. Defaults to None. + """ + if vae_path is None: + vae_path = VAE_PATH[vae_type] + + if logger is not None: + logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}") + config = AutoencoderKLCausal3D.load_config(vae_path) + if sample_size: + vae = AutoencoderKLCausal3D.from_config(config, sample_size=sample_size) + else: + vae = AutoencoderKLCausal3D.from_config(config) + + vae_ckpt = Path(vae_path) / "pytorch_model.pt" + assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}" + + ckpt = torch.load(vae_ckpt, map_location=vae.device) + if "state_dict" in ckpt: + ckpt = ckpt["state_dict"] + if any(k.startswith("vae.") for k in ckpt.keys()): + ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")} + vae.load_state_dict(ckpt) + + spatial_compression_ratio = vae.config.spatial_compression_ratio + time_compression_ratio = vae.config.time_compression_ratio + + if vae_precision is not None: + vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision]) + + vae.requires_grad_(False) + + if logger is not None: + logger.info(f"VAE to dtype: {vae.dtype}") + + if device is not None: + vae = vae.to(device) + + vae.eval() + + return vae, vae_path, spatial_compression_ratio, time_compression_ratio diff --git a/fastvideo/v1/models/hunyuan/vae/autoencoder_kl_causal_3d.py b/fastvideo/v1/models/hunyuan/vae/autoencoder_kl_causal_3d.py new file mode 100644 index 000000000..5a63783cd --- /dev/null +++ b/fastvideo/v1/models/hunyuan/vae/autoencoder_kl_causal_3d.py @@ -0,0 +1,776 @@ +# Copyright 2024 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. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== +from dataclasses import dataclass +from math import prod +from typing import Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.distributed as dist +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from fastvideo.v1.utils.parallel_states import nccl_info +from fastvideo.v1.distributed.parallel_state import get_sp_group +from fastvideo.v1.logger import init_logger + +logger = init_logger(__name__) + +try: + # This diffusers is modified and packed in the mirror. + from diffusers.loaders import FromOriginalVAEMixin +except ImportError: + # Use this to be compatible with the original diffusers. + from diffusers.loaders.single_file_model import ( + FromOriginalModelMixin as FromOriginalVAEMixin, ) + +from diffusers.models.attention_processor import (ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, Attention, + AttentionProcessor, AttnAddedKVProcessor, AttnProcessor) +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils.accelerate_utils import apply_forward_hook + +from .vae import BaseOutput, DecoderCausal3D, DecoderOutput, DiagonalGaussianDistribution, EncoderCausal3D + + +@dataclass +class DecoderOutput2(BaseOutput): + sample: torch.FloatTensor + posterior: Optional[DiagonalGaussianDistribution] = None + + +class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): + r""" + A VAE model with KL loss for encoding images/videos into latents and decoding latent representations into images/videos. + + 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 + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str] = ("DownEncoderBlockCausal3D", ), + up_block_types: Tuple[str] = ("UpDecoderBlockCausal3D", ), + block_out_channels: Tuple[int] = (64, ), + layers_per_block: int = 1, + act_fn: str = "silu", + latent_channels: int = 4, + norm_num_groups: int = 32, + sample_size: int = 32, + sample_tsize: int = 64, + scaling_factor: float = 0.18215, + force_upcast: float = True, + spatial_compression_ratio: int = 8, + time_compression_ratio: int = 4, + mid_block_add_attention: bool = True, + ): + super().__init__() + + self.time_compression_ratio = time_compression_ratio + + self.encoder = EncoderCausal3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + double_z=True, + time_compression_ratio=time_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + + self.decoder = DecoderCausal3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + time_compression_ratio=time_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + + self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) + self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) + + self.use_slicing = False + self.use_spatial_tiling = False + self.use_temporal_tiling = False + self.use_parallel = False + + # only relevant if vae tiling is enabled + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // time_compression_ratio + + self.tile_sample_min_size = self.config.sample_size + sample_size = (self.config.sample_size[0] if isinstance(self.config.sample_size, + (list, tuple)) else self.config.sample_size) + self.tile_latent_min_size = int(sample_size / (2**(len(self.config.block_out_channels) - 1))) + self.tile_overlap_factor = 0.25 + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (EncoderCausal3D, DecoderCausal3D)): + module.gradient_checkpointing = value + + def enable_temporal_tiling(self, use_tiling: bool = True): + self.use_temporal_tiling = use_tiling + + def disable_temporal_tiling(self): + self.enable_temporal_tiling(False) + + def enable_spatial_tiling(self, use_tiling: bool = True): + self.use_spatial_tiling = use_tiling + + def disable_spatial_tiling(self): + self.enable_spatial_tiling(False) + + def enable_tiling(self, use_tiling: bool = True): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger videos. + """ + self.enable_spatial_tiling(use_tiling) + self.enable_temporal_tiling(use_tiling) + + def disable_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.disable_spatial_tiling() + self.disable_temporal_tiling() + + def enable_parallel(self): + r""" + Enable sequence parallelism for the model. This will allow the vae to decode (with tiling) in parallel. + """ + self.use_parallel = True + + def enable_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.use_slicing = True + + def disable_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_slicing = False + + @property + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors( + name: str, + module: torch.nn.Module, + processors: Dict[str, AttentionProcessor], + ): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor( + self, + processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], + _remove_lora=False, + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes.") + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor, _remove_lora=_remove_lora) + else: + module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnAddedKVProcessor() + elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor, _remove_lora=True) + + @apply_forward_hook + def encode(self, + x: torch.FloatTensor, + return_dict: bool = True) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + """ + Encode a batch of images/videos into latents. + + Args: + x (`torch.FloatTensor`): Input batch of images/videos. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded images/videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + assert len(x.shape) == 5, "The input tensor should have 5 dimensions." + + if self.use_temporal_tiling and x.shape[2] > self.tile_sample_min_tsize: + return self.temporal_tiled_encode(x, return_dict=return_dict) + + if self.use_spatial_tiling and (x.shape[-1] > self.tile_sample_min_size + or x.shape[-2] > self.tile_sample_min_size): + return self.spatial_tiled_encode(x, return_dict=return_dict) + + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self.encoder(x) + + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + assert len(z.shape) == 5, "The input tensor should have 5 dimensions." + + logger.info(f"use_parallel: {self.use_parallel}") + if self.use_parallel: + return self.parallel_tiled_decode(z, return_dict=return_dict) + + if self.use_temporal_tiling and z.shape[2] > self.tile_latent_min_tsize: + return self.temporal_tiled_decode(z, return_dict=return_dict) + + if self.use_spatial_tiling and (z.shape[-1] > self.tile_latent_min_size + or z.shape[-2] > self.tile_latent_min_size): + return self.spatial_tiled_decode(z, return_dict=return_dict) + + z = self.post_quant_conv(z) + dec = self.decoder(z) + + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + @apply_forward_hook + def decode(self, + z: torch.FloatTensor, + return_dict: bool = True, + generator=None) -> Union[DecoderOutput, torch.FloatTensor]: + """ + Decode a batch of images/videos. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + + if not return_dict: + return (decoded, ) + + return DecoderOutput(sample=decoded) + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, + y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, + x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, + x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, + x, :, :] * (x / blend_extent) + return b + + def spatial_tiled_encode( + self, + x: torch.FloatTensor, + return_dict: bool = True, + return_moments: bool = False, + ) -> AutoencoderKLOutput: + r"""Encode a batch of images/videos using a tiled encoder. + + When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several + steps. This is useful to keep memory use constant regardless of image/videos size. The end result of tiled encoding is + different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the + tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the + output, but they should be much less noticeable. + + Args: + x (`torch.FloatTensor`): Input batch of images/videos. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`: + If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + # Split video into tiles and encode them separately. + rows = [] + for i in range(0, x.shape[-2], overlap_size): + row = [] + for j in range(0, x.shape[-1], overlap_size): + tile = x[:, :, :, i:i + self.tile_sample_min_size, j:j + self.tile_sample_min_size, ] + tile = self.encoder(tile) + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + moments = torch.cat(result_rows, dim=-2) + if return_moments: + return moments + + posterior = DiagonalGaussianDistribution(moments) + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def spatial_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + r""" + Decode a batch of images/videos using a tiled decoder. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, z.shape[-2], overlap_size): + row = [] + for j in range(0, z.shape[-1], overlap_size): + tile = z[:, :, :, i:i + self.tile_latent_min_size, j:j + self.tile_latent_min_size, ] + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2) + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + def temporal_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput: + + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_latent_min_tsize - blend_extent + + # Split the video into tiles and encode them separately. + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i:i + self.tile_sample_min_tsize + 1, :, :] + if self.use_spatial_tiling and (tile.shape[-1] > self.tile_sample_min_size + or tile.shape[-2] > self.tile_sample_min_size): + tile = self.spatial_tiled_encode(tile, return_moments=True) + else: + tile = self.encoder(tile) + tile = self.quant_conv(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit + 1, :, :]) + + moments = torch.cat(result_row, dim=2) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def temporal_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + # Split z into overlapping tiles and decode them separately. + + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_sample_min_tsize - blend_extent + + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i:i + self.tile_latent_min_tsize + 1, :, :] + if self.use_spatial_tiling and (tile.shape[-1] > self.tile_latent_min_size + or tile.shape[-2] > self.tile_latent_min_size): + decoded = self.spatial_tiled_decode(tile, return_dict=True).sample + else: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit + 1, :, :]) + + dec = torch.cat(result_row, dim=2) + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + def _parallel_data_generator(self, gathered_results, gathered_dim_metadata): + global_idx = 0 + for i, per_rank_metadata in enumerate(gathered_dim_metadata): + _start_shape = 0 + for shape in per_rank_metadata: + mul_shape = prod(shape) + yield (gathered_results[i, _start_shape:_start_shape + mul_shape].reshape(shape), global_idx) + _start_shape += mul_shape + global_idx += 1 + + def parallel_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + """ + Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs + """ + # world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + + sp_group = get_sp_group() + world_size = sp_group.world_size + rank = sp_group.rank_in_group + + B, C, T, H, W = z.shape + + # Calculate parameters + t_overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) + t_blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_sample_min_tsize - t_blend_extent + + s_overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + s_blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + s_row_limit = self.tile_sample_min_size - s_blend_extent + + # Calculate tile dimensions + num_t_tiles = (T + t_overlap_size - 1) // t_overlap_size + num_h_tiles = (H + s_overlap_size - 1) // s_overlap_size + num_w_tiles = (W + s_overlap_size - 1) // s_overlap_size + total_spatial_tiles = num_h_tiles * num_w_tiles + total_tiles = num_t_tiles * total_spatial_tiles + + # Calculate tiles per rank and padding + tiles_per_rank = (total_tiles + world_size - 1) // world_size + start_tile_idx = rank * tiles_per_rank + end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles) + + local_results = [] + local_dim_metadata = [] + # Process assigned tiles + for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)): + # Convert flat index to 3D indices + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + + # Calculate positions + t_start = t_idx * t_overlap_size + h_start = h_idx * s_overlap_size + w_start = w_idx * s_overlap_size + + # Extract and process tile + tile = z[:, :, t_start:t_start + self.tile_latent_min_tsize + 1, + h_start:h_start + self.tile_latent_min_size, w_start:w_start + self.tile_latent_min_size] + + # Process tile + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + + if t_start > 0: + decoded = decoded[:, :, 1:, :, :] + + # Store metadata + shape = decoded.shape + # Store decoded data (flattened) + decoded_flat = decoded.reshape(-1) + local_results.append(decoded_flat) + local_dim_metadata.append(shape) + + results = torch.cat(local_results, dim=0).contiguous() + del local_results + torch.cuda.empty_cache() + # first gather size to pad the results + local_size = torch.tensor([results.size(0)], device=results.device, dtype=torch.int64) + all_sizes = [torch.zeros(1, device=results.device, dtype=torch.int64) for _ in range(world_size)] + dist.all_gather(all_sizes, local_size) + max_size = max(size.item() for size in all_sizes) + padded_results = torch.zeros(max_size, device=results.device) + padded_results[:results.size(0)] = results + del results + torch.cuda.empty_cache() + # Gather all results + gathered_dim_metadata = [None] * world_size + gathered_results = torch.zeros_like(padded_results).repeat( + world_size, *[1] * len(padded_results.shape)).contiguous( + ) # use contiguous to make sure it won't copy data in the following operations + # gathered_results = sp_group.all_gather(padded_results) + # gathered_dim_metadata = sp_group.broadcast_object_list(gathered_dim_metadata, local_dim_metadata) + dist.all_gather_into_tensor(gathered_results, padded_results) + dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) + # Process gathered results + data = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] for _ in range(num_t_tiles)] + for current_data, global_idx in self._parallel_data_generator(gathered_results, gathered_dim_metadata): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + data[t_idx][h_idx][w_idx] = current_data + # Merge results + result_slices = [] + last_slice_data = None + for i, tem_data in enumerate(data): + slice_data = self._merge_spatial_tiles(tem_data, s_blend_extent, s_row_limit) + if i > 0: + slice_data = self.blend_t(last_slice_data, slice_data, t_blend_extent) + result_slices.append(slice_data[:, :, :t_limit, :, :]) + else: + result_slices.append(slice_data[:, :, :t_limit + 1, :, :]) + last_slice_data = slice_data + dec = torch.cat(result_slices, dim=2) + + if not return_dict: + return (dec, ) + return DecoderOutput(sample=dec) + + def _merge_spatial_tiles(self, spatial_rows, blend_extent, row_limit): + """Helper function to merge spatial tiles with blending""" + result_rows = [] + for i, row in enumerate(spatial_rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(spatial_rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def forward( + self, + sample: torch.FloatTensor, + sample_posterior: bool = False, + return_dict: bool = True, + return_posterior: bool = False, + generator: Optional[torch.Generator] = None, + ) -> Union[DecoderOutput2, torch.FloatTensor]: + r""" + Args: + sample (`torch.FloatTensor`): 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).sample + + if not return_dict: + if return_posterior: + return (dec, posterior) + else: + return (dec, ) + if return_posterior: + return DecoderOutput2(sample=dec, posterior=posterior) + else: + return DecoderOutput2(sample=dec) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections + def fuse_qkv_projections(self): + """ + Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, + key, value) are fused. For cross-attention modules, key and value projection matrices are fused. + + + + This API is 🧪 experimental. + + + """ + self.original_attn_processors = None + + for _, attn_processor in self.attn_processors.items(): + if "Added" in str(attn_processor.__class__.__name__): + raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") + + self.original_attn_processors = self.attn_processors + + for module in self.modules(): + if isinstance(module, Attention): + module.fuse_projections(fuse=True) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections + def unfuse_qkv_projections(self): + """Disables the fused QKV projection if enabled. + + + + This API is 🧪 experimental. + + + + """ + if self.original_attn_processors is not None: + self.set_attn_processor(self.original_attn_processors) diff --git a/fastvideo/v1/models/hunyuan/vae/unet_causal_3d_blocks.py b/fastvideo/v1/models/hunyuan/vae/unet_causal_3d_blocks.py new file mode 100644 index 000000000..3c466f55f --- /dev/null +++ b/fastvideo/v1/models/hunyuan/vae/unet_causal_3d_blocks.py @@ -0,0 +1,760 @@ +# Copyright 2024 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. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== + +from typing import Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import Attention, SpatialNorm +from diffusers.models.normalization import AdaGroupNorm, RMSNorm +from diffusers.utils import logging +from einops import rearrange +from torch import nn + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None): + 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 + + +class CausalConv3d(nn.Module): + """ + Implements a causal 3D convolution layer where each position only depends on previous timesteps and current spatial locations. + This maintains temporal causality in video generation tasks. + """ + + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]] = 1, + dilation: Union[int, Tuple[int, int, int]] = 1, + pad_mode="replicate", + **kwargs, + ): + super().__init__() + + self.pad_mode = pad_mode + padding = ( + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size - 1, + 0, + ) # W, H, T + self.time_causal_padding = padding + + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) + return self.conv(x) + + +class UpsampleCausal3D(nn.Module): + """ + A 3D upsampling layer with an optional convolution. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + use_conv_transpose: bool = False, + out_channels: Optional[int] = None, + name: str = "conv", + kernel_size: Optional[int] = None, + padding=1, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + interpolate=True, + upsample_factor=(2, 2, 2), + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + self.interpolate = interpolate + self.upsample_factor = upsample_factor + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + conv = None + if use_conv_transpose: + raise NotImplementedError + elif use_conv: + if kernel_size is None: + kernel_size = 3 + conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, bias=bias) + + if name == "conv": + self.conv = conv + else: + self.Conv2d_0 = conv + + def forward( + self, + hidden_states: torch.FloatTensor, + output_size: Optional[int] = None, + scale: float = 1.0, + ) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + raise NotImplementedError + + if self.use_conv_transpose: + return self.conv(hidden_states) + + # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 + dtype = hidden_states.dtype + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.float32) + + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + hidden_states = hidden_states.contiguous() + + # if `output_size` is passed we force the interpolation output + # size and do not make use of `scale_factor=2` + if self.interpolate: + B, C, T, H, W = hidden_states.shape + first_h, other_h = hidden_states.split((1, T - 1), dim=2) + if output_size is None: + if T > 1: + other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest") + + first_h = first_h.squeeze(2) + first_h = F.interpolate(first_h, scale_factor=self.upsample_factor[1:], mode="nearest") + first_h = first_h.unsqueeze(2) + else: + raise NotImplementedError + + if T > 1: + hidden_states = torch.cat((first_h, other_h), dim=2) + else: + hidden_states = first_h + + # If the input is bfloat16, we cast back to bfloat16 + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(dtype) + + if self.use_conv: + if self.name == "conv": + hidden_states = self.conv(hidden_states) + else: + hidden_states = self.Conv2d_0(hidden_states) + + return hidden_states + + +class DownsampleCausal3D(nn.Module): + """ + A 3D downsampling layer with an optional convolution. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + padding: int = 1, + name: str = "conv", + kernel_size=3, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + stride=2, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = stride + self.name = name + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + if use_conv: + conv = CausalConv3d( + self.channels, + self.out_channels, + kernel_size=kernel_size, + stride=stride, + bias=bias, + ) + else: + raise NotImplementedError + + if name == "conv": + self.Conv2d_0 = conv + self.conv = conv + elif name == "Conv2d_0": + self.conv = conv + else: + self.conv = conv + + def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + assert hidden_states.shape[1] == self.channels + + hidden_states = self.conv(hidden_states) + + return hidden_states + + +class ResnetBlockCausal3D(nn.Module): + r""" + A Resnet block. + """ + + def __init__( + self, + *, + in_channels: int, + out_channels: Optional[int] = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + groups: int = 32, + groups_out: Optional[int] = None, + pre_norm: bool = True, + eps: float = 1e-6, + non_linearity: str = "swish", + skip_time_act: bool = False, + # default, scale_shift, ada_group, spatial + time_embedding_norm: str = "default", + kernel: Optional[torch.FloatTensor] = None, + output_scale_factor: float = 1.0, + use_in_shortcut: Optional[bool] = None, + up: bool = False, + down: bool = False, + conv_shortcut_bias: bool = True, + conv_3d_out_channels: Optional[int] = None, + ): + super().__init__() + self.pre_norm = pre_norm + self.pre_norm = True + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + self.up = up + self.down = down + self.output_scale_factor = output_scale_factor + self.time_embedding_norm = time_embedding_norm + self.skip_time_act = skip_time_act + + linear_cls = nn.Linear + + if groups_out is None: + groups_out = groups + + if self.time_embedding_norm == "ada_group": + self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps) + elif self.time_embedding_norm == "spatial": + self.norm1 = SpatialNorm(in_channels, temb_channels) + else: + self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + + self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, stride=1) + + if temb_channels is not None: + if self.time_embedding_norm == "default": + self.time_emb_proj = linear_cls(temb_channels, out_channels) + elif self.time_embedding_norm == "scale_shift": + self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels) + elif (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"): + self.time_emb_proj = None + else: + raise ValueError(f"Unknown time_embedding_norm : {self.time_embedding_norm} ") + else: + self.time_emb_proj = None + + if self.time_embedding_norm == "ada_group": + self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps) + elif self.time_embedding_norm == "spatial": + self.norm2 = SpatialNorm(out_channels, temb_channels) + else: + self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) + + self.dropout = torch.nn.Dropout(dropout) + conv_3d_out_channels = conv_3d_out_channels or out_channels + self.conv2 = CausalConv3d(out_channels, conv_3d_out_channels, kernel_size=3, stride=1) + + self.nonlinearity = get_activation(non_linearity) + + self.upsample = self.downsample = None + if self.up: + self.upsample = UpsampleCausal3D(in_channels, use_conv=False) + elif self.down: + self.downsample = DownsampleCausal3D(in_channels, use_conv=False, name="op") + + self.use_in_shortcut = (self.in_channels != conv_3d_out_channels + if use_in_shortcut is None else use_in_shortcut) + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = CausalConv3d( + in_channels, + conv_3d_out_channels, + kernel_size=1, + stride=1, + bias=conv_shortcut_bias, + ) + + def forward( + self, + input_tensor: torch.FloatTensor, + temb: torch.FloatTensor, + scale: float = 1.0, + ) -> torch.FloatTensor: + hidden_states = input_tensor + + if (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"): + hidden_states = self.norm1(hidden_states, temb) + else: + hidden_states = self.norm1(hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + if self.upsample is not None: + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + input_tensor = input_tensor.contiguous() + hidden_states = hidden_states.contiguous() + input_tensor = self.upsample(input_tensor, scale=scale) + hidden_states = self.upsample(hidden_states, scale=scale) + elif self.downsample is not None: + input_tensor = self.downsample(input_tensor, scale=scale) + hidden_states = self.downsample(hidden_states, scale=scale) + + hidden_states = self.conv1(hidden_states) + + if self.time_emb_proj is not None: + if not self.skip_time_act: + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb, scale)[:, :, None, None] + + if temb is not None and self.time_embedding_norm == "default": + hidden_states = hidden_states + temb + + if (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"): + hidden_states = self.norm2(hidden_states, temb) + else: + hidden_states = self.norm2(hidden_states) + + if temb is not None and self.time_embedding_norm == "scale_shift": + scale, shift = torch.chunk(temb, 2, dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + return output_tensor + + +def get_down_block3d( + down_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + temb_channels: int, + add_downsample: bool, + downsample_stride: int, + resnet_eps: float, + resnet_act_fn: str, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + downsample_padding: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + downsample_type: Optional[str] = None, + dropout: float = 0.0, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + down_block_type = (down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type) + if down_block_type == "DownEncoderBlockCausal3D": + return DownEncoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + downsample_stride=downsample_stride, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + raise ValueError(f"{down_block_type} does not exist.") + + +def get_up_block3d( + up_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + add_upsample: bool, + upsample_scale_factor: Tuple, + resnet_eps: float, + resnet_act_fn: str, + resolution_idx: Optional[int] = None, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + upsample_type: Optional[str] = None, + dropout: float = 0.0, +) -> nn.Module: + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + up_block_type = (up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type) + if up_block_type == "UpDecoderBlockCausal3D": + return UpDecoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + resolution_idx=resolution_idx, + dropout=dropout, + add_upsample=add_upsample, + upsample_scale_factor=upsample_scale_factor, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + ) + raise ValueError(f"{up_block_type} does not exist.") + + +class UNetMidBlockCausal3D(nn.Module): + """ + A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks. + """ + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + attn_groups: Optional[int] = None, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + ): + super().__init__() + resnet_groups = (resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)) + self.add_attention = add_attention + + if attn_groups is None: + attn_groups = (resnet_groups if resnet_time_scale_shift == "default" else None) + + # there is always at least one resnet + resnets = [ + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=attn_groups, + spatial_norm_dim=(temb_channels if resnet_time_scale_shift == "spatial" else None), + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + else: + attentions.append(None) + + resnets.append( + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor: + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + B, C, T, H, W = hidden_states.shape + hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c") + attention_mask = prepare_causal_attention_mask(T, + H * W, + hidden_states.dtype, + hidden_states.device, + batch_size=B) + hidden_states = attn(hidden_states, temb=temb, attention_mask=attention_mask) + hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class DownEncoderBlockCausal3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_stride: int = 2, + downsample_padding: int = 1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + DownsampleCausal3D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + stride=downsample_stride, + ) + ]) + else: + self.downsamplers = None + + def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None, scale=scale) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale) + + return hidden_states + + +class UpDecoderBlockCausal3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + upsample_scale_factor=(2, 2, 2), + temb_channels: Optional[int] = None, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlockCausal3D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + UpsampleCausal3D( + out_channels, + use_conv=True, + out_channels=out_channels, + upsample_factor=upsample_scale_factor, + ) + ]) + else: + self.upsamplers = None + + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states diff --git a/fastvideo/v1/models/hunyuan/vae/vae.py b/fastvideo/v1/models/hunyuan/vae/vae.py new file mode 100644 index 000000000..1d0f75e51 --- /dev/null +++ b/fastvideo/v1/models/hunyuan/vae/vae.py @@ -0,0 +1,342 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.nn as nn +from diffusers.models.attention_processor import SpatialNorm +from diffusers.utils import BaseOutput, is_torch_version +from diffusers.utils.torch_utils import randn_tensor + +from .unet_causal_3d_blocks import CausalConv3d, UNetMidBlockCausal3D, get_down_block3d, get_up_block3d + + +@dataclass +class DecoderOutput(BaseOutput): + r""" + Output of decoding method. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The decoded output sample from the last layer of the model. + """ + + sample: torch.FloatTensor + + +class EncoderCausal3D(nn.Module): + r""" + The `EncoderCausal3D` layer of a variational autoencoder that encodes its input into a latent representation. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D", ), + block_out_channels: Tuple[int, ...] = (64, ), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i >= (len(block_out_channels) - 1 - num_time_downsample_layers) + and not is_final_block) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.") + + downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) + downsample_stride_T = (2, ) if add_time_downsample else (1, ) + downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) + down_block = get_down_block3d( + down_block_type, + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=bool(add_spatial_downsample or add_time_downsample), + downsample_stride=downsample_stride, + resnet_eps=1e-6, + downsample_padding=0, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=None, + ) + self.down_blocks.append(down_block) + + # mid + self.mid_block = UNetMidBlockCausal3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention, + ) + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3) + + def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor: + r"""The forward method of the `EncoderCausal3D` class.""" + assert len(sample.shape) == 5, "The input tensor should have 5 dimensions" + + sample = self.conv_in(sample) + + # down + for down_block in self.down_blocks: + sample = down_block(sample) + + # middle + sample = self.mid_block(sample) + + # post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class DecoderCausal3D(nn.Module): + r""" + The `DecoderCausal3D` layer of a variational autoencoder that decodes its latent representation into an output sample. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D", ), + block_out_channels: Tuple[int, ...] = (64, ), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + norm_type: str = "group", # group, spatial + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = CausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1) + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + temb_channels = in_channels if norm_type == "spatial" else None + + # mid + self.mid_block = UNetMidBlockCausal3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default" if norm_type == "group" else norm_type, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=temb_channels, + add_attention=mid_block_add_attention, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_upsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_upsample = bool(i < num_spatial_upsample_layers) + add_time_upsample = bool(i >= len(block_out_channels) - 1 - num_time_upsample_layers + and not is_final_block) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.") + + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) + upsample_scale_factor_T = (2, ) if add_time_upsample else (1, ) + upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW) + up_block = get_up_block3d( + up_block_type, + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + prev_output_channel=None, + add_upsample=bool(add_spatial_upsample or add_time_upsample), + upsample_scale_factor=upsample_scale_factor, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=temb_channels, + resnet_time_scale_shift=norm_type, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_type == "spatial": + self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) + else: + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward( + self, + sample: torch.FloatTensor, + latent_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + r"""The forward method of the `DecoderCausal3D` class.""" + assert len(sample.shape) == 5, "The input tensor should have 5 dimensions." + + sample = self.conv_in(sample) + + upscale_dtype = next(iter(self.up_blocks.parameters())).dtype + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + # middle + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.mid_block), + sample, + latent_embeds, + use_reentrant=False, + ) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(up_block), + sample, + latent_embeds, + use_reentrant=False, + ) + else: + # middle + sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample, latent_embeds) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds) + else: + # middle + sample = self.mid_block(sample, latent_embeds) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = up_block(sample, latent_embeds) + + # post-process + if latent_embeds is None: + sample = self.conv_norm_out(sample) + else: + sample = self.conv_norm_out(sample, latent_embeds) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class DiagonalGaussianDistribution(object): + + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + if parameters.ndim == 3: + dim = 2 # (B, L, C) + elif parameters.ndim == 5 or parameters.ndim == 4: + dim = 1 # (B, C, T, H ,W) / (B, C, H, W) + else: + raise NotImplementedError + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean, + device=self.parameters.device, + dtype=self.parameters.dtype) + + def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = randn_tensor( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + else: + reduce_dim = list(range(1, self.mean.ndim)) + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=reduce_dim, + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar + + other.logvar, + dim=reduce_dim, + ) + + def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + + def mode(self) -> torch.Tensor: + return self.mean diff --git a/fastvideo/models/loader.py b/fastvideo/v1/models/loader.py similarity index 94% rename from fastvideo/models/loader.py rename to fastvideo/v1/models/loader.py index ee3be6425..f74ef62d0 100644 --- a/fastvideo/models/loader.py +++ b/fastvideo/v1/models/loader.py @@ -1,17 +1,17 @@ from abc import ABC, abstractmethod import torch -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger import os -from fastvideo.distributed.parallel_state import initialize_sequence_parallel_group +from fastvideo.v1.distributed.parallel_state import initialize_sequence_parallel_group import glob -from fastvideo.loader.fsdp_load import load_fsdp_model -from fastvideo.loader.loader import get_model_loader +from fastvideo.v1.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.loader.loader import get_model_loader from transformers import PretrainedConfig, AutoTokenizer -from fastvideo.models.hf_transformer_utils import get_hf_config, get_diffusers_config -from fastvideo.models import get_scheduler -from fastvideo.models.registry import ModelRegistry +from fastvideo.v1.models.hf_transformer_utils import get_hf_config, get_diffusers_config +from fastvideo.v1.models import get_scheduler +from fastvideo.v1.models.registry import ModelRegistry from safetensors.torch import load_file as safetensors_load_file import json @@ -100,8 +100,8 @@ class TextEncoderLoader(ComponentLoader): return model def load_v0(self, model_path: str, architecture: str, inference_args: InferenceArgs): - from fastvideo.pipelines.hunyuan.constants import PROMPT_TEMPLATE - from fastvideo.models.hunyuan.text_encoder import TextEncoder + from fastvideo.v1.pipelines.hunyuan.constants import PROMPT_TEMPLATE + from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder # Text encoder if inference_args.prompt_template_video is not None: crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0) @@ -156,7 +156,7 @@ class VAELoader(ComponentLoader): def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): """Load the VAE based on the model path, architecture, and inference args.""" # TODO(will): move this to a constants file - from fastvideo.pipelines.hunyuan.constants import PRECISION_TO_TYPE + from fastvideo.v1.pipelines.hunyuan.constants import PRECISION_TO_TYPE config = get_diffusers_config(model=model_path) @@ -264,7 +264,7 @@ class GenericComponentLoader(ComponentLoader): return model elif self.library == "diffusers": logger.warning(f"Generic loading for diffusers components is not fully implemented") - from fastvideo.models.hf_transformer_utils import get_diffusers_config + from fastvideo.v1.models.hf_transformer_utils import get_diffusers_config model_config = get_diffusers_config(model=model_path) logger.info(f"Diffusers Model config: {model_config}") diff --git a/fastvideo/models/parameter.py b/fastvideo/v1/models/parameter.py similarity index 98% rename from fastvideo/models/parameter.py rename to fastvideo/v1/models/parameter.py index f4f8b148a..11fc6f990 100644 --- a/fastvideo/models/parameter.py +++ b/fastvideo/v1/models/parameter.py @@ -7,9 +7,9 @@ from typing import Callable, Optional, Union import torch from torch.nn import Parameter -from fastvideo.distributed import get_tensor_model_parallel_rank -from fastvideo.logger import init_logger -from fastvideo.models.utils import _make_synced_weight_loader +from fastvideo.v1.distributed import get_tensor_model_parallel_rank +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.utils import _make_synced_weight_loader __all__ = [ "BasevLLMParameter", "PackedvLLMParameter", "PerTensorScaleParameter", diff --git a/fastvideo/models/registry.py b/fastvideo/v1/models/registry.py similarity index 97% rename from fastvideo/models/registry.py rename to fastvideo/v1/models/registry.py index cc02f9e6d..b665a6ccc 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/v1/models/registry.py @@ -11,7 +11,7 @@ import importlib from functools import lru_cache import cloudpickle from torch import nn -from fastvideo.logger import logger +from fastvideo.v1.logger import logger # huggingface class name: (component_name, fastvideo module name, fastvideo class name) _TEXT_TO_VIDEO_DIT_MODELS = { @@ -44,7 +44,7 @@ _FAST_VIDEO_MODELS = { } _SUBPROCESS_COMMAND = [ - sys.executable, "-m", "fastvideo.models.dits.registry" + sys.executable, "-m", "fastvideo.v1.models.dits.registry" ] @@ -145,7 +145,7 @@ def _try_load_model_cls( model_arch: str, model: _BaseRegisteredModel, ) -> Optional[Type[nn.Module]]: - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform current_platform.verify_model_arch(model_arch) try: return model.load_model_cls() @@ -283,7 +283,7 @@ class _ModelRegistry: ModelRegistry = _ModelRegistry({ model_arch: _LazyRegisteredModel( - module_name=f"fastvideo.models.{component_name}.{mod_relname}", + module_name=f"fastvideo.v1.models.{component_name}.{mod_relname}", component_name=component_name, class_name=cls_name, ) diff --git a/fastvideo/v1/models/schedulers/base.py b/fastvideo/v1/models/schedulers/base.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py similarity index 100% rename from fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py rename to fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py diff --git a/fastvideo/models/text_encoder.py b/fastvideo/v1/models/text_encoder.py similarity index 99% rename from fastvideo/models/text_encoder.py rename to fastvideo/v1/models/text_encoder.py index fa582d2c8..4449cfd9b 100644 --- a/fastvideo/models/text_encoder.py +++ b/fastvideo/v1/models/text_encoder.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn from transformers.utils import ModelOutput -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/models/utils.py b/fastvideo/v1/models/utils.py similarity index 96% rename from fastvideo/models/utils.py rename to fastvideo/v1/models/utils.py index 8f639c018..75ea7818b 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/v1/models/utils.py @@ -50,7 +50,7 @@ def auto_attributes(init_func): def set_random_seed(seed: int) -> None: - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform current_platform.seed_everything(seed) @@ -82,7 +82,7 @@ def set_weight_attrs( # This sometimes causes OOM errors during model loading. To avoid this, # we sync the param tensor after its weight loader is called. # TODO(woosuk): Remove this hack once we have a better solution. - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform if current_platform.is_tpu() and key == "weight_loader": value = _make_synced_weight_loader(value) setattr(weight, key, value) diff --git a/fastvideo/models/vaes/common.py b/fastvideo/v1/models/vaes/common.py similarity index 99% rename from fastvideo/models/vaes/common.py rename to fastvideo/v1/models/vaes/common.py index 125b21e01..f3ecec89f 100644 --- a/fastvideo/models/vaes/common.py +++ b/fastvideo/v1/models/vaes/common.py @@ -4,7 +4,7 @@ import numpy as np import torch.nn as nn from abc import abstractmethod, ABC import torch.distributed as dist -from fastvideo.distributed import tensor_model_parallel_all_gather, get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size +from fastvideo.v1.distributed import tensor_model_parallel_all_gather, get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size from math import prod class ParallelTiledVAE(ABC): def __init__(self, *args, **kwargs): diff --git a/fastvideo/models/vaes/hunyuanvae.py b/fastvideo/v1/models/vaes/hunyuanvae.py similarity index 99% rename from fastvideo/models/vaes/hunyuanvae.py rename to fastvideo/v1/models/vaes/hunyuanvae.py index f62d89dc5..32df14429 100644 --- a/fastvideo/models/vaes/hunyuanvae.py +++ b/fastvideo/v1/models/vaes/hunyuanvae.py @@ -20,10 +20,10 @@ import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint -from fastvideo.models.vaes.common import DiagonalGaussianDistribution, ParallelTiledVAE -from fastvideo.layers.activation import get_act_fn +from fastvideo.v1.models.vaes.common import DiagonalGaussianDistribution, ParallelTiledVAE +from fastvideo.v1.layers.activation import get_act_fn -from fastvideo.models.utils import auto_attributes +from fastvideo.v1.models.utils import auto_attributes def prepare_causal_attention_mask( num_frames: int, height_width: int, dtype: torch.dtype, device: torch.device, batch_size: int = None diff --git a/fastvideo/pipelines/__init__.py b/fastvideo/v1/pipelines/__init__.py similarity index 93% rename from fastvideo/pipelines/__init__.py rename to fastvideo/v1/pipelines/__init__.py index a1919ea27..2aefab8fd 100644 --- a/fastvideo/pipelines/__init__.py +++ b/fastvideo/v1/pipelines/__init__.py @@ -1,5 +1,5 @@ """ -Diffusion pipelines for FastVideo. +Diffusion pipelines for fastvideo.v1. This package contains diffusion pipelines for generating videos and images. """ @@ -9,22 +9,22 @@ from copy import deepcopy from typing import Dict, Optional, Type, Any -from fastvideo.pipelines.pipeline_registry import PipelineRegistry -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_registry import PipelineRegistry +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger from huggingface_hub import snapshot_download from transformers import PretrainedConfig -from fastvideo.models.hf_transformer_utils import get_hf_config, get_diffusers_config -from fastvideo.models import get_scheduler -from fastvideo.loader.loader import get_model_loader +from fastvideo.v1.models.hf_transformer_utils import get_hf_config, get_diffusers_config +from fastvideo.v1.models import get_scheduler +from fastvideo.v1.loader.loader import get_model_loader import glob -from fastvideo.loader.fsdp_load import load_fsdp_model -from fastvideo.models.loader import PipelineComponentLoader +from fastvideo.v1.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.models.loader import PipelineComponentLoader logger = init_logger(__name__) # Then import the base classes -from fastvideo.pipelines.composed.composed_pipeline_base import ( +from fastvideo.v1.pipelines.composed.composed_pipeline_base import ( ComposedPipelineBase, DiffusionPipelineOutput ) diff --git a/fastvideo/pipelines/composed/__init__.py b/fastvideo/v1/pipelines/composed/__init__.py similarity index 84% rename from fastvideo/pipelines/composed/__init__.py rename to fastvideo/v1/pipelines/composed/__init__.py index d451c39df..1ba564a34 100644 --- a/fastvideo/pipelines/composed/__init__.py +++ b/fastvideo/v1/pipelines/composed/__init__.py @@ -4,7 +4,7 @@ Composed pipelines for diffusion models. This package contains pipelines that are composed of multiple stages. """ -from fastvideo.pipelines.composed.composed_pipeline_base import ( +from fastvideo.v1.pipelines.composed.composed_pipeline_base import ( ComposedPipelineBase, DiffusionPipelineOutput, ) diff --git a/fastvideo/pipelines/composed/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed/composed_pipeline_base.py similarity index 95% rename from fastvideo/pipelines/composed/composed_pipeline_base.py rename to fastvideo/v1/pipelines/composed/composed_pipeline_base.py index 284eddd0c..669b4ed4f 100644 --- a/fastvideo/pipelines/composed/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed/composed_pipeline_base.py @@ -11,10 +11,10 @@ import torch import numpy as np from tqdm.auto import tqdm -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.stages import PipelineStage -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages import PipelineStage +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/implementations/hunyuan/constants.py b/fastvideo/v1/pipelines/hunyuan/constants.py similarity index 100% rename from fastvideo/pipelines/implementations/hunyuan/constants.py rename to fastvideo/v1/pipelines/hunyuan/constants.py diff --git a/fastvideo/pipelines/hunyuan/hunyuan_loader.py b/fastvideo/v1/pipelines/hunyuan/hunyuan_loader.py similarity index 95% rename from fastvideo/pipelines/hunyuan/hunyuan_loader.py rename to fastvideo/v1/pipelines/hunyuan/hunyuan_loader.py index b2d3d76e6..fd157b92e 100644 --- a/fastvideo/pipelines/hunyuan/hunyuan_loader.py +++ b/fastvideo/v1/pipelines/hunyuan/hunyuan_loader.py @@ -4,19 +4,19 @@ from pathlib import Path from safetensors.torch import load_file as safetensors_load_file from ..loader import PipelineLoader -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE -from fastvideo.logger import init_logger +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE +from fastvideo.v1.logger import init_logger from typing import Tuple, Optional import os import glob import json -from fastvideo.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.loader.fsdp_load import load_fsdp_model -from fastvideo.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder -from fastvideo.platforms import current_platform +from fastvideo.v1.platforms import current_platform logger = init_logger(__name__) @@ -157,7 +157,7 @@ class HunyuanPipelineLoader(PipelineLoader): def load_transformer_v0(self, inference_args: InferenceArgs): """Custom transformer loading for Hunyuan""" # TODO(will): replace this with abstracted model - from fastvideo.models.hunyuan.modules import load_model + from fastvideo.v1.models.hunyuan.modules import load_model # Disable gradient torch.set_grad_enabled(False) @@ -183,7 +183,7 @@ class HunyuanPipelineLoader(PipelineLoader): def load_vae(self, inference_args: InferenceArgs): """Custom VAE loading for Hunyuan""" # TODO(will): replace this with abstracted model - from fastvideo.models.hunyuan.vae import load_vae + from fastvideo.v1.models.hunyuan.vae import load_vae vae, _, s_ratio, t_ratio = load_vae( inference_args.vae, inference_args.vae_precision, diff --git a/fastvideo/pipelines/hunyuan/hunyuan_video.py b/fastvideo/v1/pipelines/hunyuan/hunyuan_video.py similarity index 93% rename from fastvideo/pipelines/hunyuan/hunyuan_video.py rename to fastvideo/v1/pipelines/hunyuan/hunyuan_video.py index 3bcae1e33..5a807b2e0 100644 --- a/fastvideo/pipelines/hunyuan/hunyuan_video.py +++ b/fastvideo/v1/pipelines/hunyuan/hunyuan_video.py @@ -6,19 +6,19 @@ from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Tuple, Union from diffusers.image_processor import VaeImageProcessor -from fastvideo.pipelines.pipeline_base import DiffusionPipelineBase +from fastvideo.v1.pipelines.pipeline_base import DiffusionPipelineBase # from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.logger import init_logger # TODO(will): temporary import from diffusers.models import AutoencoderKL -# from fastvideo.models.encoders.encoder import TextEncoder -from fastvideo.models.hunyuan.text_encoder import TextEncoder +# from fastvideo.v1.models.encoders.encoder import TextEncoder +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder from diffusers.schedulers import KarrasDiffusionSchedulers -from fastvideo.models.hunyuan.modules import HYVideoDiffusionTransformer -from fastvideo.inference_args import InferenceArgs +from fastvideo.v1.models.hunyuan.modules import HYVideoDiffusionTransformer +from fastvideo.v1.inference_args import InferenceArgs logger = init_logger(__name__) diff --git a/fastvideo/pipelines/implementations/hunyuan/__init__.py b/fastvideo/v1/pipelines/implementations/hunyuan/__init__.py similarity index 69% rename from fastvideo/pipelines/implementations/hunyuan/__init__.py rename to fastvideo/v1/pipelines/implementations/hunyuan/__init__.py index 9e90d55d8..70c8c01d6 100644 --- a/fastvideo/pipelines/implementations/hunyuan/__init__.py +++ b/fastvideo/v1/pipelines/implementations/hunyuan/__init__.py @@ -4,7 +4,7 @@ HunYuan pipeline implementations. This package contains implementations of diffusion pipelines for HunYuan models. """ -from fastvideo.pipelines.implementations.hunyuan.hunyuan_pipeline import ( +from fastvideo.v1.pipelines.implementations.hunyuan.hunyuan_pipeline import ( EntryClass, ) diff --git a/fastvideo/v1/pipelines/implementations/hunyuan/constants.py b/fastvideo/v1/pipelines/implementations/hunyuan/constants.py new file mode 100644 index 000000000..3dd67b248 --- /dev/null +++ b/fastvideo/v1/pipelines/implementations/hunyuan/constants.py @@ -0,0 +1,89 @@ +import os + +import torch + +__all__ = [ + "C_SCALE", + "PROMPT_TEMPLATE", + "MODEL_BASE", + "PRECISIONS", + "NORMALIZATION_TYPE", + "ACTIVATION_TYPE", + "VAE_PATH", + "TEXT_ENCODER_PATH", + "TOKENIZER_PATH", + "TEXT_PROJECTION", + "DATA_TYPE", + "NEGATIVE_PROMPT", +] + +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +# =================== Constant Values ===================== +# Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid +# overflow error when tensorboard logging values. +C_SCALE = 1_000_000_000_000_000 + +# When using decoder-only models, we must provide a prompt template to instruct the text encoder +# on how to generate the text. +# -------------------------------------------------------------------- +PROMPT_TEMPLATE_ENCODE = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe 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:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") + +NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion" + +PROMPT_TEMPLATE = { + "dit-llm-encode": { + "template": PROMPT_TEMPLATE_ENCODE, + "crop_start": 36, + }, + "dit-llm-encode-video": { + "template": PROMPT_TEMPLATE_ENCODE_VIDEO, + "crop_start": 95, + }, +} + +# ======================= Model ====================== +PRECISIONS = {"fp32", "fp16", "bf16"} +NORMALIZATION_TYPE = {"layer", "rms"} +ACTIVATION_TYPE = {"relu", "silu", "gelu", "gelu_tanh"} + +# =================== Model Path ===================== +MODEL_BASE = os.getenv("MODEL_BASE", "./data/hunyuan") + +# =================== Data ======================= +DATA_TYPE = {"image", "video", "image_video"} + +# 3D VAE +VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"} + +# Text Encoder +TEXT_ENCODER_PATH = { + "clipL": f"{MODEL_BASE}/text_encoder_2", + "llm": f"{MODEL_BASE}/text_encoder", +} + +# Tokenizer +TOKENIZER_PATH = { + "clipL": f"{MODEL_BASE}/text_encoder_2", + "llm": f"{MODEL_BASE}/text_encoder", +} + +TEXT_PROJECTION = { + "linear", # Default, an nn.Linear() layer + "single_refiner", # Single TokenRefiner. Refer to LI-DiT +} diff --git a/fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_loader.py similarity index 94% rename from fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py rename to fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_loader.py index e606f1d01..1e8bd13b0 100644 --- a/fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py +++ b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_loader.py @@ -3,20 +3,20 @@ import torch.nn as nn from pathlib import Path from safetensors.torch import load_file as safetensors_load_file -from fastvideo.pipelines.loader import PipelineLoader -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.implementations.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.loader import PipelineLoader +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.implementations.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE +from fastvideo.v1.logger import init_logger from typing import Tuple, Optional -from fastvideo.models.dits import HunyuanVideoTransformer3DModel +from fastvideo.v1.models.dits import HunyuanVideoTransformer3DModel import os import glob import json -from fastvideo.loader.fsdp_load import load_fsdp_model -from fastvideo.models.hunyuan.text_encoder import TextEncoder -from fastvideo.models.vaes import AutoencoderKLHunyuanVideo +from fastvideo.v1.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.vaes import AutoencoderKLHunyuanVideo from safetensors.torch import load_file as safetensors_load_file -from fastvideo.platforms import current_platform +from fastvideo.v1.platforms import current_platform logger = init_logger(__name__) @@ -158,7 +158,7 @@ class HunyuanPipelineLoader(PipelineLoader): def load_transformer_v0(self, inference_args: InferenceArgs): """Custom transformer loading for Hunyuan""" # TODO(will): replace this with abstracted model - from fastvideo.models.hunyuan.modules import load_model + from fastvideo.v1.models.hunyuan.modules import load_model # Disable gradient torch.set_grad_enabled(False) @@ -227,7 +227,7 @@ class HunyuanPipelineLoader(PipelineLoader): def load_vae_v0(self, inference_args: InferenceArgs): """Custom VAE loading for Hunyuan""" # TODO(will): replace this with abstracted model - from fastvideo.models.hunyuan.vae import load_vae + from fastvideo.v1.models.hunyuan.vae import load_vae vae, _, s_ratio, t_ratio = load_vae( inference_args.vae, inference_args.vae_precision, diff --git a/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py similarity index 92% rename from fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py rename to fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py index 77a643fdc..4b1c39dfc 100644 --- a/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py +++ b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py @@ -9,8 +9,8 @@ from typing import Union, Any, Dict import torch from diffusers.image_processor import VaeImageProcessor -from fastvideo.pipelines.composed import ComposedPipelineBase -from fastvideo.pipelines.stages import ( +from fastvideo.v1.pipelines.composed import ComposedPipelineBase +from fastvideo.v1.pipelines.stages import ( InputValidationStage, PromptEncodingStage, TimestepPreparationStage, @@ -20,12 +20,12 @@ from fastvideo.pipelines.stages import ( DecodingStage, PostProcessingStage, ) -from fastvideo.pipelines.stages.prompt_encoding import PromptEncodingStage -# from fastvideo.pipelines.stages.timestep_preparation import FlowMatchingTimestepPreparationStage -from fastvideo.inference_args import InferenceArgs -# from fastvideo.pipelines.composed.composed_pipeline_base import DiffusionPipelineOutput -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.models.text_encoder import TextEncoder +from fastvideo.v1.pipelines.stages.prompt_encoding import PromptEncodingStage +# from fastvideo.v1.pipelines.stages.timestep_preparation import FlowMatchingTimestepPreparationStage +from fastvideo.v1.inference_args import InferenceArgs +# from fastvideo.v1.pipelines.composed.composed_pipeline_base import DiffusionPipelineOutput +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.models.text_encoder import TextEncoder # TODO(will): move PRECISION_TO_TYPE to better place from .constants import PROMPT_TEMPLATE, PRECISION_TO_TYPE @@ -33,7 +33,7 @@ from diffusers.utils import BaseOutput import numpy as np from dataclasses import dataclass -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) # class HunyuanLatentPreparationStage(LatentPreparationStage): diff --git a/fastvideo/pipelines/loader.py b/fastvideo/v1/pipelines/loader.py similarity index 93% rename from fastvideo/pipelines/loader.py rename to fastvideo/v1/pipelines/loader.py index a65a5ad2a..e1f8fa323 100644 --- a/fastvideo/pipelines/loader.py +++ b/fastvideo/v1/pipelines/loader.py @@ -1,12 +1,12 @@ from abc import ABC, abstractmethod import torch -from fastvideo.inference_args import InferenceArgs -from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler -from fastvideo.logger import init_logger -from fastvideo.platforms import current_platform -from fastvideo.distributed.parallel_state import get_sp_group -from fastvideo.pipelines.composed import ComposedPipelineBase +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler +from fastvideo.v1.logger import init_logger +from fastvideo.v1.platforms import current_platform +from fastvideo.v1.distributed.parallel_state import get_sp_group +from fastvideo.v1.pipelines.composed import ComposedPipelineBase import os import json @@ -26,14 +26,14 @@ def resolve_pipeline_cls(inference_args: InferenceArgs): """Resolve the pipeline class based on the inference args.""" # TODO(will): resolve the pipeline class based on the model path # read from hf config file _class_name - from fastvideo.pipelines.hunyuan.hunyuan_video import HunyuanVideoPipeline + from fastvideo.v1.pipelines.hunyuan.hunyuan_video import HunyuanVideoPipeline return HunyuanVideoPipeline def resolve_pipeline_cls_v2(inference_args: InferenceArgs): """Resolve the pipeline class based on the inference args.""" # TODO(will): resolve the pipeline class based on the model path # read from hf config file _class_name - from fastvideo.pipelines.implementations.hunyuan import HunyuanVideoPipeline + from fastvideo.v1.pipelines.implementations.hunyuan import HunyuanVideoPipeline return HunyuanVideoPipeline @@ -258,7 +258,7 @@ class PipelineLoader(ABC): def get_pipeline_loader(inference_args: InferenceArgs) -> PipelineLoader: """Get a pipeline loader based on the inference args.""" - from fastvideo.pipelines.implementations.hunyuan.hunyuan_loader import HunyuanPipelineLoader + from fastvideo.v1.pipelines.implementations.hunyuan.hunyuan_loader import HunyuanPipelineLoader # Map model types to their specific loaders loader_map = { "hunyuan_video": HunyuanPipelineLoader, diff --git a/fastvideo/pipelines/pipeline_base.py b/fastvideo/v1/pipelines/pipeline_base.py similarity index 97% rename from fastvideo/pipelines/pipeline_base.py rename to fastvideo/v1/pipelines/pipeline_base.py index fbe75cfc1..0349c3b76 100644 --- a/fastvideo/pipelines/pipeline_base.py +++ b/fastvideo/v1/pipelines/pipeline_base.py @@ -11,13 +11,13 @@ from tqdm.auto import tqdm from diffusers.utils.torch_utils import randn_tensor -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.models.hunyuan.text_encoder import TextEncoder -from fastvideo.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D -from fastvideo.models.hunyuan.constants import PRECISION_TO_TYPE -from fastvideo.distributed.parallel_state import get_sp_group -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from fastvideo.v1.models.hunyuan.constants import PRECISION_TO_TYPE +from fastvideo.v1.distributed.parallel_state import get_sp_group +from fastvideo.v1.logger import init_logger from einops import rearrange from diffusers.utils import BaseOutput import numpy as np @@ -420,7 +420,7 @@ class DiffusionPipelineBase(ABC): # 5. Prepare latent variables batch = self.prepare_latents(batch) - from fastvideo.utils.parallel_states import nccl_info + from fastvideo.v1.utils.parallel_states import nccl_info world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group # world_size = 1 # rank = 0 @@ -541,8 +541,8 @@ class DiffusionPipelineBase(ABC): # if sp_group: # latents = sp_group.all_gather(latents, dim=2) - from fastvideo.utils.communications import all_gather - from fastvideo.utils.parallel_states import get_sequence_parallel_state + from fastvideo.v1.utils.communications import all_gather + from fastvideo.v1.utils.parallel_states import get_sequence_parallel_state if get_sequence_parallel_state(): latents = all_gather(latents, dim=2) diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py similarity index 100% rename from fastvideo/pipelines/pipeline_batch_info.py rename to fastvideo/v1/pipelines/pipeline_batch_info.py diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/v1/pipelines/pipeline_registry.py similarity index 96% rename from fastvideo/pipelines/pipeline_registry.py rename to fastvideo/v1/pipelines/pipeline_registry.py index 9f61a37b0..10d601093 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/v1/pipelines/pipeline_registry.py @@ -7,7 +7,7 @@ from dataclasses import dataclass, field from functools import lru_cache from typing import AbstractSet, Dict, List, Optional, Tuple, Type, Union -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger import torch.nn as nn @@ -59,7 +59,7 @@ class _PipelineRegistry: @lru_cache() def import_pipeline_classes(): pipeline_arch_name_to_cls = {} - package_name = "fastvideo.pipelines.implementations" + package_name = "fastvideo.v1.pipelines.implementations" package = importlib.import_module(package_name) for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."): if ispkg: diff --git a/fastvideo/v1/pipelines/stages/__init__.py b/fastvideo/v1/pipelines/stages/__init__.py new file mode 100644 index 000000000..696f1b337 --- /dev/null +++ b/fastvideo/v1/pipelines/stages/__init__.py @@ -0,0 +1,28 @@ +""" +Pipeline stages for diffusion models. + +This package contains the various stages that can be composed to create +complete diffusion pipelines. +""" + +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.stages.input_validation import InputValidationStage +from fastvideo.v1.pipelines.stages.prompt_encoding import PromptEncodingStage +from fastvideo.v1.pipelines.stages.timestep_preparation import TimestepPreparationStage +from fastvideo.v1.pipelines.stages.latent_preparation import LatentPreparationStage +from fastvideo.v1.pipelines.stages.conditioning import ConditioningStage +from fastvideo.v1.pipelines.stages.denoising import DenoisingStage +from fastvideo.v1.pipelines.stages.decoding import DecodingStage +from fastvideo.v1.pipelines.stages.post_processing import PostProcessingStage + +__all__ = [ + "PipelineStage", + "InputValidationStage", + "PromptEncodingStage", + "TimestepPreparationStage", + "LatentPreparationStage", + "ConditioningStage", + "DenoisingStage", + "DecodingStage", + "PostProcessingStage", +] \ No newline at end of file diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/v1/pipelines/stages/base.py similarity index 94% rename from fastvideo/pipelines/stages/base.py rename to fastvideo/v1/pipelines/stages/base.py index 83ee6da17..f858725f6 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/v1/pipelines/stages/base.py @@ -11,9 +11,9 @@ import time import traceback from typing import Dict, Any -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) @@ -36,7 +36,7 @@ class PipelineStage(ABC): """ self._enable_logging = enable_logging self._stage_name = self.__class__.__name__ - self._logger = init_logger(f"fastvideo.pipelines.stages.{self._stage_name}") + self._logger = init_logger(f"fastvideo.v1.pipelines.stages.{self._stage_name}") @property def device(self) -> torch.device: diff --git a/fastvideo/pipelines/stages/conditioning.py b/fastvideo/v1/pipelines/stages/conditioning.py similarity index 90% rename from fastvideo/pipelines/stages/conditioning.py rename to fastvideo/v1/pipelines/stages/conditioning.py index c907e92d8..9765a2f13 100644 --- a/fastvideo/pipelines/stages/conditioning.py +++ b/fastvideo/v1/pipelines/stages/conditioning.py @@ -5,10 +5,10 @@ Conditioning stage for diffusion pipelines. import torch from typing import Optional -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py similarity index 88% rename from fastvideo/pipelines/stages/decoding.py rename to fastvideo/v1/pipelines/stages/decoding.py index 6863cc327..114dc6a82 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/v1/pipelines/stages/decoding.py @@ -3,11 +3,11 @@ Decoding stage for diffusion pipelines. """ import torch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.models.hunyuan.constants import PRECISION_TO_TYPE -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.models.hunyuan.constants import PRECISION_TO_TYPE +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py similarity index 95% rename from fastvideo/pipelines/stages/denoising.py rename to fastvideo/v1/pipelines/stages/denoising.py index a7101b32b..e9285cb14 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -8,13 +8,13 @@ from typing import Optional, Dict, Any, List from tqdm.auto import tqdm from einops import rearrange -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.models.hunyuan.constants import PRECISION_TO_TYPE -from fastvideo.utils.parallel_states import nccl_info, get_sequence_parallel_state -from fastvideo.utils.communications import all_gather -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.models.hunyuan.constants import PRECISION_TO_TYPE +from fastvideo.v1.utils.parallel_states import nccl_info, get_sequence_parallel_state +from fastvideo.v1.utils.communications import all_gather +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/v1/pipelines/stages/input_validation.py similarity index 94% rename from fastvideo/pipelines/stages/input_validation.py rename to fastvideo/v1/pipelines/stages/input_validation.py index 0cf5e8929..2108c7279 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/v1/pipelines/stages/input_validation.py @@ -5,10 +5,10 @@ Input validation stage for diffusion pipelines. from typing import Optional, Union, List import torch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger import random logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/v1/pipelines/stages/latent_preparation.py similarity index 94% rename from fastvideo/pipelines/stages/latent_preparation.py rename to fastvideo/v1/pipelines/stages/latent_preparation.py index ea83f1116..f54817227 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/v1/pipelines/stages/latent_preparation.py @@ -5,11 +5,11 @@ Latent preparation stage for diffusion pipelines. import torch from typing import Optional, Tuple -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs from diffusers.utils.torch_utils import randn_tensor -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/post_processing.py b/fastvideo/v1/pipelines/stages/post_processing.py similarity index 85% rename from fastvideo/pipelines/stages/post_processing.py rename to fastvideo/v1/pipelines/stages/post_processing.py index 4de05674f..8bce92fba 100644 --- a/fastvideo/pipelines/stages/post_processing.py +++ b/fastvideo/v1/pipelines/stages/post_processing.py @@ -6,11 +6,11 @@ import torch import numpy as np from typing import Optional, Union, List, Tuple -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs from diffusers.utils import BaseOutput -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/prompt_encoding.py b/fastvideo/v1/pipelines/stages/prompt_encoding.py similarity index 94% rename from fastvideo/pipelines/stages/prompt_encoding.py rename to fastvideo/v1/pipelines/stages/prompt_encoding.py index f3205fbc9..3234e9d88 100644 --- a/fastvideo/pipelines/stages/prompt_encoding.py +++ b/fastvideo/v1/pipelines/stages/prompt_encoding.py @@ -7,11 +7,11 @@ This module contains implementations of prompt encoding stages for diffusion pip from typing import Any, Dict, List, Optional, Union, Tuple import torch -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.stages import PipelineStage -from fastvideo.models.text_encoder import TextEncoder -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages import PipelineStage +from fastvideo.v1.models.text_encoder import TextEncoder +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/timestep_preparation.py b/fastvideo/v1/pipelines/stages/timestep_preparation.py similarity index 93% rename from fastvideo/pipelines/stages/timestep_preparation.py rename to fastvideo/v1/pipelines/stages/timestep_preparation.py index f16ae3224..e5b87bbd5 100644 --- a/fastvideo/pipelines/stages/timestep_preparation.py +++ b/fastvideo/v1/pipelines/stages/timestep_preparation.py @@ -8,10 +8,10 @@ from typing import Any, Dict, List, Optional, Union, Tuple import torch import inspect -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/platforms/__init__.py b/fastvideo/v1/platforms/__init__.py similarity index 100% rename from fastvideo/platforms/__init__.py rename to fastvideo/v1/platforms/__init__.py diff --git a/fastvideo/platforms/cuda.py b/fastvideo/v1/platforms/cuda.py similarity index 97% rename from fastvideo/platforms/cuda.py rename to fastvideo/v1/platforms/cuda.py index f02a4b415..f9bb18a5b 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/v1/platforms/cuda.py @@ -12,9 +12,9 @@ import torch from typing_extensions import ParamSpec # import custom ops, trigger op registration -import fastvideo.envs as envs -from fastvideo.logger import init_logger -from fastvideo.utils import import_pynvml +import fastvideo.v1.envs as envs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils import import_pynvml from .interface import DeviceCapability, Platform, PlatformEnum, _Backend @@ -115,7 +115,7 @@ class CudaPlatformBase(Platform): @classmethod def get_device_communicator_cls(cls) -> str: - return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa + return "fastvideo.v1.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa # NVML utils diff --git a/fastvideo/platforms/interface.py b/fastvideo/v1/platforms/interface.py similarity index 97% rename from fastvideo/platforms/interface.py rename to fastvideo/v1/platforms/interface.py index 8b061535f..d92e563f4 100644 --- a/fastvideo/platforms/interface.py +++ b/fastvideo/v1/platforms/interface.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Union import numpy as np import torch -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) @@ -186,7 +186,7 @@ class Platform: """ Get device specific communicator class for distributed communication. """ - return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa + return "fastvideo.v1.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa class UnspecifiedPlatform(Platform): diff --git a/fastvideo/sample/v1_fastvideo_inference.py b/fastvideo/v1/sample/v1_fastvideo_inference.py similarity index 87% rename from fastvideo/sample/v1_fastvideo_inference.py rename to fastvideo/v1/sample/v1_fastvideo_inference.py index a86105003..b0ae9842a 100644 --- a/fastvideo/sample/v1_fastvideo_inference.py +++ b/fastvideo/v1/sample/v1_fastvideo_inference.py @@ -9,12 +9,11 @@ import torchvision from einops import rearrange # Fix the import path -from fastvideo.inference_engine import InferenceEngine -from fastvideo.inference_args import InferenceArgs -from fastvideo.utils.utils import FlexibleArgumentParser -from fastvideo.distributed import init_distributed_environment, initialize_model_parallel -from fastvideo.logger import init_logger -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state +from fastvideo.v1.inference_engine import InferenceEngine +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.utils.utils import FlexibleArgumentParser +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) def main(inference_args: InferenceArgs): diff --git a/fastvideo/v1/distributed_example.py b/fastvideo/v1/tests/distributed_example.py similarity index 90% rename from fastvideo/v1/distributed_example.py rename to fastvideo/v1/tests/distributed_example.py index 2e15baa6c..c37d8a0f0 100644 --- a/fastvideo/v1/distributed_example.py +++ b/fastvideo/v1/tests/distributed_example.py @@ -1,4 +1,4 @@ -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) import os diff --git a/fastvideo/v1/distributed_run.sh b/fastvideo/v1/tests/distributed_run.sh similarity index 100% rename from fastvideo/v1/distributed_run.sh rename to fastvideo/v1/tests/distributed_run.sh diff --git a/fastvideo/v1/tests/test_clip_encoder.py b/fastvideo/v1/tests/test_clip_encoder.py new file mode 100644 index 000000000..26e71432d --- /dev/null +++ b/fastvideo/v1/tests/test_clip_encoder.py @@ -0,0 +1,168 @@ +from fastvideo.v1.models.encoders.clip import CLIPTextModel +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +import os +import torch +import torch.nn as nn +import argparse +import numpy as np +from transformers import CLIPTextModel as HFCLIPTextModel +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer + +logger = init_logger(__name__) + +def initialize_identical_weights(model1, model2, seed=42): + """Initialize both models with identical weights using a fixed seed for reproducibility.""" + # Get all parameters from both models + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Initialize each layer with identical values + with torch.no_grad(): + # Initialize weights + for name1, param1 in params1.items(): + if 'weight' in name1: + # Set seed before each weight initialization + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + + for name2, param2 in params2.items(): + if 'weight' in name2: + # Reset seed to get same initialization + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + + # Initialize biases + for name1, param1 in params1.items(): + if 'bias' in name1: + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + param1.data = param1.data.to(torch.bfloat16) + + for name2, param2 in params2.items(): + if 'bias' in name2: + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + param2.data = param2.data.to(torch.bfloat16) + + logger.info("Both models initialized with identical weights in bfloat16") + return model1, model2 + +def setup_args(): + parser = argparse.ArgumentParser(description='CLIP Text Encoder Test') + parser.add_argument('--model-path', type=str, default="openai/clip-vit-large-patch14", + help='Path to the CLIP model') + parser.add_argument('--precision', type=str, default="bfloat16", + help='Precision to use for the model (float32, float16, bfloat16)') + return parser.parse_args() + +def test_clip_encoder(): + args = setup_args() + + # Set fixed random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + # Initialize the two model implementations + logger.info(f"Loading models from {args.model_path}") + + # Load our implementation using the loader from text_encoder/__init__.py + model1, _ = load_text_encoder( + text_encoder_type="clipL", + text_encoder_precision=args.precision, + text_encoder_path=args.model_path, + logger=logger, + device=device + ) + + # Load the HuggingFace implementation directly + model2 = HFCLIPTextModel.from_pretrained(args.model_path) + if args.precision == "bfloat16": + model2 = model2.to(torch.bfloat16) + elif args.precision == "float16": + model2 = model2.to(torch.float16) + model2 = model2.to(device) + model2.eval() + + # Load tokenizer + tokenizer, _ = load_tokenizer( + tokenizer_type="clipL", + tokenizer_path=args.model_path, + logger=logger + ) + + # Test with some sample prompts + prompts = [ + "a photo of a cat", + "a beautiful landscape with mountains", + "an astronaut riding a horse on the moon" + ] + + logger.info("Testing CLIP text encoder with sample prompts") + + with torch.no_grad(): + for prompt in prompts: + logger.info(f"Testing prompt: '{prompt}'") + + # Tokenize the prompt + tokens = tokenizer( + prompt, + padding="max_length", + max_length=77, + truncation=True, + return_tensors="pt" + ).to(device) + + # Get embeddings from our implementation + outputs1 = model1( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True + ) + + # Get embeddings from HuggingFace implementation + outputs2 = model2( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True + ) + + # Compare last hidden states + last_hidden_state1 = outputs1.last_hidden_state + last_hidden_state2 = outputs2.last_hidden_state + + 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(f"Maximum difference in last hidden states: {max_diff_hidden.item()}") + logger.info(f"Mean difference in last hidden states: {mean_diff_hidden.item()}") + + # Compare pooler outputs + pooler_output1 = outputs1.pooler_output + pooler_output2 = outputs2.pooler_output + + assert pooler_output1.shape == pooler_output2.shape, \ + f"Pooler output shapes don't match: {pooler_output1.shape} vs {pooler_output2.shape}" + + max_diff_pooler = torch.max(torch.abs(pooler_output1 - pooler_output2)) + mean_diff_pooler = torch.mean(torch.abs(pooler_output1 - pooler_output2)) + + logger.info(f"Maximum difference in pooler outputs: {max_diff_pooler.item()}") + logger.info(f"Mean difference in pooler outputs: {mean_diff_pooler.item()}") + + # Check if outputs are similar (allowing for small numerical differences) + assert max_diff_hidden < 1e-4, \ + f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" + assert max_diff_pooler < 1e-4, \ + f"Pooler outputs differ significantly: max diff = {max_diff_pooler.item()}" + + logger.info("Test passed! Both CLIP text encoder implementations produce similar outputs.") + logger.info("Test completed successfully") + +if __name__ == "__main__": + test_clip_encoder() \ No newline at end of file diff --git a/fastvideo/v1/test_hunyuan_vae.py b/fastvideo/v1/tests/test_hunyuan_vae.py similarity index 98% rename from fastvideo/v1/test_hunyuan_vae.py rename to fastvideo/v1/tests/test_hunyuan_vae.py index 596607138..cb3f0e88b 100644 --- a/fastvideo/v1/test_hunyuan_vae.py +++ b/fastvideo/v1/tests/test_hunyuan_vae.py @@ -1,4 +1,4 @@ -from fastvideo.models.vaes import AutoencoderKLHunyuanVideo as MyHunyuanVAE +from fastvideo.v1.models.vaes import AutoencoderKLHunyuanVideo as MyHunyuanVAE from diffusers import AutoencoderKLHunyuanVideo as DiffusersHunyuanVAE import os @@ -7,7 +7,7 @@ import torch.nn as nn import argparse import numpy as np import json -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger from safetensors.torch import load_file logger = init_logger(__name__) diff --git a/fastvideo/v1/test_hunyuanvideo.py b/fastvideo/v1/tests/test_hunyuanvideo.py similarity index 94% rename from fastvideo/v1/test_hunyuanvideo.py rename to fastvideo/v1/tests/test_hunyuanvideo.py index 8921a6811..1194510ab 100644 --- a/fastvideo/v1/test_hunyuanvideo.py +++ b/fastvideo/v1/tests/test_hunyuanvideo.py @@ -5,8 +5,8 @@ import torch.nn as nn import argparse import numpy as np -from fastvideo.logger import init_logger -from fastvideo.distributed.parallel_state import ( +from fastvideo.v1.logger import init_logger +from fastvideo.v1.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, get_sequence_model_parallel_rank, @@ -16,13 +16,13 @@ from fastvideo.distributed.parallel_state import ( cleanup_dist_env_and_memory ) -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state +from fastvideo.v1.utils.parallel_states import initialize_sequence_parallel_state from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed.device_mesh import init_device_mesh -from fastvideo.loader.fsdp_load import shard_model -from fastvideo.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit -from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer -from fastvideo.models.hunyuan_hf.modeling_hunyuan import HunyuanVideoTransformer3DModel +from fastvideo.v1.loader.fsdp_load import shard_model +from fastvideo.v1.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit +from fastvideo.v1.models.hunyuan.modules.models import HYVideoDiffusionTransformer +from fastvideo.v1.models.hunyuan_hf.modeling_hunyuan import HunyuanVideoTransformer3DModel logger = init_logger(__name__) def initialize_identical_weights(model1, model2, seed=42): diff --git a/fastvideo/v1/test_hunyuanvideo_load.py b/fastvideo/v1/tests/test_hunyuanvideo_load.py similarity index 94% rename from fastvideo/v1/test_hunyuanvideo_load.py rename to fastvideo/v1/tests/test_hunyuanvideo_load.py index b56f38859..bb351e00d 100644 --- a/fastvideo/v1/test_hunyuanvideo_load.py +++ b/fastvideo/v1/tests/test_hunyuanvideo_load.py @@ -4,8 +4,8 @@ import torch.nn as nn import argparse import numpy as np -from fastvideo.logger import init_logger -from fastvideo.distributed.parallel_state import ( +from fastvideo.v1.logger import init_logger +from fastvideo.v1.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, get_sequence_model_parallel_rank, @@ -15,11 +15,11 @@ from fastvideo.distributed.parallel_state import ( cleanup_dist_env_and_memory ) import json -from fastvideo.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit -from fastvideo.models.hunyuan.modules.models import HUNYUAN_VIDEO_CONFIG -from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer -from fastvideo.loader.fsdp_load import load_fsdp_model -from fastvideo.models.hunyuan.inference import Inference +from fastvideo.v1.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit +from fastvideo.v1.models.hunyuan.modules.models import HUNYUAN_VIDEO_CONFIG +from fastvideo.v1.models.hunyuan.modules.models import HYVideoDiffusionTransformer +from fastvideo.v1.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.models.hunyuan.inference import Inference import glob logger = init_logger(__name__) diff --git a/fastvideo/v1/tests/test_llama_encoder.py b/fastvideo/v1/tests/test_llama_encoder.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/v1/tp_example.py b/fastvideo/v1/tests/tp_example.py similarity index 96% rename from fastvideo/v1/tp_example.py rename to fastvideo/v1/tests/tp_example.py index 88377a188..c6b27fe4e 100644 --- a/fastvideo/v1/tp_example.py +++ b/fastvideo/v1/tests/tp_example.py @@ -4,8 +4,8 @@ import torch.nn as nn import torch.distributed as dist import argparse -from fastvideo.logger import init_logger -from fastvideo.distributed.parallel_state import ( +from fastvideo.v1.logger import init_logger +from fastvideo.v1.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, get_tensor_model_parallel_rank, @@ -14,11 +14,11 @@ from fastvideo.distributed.parallel_state import ( destroy_distributed_environment, cleanup_dist_env_and_memory ) -from fastvideo.layers.linear import ( +from fastvideo.v1.layers.linear import ( ColumnParallelLinear, RowParallelLinear ) -from fastvideo.distributed.communication_op import ( +from fastvideo.v1.distributed.communication_op import ( tensor_model_parallel_all_reduce, tensor_model_parallel_all_gather ) diff --git a/fastvideo/v1/utils/checkpoint.py b/fastvideo/v1/utils/checkpoint.py new file mode 100644 index 000000000..5a4300c13 --- /dev/null +++ b/fastvideo/v1/utils/checkpoint.py @@ -0,0 +1,286 @@ +# import +import json +import os + +import torch +import torch.distributed.checkpoint as dist_cp +from peft import get_peft_model_state_dict +from safetensors.torch import load_file, save_file +from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner +from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict +from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import StateDictType + +from fastvideo.v1.utils.logging_ import main_print + + +def save_checkpoint_optimizer(model, optimizer, rank, output_dir, step, discriminator=False): + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + cpu_state = model.state_dict() + optim_state = FSDP.optim_state_dict( + model, + optimizer, + ) + + # todo move to get_state_dict + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + # save using safetensors + if rank <= 0 and not discriminator: + weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors") + save_file(cpu_state, weight_path) + config_dict = dict(model.config) + config_dict.pop('dtype') + config_path = os.path.join(save_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + optimizer_path = os.path.join(save_dir, "optimizer.pt") + torch.save(optim_state, optimizer_path) + else: + weight_path = os.path.join(save_dir, "discriminator_pytorch_model.safetensors") + save_file(cpu_state, weight_path) + optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt") + torch.save(optim_state, optimizer_path) + main_print(f"--> checkpoint saved at step {step}") + + +def save_checkpoint(transformer, rank, output_dir, step): + main_print(f"--> saving checkpoint at step {step}") + with FSDP.state_dict_type( + transformer, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + cpu_state = transformer.state_dict() + # todo move to get_state_dict + if rank <= 0: + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + # save using safetensors + weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors") + save_file(cpu_state, weight_path) + config_dict = dict(transformer.config) + if "dtype" in config_dict: + del config_dict["dtype"] # TODO + config_path = os.path.join(save_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + main_print(f"--> checkpoint saved at step {step}") + + +def save_checkpoint_generator_discriminator( + model, + optimizer, + discriminator, + discriminator_optimizer, + rank, + output_dir, + step, +): + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + cpu_state = model.state_dict() + + # todo move to get_state_dict + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + hf_weight_dir = os.path.join(save_dir, "hf_weights") + os.makedirs(hf_weight_dir, exist_ok=True) + # save using safetensors + if rank <= 0: + config_dict = dict(model.config) + config_path = os.path.join(hf_weight_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + weight_path = os.path.join(hf_weight_dir, "diffusion_pytorch_model.safetensors") + save_file(cpu_state, weight_path) + + main_print(f"--> saved HF weight checkpoint at path {hf_weight_dir}") + model_weight_dir = os.path.join(save_dir, "model_weights_state") + os.makedirs(model_weight_dir, exist_ok=True) + model_optimizer_dir = os.path.join(save_dir, "model_optimizer_state") + os.makedirs(model_optimizer_dir, exist_ok=True) + with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): + optim_state = FSDP.optim_state_dict(model, optimizer) + model_state = model.state_dict() + weight_state_dict = {"model": model_state} + dist_cp.save_state_dict( + state_dict=weight_state_dict, + storage_writer=dist_cp.FileSystemWriter(model_weight_dir), + planner=DefaultSavePlanner(), + ) + optimizer_state_dict = {"optimizer": optim_state} + dist_cp.save_state_dict( + state_dict=optimizer_state_dict, + storage_writer=dist_cp.FileSystemWriter(model_optimizer_dir), + planner=DefaultSavePlanner(), + ) + + discriminator_fsdp_state_dir = os.path.join(save_dir, "discriminator_fsdp_state") + os.makedirs(discriminator_fsdp_state_dir, exist_ok=True) + with FSDP.state_dict_type( + discriminator, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + optim_state = FSDP.optim_state_dict(discriminator, discriminator_optimizer) + model_state = discriminator.state_dict() + state_dict = {"optimizer": optim_state, "model": model_state} + if rank <= 0: + discriminator_fsdp_state_fil = os.path.join(discriminator_fsdp_state_dir, "discriminator_state.pt") + torch.save(state_dict, discriminator_fsdp_state_fil) + + main_print("--> saved FSDP state checkpoint") + + +def load_sharded_model(model, optimizer, model_dir, optimizer_dir): + with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): + weight_state_dict = {"model": model.state_dict()} + + optim_state = load_sharded_optimizer_state_dict( + model_state_dict=weight_state_dict["model"], + optimizer_key="optimizer", + storage_reader=dist_cp.FileSystemReader(optimizer_dir), + ) + optim_state = optim_state["optimizer"] + flattened_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=optim_state) + optimizer.load_state_dict(flattened_osd) + dist_cp.load_state_dict( + state_dict=weight_state_dict, + storage_reader=dist_cp.FileSystemReader(model_dir), + planner=DefaultLoadPlanner(), + ) + model_state = weight_state_dict["model"] + model.load_state_dict(model_state) + main_print(f"--> loaded model and optimizer from path {model_dir}") + return model, optimizer + + +def load_full_state_model(model, optimizer, checkpoint_file, rank): + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + discriminator_state = torch.load(checkpoint_file) + model_state = discriminator_state["model"] + if rank <= 0: + optim_state = discriminator_state["optimizer"] + else: + optim_state = None + model.load_state_dict(model_state) + discriminator_optim_state = FSDP.optim_state_dict_to_load(model=model, + optim=optimizer, + optim_state_dict=optim_state) + optimizer.load_state_dict(discriminator_optim_state) + main_print(f"--> loaded discriminator and discriminator optimizer from path {checkpoint_file}") + return model, optimizer + + +def resume_training_generator_discriminator(model, optimizer, discriminator, discriminator_optimizer, checkpoint_dir, + rank): + step = int(checkpoint_dir.split("-")[-1]) + model_weight_dir = os.path.join(checkpoint_dir, "model_weights_state") + model_optimizer_dir = os.path.join(checkpoint_dir, "model_optimizer_state") + model, optimizer = load_sharded_model(model, optimizer, model_weight_dir, model_optimizer_dir) + discriminator_ckpt_file = os.path.join(checkpoint_dir, "discriminator_fsdp_state", "discriminator_state.pt") + discriminator, discriminator_optimizer = load_full_state_model(discriminator, discriminator_optimizer, + discriminator_ckpt_file, rank) + return model, optimizer, discriminator, discriminator_optimizer, step + + +def resume_training(model, optimizer, checkpoint_dir, discriminator=False): + weight_path = os.path.join(checkpoint_dir, "diffusion_pytorch_model.safetensors") + if discriminator: + weight_path = os.path.join(checkpoint_dir, "discriminator_pytorch_model.safetensors") + model_weights = load_file(weight_path) + + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + current_state = model.state_dict() + current_state.update(model_weights) + model.load_state_dict(current_state, strict=False) + if discriminator: + optim_path = os.path.join(checkpoint_dir, "discriminator_optimizer.pt") + else: + optim_path = os.path.join(checkpoint_dir, "optimizer.pt") + optimizer_state_dict = torch.load(optim_path, weights_only=False) + optim_state = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=optimizer_state_dict) + optimizer.load_state_dict(optim_state) + step = int(checkpoint_dir.split("-")[-1]) + return model, optimizer, step + + +def save_lora_checkpoint(transformer, optimizer, rank, output_dir, step, pipeline): + with FSDP.state_dict_type( + transformer, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + full_state_dict = transformer.state_dict() + lora_optim_state = FSDP.optim_state_dict( + transformer, + optimizer, + ) + + if rank <= 0: + save_dir = os.path.join(output_dir, f"lora-checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + + # save optimizer + optim_path = os.path.join(save_dir, "lora_optimizer.pt") + torch.save(lora_optim_state, optim_path) + # save lora weight + main_print(f"--> saving LoRA checkpoint at step {step}") + transformer_lora_layers = get_peft_model_state_dict(model=transformer, state_dict=full_state_dict) + pipeline.save_lora_weights( + save_directory=save_dir, + transformer_lora_layers=transformer_lora_layers, + is_main_process=True, + ) + # save config + lora_config = { + "step": step, + "lora_params": { + "lora_rank": transformer.config.lora_rank, + "lora_alpha": transformer.config.lora_alpha, + "target_modules": transformer.config.lora_target_modules, + }, + } + config_path = os.path.join(save_dir, "lora_config.json") + with open(config_path, "w") as f: + json.dump(lora_config, f, indent=4) + main_print(f"--> LoRA checkpoint saved at step {step}") + + +def resume_lora_optimizer(transformer, checkpoint_dir, optimizer): + config_path = os.path.join(checkpoint_dir, "lora_config.json") + with open(config_path, "r") as f: + config_dict = json.load(f) + optim_path = os.path.join(checkpoint_dir, "lora_optimizer.pt") + optimizer_state_dict = torch.load(optim_path, weights_only=False) + optim_state = FSDP.optim_state_dict_to_load(model=transformer, + optim=optimizer, + optim_state_dict=optimizer_state_dict) + optimizer.load_state_dict(optim_state) + step = config_dict["step"] + main_print(f"--> Successfully resuming LoRA optimizer from step {step}") + return transformer, optimizer, step diff --git a/fastvideo/v1/utils/communications.py b/fastvideo/v1/utils/communications.py new file mode 100644 index 000000000..04fdc8194 --- /dev/null +++ b/fastvideo/v1/utils/communications.py @@ -0,0 +1,308 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Any, Tuple + +import torch +import torch.distributed as dist +from torch import Tensor + +from fastvideo.v1.utils.parallel_states import nccl_info + + +def broadcast(input_: torch.Tensor): + src = nccl_info.group_id * nccl_info.sp_size + dist.broadcast(input_, src=src, group=nccl_info.group) + + +def _all_to_all_4D(input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None) -> torch.tensor: + """ + all-to-all for QKV + + Args: + input (torch.tensor): a tensor sharded along dim scatter dim + scatter_idx (int): default 1 + gather_idx (int): default 2 + group : torch process group + + Returns: + torch.tensor: resharded tensor (bs, seqlen/P, hc, hs) + """ + assert (input.dim() == 4), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}" + + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs) + bs, shard_seqlen, hc, hs = input.shape + seqlen = shard_seqlen * seq_world_size + shard_hc = hc // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs) + input_t = (input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs).transpose(0, 2).contiguous()) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, seq_len/P, bs, hc/P, hs) scatter seqlen -all2all-> (P, seq_len/P, bs, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + torch.cuda.synchronize() + else: + output = input_t + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(seqlen, bs, shard_hc, hs) + + # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs) + output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs) + + return output + + elif scatter_idx == 1 and gather_idx == 2: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs) + bs, seqlen, shard_hc, hs = input.shape + hc = shard_hc * seq_world_size + shard_seqlen = seqlen // seq_world_size + seq_world_size = dist.get_world_size(group) + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs) + input_t = (input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, + hs).transpose(0, + 3).transpose(0, + 1).contiguous().reshape(seq_world_size, shard_hc, + shard_seqlen, bs, hs)) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, bs x hc/P, seqlen/P, hs) scatter seqlen -all2all-> (P, bs x seq_len/P, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + torch.cuda.synchronize() + else: + output = input_t + + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(hc, shard_seqlen, bs, hs) + + # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs) + output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs) + + return output + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + +class SeqAllToAll4D(torch.autograd.Function): + + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + input: Tensor, + scatter_idx: int, + gather_idx: int, + ) -> Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + return _all_to_all_4D(input, scatter_idx, gather_idx, group=group) + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]: + return ( + None, + SeqAllToAll4D.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx), + None, + None, + ) + + +def all_to_all_4D( + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return SeqAllToAll4D.apply(nccl_info.group, input_, scatter_dim, gather_dim) + + +def _all_to_all( + input_: torch.Tensor, + world_size: int, + group: dist.ProcessGroup, + scatter_dim: int, + gather_dim: int, +): + input_list = [t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +class _AllToAll(torch.autograd.Function): + """All-to-all communication. + + Args: + input_: input matrix + process_group: communication group + scatter_dim: scatter dimension + gather_dim: gather dimension + """ + + @staticmethod + def forward(ctx, input_, process_group, scatter_dim, gather_dim): + ctx.process_group = process_group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.world_size = dist.get_world_size(process_group) + output = _all_to_all(input_, ctx.world_size, process_group, scatter_dim, gather_dim) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_output = _all_to_all( + grad_output, + ctx.world_size, + ctx.process_group, + ctx.gather_dim, + ctx.scatter_dim, + ) + return ( + grad_output, + None, + None, + None, + ) + + +def all_to_all( + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return _AllToAll.apply(input_, nccl_info.group, scatter_dim, gather_dim) + + +class _AllGather(torch.autograd.Function): + """All-gather communication with autograd support. + + Args: + input_: input tensor + dim: dimension along which to concatenate + """ + + @staticmethod + def forward(ctx, input_, dim): + ctx.dim = dim + world_size = nccl_info.sp_size + group = nccl_info.group + input_size = list(input_.size()) + + ctx.input_size = input_size[dim] + + tensor_list = [torch.empty_like(input_) for _ in range(world_size)] + input_ = input_.contiguous() + dist.all_gather(tensor_list, input_, group=group) + + output = torch.cat(tensor_list, dim=dim) + return output + + @staticmethod + def backward(ctx, grad_output): + world_size = nccl_info.sp_size + rank = nccl_info.rank_within_group + dim = ctx.dim + input_size = ctx.input_size + + sizes = [input_size] * world_size + + grad_input_list = torch.split(grad_output, sizes, dim=dim) + grad_input = grad_input_list[rank] + + return grad_input, None + + +def all_gather(input_: torch.Tensor, dim: int = 1): + """Performs an all-gather operation on the input tensor along the specified dimension. + + Args: + input_ (torch.Tensor): Input tensor of shape [B, H, S, D]. + dim (int, optional): Dimension along which to concatenate. Defaults to 1. + + Returns: + torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'. + """ + return _AllGather.apply(input_, dim) + + +def prepare_sequence_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask): + if nccl_info.sp_size == 1: + return ( + hidden_states, + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + ) + + def prepare(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask): + hidden_states = all_to_all(hidden_states, scatter_dim=2, gather_dim=0) + encoder_hidden_states = all_to_all(encoder_hidden_states, scatter_dim=1, gather_dim=0) + attention_mask = all_to_all(attention_mask, scatter_dim=1, gather_dim=0) + encoder_attention_mask = all_to_all(encoder_attention_mask, scatter_dim=1, gather_dim=0) + return ( + hidden_states, + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + ) + + sp_size = nccl_info.sp_size + frame = hidden_states.shape[2] + assert frame % sp_size == 0, "frame should be a multiple of sp_size" + + ( + hidden_states, + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + ) = prepare( + hidden_states, + encoder_hidden_states.repeat(1, sp_size, 1), + attention_mask.repeat(1, sp_size, 1, 1), + encoder_attention_mask.repeat(1, sp_size), + ) + + return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask + + +def sp_parallel_dataloader_wrapper(dataloader, device, train_batch_size, sp_size, train_sp_batch_size): + while True: + for data_item in dataloader: + latents, cond, attn_mask, cond_mask = data_item + latents = latents.to(device) + cond = cond.to(device) + attn_mask = attn_mask.to(device) + cond_mask = cond_mask.to(device) + frame = latents.shape[2] + if frame == 1: + yield latents, cond, attn_mask, cond_mask + else: + latents, cond, attn_mask, cond_mask = prepare_sequence_parallel_data( + latents, cond, attn_mask, cond_mask) + assert (train_batch_size * sp_size >= + train_sp_batch_size), "train_batch_size * sp_size should be greater than train_sp_batch_size" + for iter in range(train_batch_size * sp_size // train_sp_batch_size): + st_idx = iter * train_sp_batch_size + ed_idx = (iter + 1) * train_sp_batch_size + encoder_hidden_states = cond[st_idx:ed_idx] + attention_mask = attn_mask[st_idx:ed_idx] + encoder_attention_mask = cond_mask[st_idx:ed_idx] + yield ( + latents[st_idx:ed_idx], + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + ) diff --git a/fastvideo/v1/utils/dataset_utils.py b/fastvideo/v1/utils/dataset_utils.py new file mode 100644 index 000000000..3c8c9e5a6 --- /dev/null +++ b/fastvideo/v1/utils/dataset_utils.py @@ -0,0 +1,342 @@ +import math +import random +from collections import Counter +from typing import List, Optional + +import decord +import torch +import torch.utils +import torch.utils.data +from torch.nn import functional as F +from torch.utils.data import Sampler + +IMG_EXTENSIONS = [".jpg", ".JPG", ".jpeg", ".JPEG", ".png", ".PNG"] + + +def is_image_file(filename): + return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) + + +class DecordInit(object): + """Using Decord(https://github.com/dmlc/decord) to initialize the video_reader.""" + + def __init__(self, num_threads=1): + self.num_threads = num_threads + self.ctx = decord.cpu(0) + + def __call__(self, filename): + """Perform the Decord initialization. + Args: + results (dict): The resulting dict to be modified and passed + to the next transform in pipeline. + """ + reader = decord.VideoReader(filename, ctx=self.ctx, num_threads=self.num_threads) + return reader + + def __repr__(self): + repr_str = (f"{self.__class__.__name__}(" + f"sr={self.sr}," + f"num_threads={self.num_threads})") + return repr_str + + +def pad_to_multiple(number, ds_stride): + remainder = number % ds_stride + if remainder == 0: + return number + else: + padding = ds_stride - remainder + return number + padding + + +# TODO +class Collate: + + def __init__(self, args): + self.batch_size = args.train_batch_size + self.group_frame = args.group_frame + self.group_resolution = args.group_resolution + + self.max_height = args.max_height + self.max_width = args.max_width + self.ae_stride = args.ae_stride + + self.ae_stride_t = args.ae_stride_t + self.ae_stride_thw = (self.ae_stride_t, self.ae_stride, self.ae_stride) + + self.patch_size = args.patch_size + self.patch_size_t = args.patch_size_t + + self.num_frames = args.num_frames + self.use_image_num = args.use_image_num + self.max_thw = (self.num_frames, self.max_height, self.max_width) + + def package(self, batch): + batch_tubes = [i["pixel_values"] for i in batch] # b [c t h w] + input_ids = [i["input_ids"] for i in batch] # b [1 l] + cond_mask = [i["cond_mask"] for i in batch] # b [1 l] + return batch_tubes, input_ids, cond_mask + + def __call__(self, batch): + batch_tubes, input_ids, cond_mask = self.package(batch) + + ds_stride = self.ae_stride * self.patch_size + t_ds_stride = self.ae_stride_t * self.patch_size_t + + pad_batch_tubes, attention_mask, input_ids, cond_mask = self.process( + batch_tubes, + input_ids, + cond_mask, + t_ds_stride, + ds_stride, + self.max_thw, + self.ae_stride_thw, + ) + assert not torch.any(torch.isnan(pad_batch_tubes)), "after pad_batch_tubes" + return pad_batch_tubes, attention_mask, input_ids, cond_mask + + def process( + self, + batch_tubes, + input_ids, + cond_mask, + t_ds_stride, + ds_stride, + max_thw, + ae_stride_thw, + ): + # pad to max multiple of ds_stride + batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)] + assert len(batch_input_size) == self.batch_size + if self.group_frame or self.group_resolution or self.batch_size == 1: # + len_each_batch = batch_input_size + idx_length_dict = dict([*zip(list(range(self.batch_size)), len_each_batch)]) + count_dict = Counter(len_each_batch) + if len(count_dict) != 1: + sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1]) + pick_length = sorted_by_value[-1][0] # the highest frequency + candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length] + random_select_batch = [ + random.choice(candidate_batch) for _ in range(len(len_each_batch) - len(candidate_batch)) + ] + print( + batch_input_size, + idx_length_dict, + count_dict, + sorted_by_value, + pick_length, + candidate_batch, + random_select_batch, + ) + pick_idx = candidate_batch + random_select_batch + + batch_tubes = [batch_tubes[i] for i in pick_idx] + batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)] + input_ids = [input_ids[i] for i in pick_idx] # b [1, l] + cond_mask = [cond_mask[i] for i in pick_idx] # b [1, l] + + for i in range(1, self.batch_size): + assert batch_input_size[0] == batch_input_size[i] + max_t = max([i[1] for i in batch_input_size]) + max_h = max([i[2] for i in batch_input_size]) + max_w = max([i[3] for i in batch_input_size]) + else: + max_t, max_h, max_w = max_thw + pad_max_t, pad_max_h, pad_max_w = ( + pad_to_multiple(max_t - 1 + self.ae_stride_t, t_ds_stride), + pad_to_multiple(max_h, ds_stride), + pad_to_multiple(max_w, ds_stride), + ) + pad_max_t = pad_max_t + 1 - self.ae_stride_t + each_pad_t_h_w = [[pad_max_t - i.shape[1], pad_max_h - i.shape[2], pad_max_w - i.shape[3]] for i in batch_tubes] + pad_batch_tubes = [ + F.pad(im, (0, pad_w, 0, pad_h, 0, pad_t), value=0) + for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes) + ] + pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0) + + max_tube_size = [pad_max_t, pad_max_h, pad_max_w] + max_latent_size = [ + ((max_tube_size[0] - 1) // ae_stride_thw[0] + 1), + max_tube_size[1] // ae_stride_thw[1], + max_tube_size[2] // ae_stride_thw[2], + ] + valid_latent_size = [[ + int(math.ceil((i[1] - 1) / ae_stride_thw[0])) + 1, + int(math.ceil(i[2] / ae_stride_thw[1])), + int(math.ceil(i[3] / ae_stride_thw[2])), + ] for i in batch_input_size] + attention_mask = [ + F.pad( + torch.ones(i, dtype=pad_batch_tubes.dtype), + ( + 0, + max_latent_size[2] - i[2], + 0, + max_latent_size[1] - i[1], + 0, + max_latent_size[0] - i[0], + ), + value=0, + ) for i in valid_latent_size + ] + attention_mask = torch.stack(attention_mask) # b t h w + if self.batch_size == 1 or self.group_frame or self.group_resolution: + assert torch.all(attention_mask.bool()) + + input_ids = torch.stack(input_ids) # b 1 l + cond_mask = torch.stack(cond_mask) # b 1 l + + return pad_batch_tubes, attention_mask, input_ids, cond_mask + + +def split_to_even_chunks(indices, lengths, num_chunks, batch_size): + """ + Split a list of indices into `chunks` chunks of roughly equal lengths. + """ + + if len(indices) % num_chunks != 0: + chunks = [indices[i::num_chunks] for i in range(num_chunks)] + else: + num_indices_per_chunk = len(indices) // num_chunks + + chunks = [[] for _ in range(num_chunks)] + chunks_lengths = [0 for _ in range(num_chunks)] + for index in indices: + shortest_chunk = chunks_lengths.index(min(chunks_lengths)) + chunks[shortest_chunk].append(index) + chunks_lengths[shortest_chunk] += lengths[index] + if len(chunks[shortest_chunk]) == num_indices_per_chunk: + chunks_lengths[shortest_chunk] = float("inf") + # return chunks + + pad_chunks = [] + for idx, chunk in enumerate(chunks): + if batch_size != len(chunk): + assert batch_size > len(chunk) + if len(chunk) != 0: + chunk = chunk + [random.choice(chunk) for _ in range(batch_size - len(chunk))] + else: + chunk = random.choice(pad_chunks) + print(chunks[idx], "->", chunk) + pad_chunks.append(chunk) + return pad_chunks + + +def group_frame_fun(indices, lengths): + # sort by num_frames + indices.sort(key=lambda i: lengths[i], reverse=True) + return indices + + +def megabatch_frame_alignment(megabatches, lengths): + aligned_magabatches = [] + for _, megabatch in enumerate(megabatches): + assert len(megabatch) != 0 + len_each_megabatch = [lengths[i] for i in megabatch] + idx_length_dict = dict([*zip(megabatch, len_each_megabatch)]) + count_dict = Counter(len_each_megabatch) + + # mixed frame length, align megabatch inside + if len(count_dict) != 1: + sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1]) + pick_length = sorted_by_value[-1][0] # the highest frequency + candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length] + random_select_batch = [ + random.choice(candidate_batch) for i in range(len(idx_length_dict) - len(candidate_batch)) + ] + aligned_magabatch = candidate_batch + random_select_batch + aligned_magabatches.append(aligned_magabatch) + # already aligned megabatches + else: + aligned_magabatches.append(megabatch) + + return aligned_magabatches + + +def get_length_grouped_indices( + lengths, + batch_size, + world_size, + generator=None, + group_frame=False, + group_resolution=False, + seed=42, +): + # We need to use torch for the random part as a distributed sampler will set the random seed for torch. + if generator is None: + generator = torch.Generator().manual_seed(seed) # every rank will generate a fixed order but random index + + indices = torch.randperm(len(lengths), generator=generator).tolist() + + # sort dataset according to frame + indices = group_frame_fun(indices, lengths) + + # chunk dataset to megabatches + megabatch_size = world_size * batch_size + megabatches = [indices[i:i + megabatch_size] for i in range(0, len(lengths), megabatch_size)] + + # make sure the length in each magabatch is align with each other + megabatches = megabatch_frame_alignment(megabatches, lengths) + + # aplit aligned megabatch into batches + megabatches = [split_to_even_chunks(megabatch, lengths, world_size, batch_size) for megabatch in megabatches] + + # random megabatches to do video-image mix training + indices = torch.randperm(len(megabatches), generator=generator).tolist() + shuffled_megabatches = [megabatches[i] for i in indices] + + # expand indices and return + return [i for megabatch in shuffled_megabatches for batch in megabatch for i in batch] + + +class LengthGroupedSampler(Sampler): + r""" + Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while + keeping a bit of randomness. + """ + + def __init__( + self, + batch_size: int, + rank: int, + world_size: int, + lengths: Optional[List[int]] = None, + group_frame=False, + group_resolution=False, + generator=None, + ): + if lengths is None: + raise ValueError("Lengths must be provided.") + + self.batch_size = batch_size + self.rank = rank + self.world_size = world_size + self.lengths = lengths + self.group_frame = group_frame + self.group_resolution = group_resolution + self.generator = generator + + def __len__(self): + return len(self.lengths) + + def __iter__(self): + indices = get_length_grouped_indices( + self.lengths, + self.batch_size, + self.world_size, + group_frame=self.group_frame, + group_resolution=self.group_resolution, + generator=self.generator, + ) + + def distributed_sampler(lst, rank, batch_size, world_size): + result = [] + index = rank * batch_size + while index < len(lst): + result.extend(lst[index:index + batch_size]) + index += batch_size * world_size + return result + + indices = distributed_sampler(indices, self.rank, self.batch_size, self.world_size) + return iter(indices) diff --git a/fastvideo/v1/utils/env_utils.py b/fastvideo/v1/utils/env_utils.py new file mode 100644 index 000000000..13591e9aa --- /dev/null +++ b/fastvideo/v1/utils/env_utils.py @@ -0,0 +1,38 @@ +import platform + +import accelerate +import peft +import torch +import transformers +from transformers.utils import is_torch_cuda_available, is_torch_npu_available + +VERSION = "1.2.0" + +if __name__ == "__main__": + info = { + "FastVideo version": VERSION, + "Platform": platform.platform(), + "Python version": platform.python_version(), + "PyTorch version": torch.__version__, + "Transformers version": transformers.__version__, + "Accelerate version": accelerate.__version__, + "PEFT version": peft.__version__, + } + + if is_torch_cuda_available(): + info["PyTorch version"] += " (GPU)" + info["GPU type"] = torch.cuda.get_device_name() + + if is_torch_npu_available(): + info["PyTorch version"] += " (NPU)" + info["NPU type"] = torch.npu.get_device_name() + info["CANN version"] = torch.version.cann # codespell:ignore + + try: + import bitsandbytes + + info["Bitsandbytes version"] = bitsandbytes.__version__ + except Exception: + pass + + print("\n" + "\n".join([f"- {key}: {value}" for key, value in info.items()]) + "\n") diff --git a/fastvideo/v1/utils/fsdp_util.py b/fastvideo/v1/utils/fsdp_util.py new file mode 100644 index 000000000..96a1e11ce --- /dev/null +++ b/fastvideo/v1/utils/fsdp_util.py @@ -0,0 +1,132 @@ +# ruff: noqa: E731 +import functools +from functools import partial + +import torch +from peft.utils.other import fsdp_auto_wrap_policy +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (CheckpointImpl, apply_activation_checkpointing, + checkpoint_wrapper) +from torch.distributed.fsdp import MixedPrecision, ShardingStrategy +from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy + +from fastvideo.v1.models.mochi_hf.modeling_mochi import MochiTransformerBlock +from fastvideo.v1.utils.load import get_no_split_modules + +non_reentrant_wrapper = partial( + checkpoint_wrapper, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, +) + +check_fn = lambda submodule: isinstance(submodule, MochiTransformerBlock) + + +def apply_fsdp_checkpointing(model, no_split_modules, p=1): + # https://github.com/foundation-model-stack/fms-fsdp/blob/408c7516d69ea9b6bcd4c0f5efab26c0f64b3c2d/fms_fsdp/policies/ac_handler.py#L16 + """apply activation checkpointing to model + returns None as model is updated directly + """ + print("--> applying fdsp activation checkpointing...") + block_idx = 0 + cut_off = 1 / 2 + # when passing p as a fraction number (e.g. 1/3), it will be interpreted + # as a string in argv, thus we need eval("1/3") here for fractions. + p = eval(p) if isinstance(p, str) else p + + def selective_checkpointing(submodule): + nonlocal block_idx + nonlocal cut_off + + if isinstance(submodule, no_split_modules): + block_idx += 1 + if block_idx * p >= cut_off: + cut_off += 1 + return True + return False + + apply_activation_checkpointing( + model, + checkpoint_wrapper_fn=non_reentrant_wrapper, + check_fn=selective_checkpointing, + ) + + +def get_mixed_precision(master_weight_type="fp32"): + weight_type = torch.float32 if master_weight_type == "fp32" else torch.bfloat16 + mixed_precision = MixedPrecision( + param_dtype=weight_type, + # Gradient communication precision. + reduce_dtype=weight_type, + # Buffer precision. + buffer_dtype=weight_type, + cast_forward_inputs=False, + ) + return mixed_precision + + +def get_dit_fsdp_kwargs( + transformer, + sharding_strategy, + use_lora=False, + cpu_offload=False, + master_weight_type="fp32", +): + no_split_modules = get_no_split_modules(transformer) + if use_lora: + auto_wrap_policy = fsdp_auto_wrap_policy + else: + auto_wrap_policy = functools.partial( + transformer_auto_wrap_policy, + transformer_layer_cls=no_split_modules, + ) + + # we use float32 for fsdp but autocast during training + mixed_precision = get_mixed_precision(master_weight_type) + + if sharding_strategy == "full": + sharding_strategy = ShardingStrategy.FULL_SHARD + elif sharding_strategy == "hybrid_full": + sharding_strategy = ShardingStrategy.HYBRID_SHARD + elif sharding_strategy == "none": + sharding_strategy = ShardingStrategy.NO_SHARD + auto_wrap_policy = None + elif sharding_strategy == "hybrid_zero2": + sharding_strategy = ShardingStrategy._HYBRID_SHARD_ZERO2 + + device_id = torch.cuda.current_device() + cpu_offload = (torch.distributed.fsdp.CPUOffload(offload_params=True) if cpu_offload else None) + fsdp_kwargs = { + "auto_wrap_policy": auto_wrap_policy, + "mixed_precision": mixed_precision, + "sharding_strategy": sharding_strategy, + "device_id": device_id, + "limit_all_gathers": True, + "cpu_offload": cpu_offload, + } + + # Add LoRA-specific settings when LoRA is enabled + if use_lora: + fsdp_kwargs.update({ + "use_orig_params": False, # Required for LoRA memory savings + "sync_module_states": True, + }) + + return fsdp_kwargs, no_split_modules + + +def get_discriminator_fsdp_kwargs(master_weight_type="fp32"): + auto_wrap_policy = None + + # Use existing mixed precision settings + + mixed_precision = get_mixed_precision(master_weight_type) + sharding_strategy = ShardingStrategy.NO_SHARD + device_id = torch.cuda.current_device() + fsdp_kwargs = { + "auto_wrap_policy": auto_wrap_policy, + "mixed_precision": mixed_precision, + "sharding_strategy": sharding_strategy, + "device_id": device_id, + "limit_all_gathers": True, + } + + return fsdp_kwargs diff --git a/fastvideo/v1/utils/load.py b/fastvideo/v1/utils/load.py new file mode 100644 index 000000000..d305e0106 --- /dev/null +++ b/fastvideo/v1/utils/load.py @@ -0,0 +1,339 @@ +import os +from pathlib import Path + +import torch +import torch.nn.functional as F +from diffusers import AutoencoderKLHunyuanVideo, AutoencoderKLMochi +from torch import nn +from transformers import AutoTokenizer, T5EncoderModel + +from fastvideo.v1.models.hunyuan.modules.models import (HYVideoDiffusionTransformer, MMDoubleStreamBlock, + MMSingleStreamBlock) +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from fastvideo.v1.models.hunyuan_hf.modeling_hunyuan import (HunyuanVideoSingleTransformerBlock, + HunyuanVideoTransformer3DModel, HunyuanVideoTransformerBlock) +from fastvideo.v1.models.mochi_hf.modeling_mochi import MochiTransformer3DModel, MochiTransformerBlock +from fastvideo.v1.utils.logging_ import main_print + +hunyuan_config = { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + "guidance_embed": True, +} + +PROMPT_TEMPLATE_ENCODE = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe 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:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") + +NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion" + +PROMPT_TEMPLATE = { + "dit-llm-encode": { + "template": PROMPT_TEMPLATE_ENCODE, + "crop_start": 36, + }, + "dit-llm-encode-video": { + "template": PROMPT_TEMPLATE_ENCODE_VIDEO, + "crop_start": 95, + }, +} + + +class HunyuanTextEncoderWrapper(nn.Module): + + def __init__(self, pretrained_model_name_or_path, device): + super().__init__() + + text_len = 256 + crop_start = PROMPT_TEMPLATE["dit-llm-encode-video"].get("crop_start", 0) + + max_length = text_len + crop_start + + # prompt_template + prompt_template = PROMPT_TEMPLATE["dit-llm-encode"] + + # prompt_template_video + prompt_template_video = PROMPT_TEMPLATE["dit-llm-encode-video"] + text_encoder_path = os.path.join(pretrained_model_name_or_path, "text_encoder") + self.text_encoder = TextEncoder( + text_encoder_type="llm", + text_encoder_path=text_encoder_path, + max_length=max_length, + text_encoder_precision="fp16", + tokenizer_type="llm", + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=2, + apply_final_norm=False, + reproduce=False, + logger=None, + device=device, + ) + text_encoder_path_2 = os.path.join(pretrained_model_name_or_path, "text_encoder_2") + self.text_encoder_2 = TextEncoder( + text_encoder_type="clipL", + text_encoder_path=text_encoder_path_2, + max_length=77, + text_encoder_precision="fp16", + tokenizer_type="clipL", + reproduce=False, + logger=None, + device=device, + ) + + def encode_(self, prompt, text_encoder, clip_skip=None): + # TODO + device = self.text_encoder.device + data_type = "video" + num_videos_per_prompt = 1 + + text_inputs = text_encoder.text2tokens(prompt, data_type=data_type) + + if clip_skip is None: + prompt_outputs = text_encoder.encode(text_inputs, data_type="video", device=device) + prompt_embeds = prompt_outputs.hidden_state + else: + prompt_outputs = text_encoder.encode( + text_inputs, + output_hidden_states=True, + data_type=data_type, + device=device, + ) + prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)] + + prompt_embeds = text_encoder.model.text_model.final_layer_norm(prompt_embeds) + + attention_mask = prompt_outputs.attention_mask + if attention_mask is not None: + attention_mask = attention_mask.to(device) + bs_embed, seq_len = attention_mask.shape + attention_mask = attention_mask.repeat(1, num_videos_per_prompt) + attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len) + + if text_encoder is not None: + prompt_embeds_dtype = text_encoder.dtype + elif self.transformer is not None: + prompt_embeds_dtype = self.transformer.dtype + else: + prompt_embeds_dtype = prompt_embeds.dtype + + prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) + + if prompt_embeds.ndim == 2: + bs_embed, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, -1) + else: + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) + return (prompt_embeds, attention_mask) + + def encode_prompt(self, prompt): + prompt_embeds, attention_mask = self.encode_(prompt, self.text_encoder) + prompt_embeds_2, attention_mask_2 = self.encode_(prompt, self.text_encoder_2) + prompt_embeds_2 = F.pad( + prompt_embeds_2, + (0, prompt_embeds.shape[2] - prompt_embeds_2.shape[1]), + value=0, + ).unsqueeze(1) + prompt_embeds = torch.cat([prompt_embeds_2, prompt_embeds], dim=1) + return prompt_embeds, attention_mask + + +class MochiTextEncoderWrapper(nn.Module): + + def __init__(self, pretrained_model_name_or_path, device): + super().__init__() + self.text_encoder = T5EncoderModel.from_pretrained(os.path.join(pretrained_model_name_or_path, + "text_encoder")).to(device) + self.tokenizer = AutoTokenizer.from_pretrained(os.path.join(pretrained_model_name_or_path, "tokenizer")) + self.max_sequence_length = 256 + + def encode_prompt(self, prompt): + device = self.text_encoder.device + dtype = self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.bool().to(device) + + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.max_sequence_length - 1:-1]) + main_print(f"Truncated text input: {prompt} to: {removed_text} for model input.") + prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.view(batch_size, seq_len, -1) + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + + return prompt_embeds, prompt_attention_mask + + +def load_hunyuan_state_dict(model, dit_model_name_or_path): + load_key = "module" + model_path = dit_model_name_or_path + bare_model = "unknown" + + state_dict = torch.load(model_path, map_location=lambda storage, loc: storage, weights_only=True) + + if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict): + bare_model = False + if bare_model is False: + if load_key in state_dict: + state_dict = state_dict[load_key] + else: + raise KeyError(f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint " + f"are: {list(state_dict.keys())}.") + model.load_state_dict(state_dict, strict=True) + return model + + +def load_transformer( + model_type, + dit_model_name_or_path, + pretrained_model_name_or_path, + master_weight_type, +): + if model_type == "mochi": + if dit_model_name_or_path: + transformer = MochiTransformer3DModel.from_pretrained( + dit_model_name_or_path, + torch_dtype=master_weight_type, + # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, + ) + else: + transformer = MochiTransformer3DModel.from_pretrained( + pretrained_model_name_or_path, + subfolder="transformer", + torch_dtype=master_weight_type, + # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, + ) + elif model_type == "hunyuan_hf": + if dit_model_name_or_path: + transformer = HunyuanVideoTransformer3DModel.from_pretrained( + dit_model_name_or_path, + torch_dtype=master_weight_type, + # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, + ) + else: + transformer = HunyuanVideoTransformer3DModel.from_pretrained( + pretrained_model_name_or_path, + subfolder="transformer", + torch_dtype=master_weight_type, + # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, + ) + elif model_type == "hunyuan": + transformer = HYVideoDiffusionTransformer( + in_channels=16, + out_channels=16, + **hunyuan_config, + dtype=master_weight_type, + ) + transformer = load_hunyuan_state_dict(transformer, dit_model_name_or_path) + if master_weight_type == torch.bfloat16: + transformer = transformer.bfloat16() + else: + raise ValueError(f"Unsupported model type: {model_type}") + return transformer + + +def load_vae(model_type, pretrained_model_name_or_path): + weight_dtype = torch.float32 + if model_type == "mochi": + vae = AutoencoderKLMochi.from_pretrained(pretrained_model_name_or_path, + subfolder="vae", + torch_dtype=weight_dtype).to("cuda") + autocast_type = torch.bfloat16 + fps = 30 + elif model_type == "hunyuan_hf": + vae = AutoencoderKLHunyuanVideo.from_pretrained(pretrained_model_name_or_path, + subfolder="vae", + torch_dtype=weight_dtype).to("cuda") + autocast_type = torch.bfloat16 + fps = 24 + elif model_type == "hunyuan": + vae_precision = torch.float32 + vae_path = os.path.join(pretrained_model_name_or_path, "hunyuan-video-t2v-720p/vae") + + config = AutoencoderKLCausal3D.load_config(vae_path) + vae = AutoencoderKLCausal3D.from_config(config) + + vae_ckpt = Path(vae_path) / "pytorch_model.pt" + assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}" + + ckpt = torch.load(vae_ckpt, map_location=vae.device, weights_only=True) + if "state_dict" in ckpt: + ckpt = ckpt["state_dict"] + if any(k.startswith("vae.") for k in ckpt.keys()): + ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")} + vae.load_state_dict(ckpt) + vae = vae.to(dtype=vae_precision) + vae.requires_grad_(False) + vae = vae.to("cuda") + vae.eval() + autocast_type = torch.float32 + fps = 24 + return vae, autocast_type, fps + + +def load_text_encoder(model_type, pretrained_model_name_or_path, device): + if model_type == "mochi": + text_encoder = MochiTextEncoderWrapper(pretrained_model_name_or_path, device) + elif model_type == "hunyuan" or "hunyuan_hf": + text_encoder = HunyuanTextEncoderWrapper(pretrained_model_name_or_path, device) + else: + raise ValueError(f"Unsupported model type: {model_type}") + return text_encoder + + +def get_no_split_modules(transformer): + # if of type MochiTransformer3DModel + if isinstance(transformer, MochiTransformer3DModel): + return (MochiTransformerBlock, ) + elif isinstance(transformer, HunyuanVideoTransformer3DModel): + return (HunyuanVideoSingleTransformerBlock, HunyuanVideoTransformerBlock) + elif isinstance(transformer, HYVideoDiffusionTransformer): + return (MMDoubleStreamBlock, MMSingleStreamBlock) + else: + raise ValueError(f"Unsupported transformer type: {type(transformer)}") + + +if __name__ == "__main__": + # test encode prompt + device = torch.cuda.current_device() + pretrained_model_name_or_path = "data/hunyuan" + text_encoder = load_text_encoder("hunyuan", pretrained_model_name_or_path, device) + prompt = "A man on stage claps his hands together while facing the audience. The audience, visible in the foreground, holds up mobile devices to record the event, capturing the moment from various angles. The background features a large banner with text identifying the man on stage. Throughout the sequence, the man's expression remains engaged and directed towards the audience. The camera angle remains constant, focusing on capturing the interaction between the man on stage and the audience." + prompt_embeds, attention_mask = text_encoder.encode_prompt(prompt) diff --git a/fastvideo/v1/utils/logging_.py b/fastvideo/v1/utils/logging_.py new file mode 100644 index 000000000..ee1523f87 --- /dev/null +++ b/fastvideo/v1/utils/logging_.py @@ -0,0 +1,24 @@ +import os +import pdb +import sys + + +def main_print(content): + if int(os.environ["LOCAL_RANK"]) <= 0: + print(content) + + +# ForkedPdb().set_trace() +class ForkedPdb(pdb.Pdb): + """A Pdb subclass that may be used + from a forked multiprocessing child + + """ + + def interaction(self, *args, **kwargs): + _stdin = sys.stdin + try: + sys.stdin = open("/dev/stdin") + pdb.Pdb.interaction(self, *args, **kwargs) + finally: + sys.stdin = _stdin diff --git a/fastvideo/v1/utils/optimizer.py b/fastvideo/v1/utils/optimizer.py new file mode 100644 index 000000000..661f30379 --- /dev/null +++ b/fastvideo/v1/utils/optimizer.py @@ -0,0 +1,68 @@ +import torch +from accelerate.logging import get_logger + +logger = get_logger(__name__) + + +def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False): + # Optimizer creation + supported_optimizers = ["adam", "adamw", "prodigy"] + if args.optimizer not in supported_optimizers: + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not (args.optimizer.lower() not in ["adam", "adamw"]): + logger.warning(f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}") + + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError("To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.") + + if args.optimizer.lower() == "adamw": + optimizer_class = (bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW) + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + eps=args.adam_epsilon, + weight_decay=args.adam_weight_decay, + ) + elif args.optimizer.lower() == "adam": + optimizer_class = bnb.optim.Adam8bit if args.use_8bit_adam else torch.optim.Adam + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + eps=args.adam_epsilon, + weight_decay=args.adam_weight_decay, + ) + elif args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0") + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + + return optimizer diff --git a/fastvideo/v1/utils/parallel_states.py b/fastvideo/v1/utils/parallel_states.py new file mode 100644 index 000000000..c1f4cd751 --- /dev/null +++ b/fastvideo/v1/utils/parallel_states.py @@ -0,0 +1,63 @@ +import os + +import torch.distributed as dist + + +class COMM_INFO: + + def __init__(self): + self.group = None + self.sp_size = 1 + self.global_rank = 0 + self.rank_within_group = 0 + self.group_id = 0 + + +nccl_info = COMM_INFO() +_SEQUENCE_PARALLEL_STATE = False + + +def initialize_sequence_parallel_state(sequence_parallel_size): + global _SEQUENCE_PARALLEL_STATE + if sequence_parallel_size > 1: + _SEQUENCE_PARALLEL_STATE = True + initialize_sequence_parallel_group(sequence_parallel_size) + else: + nccl_info.sp_size = 1 + nccl_info.global_rank = int(os.getenv("RANK", "0")) + nccl_info.rank_within_group = 0 + nccl_info.group_id = int(os.getenv("RANK", "0")) + + +def set_sequence_parallel_state(state): + global _SEQUENCE_PARALLEL_STATE + _SEQUENCE_PARALLEL_STATE = state + + +def get_sequence_parallel_state(): + return _SEQUENCE_PARALLEL_STATE + + +def initialize_sequence_parallel_group(sequence_parallel_size): + """Initialize the sequence parallel group.""" + rank = int(os.getenv("RANK", "0")) + world_size = int(os.getenv("WORLD_SIZE", "1")) + assert ( + world_size % sequence_parallel_size == 0 + ), "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format( + world_size, sequence_parallel_size) + nccl_info.sp_size = sequence_parallel_size + nccl_info.global_rank = rank + num_sequence_parallel_groups: int = world_size // sequence_parallel_size + for i in range(num_sequence_parallel_groups): + ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) + group = dist.new_group(ranks) + if rank in ranks: + nccl_info.group = group + nccl_info.rank_within_group = rank - i * sequence_parallel_size + nccl_info.group_id = i + + +def destroy_sequence_parallel_group(): + """Destroy the sequence parallel group.""" + dist.destroy_process_group() diff --git a/fastvideo/utils/utils.py b/fastvideo/v1/utils/utils.py similarity index 98% rename from fastvideo/utils/utils.py rename to fastvideo/v1/utils/utils.py index 7c9896432..2b5c5ee68 100644 --- a/fastvideo/utils/utils.py +++ b/fastvideo/v1/utils/utils.py @@ -1,7 +1,7 @@ import torch -import fastvideo.envs as envs +import fastvideo.v1.envs as envs import inspect -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger import argparse import sys from typing import List, Dict, Union, Type, Any, TypeVar @@ -63,7 +63,7 @@ def current_stream() -> torch.cuda.Stream: the underlying hypothesis is that we do not call `torch._C._cuda_setStream` from C/C++ code. """ - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform global _current_stream if _current_stream is None: # when this function is called before any stream is set, diff --git a/fastvideo/v1/utils/validation.py b/fastvideo/v1/utils/validation.py new file mode 100644 index 000000000..6d815a6ab --- /dev/null +++ b/fastvideo/v1/utils/validation.py @@ -0,0 +1,308 @@ +# isort: skip_file +import gc +import os +from typing import List, Optional, Union + +import numpy as np +import torch +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import export_to_video +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from einops import rearrange +from tqdm import tqdm + +import wandb +from fastvideo.v1.distill.solver import PCMFMScheduler +from fastvideo.v1.models.mochi_hf.pipeline_mochi import (linear_quadratic_schedule, retrieve_timesteps) +from fastvideo.v1.utils.communications import all_gather +from fastvideo.v1.utils.load import load_vae +from fastvideo.v1.utils.parallel_states import (get_sequence_parallel_state, nccl_info) + + +def prepare_latents( + batch_size, + num_channels_latents, + height, + width, + num_frames, + dtype, + device, + generator, + vae_spatial_scale_factor, + vae_temporal_scale_factor, +): + height = height // vae_spatial_scale_factor + width = width // vae_spatial_scale_factor + num_frames = (num_frames - 1) // vae_temporal_scale_factor + 1 + + shape = (batch_size, num_channels_latents, num_frames, height, width) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + +def sample_validation_video( + model_type, + transformer, + vae, + scheduler, + scheduler_type="euler", + height: Optional[int] = None, + width: Optional[int] = None, + num_frames: int = 16, + num_inference_steps: int = 28, + timesteps: List[int] = None, + guidance_scale: float = 4.5, + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + vae_spatial_scale_factor=8, + vae_temporal_scale_factor=6, + num_channels_latents=12, +): + device = vae.device + + batch_size = prompt_embeds.shape[0] + + do_classifier_free_guidance = guidance_scale > 1.0 + if do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + # 4. Prepare latent variables + # TODO: Remove hardcore + latents = prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + prompt_embeds.dtype, + device, + generator, + vae_spatial_scale_factor, + vae_temporal_scale_factor, + ) + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + if get_sequence_parallel_state(): + latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous() + latents = latents[:, :, rank, :, :, :] + + # 5. Prepare timestep + # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77 + threshold_noise = 0.025 + sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise) + sigmas = np.array(sigmas) + if scheduler_type == "euler" and model_type == "mochi": #todo + timesteps, num_inference_steps = retrieve_timesteps( + scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + ) + else: + timesteps, num_inference_steps = retrieve_timesteps( + scheduler, + num_inference_steps, + device, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * scheduler.order, 0) + + # 6. Denoising loop + # with self.progress_bar(total=num_inference_steps) as progress_bar: + # write with tqdm instead + # only enable if nccl_info.global_rank == 0 + + with tqdm( + total=num_inference_steps, + disable=nccl_info.rank_within_group != 0, + desc="Validation sampling...", + ) as progress_bar: + for i, t in enumerate(timesteps): + latent_model_input = (torch.cat([latents] * 2) if do_classifier_free_guidance else latents) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + with torch.autocast("cuda", dtype=torch.bfloat16): + noise_pred = transformer( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + timestep=timestep, + encoder_attention_mask=prompt_attention_mask, + return_dict=False, + )[0] + + # Mochi CFG + Sampling runs in FP32 + noise_pred = noise_pred.to(torch.float32) + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0] + latents = latents.to(latents_dtype) + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0): + progress_bar.update() + + if get_sequence_parallel_state(): + latents = all_gather(latents, dim=2) + + if output_type == "latent": + video = latents + else: + # unscale/denormalize the latents + # denormalize with the mean and std if available and not None + has_latents_mean = (hasattr(vae.config, "latents_mean") and vae.config.latents_mean is not None) + has_latents_std = (hasattr(vae.config, "latents_std") and vae.config.latents_std is not None) + if has_latents_mean and has_latents_std: + latents_mean = (torch.tensor(vae.config.latents_mean).view(1, 12, 1, 1, + 1).to(latents.device, latents.dtype)) + latents_std = (torch.tensor(vae.config.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)) + latents = latents * latents_std / vae.config.scaling_factor + latents_mean + else: + latents = latents / vae.config.scaling_factor + with torch.autocast("cuda", dtype=vae.dtype): + video = vae.decode(latents, return_dict=False)[0] + video_processor = VideoProcessor(vae_scale_factor=vae_spatial_scale_factor) + video = video_processor.postprocess_video(video, output_type=output_type) + + return (video, ) + + +@torch.no_grad() +@torch.autocast("cuda", dtype=torch.bfloat16) +def log_validation( + args, + transformer, + device, + weight_dtype, # TODO + global_step, + scheduler_type="euler", + shift=1.0, + num_euler_timesteps=100, + linear_quadratic_threshold=0.025, + linear_range=0.5, + ema=False, +): + # TODO + print("Running validation....\n") + if args.model_type == "mochi": + vae_spatial_scale_factor = 8 + vae_temporal_scale_factor = 6 + num_channels_latents = 12 + elif args.model_type == "hunyuan" or "hunyuan_hf": + vae_spatial_scale_factor = 8 + vae_temporal_scale_factor = 4 + num_channels_latents = 16 + else: + raise ValueError(f"Model type {args.model_type} not supported") + vae, autocast_type, fps = load_vae(args.model_type, args.pretrained_model_name_or_path) + vae.enable_tiling() + if scheduler_type == "euler": + scheduler = FlowMatchEulerDiscreteScheduler(shift=shift) + else: + linear_quadraic = True if scheduler_type == "pcm_linear_quadratic" else False + scheduler = PCMFMScheduler( + 1000, + shift, + num_euler_timesteps, + linear_quadraic, + linear_quadratic_threshold, + linear_range, + ) + # args.validation_prompt_dir + + validation_guidance_scale_ls = args.validation_guidance_scale.split(",") + validation_guidance_scale_ls = [float(scale) for scale in validation_guidance_scale_ls] + for validation_sampling_step in args.validation_sampling_steps.split(","): + validation_sampling_step = int(validation_sampling_step) + for validation_guidance_scale in validation_guidance_scale_ls: + videos = [] + # prompt_embed are named embed0 to embedN + # check how many embeds are there + embe_dir = os.path.join(args.validation_prompt_dir, "prompt_embed") + mask_dir = os.path.join(args.validation_prompt_dir, "prompt_attention_mask") + embeds = sorted([f for f in os.listdir(embe_dir)]) + masks = sorted([f for f in os.listdir(mask_dir)]) + num_embeds = len(embeds) + validation_prompt_ids = list(range(num_embeds)) + num_sp_groups = int(os.getenv("WORLD_SIZE", "1")) // nccl_info.sp_size + # pad to multiple of groups + if num_embeds % num_sp_groups != 0: + validation_prompt_ids += [0] * (num_sp_groups - num_embeds % num_sp_groups) + num_embeds_per_group = len(validation_prompt_ids) // num_sp_groups + local_prompt_ids = validation_prompt_ids[nccl_info.group_id * + num_embeds_per_group:(nccl_info.group_id + 1) * + num_embeds_per_group] + + for i in local_prompt_ids: + prompt_embed_path = os.path.join(embe_dir, f"{embeds[i]}") + prompt_mask_path = os.path.join(mask_dir, f"{masks[i]}") + prompt_embeds = (torch.load(prompt_embed_path, map_location="cpu", + weights_only=True).to(device).unsqueeze(0)) + prompt_attention_mask = (torch.load(prompt_mask_path, map_location="cpu", + weights_only=True).to(device).unsqueeze(0)) + negative_prompt_embeds = torch.zeros(256, 4096).to(device).unsqueeze(0) + negative_prompt_attention_mask = (torch.zeros(256).bool().to(device).unsqueeze(0)) + generator = torch.Generator(device="cpu").manual_seed(12345) + video = sample_validation_video( + args.model_type, + transformer, + vae, + scheduler, + scheduler_type=scheduler_type, + num_frames=args.num_frames, + height=args.num_height, + width=args.num_width, + num_inference_steps=validation_sampling_step, + guidance_scale=validation_guidance_scale, + generator=generator, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_attention_mask=negative_prompt_attention_mask, + vae_spatial_scale_factor=vae_spatial_scale_factor, + vae_temporal_scale_factor=vae_temporal_scale_factor, + num_channels_latents=num_channels_latents, + )[0] + if nccl_info.rank_within_group == 0: + videos.append(video[0]) + # collect videos from all process to process zero + + gc.collect() + torch.cuda.empty_cache() + # log if main process + torch.distributed.barrier() + all_videos = [None for i in range(int(os.getenv("WORLD_SIZE", "1")))] # remove padded videos + torch.distributed.all_gather_object(all_videos, videos) + if nccl_info.global_rank == 0: + # remove padding + videos = [video for videos in all_videos for video in videos] + videos = videos[:num_embeds] + # linearize all videos + video_filenames = [] + for i, video in enumerate(videos): + filename = os.path.join( + args.output_dir, + f"validation_step_{global_step}_sample_{validation_sampling_step}_guidance_{validation_guidance_scale}_video_{i}.mp4", + ) + export_to_video(video, filename, fps=fps) + video_filenames.append(filename) + + logs = { + f"{'ema_' if ema else ''}validation_sample_{validation_sampling_step}_guidance_{validation_guidance_scale}": + [wandb.Video(filename) for i, filename in enumerate(video_filenames)] + } + wandb.log(logs, step=global_step) diff --git a/scripts/inference/inference_hunyuan.sh b/scripts/inference/inference_hunyuan.sh index d42d008b8..2bf74a5e7 100644 --- a/scripts/inference/inference_hunyuan.sh +++ b/scripts/inference/inference_hunyuan.sh @@ -7,34 +7,34 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ --height 720 \ --width 1280 \ --num_frames 125 \ - --num_inference_steps 1 \ + --num_inference_steps 6 \ --guidance_scale 1 \ --embedded_cfg_scale 6 \ --flow_shift 17 \ --flow-reverse \ --prompt ./assets/prompt.txt \ --seed 1024 \ - --output_path outputs_video/hunyuan/vae_sp_v0_1/ \ + --output_path outputs_video/hunyuan/vae_sp/ \ --model_path $MODEL_BASE \ --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ --vae-sp # Inference original Hunyuan -# num_gpus=4 -# export MODEL_BASE=data/hunyuan -# torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ -# fastvideo/sample/sample_t2v_hunyuan.py \ -# --height 720 \ -# --width 1280 \ -# --num_frames 125 \ -# --num_inference_steps 50 \ -# --guidance_scale 1 \ -# --embedded_cfg_scale 6 \ -# --flow_shift 7 \ -# --flow-reverse \ -# --prompt ./assets/prompt.txt \ -# --seed 1024 \ -# --output_path outputs_video/hunyuan/vae_sp/ \ -# --model_path $MODEL_BASE \ -# --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ -# --vae-sp +num_gpus=4 +export MODEL_BASE=data/hunyuan +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/sample/sample_t2v_hunyuan.py \ + --height 720 \ + --width 1280 \ + --num_frames 125 \ + --num_inference_steps 50 \ + --guidance_scale 1 \ + --embedded_cfg_scale 6 \ + --flow_shift 7 \ + --flow-reverse \ + --prompt ./assets/prompt.txt \ + --seed 1024 \ + --output_path outputs_video/hunyuan/vae_sp/ \ + --model_path $MODEL_BASE \ + --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ + --vae-sp diff --git a/scripts/inference/v1_inference_hunyuan.sh b/scripts/inference/v1_inference_hunyuan.sh index 5d3293c64..4d650b432 100644 --- a/scripts/inference/v1_inference_hunyuan.sh +++ b/scripts/inference/v1_inference_hunyuan.sh @@ -4,7 +4,7 @@ num_gpus=4 export MODEL_BASE=FastVideo/FastHunyuan-diffusers # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ - fastvideo/sample/v1_fastvideo_inference.py \ + fastvideo/v1/sample/v1_fastvideo_inference.py \ --sp_size 4 \ --height 720 \ --width 1280 \ From bf1fc279892a93f203b02fec2a99a6c6cb670eb1 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Fri, 14 Mar 2025 19:17:53 +0000 Subject: [PATCH 03/12] remove unneeded file --- fastvideo/sample/new_sample_hunyuan.py | 103 ------------------------- 1 file changed, 103 deletions(-) delete mode 100644 fastvideo/sample/new_sample_hunyuan.py diff --git a/fastvideo/sample/new_sample_hunyuan.py b/fastvideo/sample/new_sample_hunyuan.py deleted file mode 100644 index 0ac839bea..000000000 --- a/fastvideo/sample/new_sample_hunyuan.py +++ /dev/null @@ -1,103 +0,0 @@ -import argparse -import os -from pathlib import Path - -import imageio -import numpy as np -import torch -import torchvision -from einops import rearrange - -# Fix the import path -from fastvideo.inference import DiffusionInference -from fastvideo.inference_args import InferenceArgs -from fastvideo.utils.utils import FlexibleArgumentParser -from fastvideo.distributed import init_distributed_environment, initialize_model_parallel -from fastvideo.logger import init_logger -import torch.distributed as dist -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state -logger = init_logger(__name__) - -def initialize_distributed(): - local_rank = int(os.getenv("RANK", 0)) - world_size = int(os.getenv("WORLD_SIZE", 1)) - print("world_size", world_size) - torch.cuda.set_device(local_rank) - dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank) - initialize_sequence_parallel_state(world_size) - -def main(inference_args: InferenceArgs): - # initialize_distributed() - # print(nccl_info.sp_size) - # local_rank = int(os.environ.get("LOCAL_RANK", 0)) - # rank = int(os.environ.get("RANK", 0)) - # world_size = int(os.environ.get("WORLD_SIZE", 1)) - - # logger.info(f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}") - # init_distributed_environment( - # world_size=world_size, - # rank=rank, - # local_rank=local_rank - # ) - # print(inference_args.sp_size) - - # Initialize tensor model parallel groups - # initialize_model_parallel( - # sequence_model_parallel_size=inference_args.sp_size - # ) - initialize_distributed() - - - print('Loading pipeline') - # Create inference object using the updated API - diffusion = DiffusionInference.load_pipeline( - # model_path=models_root_path, - inference_args, - # model_loader_cls=ModelLoader, - # pipeline_cls=HunyuanVideoPipeline - ) - print('Pipeline loaded') - # return - - # Load prompts - if inference_args.prompt.endswith('.txt'): - with open(inference_args.prompt) as f: - prompts = [line.strip() for line in f.readlines()] - else: - prompts = [inference_args.prompt] - - # Process each prompt - for prompt in prompts: - outputs = diffusion.predict( - prompt=prompt, - inference_args=inference_args, - ) - - # Process outputs - videos = rearrange(outputs["samples"], "b c t h w -> t b c h w") - frames = [] - for x in videos: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) - - # Save video - os.makedirs(os.path.dirname(inference_args.output_path), exist_ok=True) - imageio.mimsave( - os.path.join(inference_args.output_path, f"{prompt[:100]}.mp4"), - frames, - fps=inference_args.fps - ) - - -if __name__ == "__main__": - parser = FlexibleArgumentParser() - InferenceArgs.add_cli_args(parser) - args = parser.parse_args() - inference_args = InferenceArgs.from_cli_args(args) - - # Validate arguments - if inference_args.vae_sp and not inference_args.vae_tiling: - raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.") - - main(inference_args) From eafeea4a3fad6cd9a9ff9d90cda19e6809e7c87e Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Fri, 14 Mar 2025 19:19:35 +0000 Subject: [PATCH 04/12] revert pyproject.toml --- pyproject.toml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ac44c82c0..c59afaa69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,14 +18,14 @@ dependencies = [ "scipy==1.14.1", "six==1.16.0", "h5py==3.12.1", # Machine Learning & Transformers - "transformers", "tokenizers", "sentencepiece==0.2.0", + "transformers==4.46.1", "tokenizers==0.20.1", "sentencepiece==0.2.0", "timm==1.0.11", "peft==0.13.2", "diffusers==0.32.0", "bitsandbytes", # Acceleration & Optimization "accelerate==1.0.1", "torch", # Computer Vision & Image Processing - "opencv-python==4.10.0.84", "pillow", "imageio==2.36.0", + "opencv-python==4.10.0.84", "pillow==10.2.0", "imageio==2.36.0", "imageio-ffmpeg==0.5.1", "decord==0.6.0", "einops", # Experiment Tracking & Logging @@ -33,7 +33,7 @@ dependencies = [ # Miscellaneous Utilities "tqdm==4.66.5", "PyYAML==6.0.1", "idna==3.6", "protobuf==5.28.3", - "gradio", "huggingface_hub==0.26.1", "moviepy==1.0.3", "flask", "flask_restful", "aiohttp", + "gradio==5.3.0", "huggingface_hub==0.26.1", "moviepy==1.0.3", "flask", "flask_restful", "aiohttp", # System & Monitoring Tools "gpustat", "watch", @@ -76,4 +76,4 @@ skip ="./data,./wandb,./csrc/sliding_tile_attention/tk" column_limit = 120 [tool.isort] -line_length = 120 +line_length = 120 \ No newline at end of file From 1976b23121139ffba817bffbe7c2b516857e7ded Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Fri, 14 Mar 2025 19:29:29 +0000 Subject: [PATCH 05/12] move v1's v0 code into v1/v0_reference_src --- fastvideo/v1/{pipelines => v0_reference_src}/loader.py | 0 fastvideo/v1/{ => v0_reference_src}/models/hunyuan/__init__.py | 0 fastvideo/v1/{ => v0_reference_src}/models/hunyuan/constants.py | 0 .../{ => v0_reference_src}/models/hunyuan/diffusion/__init__.py | 0 .../models/hunyuan/diffusion/pipelines/__init__.py | 0 .../models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py | 0 .../models/hunyuan/diffusion/schedulers/__init__.py | 0 .../diffusion/schedulers/scheduling_flow_match_discrete.py | 0 fastvideo/v1/{ => v0_reference_src}/models/hunyuan/idle_config.py | 0 fastvideo/v1/{ => v0_reference_src}/models/hunyuan/inference.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/modules/__init__.py | 0 .../models/hunyuan/modules/activation_layers.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/modules/attenion.py | 0 .../{ => v0_reference_src}/models/hunyuan/modules/embed_layers.py | 0 .../{ => v0_reference_src}/models/hunyuan/modules/mlp_layers.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/modules/models.py | 0 .../models/hunyuan/modules/modulate_layers.py | 0 .../{ => v0_reference_src}/models/hunyuan/modules/norm_layers.py | 0 .../models/hunyuan/modules/posemb_layers.py | 0 .../models/hunyuan/modules/token_refiner.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/prompt_rewrite.py | 0 .../models/hunyuan/text_encoder/__init__.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/utils/__init__.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/utils/data_utils.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/utils/file_utils.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/utils/helpers.py | 0 .../hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py | 0 .../v1/{ => v0_reference_src}/models/hunyuan/vae/__init__.py | 0 .../models/hunyuan/vae/autoencoder_kl_causal_3d.py | 0 .../models/hunyuan/vae/unet_causal_3d_blocks.py | 0 fastvideo/v1/{ => v0_reference_src}/models/hunyuan/vae/vae.py | 0 fastvideo/v1/{pipelines => v0_reference_src}/pipeline_base.py | 0 .../v1/{ => v0_reference_src}/pipelines/hunyuan/constants.py | 0 .../v1/{ => v0_reference_src}/pipelines/hunyuan/hunyuan_loader.py | 0 .../v1/{ => v0_reference_src}/pipelines/hunyuan/hunyuan_video.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/checkpoint.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/communications.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/dataset_utils.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/env_utils.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/fsdp_util.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/load.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/logging_.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/optimizer.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/parallel_states.py | 0 fastvideo/v1/{ => v0_reference_src}/utils/validation.py | 0 45 files changed, 0 insertions(+), 0 deletions(-) rename fastvideo/v1/{pipelines => v0_reference_src}/loader.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/__init__.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/constants.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/diffusion/__init__.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/diffusion/pipelines/__init__.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/diffusion/schedulers/__init__.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/idle_config.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/inference.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/__init__.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/activation_layers.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/attenion.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/embed_layers.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/mlp_layers.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/models.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/modulate_layers.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/norm_layers.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/posemb_layers.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/modules/token_refiner.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/prompt_rewrite.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/text_encoder/__init__.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/utils/__init__.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/utils/data_utils.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/utils/file_utils.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/utils/helpers.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/vae/__init__.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/vae/autoencoder_kl_causal_3d.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/vae/unet_causal_3d_blocks.py (100%) rename fastvideo/v1/{ => v0_reference_src}/models/hunyuan/vae/vae.py (100%) rename fastvideo/v1/{pipelines => v0_reference_src}/pipeline_base.py (100%) rename fastvideo/v1/{ => v0_reference_src}/pipelines/hunyuan/constants.py (100%) rename fastvideo/v1/{ => v0_reference_src}/pipelines/hunyuan/hunyuan_loader.py (100%) rename fastvideo/v1/{ => v0_reference_src}/pipelines/hunyuan/hunyuan_video.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/checkpoint.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/communications.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/dataset_utils.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/env_utils.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/fsdp_util.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/load.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/logging_.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/optimizer.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/parallel_states.py (100%) rename fastvideo/v1/{ => v0_reference_src}/utils/validation.py (100%) diff --git a/fastvideo/v1/pipelines/loader.py b/fastvideo/v1/v0_reference_src/loader.py similarity index 100% rename from fastvideo/v1/pipelines/loader.py rename to fastvideo/v1/v0_reference_src/loader.py diff --git a/fastvideo/v1/models/hunyuan/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/__init__.py similarity index 100% rename from fastvideo/v1/models/hunyuan/__init__.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/__init__.py diff --git a/fastvideo/v1/models/hunyuan/constants.py b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py similarity index 100% rename from fastvideo/v1/models/hunyuan/constants.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/constants.py diff --git a/fastvideo/v1/models/hunyuan/diffusion/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/__init__.py similarity index 100% rename from fastvideo/v1/models/hunyuan/diffusion/__init__.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/__init__.py diff --git a/fastvideo/v1/models/hunyuan/diffusion/pipelines/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/__init__.py similarity index 100% rename from fastvideo/v1/models/hunyuan/diffusion/pipelines/__init__.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/__init__.py diff --git a/fastvideo/v1/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py similarity index 100% rename from fastvideo/v1/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py diff --git a/fastvideo/v1/models/hunyuan/diffusion/schedulers/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/__init__.py similarity index 100% rename from fastvideo/v1/models/hunyuan/diffusion/schedulers/__init__.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/__init__.py diff --git a/fastvideo/v1/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py similarity index 100% rename from fastvideo/v1/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py diff --git a/fastvideo/v1/models/hunyuan/idle_config.py b/fastvideo/v1/v0_reference_src/models/hunyuan/idle_config.py similarity index 100% rename from fastvideo/v1/models/hunyuan/idle_config.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/idle_config.py diff --git a/fastvideo/v1/models/hunyuan/inference.py b/fastvideo/v1/v0_reference_src/models/hunyuan/inference.py similarity index 100% rename from fastvideo/v1/models/hunyuan/inference.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/inference.py diff --git a/fastvideo/v1/models/hunyuan/modules/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/__init__.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py diff --git a/fastvideo/v1/models/hunyuan/modules/activation_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/activation_layers.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/activation_layers.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/activation_layers.py diff --git a/fastvideo/v1/models/hunyuan/modules/attenion.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/attenion.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py diff --git a/fastvideo/v1/models/hunyuan/modules/embed_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/embed_layers.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/embed_layers.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/embed_layers.py diff --git a/fastvideo/v1/models/hunyuan/modules/mlp_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/mlp_layers.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/mlp_layers.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/mlp_layers.py diff --git a/fastvideo/v1/models/hunyuan/modules/models.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/models.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py diff --git a/fastvideo/v1/models/hunyuan/modules/modulate_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/modulate_layers.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/modulate_layers.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/modulate_layers.py diff --git a/fastvideo/v1/models/hunyuan/modules/norm_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/norm_layers.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/norm_layers.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/norm_layers.py diff --git a/fastvideo/v1/models/hunyuan/modules/posemb_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/posemb_layers.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/posemb_layers.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/posemb_layers.py diff --git a/fastvideo/v1/models/hunyuan/modules/token_refiner.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/token_refiner.py similarity index 100% rename from fastvideo/v1/models/hunyuan/modules/token_refiner.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/modules/token_refiner.py diff --git a/fastvideo/v1/models/hunyuan/prompt_rewrite.py b/fastvideo/v1/v0_reference_src/models/hunyuan/prompt_rewrite.py similarity index 100% rename from fastvideo/v1/models/hunyuan/prompt_rewrite.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/prompt_rewrite.py diff --git a/fastvideo/v1/models/hunyuan/text_encoder/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py similarity index 100% rename from fastvideo/v1/models/hunyuan/text_encoder/__init__.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py diff --git a/fastvideo/v1/models/hunyuan/utils/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/__init__.py similarity index 100% rename from fastvideo/v1/models/hunyuan/utils/__init__.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/utils/__init__.py diff --git a/fastvideo/v1/models/hunyuan/utils/data_utils.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/data_utils.py similarity index 100% rename from fastvideo/v1/models/hunyuan/utils/data_utils.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/utils/data_utils.py diff --git a/fastvideo/v1/models/hunyuan/utils/file_utils.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/file_utils.py similarity index 100% rename from fastvideo/v1/models/hunyuan/utils/file_utils.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/utils/file_utils.py diff --git a/fastvideo/v1/models/hunyuan/utils/helpers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/helpers.py similarity index 100% rename from fastvideo/v1/models/hunyuan/utils/helpers.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/utils/helpers.py diff --git a/fastvideo/v1/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py similarity index 100% rename from fastvideo/v1/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py diff --git a/fastvideo/v1/models/hunyuan/vae/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py similarity index 100% rename from fastvideo/v1/models/hunyuan/vae/__init__.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py diff --git a/fastvideo/v1/models/hunyuan/vae/autoencoder_kl_causal_3d.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py similarity index 100% rename from fastvideo/v1/models/hunyuan/vae/autoencoder_kl_causal_3d.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py diff --git a/fastvideo/v1/models/hunyuan/vae/unet_causal_3d_blocks.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py similarity index 100% rename from fastvideo/v1/models/hunyuan/vae/unet_causal_3d_blocks.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py diff --git a/fastvideo/v1/models/hunyuan/vae/vae.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/vae.py similarity index 100% rename from fastvideo/v1/models/hunyuan/vae/vae.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/vae/vae.py diff --git a/fastvideo/v1/pipelines/pipeline_base.py b/fastvideo/v1/v0_reference_src/pipeline_base.py similarity index 100% rename from fastvideo/v1/pipelines/pipeline_base.py rename to fastvideo/v1/v0_reference_src/pipeline_base.py diff --git a/fastvideo/v1/pipelines/hunyuan/constants.py b/fastvideo/v1/v0_reference_src/pipelines/hunyuan/constants.py similarity index 100% rename from fastvideo/v1/pipelines/hunyuan/constants.py rename to fastvideo/v1/v0_reference_src/pipelines/hunyuan/constants.py diff --git a/fastvideo/v1/pipelines/hunyuan/hunyuan_loader.py b/fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_loader.py similarity index 100% rename from fastvideo/v1/pipelines/hunyuan/hunyuan_loader.py rename to fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_loader.py diff --git a/fastvideo/v1/pipelines/hunyuan/hunyuan_video.py b/fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_video.py similarity index 100% rename from fastvideo/v1/pipelines/hunyuan/hunyuan_video.py rename to fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_video.py diff --git a/fastvideo/v1/utils/checkpoint.py b/fastvideo/v1/v0_reference_src/utils/checkpoint.py similarity index 100% rename from fastvideo/v1/utils/checkpoint.py rename to fastvideo/v1/v0_reference_src/utils/checkpoint.py diff --git a/fastvideo/v1/utils/communications.py b/fastvideo/v1/v0_reference_src/utils/communications.py similarity index 100% rename from fastvideo/v1/utils/communications.py rename to fastvideo/v1/v0_reference_src/utils/communications.py diff --git a/fastvideo/v1/utils/dataset_utils.py b/fastvideo/v1/v0_reference_src/utils/dataset_utils.py similarity index 100% rename from fastvideo/v1/utils/dataset_utils.py rename to fastvideo/v1/v0_reference_src/utils/dataset_utils.py diff --git a/fastvideo/v1/utils/env_utils.py b/fastvideo/v1/v0_reference_src/utils/env_utils.py similarity index 100% rename from fastvideo/v1/utils/env_utils.py rename to fastvideo/v1/v0_reference_src/utils/env_utils.py diff --git a/fastvideo/v1/utils/fsdp_util.py b/fastvideo/v1/v0_reference_src/utils/fsdp_util.py similarity index 100% rename from fastvideo/v1/utils/fsdp_util.py rename to fastvideo/v1/v0_reference_src/utils/fsdp_util.py diff --git a/fastvideo/v1/utils/load.py b/fastvideo/v1/v0_reference_src/utils/load.py similarity index 100% rename from fastvideo/v1/utils/load.py rename to fastvideo/v1/v0_reference_src/utils/load.py diff --git a/fastvideo/v1/utils/logging_.py b/fastvideo/v1/v0_reference_src/utils/logging_.py similarity index 100% rename from fastvideo/v1/utils/logging_.py rename to fastvideo/v1/v0_reference_src/utils/logging_.py diff --git a/fastvideo/v1/utils/optimizer.py b/fastvideo/v1/v0_reference_src/utils/optimizer.py similarity index 100% rename from fastvideo/v1/utils/optimizer.py rename to fastvideo/v1/v0_reference_src/utils/optimizer.py diff --git a/fastvideo/v1/utils/parallel_states.py b/fastvideo/v1/v0_reference_src/utils/parallel_states.py similarity index 100% rename from fastvideo/v1/utils/parallel_states.py rename to fastvideo/v1/v0_reference_src/utils/parallel_states.py diff --git a/fastvideo/v1/utils/validation.py b/fastvideo/v1/v0_reference_src/utils/validation.py similarity index 100% rename from fastvideo/v1/utils/validation.py rename to fastvideo/v1/v0_reference_src/utils/validation.py From d5ac1e9beee736c3193dfadcd61b5dca5c31e98e Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Fri, 14 Mar 2025 19:31:38 +0000 Subject: [PATCH 06/12] remove unused attention --- fastvideo/v1/attention/backends/__init__.py | 0 fastvideo/v1/attention/backends/abstract.py | 300 ----- fastvideo/v1/attention/backends/flash_attn.py | 1042 ----------------- fastvideo/v1/attention/backends/utils.py | 619 ---------- fastvideo/v1/attention/layer.py | 135 --- 5 files changed, 2096 deletions(-) delete mode 100644 fastvideo/v1/attention/backends/__init__.py delete mode 100644 fastvideo/v1/attention/backends/abstract.py delete mode 100755 fastvideo/v1/attention/backends/flash_attn.py delete mode 100644 fastvideo/v1/attention/backends/utils.py delete mode 100644 fastvideo/v1/attention/layer.py diff --git a/fastvideo/v1/attention/backends/__init__.py b/fastvideo/v1/attention/backends/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/fastvideo/v1/attention/backends/abstract.py b/fastvideo/v1/attention/backends/abstract.py deleted file mode 100644 index 0cd95e074..000000000 --- a/fastvideo/v1/attention/backends/abstract.py +++ /dev/null @@ -1,300 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from abc import ABC, abstractmethod -from contextlib import contextmanager -from dataclasses import dataclass, fields -from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, - Protocol, Set, Tuple, Type, TypeVar) - -import torch - -from vllm.multimodal import MultiModalPlaceholderMap - -if TYPE_CHECKING: - from vllm.worker.model_runner_base import (ModelRunnerBase, - ModelRunnerInputBase, - ModelRunnerInputBuilderBase) - - -class AttentionType: - """ - Attention type. - Use string to be compatible with `torch.compile`. - """ - # Decoder attention between previous layer Q/K/V - DECODER = "decoder" - # Encoder attention between previous layer Q/K/V for encoder-decoder - ENCODER = "encoder" - # Encoder attention between previous layer Q/K/V - ENCODER_ONLY = "encoder_only" - # Attention between dec. Q and enc. K/V for encoder-decoder - ENCODER_DECODER = "encoder_decoder" - - -class AttentionBackend(ABC): - """Abstract class for attention backends.""" - # For some attention backends, we allocate an output tensor before - # calling the custom op. When piecewise cudagraph is enabled, this - # makes sure the output tensor is allocated inside the cudagraph. - accept_output_buffer: bool = False - - @staticmethod - @abstractmethod - def get_name() -> str: - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_impl_cls() -> Type["AttentionImpl"]: - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_metadata_cls() -> Type["AttentionMetadata"]: - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_state_cls() -> Type["AttentionState"]: - raise NotImplementedError - - @classmethod - def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata": - return cls.get_metadata_cls()(*args, **kwargs) - - @staticmethod - @abstractmethod - def get_builder_cls() -> Type["AttentionMetadataBuilder"]: - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - ) -> Tuple[int, ...]: - raise NotImplementedError - - @staticmethod - @abstractmethod - def swap_blocks( - src_kv_cache: torch.Tensor, - dst_kv_cache: torch.Tensor, - src_to_dst: torch.Tensor, - ) -> None: - raise NotImplementedError - - @staticmethod - @abstractmethod - def copy_blocks( - kv_caches: List[torch.Tensor], - src_to_dists: torch.Tensor, - ) -> None: - raise NotImplementedError - - def advance_step(self, model_input: "ModelRunnerInputBase", - sampled_token_ids: Optional[torch.Tensor], - block_size: int, num_seqs: int, num_queries: int) -> None: - raise NotImplementedError - - -@dataclass -class AttentionMetadata: - """Attention metadata for prefill and decode batched together.""" - # Total number of prefill requests. - num_prefills: int - # Number of prefill tokens. - num_prefill_tokens: int - # Number of decode tokens. Note that it is equivalent to the number of - # decode requests. - num_decode_tokens: int - # (num_tokens,). The indices of the token slots that input tokens will be - # stored into. E.g., if `slot_mapping` is [35, 2, 17] and the block size - # is 16, the three tokens are stored in the 3rd slot in block 2, 2nd slot - # in block 0, and 1st slot in block 1, respectively. - slot_mapping: torch.Tensor - - # The index maps that relate multi-modal embeddings to the corresponding - # placeholders. - # - # N.B. These aren't really related to attention and don't belong on this - # type -- this is just a temporary solution to make them available to - # `model_executable`. - multi_modal_placeholder_index_maps: Optional[Dict[ - str, MultiModalPlaceholderMap.IndexMap]] - - # Enable/disable KV scales calculation. This is so that we can disable the - # calculation until after prefill and cuda graph capture. - enable_kv_scales_calculation: bool - - @property - @abstractmethod - def prefill_metadata(self) -> Optional["AttentionMetadata"]: - """Return the attention metadata that's required to run prefill - attention.""" - pass - - @property - @abstractmethod - def decode_metadata(self) -> Optional["AttentionMetadata"]: - """Return the attention metadata that's required to run decode - attention.""" - pass - - def asdict_zerocopy(self, - skip_fields: Optional[Set[str]] = None - ) -> Dict[str, Any]: - """Similar to dataclasses.asdict, but avoids deepcopying.""" - if skip_fields is None: - skip_fields = set() - # Note that if we add dataclasses as fields, they will need - # similar handling. - return { - field.name: getattr(self, field.name) - for field in fields(self) if field.name not in skip_fields - } - - -T = TypeVar("T", bound=AttentionMetadata) - - -class AttentionState(ABC, Generic[T]): - """Holds attention backend-specific objects reused during the - lifetime of the model runner.""" - - @abstractmethod - def __init__(self, runner: "ModelRunnerBase"): - ... - - @abstractmethod - @contextmanager - def graph_capture(self, max_batch_size: int): - """Context manager used when capturing CUDA graphs.""" - yield - - @abstractmethod - def graph_clone(self, batch_size: int) -> "AttentionState[T]": - """Clone attention state to save in CUDA graph metadata.""" - ... - - @abstractmethod - def graph_capture_get_metadata_for_batch( - self, - batch_size: int, - is_encoder_decoder_model: bool = False) -> T: - """Get attention metadata for CUDA graph capture of batch_size.""" - ... - - @abstractmethod - def get_graph_input_buffers( - self, - attn_metadata: T, - is_encoder_decoder_model: bool = False) -> Dict[str, Any]: - """Get attention-specific input buffers for CUDA graph capture.""" - ... - - @abstractmethod - def prepare_graph_input_buffers( - self, - input_buffers: Dict[str, Any], - attn_metadata: T, - is_encoder_decoder_model: bool = False) -> None: - """In-place modify input buffers dict for CUDA graph replay.""" - ... - - @abstractmethod - def begin_forward(self, model_input: "ModelRunnerInputBase") -> None: - """Prepare state for forward pass.""" - ... - - -class AttentionMetadataBuilder(ABC, Generic[T]): - """Abstract class for attention metadata builders.""" - - @abstractmethod - def __init__(self, input_builder: "ModelRunnerInputBuilderBase") -> None: - """Create the builder, remember some configuration and parameters.""" - raise NotImplementedError - - @abstractmethod - def prepare(self) -> None: - """Prepare for one batch.""" - raise NotImplementedError - - @abstractmethod - def build(self, seq_lens: List[int], query_lens: List[int], - cuda_graph_pad_size: int, batch_size: int) -> T: - """Build attention metadata with on-device tensors.""" - raise NotImplementedError - - -class AttentionLayer(Protocol): - - _k_scale: torch.Tensor - _v_scale: torch.Tensor - _k_scale_float: float - _v_scale_float: float - - def forward( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> torch.Tensor: - ... - - -class AttentionImpl(ABC, Generic[T]): - - @abstractmethod - def __init__( - self, - num_heads: int, - head_size: int, - scale: float, - num_kv_heads: Optional[int] = None, - alibi_slopes: Optional[List[float]] = None, - sliding_window: Optional[int] = None, - kv_cache_dtype: str = "auto", - blocksparse_params: Optional[Dict[str, Any]] = None, - logits_soft_cap: Optional[float] = None, - attn_type: str = AttentionType.DECODER, - ) -> None: - raise NotImplementedError - - @abstractmethod - def forward( - self, - layer: AttentionLayer, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: T, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - raise NotImplementedError - - -class MLAAttentionImpl(AttentionImpl[T], Generic[T]): - - @abstractmethod - def forward( - self, - layer: AttentionLayer, - hidden_states_or_cq: torch.Tensor, - kv_c_normed: torch.Tensor, - k_pe: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: T, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - raise NotImplementedError - - -def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: - return kv_cache_dtype != "auto" diff --git a/fastvideo/v1/attention/backends/flash_attn.py b/fastvideo/v1/attention/backends/flash_attn.py deleted file mode 100755 index c9c1e0c35..000000000 --- a/fastvideo/v1/attention/backends/flash_attn.py +++ /dev/null @@ -1,1042 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Attention layer with FlashAttention.""" -from collections import defaultdict -from dataclasses import dataclass -from itertools import accumulate -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type - -import torch -import math - -from vllm import _custom_ops as ops -# yapf conflicts with isort for this block -# yapf: disable -from fastvideo.v1.attention.backends.abstract import (AttentionBackend, AttentionImpl, - AttentionLayer, - AttentionMetadata, - AttentionMetadataBuilder, - AttentionType, - is_quantized_kv_cache) -# yapf: enable -from fastvideo.v1.attention.backends.utils import ( - PAD_SLOT_ID, CommonAttentionState, compute_slot_mapping, - compute_slot_mapping_start_idx, get_flash_attn_version, - get_num_prefill_decode_query_kv_tokens, get_seq_len_block_table_args, - is_all_cross_attn_metadata_set, is_all_encoder_attn_metadata_set, - is_block_tables_empty) -from fastvideo.v1.logger import init_logger -# from vllm.multimodal import MultiModalPlaceholderMap -from vllm.utils import async_tensor_h2d, make_tensor_with_pad -from vllm.vllm_flash_attn import (flash_attn_varlen_func, - flash_attn_with_kvcache) - -if TYPE_CHECKING: - from vllm.worker.model_runner import (ModelInputForGPUBuilder, - ModelInputForGPUWithSamplingMetadata) - -logger = init_logger(__name__) - - -class FlashAttentionBackend(AttentionBackend): - - accept_output_buffer: bool = True - - @staticmethod - def get_supported_head_sizes() -> List[int]: - return [32, 64, 96, 128, 160, 192, 224, 256] - - @staticmethod - def get_name() -> str: - return "FLASH_ATTN" - - @staticmethod - def get_impl_cls() -> Type["FlashAttentionImpl"]: - return FlashAttentionImpl - - @staticmethod - def get_metadata_cls() -> Type["AttentionMetadata"]: - return FlashAttentionMetadata - - @staticmethod - def get_builder_cls() -> Type["FlashAttentionMetadataBuilder"]: - return FlashAttentionMetadataBuilder - - @staticmethod - def get_state_cls() -> Type["CommonAttentionState"]: - return CommonAttentionState - - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - ) -> Tuple[int, ...]: - if block_size % 16 != 0: - raise ValueError("Block size must be a multiple of 16.") - return (2, num_blocks, block_size, num_kv_heads, head_size) - - @staticmethod - def swap_blocks( - src_kv_cache: torch.Tensor, - dst_kv_cache: torch.Tensor, - src_to_dst: torch.Tensor, - ) -> None: - src_key_cache = src_kv_cache[0] - dst_key_cache = dst_kv_cache[0] - ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst) - src_value_cache = src_kv_cache[1] - dst_value_cache = dst_kv_cache[1] - ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst) - - @staticmethod - def copy_blocks( - kv_caches: List[torch.Tensor], - src_to_dists: torch.Tensor, - ) -> None: - key_caches = [kv_cache[0] for kv_cache in kv_caches] - value_caches = [kv_cache[1] for kv_cache in kv_caches] - - ops.copy_blocks(key_caches, value_caches, src_to_dists) - - -@dataclass -class FlashAttentionMetadata(AttentionMetadata): - """Metadata for FlashAttentionBackend. - - NOTE: Any python object stored here is not updated when it is - cuda-graph replayed. If you have values that need to be changed - dynamically, it should be stored in tensor. The tensor has to be - updated from `CUDAGraphRunner.forward` API. - """ - # (batch_size,). The sequence length per sequence. Sequence length means - # the computed tokens + new tokens None if it is a decoding. - seq_lens: Optional[List[int]] - # seq_lens stored as a tensor. - seq_lens_tensor: Optional[torch.Tensor] - - # NOTE(sang): Definition of context_len, query_len, and seq_len. - # |---------- N-1 iteration --------| - # |---------------- N iteration ---------------------| - # |- tokenA -|......................|-- newTokens ---| - # |---------- context_len ----------| - # |-------------------- seq_len ---------------------| - # |-- query_len ---| - - # Maximum sequence length among prefill batch. 0 if there are decoding - # requests only. - max_prefill_seq_len: int - # Maximum sequence length among decode batch. 0 if there are prefill - # requests only. - max_decode_seq_len: int - # (batch_size,) A tensor of context lengths (tokens that are computed - # so far). - context_lens_tensor: Optional[torch.Tensor] - - # (batch_size, max_blocks_per_seq). - # Block addresses per sequence. (Seq id -> list of physical block) - # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks - # in the kv cache. Each block can contain up to block_size tokens. - # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph - # captured. - block_tables: Optional[torch.Tensor] - - # Whether or not if cuda graph is enabled. - # Cuda-graph is currently enabled for decoding only. - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention. - - use_cuda_graph: bool - - # Maximum query length in the batch. - max_query_len: Optional[int] = None - - # Max number of query tokens among request in the batch. - max_decode_query_len: Optional[int] = None - - # (batch_size + 1,). The cumulative subquery lengths of the sequences in - # the batch, used to index into subquery. E.g., if the subquery length - # is [4, 6], it is [0, 4, 10]. - query_start_loc: Optional[torch.Tensor] = None - # (batch_size + 1,). The cumulative sequence lengths of the sequences in - # the batch, used to index into sequence. E.g., if the sequence length is - # [4, 6], it is [0, 4, 10]. - seq_start_loc: Optional[torch.Tensor] = None - - _cached_prefill_metadata: Optional["FlashAttentionMetadata"] = None - _cached_decode_metadata: Optional["FlashAttentionMetadata"] = None - - # Begin encoder attn & enc/dec cross-attn fields... - - # Encoder sequence lengths representation - encoder_seq_lens: Optional[List[int]] = None - encoder_seq_lens_tensor: Optional[torch.Tensor] = None - # (batch_size + 1,). The cumulative sequence lengths of the sequences in - # the batch, used to index into sequence. E.g., if the sequence length is - # [4, 6], it is [0, 4, 10]. - encoder_seq_start_loc: Optional[torch.Tensor] = None - # Maximum sequence length among encoder sequences - max_encoder_seq_len: Optional[int] = None - # Number of tokens input to encoder - num_encoder_tokens: Optional[int] = None - - # Cross-attention memory-mapping data structures: slot mapping - # and block tables - cross_slot_mapping: Optional[torch.Tensor] = None - cross_block_tables: Optional[torch.Tensor] = None - - @property - def is_all_encoder_attn_metadata_set(self): - ''' - All attention metadata required for encoder attention is set. - ''' - return is_all_encoder_attn_metadata_set(self) - - @property - def is_all_cross_attn_metadata_set(self): - ''' - All attention metadata required for enc/dec cross-attention is set. - - Superset of encoder attention required metadata. - ''' - return is_all_cross_attn_metadata_set(self) - - @property - def prefill_metadata(self) -> Optional["FlashAttentionMetadata"]: - if self.num_prefills == 0: - return None - - if self._cached_prefill_metadata is not None: - return self._cached_prefill_metadata - - assert ((self.seq_lens is not None) - or (self.encoder_seq_lens is not None)) - assert ((self.seq_lens_tensor is not None) - or (self.encoder_seq_lens_tensor is not None)) - - # Compute some attn_metadata fields which default to None - query_start_loc = (None if self.query_start_loc is None else - self.query_start_loc[:self.num_prefills + 1]) - slot_mapping = (None if self.slot_mapping is None else - self.slot_mapping[:self.num_prefill_tokens]) - seq_lens = (None if self.seq_lens is None else - self.seq_lens[:self.num_prefills]) - seq_lens_tensor = (None if self.seq_lens_tensor is None else - self.seq_lens_tensor[:self.num_prefills]) - seq_start_loc = (None if self.seq_start_loc is None else - self.seq_start_loc[:self.num_prefills + 1]) - context_lens_tensor = (None if self.context_lens_tensor is None else - self.context_lens_tensor[:self.num_prefills]) - block_tables = (None if self.block_tables is None else - self.block_tables[:self.num_prefills]) - - self._cached_prefill_metadata = FlashAttentionMetadata( - num_prefills=self.num_prefills, - num_prefill_tokens=self.num_prefill_tokens, - num_decode_tokens=0, - slot_mapping=slot_mapping, - multi_modal_placeholder_index_maps=self. - multi_modal_placeholder_index_maps, - enable_kv_scales_calculation=self.enable_kv_scales_calculation, - seq_lens=seq_lens, - seq_lens_tensor=seq_lens_tensor, - max_query_len=self.max_query_len, - max_prefill_seq_len=self.max_prefill_seq_len, - max_decode_query_len=0, - max_decode_seq_len=0, - query_start_loc=query_start_loc, - seq_start_loc=seq_start_loc, - context_lens_tensor=context_lens_tensor, - block_tables=block_tables, - use_cuda_graph=False, - # Begin encoder & cross attn fields below... - encoder_seq_lens=self.encoder_seq_lens, - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor, - encoder_seq_start_loc=self.encoder_seq_start_loc, - max_encoder_seq_len=self.max_encoder_seq_len, - cross_slot_mapping=self.cross_slot_mapping, - cross_block_tables=self.cross_block_tables) - return self._cached_prefill_metadata - - @property - def decode_metadata(self) -> Optional["FlashAttentionMetadata"]: - if self.num_decode_tokens == 0: - return None - - if self._cached_decode_metadata is not None: - return self._cached_decode_metadata - assert ((self.seq_lens_tensor is not None) - or (self.encoder_seq_lens_tensor is not None)) - - # Compute some attn_metadata fields which default to None - slot_mapping = (None if self.slot_mapping is None else - self.slot_mapping[self.num_prefill_tokens:]) - seq_lens_tensor = (None if self.seq_lens_tensor is None else - self.seq_lens_tensor[self.num_prefills:]) - block_tables = (None if self.block_tables is None else - self.block_tables[self.num_prefills:]) - - self._cached_decode_metadata = FlashAttentionMetadata( - num_prefills=0, - num_prefill_tokens=0, - num_decode_tokens=self.num_decode_tokens, - slot_mapping=slot_mapping, - multi_modal_placeholder_index_maps=None, - enable_kv_scales_calculation=True, - seq_lens=None, - seq_lens_tensor=seq_lens_tensor, - max_decode_query_len=self.max_decode_query_len, - max_query_len=self.max_query_len, - max_prefill_seq_len=0, - max_decode_seq_len=self.max_decode_seq_len, - # Batch may be composed of prefill|decodes, adjust query start - # indices to refer to the start of decodes. E.g. - # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6]. - query_start_loc=(self.query_start_loc[self.num_prefills:] - - self.query_start_loc[self.num_prefills]) - if self.query_start_loc is not None else None, - seq_start_loc=self.seq_start_loc[self.num_prefills:] - if self.seq_start_loc is not None else None, - context_lens_tensor=None, - block_tables=block_tables, - use_cuda_graph=self.use_cuda_graph, - # Begin encoder & cross attn fields below... - encoder_seq_lens=self.encoder_seq_lens, - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor, - encoder_seq_start_loc=self.encoder_seq_start_loc, - max_encoder_seq_len=self.max_encoder_seq_len, - cross_slot_mapping=self.cross_slot_mapping, - cross_block_tables=self.cross_block_tables) - return self._cached_decode_metadata - - def advance_step(self, - model_input: "ModelInputForGPUWithSamplingMetadata", - sampled_token_ids: Optional[torch.Tensor], - block_size: int, - num_seqs: int, - num_queries: int, - turn_prefills_into_decodes: bool = False): - """ - Update metadata in-place to advance one decode step. - """ - # When using cudagraph, the num_seqs is padded to the next captured - # batch sized, but num_queries tracks the actual number of requests in - # the batch. For --enforce-eager mode, num_seqs == num_queries - if num_seqs != num_queries: - assert num_seqs > num_queries - assert self.use_cuda_graph - - if turn_prefills_into_decodes: - # When Mutli-Step is enabled with Chunked-Prefill, prefills and - # decodes are scheduled together. In the first step, all the - # prefills turn into decodes. This update reflects that - # conversion. - assert self.num_decode_tokens + self.num_prefills == num_seqs - self.num_decode_tokens += self.num_prefills - self.num_prefills = 0 - self.num_prefill_tokens = 0 - self.max_prefill_seq_len = 0 - self.max_query_len = 1 - - self.slot_mapping = self.slot_mapping[:num_seqs] - else: - assert self.seq_lens is not None - assert self.max_decode_seq_len == max(self.seq_lens) - - assert self.num_prefills == 0 - assert self.num_prefill_tokens == 0 - assert self.num_decode_tokens == num_seqs - assert self.slot_mapping.shape == (num_seqs, ) - - assert self.seq_lens is not None - assert len(self.seq_lens) == num_seqs - assert self.seq_lens_tensor is not None - assert self.seq_lens_tensor.shape == (num_seqs, ) - assert self.max_query_len == 1 - assert self.max_prefill_seq_len == 0 - - assert self.query_start_loc is not None - assert self.query_start_loc.shape == (num_queries + 1, ) - assert self.seq_start_loc is not None - assert self.seq_start_loc.shape == (num_seqs + 1, ) - - assert self.context_lens_tensor is not None - assert self.context_lens_tensor.shape == (num_queries, ) - - assert self.block_tables is not None - assert self.block_tables.shape[0] == num_seqs - - # Update query lengths. Note that we update only queries and not seqs, - # since tensors may be padded due to captured cuda graph batch size - for i in range(num_queries): - self.seq_lens[i] += 1 - self.max_decode_seq_len = max(self.seq_lens) - - ops.advance_step_flashattn(num_seqs=num_seqs, - num_queries=num_queries, - block_size=block_size, - input_tokens=model_input.input_tokens, - sampled_token_ids=sampled_token_ids, - input_positions=model_input.input_positions, - seq_lens=self.seq_lens_tensor, - slot_mapping=self.slot_mapping, - block_tables=self.block_tables) - - -class FlashAttentionMetadataBuilder( - AttentionMetadataBuilder[FlashAttentionMetadata]): - - def __init__(self, input_builder: "ModelInputForGPUBuilder"): - self.input_builder = input_builder - self.runner = input_builder.runner - self.sliding_window = input_builder.sliding_window - self.block_size = input_builder.block_size - - def prepare(self): - self.slot_mapping: List[int] = [] - self.prefill_seq_lens: List[int] = [] - self.context_lens: List[int] = [] - self.block_tables: List[List[int]] = [] - self.curr_seq_lens: List[int] = [] - self.multimodal_placeholder_maps: Dict[ - str, - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap) - self.num_prefills = 0 - self.num_prefill_tokens = 0 - self.num_decode_tokens = 0 - self.has_prefix_cache_hit = False - - def _add_seq_group( - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup", - chunked_prefill_enabled: bool, prefix_cache_hit: bool): - """Add a sequence group to the metadata. Specifically update/append - 1. context length. - 2. block table. - 3. slot mapping. - """ - is_prompt = inter_data.is_prompt - block_tables = inter_data.block_tables - - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len, - curr_sliding_window_block) in zip( - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens], - inter_data.orig_seq_lens, inter_data.seq_lens, - inter_data.query_lens, inter_data.context_lens, - inter_data.curr_sliding_window_blocks): - self.context_lens.append(context_len) - - if is_prompt: - mm_maps = inter_data.multi_modal_placeholder_maps - if mm_maps: - for modality, placeholders in mm_maps.items(): - self.multimodal_placeholder_maps[modality].extend( - placeholders) - - self.num_prefills += 1 - self.num_prefill_tokens += token_len - self.prefill_seq_lens.append(seq_len) - else: - self.num_decode_tokens += query_len - self.curr_seq_lens.append(curr_seq_len) - - # Compute block table. - # TODO(sang): Combine chunked prefill and prefix caching by - # only allowing multiple of block_size chunk size. - # NOTE: This only works for oooooooxxx style attention. - block_table = [] - if prefix_cache_hit: - # NOTE(woosuk): For flash-attn, the block table should - # include the entries for the incoming prefill tokens. - block_table = block_tables[seq_id] - elif ((chunked_prefill_enabled or not is_prompt) - and block_tables is not None): - if curr_sliding_window_block == 0: - block_table = block_tables[seq_id] - else: - block_table = block_tables[seq_id][ - -curr_sliding_window_block:] - self.block_tables.append(block_table) - - # Compute slot mapping. - is_profile_run = is_block_tables_empty(block_tables) - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len, - context_len, - self.sliding_window) - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id, - seq_len, context_len, start_idx, - self.block_size, inter_data.block_tables) - - def _get_graph_runner_block_tables( - self, num_seqs: int, - block_tables: List[List[int]]) -> torch.Tensor: - # The shape of graph_block_tables is - # [max batch size, max context len // block size]. - max_batch_size, max_blocks = self.runner.graph_block_tables.shape - assert max_batch_size >= num_seqs - - graph_block_tables = self.runner.graph_block_tables[:num_seqs] - for i, block_table in enumerate(block_tables): - if block_table: - num_blocks = len(block_table) - if num_blocks <= max_blocks: - graph_block_tables[i, :num_blocks] = block_table - else: - # It may be possible to have more blocks allocated due - # to lookahead slots of multi-step, however, they are - # not used anyway, so can be safely ignored. - graph_block_tables[ - i, :max_blocks] = block_table[:max_blocks] - - return torch.from_numpy(graph_block_tables).to( - device=self.runner.device, non_blocking=True) - - def build(self, seq_lens: List[int], query_lens: List[int], - cuda_graph_pad_size: int, batch_size: int): - """Build attention metadata with on-device tensors. - - Args: - seq_lens: The maybe padded sequence lengths of the input sequences. - query_lens: The query lengths of the input sequences. - cuda_graph_pad_size: The padding size for cuda graph. - -1 if cuda graph is not used. - batch_size: The maybe padded batch size. - """ - prefix_cache_hit = any([ - inter_data.prefix_cache_hit - for inter_data in self.input_builder.inter_data_list - ]) - for inter_data in self.input_builder.inter_data_list: - self._add_seq_group(inter_data, - self.input_builder.chunked_prefill_enabled, - prefix_cache_hit) - - device = self.runner.device - use_captured_graph = cuda_graph_pad_size != -1 - - max_query_len = max(query_lens) - decode_query_lens = query_lens[self.num_prefills:] - if len(decode_query_lens) > 0: - max_decode_query_len = max(decode_query_lens) - else: - max_decode_query_len = 1 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0) - max_decode_seq_len = max(self.curr_seq_lens, default=0) - num_decode_tokens = self.num_decode_tokens - query_start_loc = list(accumulate(query_lens, initial=0)) - seq_start_loc = list(accumulate(seq_lens, initial=0)) - - num_seqs = len(seq_lens) - if use_captured_graph: - self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size) - self.block_tables.extend([] * cuda_graph_pad_size) - num_decode_tokens = batch_size - self.num_prefill_tokens - block_tables = self._get_graph_runner_block_tables( - num_seqs, self.block_tables) - else: - block_tables = make_tensor_with_pad( - self.block_tables, - pad=0, - dtype=torch.int, - device=device, - ) - assert max_query_len > 0, ("query_lens: {}".format(query_lens)) - - assert device is not None - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int, - device, self.runner.pin_memory) - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device, - self.runner.pin_memory) - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long, - device, self.runner.pin_memory) - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32, - device, - self.runner.pin_memory) - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32, - device, self.runner.pin_memory) - placeholder_index_maps = { - modality: placeholder_map.index_map() - for modality, placeholder_map in - self.multimodal_placeholder_maps.items() - } - - return FlashAttentionMetadata( - num_prefills=self.num_prefills, - slot_mapping=slot_mapping_tensor, - num_prefill_tokens=self.num_prefill_tokens, - num_decode_tokens=num_decode_tokens, - seq_lens=seq_lens, - multi_modal_placeholder_index_maps=placeholder_index_maps, - enable_kv_scales_calculation=True, - seq_lens_tensor=seq_lens_tensor, - max_query_len=max_query_len, - max_decode_query_len=max_decode_query_len, - max_prefill_seq_len=max_prefill_seq_len, - max_decode_seq_len=max_decode_seq_len, - query_start_loc=query_start_loc_tensor, - seq_start_loc=seq_start_loc_tensor, - context_lens_tensor=context_lens_tensor, - block_tables=block_tables, - use_cuda_graph=use_captured_graph, - ) - - -class FlashAttentionImpl(AttentionImpl): - """ - If the input tensors contain prompt tokens, the layout is as follows: - |<--------------- num_prefill_tokens ----------------->| - |<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->| - - Otherwise, the layout is as follows: - |<----------------- num_decode_tokens ------------------>| - |<--decode_0-->|..........|<--decode_M-1-->|<--padding-->| - - Generation tokens can contain padding when cuda-graph is used. - Currently, prompt tokens don't contain any padding. - - The prompts might have different lengths, while the generation tokens - always have length 1. - - If chunked prefill is enabled, prefill tokens and decode tokens can be - batched together in a flattened 1D query. - - |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->| - |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->| - - Currently, cuda graph is disabled for chunked prefill, meaning there's no - padding between prefill and decode tokens. - """ - - def __init__( - self, - num_heads: int, - head_size: int, - scale: float, - num_kv_heads: int, - alibi_slopes: Optional[List[float]], - sliding_window: Optional[int], - kv_cache_dtype: str, - blocksparse_params: Optional[Dict[str, Any]] = None, - logits_soft_cap: Optional[float] = None, - attn_type: str = AttentionType.DECODER, - ) -> None: - if blocksparse_params is not None: - raise ValueError( - "FlashAttention does not support block-sparse attention.") - self.num_heads = num_heads - self.head_size = head_size - self.scale = float(scale) - self.num_kv_heads = num_kv_heads - if alibi_slopes is not None: - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32) - self.alibi_slopes = alibi_slopes - self.sliding_window = ((sliding_window - 1, - 0) if sliding_window is not None else (-1, -1)) - self.kv_cache_dtype = kv_cache_dtype - if is_quantized_kv_cache(self.kv_cache_dtype): - raise NotImplementedError( - "FlashAttention with FP8 KV cache not yet supported") - if logits_soft_cap is None: - # In flash-attn, setting logits_soft_cap as 0 means no soft cap. - logits_soft_cap = 0 - self.logits_soft_cap = logits_soft_cap - - assert self.num_heads % self.num_kv_heads == 0 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads - - support_head_sizes = FlashAttentionBackend.get_supported_head_sizes() - if head_size not in support_head_sizes: - raise ValueError( - f"Head size {head_size} is not supported by FlashAttention. " - f"Supported head sizes are: {support_head_sizes}.") - self.attn_type = attn_type - self.vllm_flash_attn_version = get_flash_attn_version() - - def forward( - self, - layer: AttentionLayer, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: FlashAttentionMetadata, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Forward pass with FlashAttention. - - Args: - query: shape = [num_tokens, num_heads, head_size] - key: shape = [num_tokens, num_kv_heads, head_size] - value: shape = [num_tokens, num_kv_heads, head_size] - output: shape = [num_tokens, num_heads, head_size] - kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size] - NOTE: kv_cache will be an empty tensor with shape [0] - for profiling run. - attn_metadata: Metadata for attention. - NOTE: It in-place updates the output tensor. - """ - # NOTE(woosuk): FlashAttention does not support FP8 KV cache. - assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0, ( - "key/v_scale is not supported in FlashAttention.") - - assert output is not None, "Output tensor must be provided." - - attn_type = self.attn_type - if (attn_type == AttentionType.ENCODER - and (not attn_metadata.is_all_encoder_attn_metadata_set)): - raise AttributeError("Encoder attention requires setting " - "encoder metadata attributes.") - elif (attn_type == AttentionType.ENCODER_DECODER - and (not attn_metadata.is_all_cross_attn_metadata_set)): - raise AttributeError("Encoder/decoder cross-attention " - "requires setting cross-attention " - "metadata attributes.") - - kv_cache_dtype: str = self.kv_cache_dtype - softmax_scale: float = self.scale - window_size = self.sliding_window - alibi_slopes: Optional[torch.Tensor] = self.alibi_slopes - logits_soft_cap: Optional[float] = self.logits_soft_cap - - if kv_cache.numel() > 0: - key_cache = kv_cache[0] - value_cache = kv_cache[1] - # We skip updating the KV cache under two conditions: - # a. When the Attention Type is ENCODER. In this phase, we compute - # only the encoder attention without updating the cache. - # b. When both Key and Value are None. This occurs during - # cross-attention computation in the decoding phase, where the - # KV cache is already populated with the cross-attention - # tensor. Thus, we skip cache updates during this time. - if (attn_type != AttentionType.ENCODER) and (key is not None) and ( - value is not None): - if attn_type == AttentionType.ENCODER_DECODER: - # Update cross-attention KV cache (prefill-only) - updated_slot_mapping = attn_metadata.cross_slot_mapping - else: - # Update self-attention KV cache (prefill/decode) - updated_slot_mapping = attn_metadata.slot_mapping - - # Reshape the input keys and values and store them in the cache. - # If kv_cache is not provided, the new key and value tensors are - # not cached. This happens during the initial memory - # profiling run. - torch.ops._C_cache_ops.reshape_and_cache_flash( - key, - value, - kv_cache[0], - kv_cache[1], - updated_slot_mapping.flatten(), # type: ignore[union-attr] - kv_cache_dtype, - layer._k_scale, - layer._v_scale, - ) - - (num_prefill_query_tokens, num_prefill_kv_tokens, - num_decode_query_tokens) = \ - get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type) - decode_query = query[num_prefill_query_tokens:] - decode_output = output[num_prefill_query_tokens:] - # QKV for prefill. - query = query[:num_prefill_query_tokens] - prefill_output = output[:num_prefill_query_tokens] - assert query.shape[0] == num_prefill_query_tokens - assert decode_query.shape[0] == num_decode_query_tokens - - if prefill_meta := attn_metadata.prefill_metadata: - # Prompt run. - if (kv_cache.numel() == 0 or prefill_meta.block_tables is None - or prefill_meta.block_tables.numel() == 0): - # normal attention - # When block_tables are not filled, it means q and k are the - # prompt, and they have the same length. - q_seq_start_loc, q_seq_len, k_seq_start_loc, k_seq_len = \ - _get_query_key_seq_metadata(prefill_meta, True, attn_type) - - key = key[:num_prefill_kv_tokens] - value = value[:num_prefill_kv_tokens] - - flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=q_seq_start_loc, - cu_seqlens_k=k_seq_start_loc, - max_seqlen_q=q_seq_len, - max_seqlen_k=k_seq_len, - softmax_scale=softmax_scale, - causal=_get_causal_option(attn_type), - window_size=window_size, - alibi_slopes=alibi_slopes, - softcap=logits_soft_cap, - out=prefill_output, - fa_version=self.vllm_flash_attn_version, - ) - else: - # prefix-enabled attention - assert attn_type == AttentionType.DECODER, ( - "Only decoder-only models support prefix caching") - assert prefill_meta.seq_lens is not None - max_seq_len = max(prefill_meta.seq_lens) - flash_attn_varlen_func( # noqa - q=query, - k=key_cache, - v=value_cache, - cu_seqlens_q=prefill_meta.query_start_loc, - max_seqlen_q=prefill_meta.max_query_len, - seqused_k=prefill_meta.seq_lens_tensor, - max_seqlen_k=max_seq_len, - softmax_scale=softmax_scale, - causal=True, - window_size=window_size, - alibi_slopes=alibi_slopes, - block_table=prefill_meta.block_tables, - softcap=logits_soft_cap, - out=prefill_output, - fa_version=self.vllm_flash_attn_version, - ) - - if decode_meta := attn_metadata.decode_metadata: - # Decoding run. - # Use flash_attn_varlen_func kernel for speculative decoding - # because different queries might have different lengths. - - assert decode_meta.max_decode_query_len is not None - # use only for actual varlen decoding - if decode_meta.max_decode_query_len > 1: - assert attn_type == AttentionType.DECODER, ( - "Only decoder-only models support max_decode_query_len > 1" - ) - flash_attn_varlen_func( - q=decode_query, - k=key_cache, - v=value_cache, - cu_seqlens_q=decode_meta.query_start_loc, - max_seqlen_q=decode_meta.max_decode_query_len, - seqused_k=decode_meta.seq_lens_tensor, - max_seqlen_k=decode_meta.max_decode_seq_len, - softmax_scale=softmax_scale, - causal=True, - window_size=window_size, - alibi_slopes=alibi_slopes, - softcap=logits_soft_cap, - block_table=decode_meta.block_tables, - out=decode_output, - fa_version=self.vllm_flash_attn_version, - ) - else: - # Use flash_attn_with_kvcache for normal decoding. - ( - seq_lens_arg, - _, - block_tables_arg, - ) = get_seq_len_block_table_args(decode_meta, False, attn_type) - flash_attn_with_kvcache( - q=decode_query.unsqueeze(1), - k_cache=key_cache, - v_cache=value_cache, - block_table=block_tables_arg, - cache_seqlens=seq_lens_arg, - softmax_scale=softmax_scale, - causal=True, - window_size=window_size, - alibi_slopes=alibi_slopes, - softcap=logits_soft_cap, - out=decode_output.unsqueeze(1), - fa_version=self.vllm_flash_attn_version, - ) - return output - - -def _get_query_key_seq_metadata( - attn_metadata, - is_prompt: bool, - attn_type: str, -) -> tuple: - """ - Returns sequence metadata for key and query based on the specified - attention type and whether input is a prompt. - - This function computes the starting locations and maximum sequence lengths - for key and query sequences for different attention types. - - Args: - attn_metadata: The attention metadata object - is_prompt (bool): A flag indicating if the input is a prompt - attn_type (AttentionType): The type of attention being used. - - Returns: - tuple: A tuple containing four integers: - - Starting location for the query sequence. - - Maximum sequence length for the query sequence. - - Starting location for the key sequence. - - Maximum sequence length for the key sequence. - - Raises: - AttributeError: If an invalid attention type is provided. - """ - if attn_type == AttentionType.DECODER: - # Decoder self-attention - # Choose max_seq_len based on whether we are in prompt_run - if is_prompt: - max_seq_len = attn_metadata.max_prefill_seq_len - else: - max_seq_len = attn_metadata.max_decode_seq_len - return (attn_metadata.seq_start_loc, max_seq_len, - attn_metadata.seq_start_loc, max_seq_len) - - elif attn_type == AttentionType.ENCODER_DECODER: - # This is cross attention between the where the key - # is the precomputed encoder attention and query - # is the input sequence. - # Choose query max length based on whether it is prompt - # or not. - if is_prompt: - max_seq_len = attn_metadata.max_prefill_seq_len - else: - max_seq_len = attn_metadata.max_decode_seq_len - return (attn_metadata.seq_start_loc, max_seq_len, - attn_metadata.encoder_seq_start_loc, - attn_metadata.max_encoder_seq_len) - elif attn_type == AttentionType.ENCODER: - # For encoder attention both the query and the key are same i.e the - # encoder sequence. - return (attn_metadata.encoder_seq_start_loc, - attn_metadata.max_encoder_seq_len, - attn_metadata.encoder_seq_start_loc, - attn_metadata.max_encoder_seq_len) - elif attn_type == AttentionType.ENCODER_ONLY: - assert is_prompt, "Should not have decode for encoder only model." - return (attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len, - attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len) - else: - raise AttributeError(f"Invalid attention type {str(attn_type)}") - - -def _get_causal_option(attn_type: str) -> bool: - """ - Determine whether the given attention type is suitable for causal - attention mechanisms. - - Args: - attn_type (AttentionType): The type of attention being evaluated - - Returns: - bool: Returns `True` if the attention type is suitable for causal - attention (i.e., not encoder, encoder-only, or encoder-decoder), - otherwise returns `False`. - """ - return not (attn_type == AttentionType.ENCODER - or attn_type == AttentionType.ENCODER_ONLY - or attn_type == AttentionType.ENCODER_DECODER) - - -class TextEncoderFlashAttentionImpl: - """Simplified Flash Attention implementation for text encoders in diffusion models. - - This implementation removes KV caching and other generation-specific features, - focusing only on the encoder attention pattern needed for text encoders in - diffusion models like DiT. - """ - - def __init__( - self, - num_heads: int, - head_size: int, - scale: Optional[float] = None, - ) -> None: - """Initialize the text encoder flash attention implementation. - - Args: - num_heads: Number of attention heads - head_size: Size of each attention head - scale: Attention scale factor. If None, defaults to 1/sqrt(head_size) - """ - self.num_heads = num_heads - self.head_size = head_size - self.num_kv_heads = num_heads # For encoder, typically num_kv_heads = num_heads - - # Set default scale if not provided - if scale is None: - scale = 1.0 / math.sqrt(head_size) - self.scale = float(scale) - - # Set defaults for optional parameters - self.alibi_slopes = None - self.sliding_window = (-1, -1) # No sliding window by default - self.logits_soft_cap = 0 # No soft cap by default - - # Always use encoder-only attention type - self.attn_type = AttentionType.ENCODER_ONLY - self.vllm_flash_attn_version = get_flash_attn_version() - - # Verify head size is supported - support_head_sizes = FlashAttentionBackend.get_supported_head_sizes() - if head_size not in support_head_sizes: - raise ValueError( - f"Head size {head_size} is not supported by FlashAttention. " - f"Supported head sizes are: {support_head_sizes}.") - - def forward( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - seq_lens: Optional[List[int]] = None, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Forward pass with FlashAttention for text encoder. - - Args: - query: shape = [num_tokens, num_heads, head_size] - key: shape = [num_tokens, num_kv_heads, head_size] - value: shape = [num_tokens, num_kv_heads, head_size] - seq_lens: List of sequence lengths for batched input - output: shape = [num_tokens, num_heads, head_size] - Pre-allocated output tensor (optional) - - Returns: - torch.Tensor: Output tensor with shape [num_tokens, num_heads, head_size] - """ - device = query.device - batch_size = len(seq_lens) if seq_lens is not None else 1 - - if output is None: - output = torch.empty_like(query) - - if seq_lens is None: - # Single sequence case - use standard flash attention - return flash_attn_varlen_func( - q=query, - k=key, - v=value, - softmax_scale=self.scale, - causal=False, # Non-causal for encoder - window_size=self.sliding_window, - alibi_slopes=self.alibi_slopes, - softcap=self.logits_soft_cap, - out=output, - fa_version=self.vllm_flash_attn_version, - ) - else: - # Multiple sequences case - use varlen flash attention - max_seq_len = max(seq_lens) - seq_start_loc = list(accumulate(seq_lens, initial=0)) - seq_start_loc_tensor = torch.tensor(seq_start_loc, - dtype=torch.int32, - device=device) - - return flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=seq_start_loc_tensor, - cu_seqlens_k=seq_start_loc_tensor, - max_seqlen_q=max_seq_len, - max_seqlen_k=max_seq_len, - softmax_scale=self.scale, - causal=False, # Non-causal for encoder - window_size=self.sliding_window, - alibi_slopes=self.alibi_slopes, - softcap=self.logits_soft_cap, - out=output, - fa_version=self.vllm_flash_attn_version, - ) diff --git a/fastvideo/v1/attention/backends/utils.py b/fastvideo/v1/attention/backends/utils.py deleted file mode 100644 index 4374b5422..000000000 --- a/fastvideo/v1/attention/backends/utils.py +++ /dev/null @@ -1,619 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Attention backend utils""" -from collections import defaultdict -from contextlib import contextmanager -from itertools import accumulate -from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Type, TypeVar, Union - -import numpy as np -import torch - -from vllm import envs -from vllm.attention import (AttentionMetadata, AttentionMetadataBuilder, - AttentionState) -from vllm.attention.backends.abstract import AttentionType -from vllm.logger import init_logger -from vllm.multimodal import MultiModalPlaceholderMap -from vllm.platforms import current_platform -from vllm.utils import async_tensor_h2d, make_tensor_with_pad - -logger = init_logger(__name__) - -if TYPE_CHECKING: - from vllm.worker.model_runner_base import ModelRunnerBase - -# Error string(s) for encoder/decoder -# unsupported attention scenarios -STR_NOT_IMPL_ENC_DEC_ROCM_HIP = ("ROCm/HIP is not currently supported " - "with encoder/decoder models.") - -PAD_SLOT_ID = -1 - -# Switch to numpy implementation of compute_slot_mapping -# if we have at least this many elements. Could be tuned further. -_COMPUTE_SLOT_MAPPING_NUMPY_NUMEL = 256 - -if TYPE_CHECKING: - from vllm.worker.model_runner import ModelInputForGPUBuilder - - -def is_block_tables_empty(block_tables: Union[None, Dict]): - """ - Check if block_tables is None or a dictionary with all None values. - """ - if block_tables is None: - return True - return (isinstance(block_tables, dict) - and all(value is None for value in block_tables.values())) - - -def compute_slot_mapping_start_idx(is_prompt: bool, query_len: int, - context_len: int, sliding_window: int): - """ - Compute the start index of slot mapping. - """ - start_idx = 0 - if is_prompt and sliding_window is not None: - start_idx = max(0, query_len - sliding_window) - return start_idx - - -def _compute_slot_mapping_python(slot_mapping: List[int], - block_table: List[int], range_start: int, - range_end: int, block_size: int): - for i in range(range_start, range_end): - block_number = block_table[i // block_size] - block_offset = i % block_size - slot = block_number * block_size + block_offset - slot_mapping.append(slot) - - -def _compute_slot_mapping_numpy(slot_mapping: List[int], - block_table: List[int], range_start: int, - range_end: int, block_size: int): - block_table_array = np.array(block_table) - idx = np.arange(range_start, range_end) - block_offset = idx % block_size - idx //= block_size - seq_slot_mapping_array = block_table_array[idx] - seq_slot_mapping_array *= block_size - seq_slot_mapping_array += block_offset - slot_mapping.extend(seq_slot_mapping_array) - - -def compute_slot_mapping(is_profile_run: bool, slot_mapping: List[int], - seq_id: int, seq_len: int, context_len: int, - start_idx: int, block_size: int, - block_tables: Dict[int, List[int]]): - """ - Compute slot mapping. - """ - if is_profile_run: - # During memory profiling, the block tables are not - # initialized yet. In this case, we just use a dummy - # slot mapping. - # In embeddings, the block tables are {seq_id: None}. - slot_mapping.extend([PAD_SLOT_ID] * seq_len) - return - - # Mask the [0, start_idx) tokens of the prompt with - # PAD_SLOT_ID, where start_idx is max(0, seq_len - - # sliding_window). For example, if the prompt len is 10, - # sliding window is 8, and block size is 4, the first two - # tokens are masked and the slot mapping will be - # [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1]. - padding_mask_len = max(0, start_idx - context_len) - slot_mapping.extend([PAD_SLOT_ID] * padding_mask_len) - - range_start = max(start_idx, context_len) - range_end = seq_len - numel = range_end - range_start - block_table = block_tables[seq_id] - - # numpy implementation will be faster than python if we have - # many elements, otherwise it will be slower. - if numel < _COMPUTE_SLOT_MAPPING_NUMPY_NUMEL: - _compute_slot_mapping_python(slot_mapping, block_table, range_start, - range_end, block_size) - else: - _compute_slot_mapping_numpy(slot_mapping, block_table, range_start, - range_end, block_size) - - -TAttentionMetadata = TypeVar("TAttentionMetadata", bound='AttentionMetadata') - - -class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]): - - _metadata_cls: Type[TAttentionMetadata] - - def __init__(self, input_builder: "ModelInputForGPUBuilder"): - self.input_builder = input_builder - self.runner = input_builder.runner - - self.sliding_window = input_builder.sliding_window - self.block_size = input_builder.block_size - - def prepare(self): - self.slot_mapping: List[int] = [] - self.prefill_seq_lens: List[int] = [] - self.context_lens: List[int] = [] - self.block_tables: List[List[int]] = [] - self.curr_seq_lens: List[int] = [] - self.multimodal_placeholder_maps: Dict[ - str, - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap) - self.num_prefills = 0 - self.num_prefill_tokens = 0 - self.num_decode_tokens = 0 - - def _add_seq_group( - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup", - chunked_prefill_enabled: bool): - is_prompt = inter_data.is_prompt - block_tables = inter_data.block_tables - - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len, - curr_sliding_window_block) in zip( - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens], - inter_data.orig_seq_lens, inter_data.seq_lens, - inter_data.query_lens, inter_data.context_lens, - inter_data.curr_sliding_window_blocks): - self.context_lens.append(context_len) - if is_prompt: - mm_maps = inter_data.multi_modal_placeholder_maps - if mm_maps: - for modality, placeholders in mm_maps.items(): - self.multimodal_placeholder_maps[modality].extend( - placeholders) - - self.num_prefills += 1 - self.num_prefill_tokens += token_len - self.prefill_seq_lens.append(seq_len) - else: - assert query_len == 1, ( - "seq_len: {}, context_len: {}, query_len: {}".format( - seq_len, context_len, query_len)) - self.num_decode_tokens += query_len - self.curr_seq_lens.append(curr_seq_len) - - # Compute block table. - # TODO(sang): Combine chunked prefill and prefix caching by - # only allowing multiple of block_size chunk size. - # NOTE: This only works for oooooooxxx style attention. - block_table = [] - if inter_data.prefix_cache_hit: - block_table = block_tables[seq_id] - elif ((chunked_prefill_enabled or not is_prompt) - and block_tables is not None): - if curr_sliding_window_block == 0: - block_table = block_tables[seq_id] - else: - block_table = block_tables[seq_id][ - -curr_sliding_window_block:] - self.block_tables.append(block_table) - - # Compute slot mapping. - is_profile_run = is_block_tables_empty(block_tables) - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len, - context_len, - self.sliding_window) - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id, - seq_len, context_len, start_idx, - self.block_size, inter_data.block_tables) - - def build(self, seq_lens: List[int], query_lens: List[int], - cuda_graph_pad_size: int, batch_size: int): - """Build attention metadata with on-device tensors. - - Args: - seq_lens: The maybe padded sequence lengths of the input sequences. - query_lens: The query lengths of the input sequences. - cuda_graph_pad_size: The padding size for cuda graph. - -1 if cuda graph is not used. - batch_size: The maybe padded batch size. - """ - for inter_data in self.input_builder.inter_data_list: - self._add_seq_group(inter_data, - self.input_builder.chunked_prefill_enabled) - - device = self.runner.device - use_captured_graph = cuda_graph_pad_size != -1 - - max_query_len = max(query_lens) - max_prefill_seq_len = max(self.prefill_seq_lens, default=0) - max_decode_seq_len = max(self.curr_seq_lens, default=0) - num_decode_tokens = self.num_decode_tokens - query_start_loc = list(accumulate(query_lens, initial=0)) - seq_start_loc = list(accumulate(seq_lens, initial=0)) - - if use_captured_graph: - self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size) - self.block_tables.extend([] * cuda_graph_pad_size) - num_decode_tokens = batch_size - - # The shape of graph_block_tables is - # [max batch size, max context len // block size]. - input_block_tables = self.runner.graph_block_tables[:batch_size] - for i, block_table in enumerate(self.block_tables): - if block_table: - input_block_tables[i, :len(block_table)] = block_table - block_tables = torch.from_numpy(input_block_tables).to( - device, non_blocking=True) - else: - block_tables = make_tensor_with_pad( - self.block_tables, - pad=0, - dtype=torch.int, - device=device, - ) - assert max_query_len > 0, "query_lens: {}".format(query_lens) - - assert device is not None - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int, - device, self.runner.pin_memory) - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device, - self.runner.pin_memory) - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long, - device, self.runner.pin_memory) - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32, - device, - self.runner.pin_memory) - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32, - device, self.runner.pin_memory) - placeholder_index_maps = { - modality: placeholder_map.index_map() - for modality, placeholder_map in - self.multimodal_placeholder_maps.items() - } - - return self._metadata_cls( # type: ignore - num_prefills=self.num_prefills, - slot_mapping=slot_mapping_tensor, - multi_modal_placeholder_index_maps=placeholder_index_maps, - enable_kv_scales_calculation=True, - num_prefill_tokens=self.num_prefill_tokens, - num_decode_tokens=num_decode_tokens, - seq_lens=seq_lens, - seq_lens_tensor=seq_lens_tensor, - max_query_len=max_query_len, - max_prefill_seq_len=max_prefill_seq_len, - max_decode_seq_len=max_decode_seq_len, - query_start_loc=query_start_loc_tensor, - seq_start_loc=seq_start_loc_tensor, - context_lens_tensor=context_lens_tensor, - block_tables=block_tables, - use_cuda_graph=use_captured_graph, - ) - - -class CommonAttentionState(AttentionState): - - def __init__(self, runner: "ModelRunnerBase"): - self.runner = runner - self._is_graph_capturing = False - - @contextmanager - def graph_capture(self, max_batch_size: int): - - self._is_graph_capturing = True - - self._graph_slot_mapping = torch.full((max_batch_size, ), - PAD_SLOT_ID, - dtype=torch.long, - device=self.runner.device) - self._graph_seq_lens = torch.ones(max_batch_size, - dtype=torch.int32, - device=self.runner.device) - self._graph_block_tables = torch.from_numpy( - self.runner.graph_block_tables).to(device=self.runner.device) - - yield - - self._is_graph_capturing = False - del self._graph_slot_mapping - del self._graph_seq_lens - del self._graph_block_tables - - def graph_clone(self, batch_size: int) -> "CommonAttentionState": - assert self._is_graph_capturing - return self.__class__(self.runner) - - def graph_capture_get_metadata_for_batch( - self, batch_size: int, is_encoder_decoder_model: bool = False): - assert self._is_graph_capturing - attn_metadata = self.runner.attn_backend.make_metadata( - num_prefills=0, - num_prefill_tokens=0, - num_decode_tokens=batch_size, - slot_mapping=self._graph_slot_mapping[:batch_size], - multi_modal_placeholder_index_maps=None, - enable_kv_scales_calculation=True, - seq_lens=None, - seq_lens_tensor=self._graph_seq_lens[:batch_size], - max_query_len=1, - max_decode_query_len=1, - max_prefill_seq_len=0, - max_decode_seq_len=self.runner.max_seq_len_to_capture, - query_start_loc=None, - seq_start_loc=None, - context_lens_tensor=None, - block_tables=self._graph_block_tables[:batch_size], - use_cuda_graph=True, - ) - if is_encoder_decoder_model: - # The encoder decoder model works only with XFormers and - # Flash Attention backend. Assert the same. - assert self.runner.attn_backend.get_name() in\ - ["XFORMERS", "FLASH_ATTN"], \ - f"Expected attn_backend name to be either 'XFORMERS' or " \ - f"'FLASH_ATTN', but "\ - f"got '{self.runner.attn_backend.get_name()}'" - self._update_captured_metadata_for_enc_dec_model( - batch_size=batch_size, attn_metadata=attn_metadata) - - return attn_metadata - - def get_graph_input_buffers( - self, - attn_metadata, - is_encoder_decoder_model: bool = False) -> Dict[str, Any]: - input_buffers = { - "slot_mapping": attn_metadata.slot_mapping, - "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor, - "block_tables": attn_metadata.decode_metadata.block_tables, - } - if is_encoder_decoder_model: - # The encoder decoder model works only with XFormers and - # Flash Attention backend. Assert the same. - assert self.runner.attn_backend.get_name() in\ - ["XFORMERS", "FLASH_ATTN"], \ - f"Expected attn_backend name to be either 'XFORMERS' or "\ - f"'FLASH_ATTN', but "\ - f"got '{self.runner.attn_backend.get_name()}'" - self._add_additonal_input_buffers_for_enc_dec_model( - attn_metadata=attn_metadata, input_buffers=input_buffers) - return input_buffers - - def prepare_graph_input_buffers( - self, - input_buffers, - attn_metadata, - is_encoder_decoder_model: bool = False) -> None: - input_buffers["seq_lens_tensor"].copy_( - attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True) - input_buffers["block_tables"].copy_( - attn_metadata.decode_metadata.block_tables, non_blocking=True) - if is_encoder_decoder_model: - # The encoder decoder model works only with XFormers and - # Flash Attention backend. Assert the same. - assert self.runner.attn_backend.get_name() in\ - ["XFORMERS", "FLASH_ATTN"], \ - f"Expected attn_backend name to be either 'XFORMERS' or "\ - f"'FLASH_ATTN', but "\ - f"got '{self.runner.attn_backend.get_name()}'" - self._prepare_input_buffers_for_enc_dec_model( - attn_metadata, input_buffers) - - def begin_forward(self, model_input) -> None: - return - - def _update_captured_metadata_for_enc_dec_model(self, batch_size: int, - attn_metadata): - """ - Updates the attention metadata parameters for CUDA graph capture in an - encoder-decoder model. - - This method modifies attention-related tensors and metadata required - for CUDA graph capture in encoder-decoder models. Specifically, it - updates the cross-attention and encoder sequence tensors in the - AttentionMetadata object. - """ - # During decode phase the cross_slot_mapping will be empty. Hence set - # an empty tensor for CUDA Graph capture. - attn_metadata.cross_slot_mapping = torch.tensor( - [], dtype=torch.int).cuda() - attn_metadata.cross_block_tables = torch.full( - (batch_size, self.runner.get_max_block_per_batch()), - 1, - dtype=torch.int).cuda() - attn_metadata.encoder_seq_lens = torch.full((batch_size, ), - 1, - dtype=torch.int).cuda() - attn_metadata.encoder_seq_lens_tensor = torch.full( - (batch_size, ), 1, dtype=torch.int).cuda() - attn_metadata.max_encoder_seq_len = self.runner.max_seq_len_to_capture - attn_metadata.num_encoder_tokens = 0 - - def _add_additonal_input_buffers_for_enc_dec_model( - self, attn_metadata, input_buffers: Dict[str, Any]): - """ - Saves additional input buffers specific to the encoder-decoder model - from the attention metadata. - - This method extracts and stores encoder-decoder related input buffers - from the `attn_metadata` into the `input_buffers` dictionary. The - buffers include encoder sequence lengths, cross-slot mappings, and - cross-block tables, which are essential for the encoder-decoder model - during CUDA graph replay. - """ - input_buffers["encoder_seq_lens_tensor"] = ( - attn_metadata.decode_metadata.encoder_seq_lens_tensor) - input_buffers["cross_slot_mapping"] = ( - attn_metadata.decode_metadata.cross_slot_mapping) - input_buffers["cross_block_tables"] = ( - attn_metadata.decode_metadata.cross_block_tables) - - def _prepare_input_buffers_for_enc_dec_model(self, attn_metadata, - input_buffers: Dict[str, - Any]): - """ - Populates input buffers with data from the encoder-decoder model's - attention metadata. - - This method fills the input buffers with encoder-decoder specific - tensors. It copies data from the `attn_metadata` and keyword arguments - (`kwargs`) into corresponding buffers in the `input_buffers` dictionary. - The copied data includes attention-related metadata as well as input - IDs and positional information for the encoder. - """ - input_buffers["encoder_seq_lens_tensor"].copy_( - attn_metadata.decode_metadata.encoder_seq_lens_tensor, - non_blocking=True) - input_buffers["cross_slot_mapping"].copy_( - attn_metadata.decode_metadata.cross_slot_mapping, - non_blocking=True) - input_buffers["cross_block_tables"].copy_( - attn_metadata.decode_metadata.cross_block_tables, - non_blocking=True) - - -def is_all_encoder_attn_metadata_set(attn_metadata): - ''' - All attention metadata required for encoder attention is set. - ''' - return ((attn_metadata.encoder_seq_lens is not None) - and (attn_metadata.encoder_seq_lens_tensor is not None) - and (attn_metadata.max_encoder_seq_len is not None)) - - -def is_all_cross_attn_metadata_set(attn_metadata): - ''' - All attention metadata required for enc/dec cross-attention is set. - - Superset of encoder attention required metadata. - ''' - return (attn_metadata.is_all_encoder_attn_metadata_set - and (attn_metadata.cross_slot_mapping is not None) - and (attn_metadata.cross_block_tables is not None)) - - -def get_seq_len_block_table_args( - attn_metadata, - is_prompt: bool, - attn_type: str, -) -> tuple: - ''' - The particular choice of sequence-length- and block-table-related - attributes which should be extracted from attn_metadata is dependent - on the type of attention operation. - - Decoder attn -> select entirely decoder self-attention-related fields - Encoder/decoder cross-attn -> select encoder sequence lengths & - cross-attn block-tables fields - Encoder attn -> select encoder sequence lengths fields & no block tables - - Arguments: - - * attn_metadata: Attention metadata structure associated with attention op - * is_prompt: True if prefill, False otherwise - * attn_type: encoder attention, decoder self-attention, - encoder/decoder cross-attention - - Returns: - - * Appropriate sequence-lengths tensor - * Appropriate max sequence-length scalar - * Appropriate block tables (or None) - ''' - - if attn_type == AttentionType.DECODER: - # Decoder self-attention - # Choose max_seq_len based on whether we are in prompt_run - if is_prompt: - max_seq_len = attn_metadata.max_prefill_seq_len - else: - max_seq_len = attn_metadata.max_decode_seq_len - return (attn_metadata.seq_lens_tensor, max_seq_len, - attn_metadata.block_tables) - elif attn_type == AttentionType.ENCODER_DECODER: - # Enc/dec cross-attention KVs match encoder sequence length; - # cross-attention utilizes special "cross" block tables - return (attn_metadata.encoder_seq_lens_tensor, - attn_metadata.max_encoder_seq_len, - attn_metadata.cross_block_tables) - elif attn_type == AttentionType.ENCODER: - # No block tables associated with encoder attention - return (attn_metadata.encoder_seq_lens_tensor, - attn_metadata.max_encoder_seq_len, None) - else: - raise AttributeError(f"Invalid attention type {str(attn_type)}") - - -def get_num_prefill_decode_query_kv_tokens( - attn_metadata, - attn_type: str, -) -> Tuple[int, int, int]: - """ - Calculate the number of prefill and decode tokens for query, key/value - based on the attention metadata and the specified attention type. - - Args: - attn_metadata (FlashAttentionMetadata): Attention Metadata object. - attn_type (AttentionType): The type of attention being used. - Returns: - Tuple[int, int, int]: A tuple containing three integers: - - The number of prefill query tokens. - - The number of prefill key/value tokens. - - The number of decode query tokens. - - Raises: - AssertionError: If the number of encoder tokens in `attn_metadata` - is `None` when required for the calculations. - """ - num_prefill_query_tokens = 0 - num_decode_query_tokens = 0 - num_prefill_kv_tokens = 0 - if attn_type == AttentionType.ENCODER: - # Encoder attention is only invoked during prefill phase. - # The same input servers a both query and key. - assert attn_metadata.num_encoder_tokens is not None - num_prefill_query_tokens = attn_metadata.num_encoder_tokens - num_prefill_kv_tokens = attn_metadata.num_encoder_tokens - num_decode_query_tokens = 0 - elif attn_type == AttentionType.ENCODER_DECODER: - assert attn_metadata.num_encoder_tokens is not None - num_prefill_query_tokens = attn_metadata.num_prefill_tokens - # The key is the encoder/cross-attention. - num_prefill_kv_tokens = attn_metadata.num_encoder_tokens - num_decode_query_tokens = attn_metadata.num_decode_tokens - else: # attn_type == AttentionType.DECODER or - # attn_type == AttentionType.ENCODER_ONLY - num_prefill_query_tokens = attn_metadata.num_prefill_tokens - num_prefill_kv_tokens = attn_metadata.num_prefill_tokens - num_decode_query_tokens = attn_metadata.num_decode_tokens - - return (num_prefill_query_tokens, num_prefill_kv_tokens, - num_decode_query_tokens) - - -def get_flash_attn_version(): - try: - from vllm.vllm_flash_attn.flash_attn_interface import ( - fa_version_unsupported_reason, is_fa_version_supported) - - # if hopper default to FA3, otherwise stick to FA2 for now - # TODO(lucas): profile FA3 on ampere to see if it makes sense to - # use FA3 as default for both - if current_platform.get_device_capability()[0] == 9: - fa_version = 3 if is_fa_version_supported(3) else 2 - else: - fa_version = 2 - - if envs.VLLM_FLASH_ATTN_VERSION is not None: - assert envs.VLLM_FLASH_ATTN_VERSION in [2, 3] - fa_version = envs.VLLM_FLASH_ATTN_VERSION - if (current_platform.get_device_capability()[0] == 10 - and envs.VLLM_FLASH_ATTN_VERSION == 3): - logger.warning("Cannot use FA version 3 on Blackwell platform", - "defaulting to FA version 2.") - fa_version = 2 - - if not is_fa_version_supported(fa_version): - logger.error("Cannot use FA version %d is not supported due to %s", - fa_version, fa_version_unsupported_reason(fa_version)) - - assert is_fa_version_supported(fa_version) - return fa_version - except (ImportError, AssertionError): - return None diff --git a/fastvideo/v1/attention/layer.py b/fastvideo/v1/attention/layer.py deleted file mode 100644 index 2fa49c7b9..000000000 --- a/fastvideo/v1/attention/layer.py +++ /dev/null @@ -1,135 +0,0 @@ -import torch -import torch.nn as nn -from typing import Optional, Dict, Any, Type - -from fastvideo.v1.attention.flash_attn import LocalAttention, DistributedAttention - - -class AttentionBackend: - """Base class for attention backends that provides implementation classes.""" - - @staticmethod - def get_impl_cls() -> Type: - """Returns the implementation class for the attention backend.""" - raise NotImplementedError("Subclasses must implement get_impl_cls") - - -class LocalAttentionBackend(AttentionBackend): - """Backend for local attention implementation.""" - - @staticmethod - def get_impl_cls() -> Type: - """Returns the LocalAttention implementation class.""" - return LocalAttention - - -class DistributedAttentionBackend(AttentionBackend): - """Backend for distributed attention implementation.""" - - @staticmethod - def get_impl_cls() -> Type: - """Returns the DistributedAttention implementation class.""" - return DistributedAttention - - -class AttentionLayer(nn.Module): - """Base attention layer that uses an implementation class from a backend.""" - - def __init__( - self, - attn_backend: AttentionBackend, - dropout_rate: float = 0.0, - causal: bool = False, - softmax_scale: Optional[float] = None, - **extra_impl_args - ): - super().__init__() - impl_cls = attn_backend.get_impl_cls() - self.impl = impl_cls( - dropout_rate=dropout_rate, - causal=causal, - softmax_scale=softmax_scale, - **extra_impl_args - ) - - def forward(self, *args, **kwargs): - """Forward pass delegated to the implementation.""" - return self.impl(*args, **kwargs) - - -class LocalAttentionLayer(AttentionLayer): - """Wrapper for local attention implementation.""" - - def __init__( - self, - dropout_rate: float = 0.0, - causal: bool = False, - softmax_scale: Optional[float] = None, - **extra_impl_args - ): - super().__init__( - LocalAttentionBackend(), - dropout_rate=dropout_rate, - causal=causal, - softmax_scale=softmax_scale, - **extra_impl_args - ) - - def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - """Apply local attention. - - Args: - q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] - k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] - v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] - - Returns: - torch.Tensor: Output tensor after attention - """ - return self.impl(q, k, v) - - -class DistributedAttentionLayer(AttentionLayer): - """Wrapper for distributed attention implementation.""" - - def __init__( - self, - dropout_rate: float = 0.0, - causal: bool = False, - softmax_scale: Optional[float] = None, - **extra_impl_args - ): - super().__init__( - DistributedAttentionBackend(), - dropout_rate=dropout_rate, - causal=causal, - softmax_scale=softmax_scale, - **extra_impl_args - ) - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - replicated_q: Optional[torch.Tensor] = None, - replicated_k: Optional[torch.Tensor] = None, - replicated_v: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Forward pass for distributed attention. - - Args: - q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] - k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] - v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] - replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens - replicated_k (Optional[torch.Tensor]): Replicated key tensor - replicated_v (Optional[torch.Tensor]): Replicated value tensor - - Returns: - Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - - o (torch.Tensor): Output tensor after attention for the main sequence - - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided - """ - return self.impl(q, k, v, replicated_q, replicated_k, replicated_v) - From d2db0d475ba71b0f48c0bb00db5114dade5e81b8 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Fri, 14 Mar 2025 23:13:34 +0000 Subject: [PATCH 07/12] fix import paths --- .../device_communicators/pynccl.py | 2 +- .../device_communicators/pynccl_wrapper.py | 2 +- fastvideo/v1/inference_args.py | 2 +- fastvideo/v1/inference_engine.py | 14 +- fastvideo/v1/models/dits/hunyuanvideo.py | 4 +- fastvideo/v1/models/encoders/encoder.py | 323 ------------------ fastvideo/v1/models/loader.py | 2 +- fastvideo/v1/pipelines/stages/decoding.py | 2 +- fastvideo/v1/pipelines/stages/denoising.py | 7 +- fastvideo/v1/sample/v1_fastvideo_inference.py | 2 +- fastvideo/v1/{utils => }/utils.py | 22 ++ 11 files changed, 39 insertions(+), 343 deletions(-) delete mode 100644 fastvideo/v1/models/encoders/encoder.py rename fastvideo/v1/{utils => }/utils.py (95%) diff --git a/fastvideo/v1/distributed/device_communicators/pynccl.py b/fastvideo/v1/distributed/device_communicators/pynccl.py index 8d1bbd7c9..3d85a831c 100644 --- a/fastvideo/v1/distributed/device_communicators/pynccl.py +++ b/fastvideo/v1/distributed/device_communicators/pynccl.py @@ -13,7 +13,7 @@ from fastvideo.v1.distributed.device_communicators.pynccl_wrapper import ( ncclRedOpTypeEnum, ncclUniqueId) from fastvideo.v1.distributed.utils import StatelessProcessGroup from fastvideo.v1.logger import init_logger -from fastvideo.v1.utils.utils import current_stream +from fastvideo.v1.utils import current_stream logger = init_logger(__name__) diff --git a/fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py b/fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py index 6a0d8179a..cc3770a9e 100644 --- a/fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py +++ b/fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py @@ -31,7 +31,7 @@ import torch from torch.distributed import ReduceOp from fastvideo.v1.logger import init_logger -from fastvideo.v1.utils.utils import find_nccl_library +from fastvideo.v1.utils import find_nccl_library logger = init_logger(__name__) diff --git a/fastvideo/v1/inference_args.py b/fastvideo/v1/inference_args.py index 9cd9923f8..371a41d00 100644 --- a/fastvideo/v1/inference_args.py +++ b/fastvideo/v1/inference_args.py @@ -18,7 +18,7 @@ import argparse import dataclasses from fastvideo.v1.logger import init_logger -from fastvideo.v1.utils.utils import FlexibleArgumentParser +from fastvideo.v1.utils import FlexibleArgumentParser from typing import List, Optional diff --git a/fastvideo/v1/inference_engine.py b/fastvideo/v1/inference_engine.py index cd5b2da86..285bb6219 100644 --- a/fastvideo/v1/inference_engine.py +++ b/fastvideo/v1/inference_engine.py @@ -7,20 +7,16 @@ This module provides classes and functions for running inference with diffusion import os import time import torch -import numpy as np -from typing import Any, Dict, List, Optional, Union, Type -from PIL import Image -import uuid -import random +from typing import Any, Dict from fastvideo.v1.inference_args import InferenceArgs from fastvideo.v1.pipelines import ComposedPipelineBase, build_pipeline -from fastvideo.v1.pipelines.pipeline_registry import PipelineRegistry from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.logger import init_logger -# TODO(will): remove -from fastvideo.v1.models.hunyuan.utils.data_utils import align_to -from fastvideo.v1.models.hunyuan.constants import NEGATIVE_PROMPT +# TODO(will): remove, check if this is hunyuan specific +from fastvideo.v1.utils import align_to +# TODO(will): remove, move this to hunyuan stage +from fastvideo.v1.pipelines.implementations.hunyuan.constants import NEGATIVE_PROMPT logger = init_logger(__name__) diff --git a/fastvideo/v1/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py index 5f5e71f28..96431b546 100644 --- a/fastvideo/v1/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/models/dits/hunyuanvideo.py @@ -7,8 +7,8 @@ from fastvideo.v1.layers.layernorm import LayerNormScaleShift, ScaleResidual, Sc from fastvideo.v1.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection, unpatchify from fastvideo.v1.layers.rotary_embedding import _apply_rotary_emb, get_rotary_pos_embed from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_world_size, get_sequence_model_parallel_rank -# TODO: RMSNorm .... -from fastvideo.v1.models.hunyuan.modules.norm_layers import RMSNorm +# TODO(will-PY-refactor): RMSNorm .... +from fastvideo.v1.v0_reference_src.models.hunyuan.modules.norm_layers import RMSNorm from fastvideo.v1.layers.mlp import MLP from fastvideo.v1.models.dits.base import BaseDiT diff --git a/fastvideo/v1/models/encoders/encoder.py b/fastvideo/v1/models/encoders/encoder.py deleted file mode 100644 index aadfaa87f..000000000 --- a/fastvideo/v1/models/encoders/encoder.py +++ /dev/null @@ -1,323 +0,0 @@ -from dataclasses import dataclass -from typing import Optional, Tuple - -import torch -import torch.nn as nn -from transformers import AutoModel, AutoTokenizer, CLIPTextModel, CLIPTokenizer -from transformers.utils import ModelOutput - -from ..constants import PRECISION_TO_TYPE, TEXT_ENCODER_PATH, TOKENIZER_PATH - - -def use_default(value, default): - return value if value is not None else default - - -def load_text_encoder( - text_encoder_type, - text_encoder_precision=None, - text_encoder_path=None, - logger=None, - device=None, -): - if text_encoder_path is None: - text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type] - if logger is not None: - logger.info(f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}") - - if text_encoder_type == "clipL": - text_encoder = CLIPTextModel.from_pretrained(text_encoder_path) - text_encoder.final_layer_norm = text_encoder.text_model.final_layer_norm - elif text_encoder_type == "llm": - text_encoder = AutoModel.from_pretrained(text_encoder_path, low_cpu_mem_usage=True) - text_encoder.final_layer_norm = text_encoder.norm - else: - raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") - # from_pretrained will ensure that the model is in eval mode. - - if text_encoder_precision is not None: - text_encoder = text_encoder.to(dtype=PRECISION_TO_TYPE[text_encoder_precision]) - - text_encoder.requires_grad_(False) - - if logger is not None: - logger.info(f"Text encoder to dtype: {text_encoder.dtype}") - - if device is not None: - text_encoder = text_encoder.to(device) - - return text_encoder, text_encoder_path - - -def load_tokenizer(tokenizer_type, tokenizer_path=None, padding_side="right", logger=None): - if tokenizer_path is None: - tokenizer_path = TOKENIZER_PATH[tokenizer_type] - if logger is not None: - logger.info(f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}") - - if tokenizer_type == "clipL": - tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path, max_length=77) - elif tokenizer_type == "llm": - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, padding_side=padding_side) - else: - raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}") - - return tokenizer, tokenizer_path - - -@dataclass -class TextEncoderModelOutput(ModelOutput): - """ - Base class for model's outputs that also contains a pooling of the last hidden states. - - Args: - hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): - Sequence of hidden-states at the output of the last layer of the model. - attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: - hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): - Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + - one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. - Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. - text_outputs (`list`, *optional*, returned when `return_texts=True` is passed): - List of decoded texts. - """ - - hidden_state: torch.FloatTensor = None - attention_mask: Optional[torch.LongTensor] = None - hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None - text_outputs: Optional[list] = None - - -class TextEncoder(nn.Module): - - def __init__( - self, - text_encoder_type: str, - max_length: int, - text_encoder_precision: Optional[str] = None, - text_encoder_path: Optional[str] = None, - tokenizer_type: Optional[str] = None, - tokenizer_path: Optional[str] = None, - output_key: Optional[str] = None, - use_attention_mask: bool = True, - input_max_length: Optional[int] = None, - prompt_template: Optional[dict] = None, - prompt_template_video: Optional[dict] = None, - hidden_state_skip_layer: Optional[int] = None, - apply_final_norm: bool = False, - reproduce: bool = False, - logger=None, - device=None, - ): - super().__init__() - self.text_encoder_type = text_encoder_type - self.max_length = max_length - self.precision = text_encoder_precision - self.model_path = text_encoder_path - self.tokenizer_type = (tokenizer_type if tokenizer_type is not None else text_encoder_type) - self.tokenizer_path = (tokenizer_path if tokenizer_path is not None else text_encoder_path) - self.use_attention_mask = use_attention_mask - if prompt_template_video is not None: - assert (use_attention_mask is True), "Attention mask is True required when training videos." - self.input_max_length = (input_max_length if input_max_length is not None else max_length) - self.prompt_template = prompt_template - self.prompt_template_video = prompt_template_video - self.hidden_state_skip_layer = hidden_state_skip_layer - self.apply_final_norm = apply_final_norm - self.reproduce = reproduce - self.logger = logger - - self.use_template = self.prompt_template is not None - if self.use_template: - assert (isinstance(self.prompt_template, dict) and "template" in self.prompt_template - ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}" - assert "{}" in str(self.prompt_template["template"]), ( - "`prompt_template['template']` must contain a placeholder `{}` for the input text, " - f"got {self.prompt_template['template']}") - - self.use_video_template = self.prompt_template_video is not None - if self.use_video_template: - if self.prompt_template_video is not None: - assert ( - isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video - ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}" - assert "{}" in str(self.prompt_template_video["template"]), ( - "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, " - f"got {self.prompt_template_video['template']}") - - if "t5" in text_encoder_type: - self.output_key = output_key or "last_hidden_state" - elif "clip" in text_encoder_type: - self.output_key = output_key or "pooler_output" - elif "llm" in text_encoder_type or "glm" in text_encoder_type: - self.output_key = output_key or "last_hidden_state" - else: - raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") - - self.model, self.model_path = load_text_encoder( - text_encoder_type=self.text_encoder_type, - text_encoder_precision=self.precision, - text_encoder_path=self.model_path, - logger=self.logger, - device=device, - ) - self.dtype = self.model.dtype - self.device = self.model.device - - self.tokenizer, self.tokenizer_path = load_tokenizer( - tokenizer_type=self.tokenizer_type, - tokenizer_path=self.tokenizer_path, - padding_side="right", - logger=self.logger, - ) - - def __repr__(self): - return f"{self.text_encoder_type} ({self.precision} - {self.model_path})" - - @staticmethod - def apply_text_to_template(text, template, prevent_empty_text=True): - """ - Apply text to template. - - Args: - text (str): Input text. - template (str or list): Template string or list of chat conversation. - prevent_empty_text (bool): If True, we will prevent the user text from being empty - by adding a space. Defaults to True. - """ - if isinstance(template, str): - # Will send string to tokenizer. Used for llm - return template.format(text) - else: - raise TypeError(f"Unsupported template type: {type(template)}") - - def text2tokens(self, text, data_type="image"): - """ - Tokenize the input text. - - Args: - text (str or list): Input text. - """ - tokenize_input_type = "str" - if self.use_template: - if data_type == "image": - prompt_template = self.prompt_template["template"] - elif data_type == "video": - prompt_template = self.prompt_template_video["template"] - else: - raise ValueError(f"Unsupported data type: {data_type}") - if isinstance(text, (list, tuple)): - text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text] - if isinstance(text[0], list): - tokenize_input_type = "list" - elif isinstance(text, str): - text = self.apply_text_to_template(text, prompt_template) - if isinstance(text, list): - tokenize_input_type = "list" - else: - raise TypeError(f"Unsupported text type: {type(text)}") - - kwargs = dict( - truncation=True, - max_length=self.max_length, - padding="max_length", - return_tensors="pt", - ) - if tokenize_input_type == "str": - return self.tokenizer( - text, - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=True, - **kwargs, - ) - elif tokenize_input_type == "list": - return self.tokenizer.apply_chat_template( - text, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - **kwargs, - ) - else: - raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}") - - def encode( - self, - batch_encoding, - use_attention_mask=None, - output_hidden_states=False, - do_sample=None, - hidden_state_skip_layer=None, - return_texts=False, - data_type="image", - device=None, - ): - """ - Args: - batch_encoding (dict): Batch encoding from tokenizer. - use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask. - Defaults to None. - output_hidden_states (bool): Whether to output hidden states. If False, return the value of - self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer, - output_hidden_states will be set True. Defaults to False. - do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None. - When self.produce is False, do_sample is set to True by default. - hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer. - If None, self.output_key will be used. Defaults to None. - return_texts (bool): Whether to return the decoded texts. Defaults to False. - """ - device = self.model.device if device is None else device - use_attention_mask = use_default(use_attention_mask, self.use_attention_mask) - hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer) - do_sample = use_default(do_sample, not self.reproduce) - attention_mask = (batch_encoding["attention_mask"].to(device) if use_attention_mask else None) - outputs = self.model( - input_ids=batch_encoding["input_ids"].to(device), - attention_mask=attention_mask, - output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None, - ) - if hidden_state_skip_layer is not None: - last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)] - # Real last hidden state already has layer norm applied. So here we only apply it - # for intermediate layers. - if hidden_state_skip_layer > 0 and self.apply_final_norm: - last_hidden_state = self.model.final_layer_norm(last_hidden_state) - else: - last_hidden_state = outputs[self.output_key] - - # Remove hidden states of instruction tokens, only keep prompt tokens. - if self.use_template: - if data_type == "image": - crop_start = self.prompt_template.get("crop_start", -1) - elif data_type == "video": - crop_start = self.prompt_template_video.get("crop_start", -1) - else: - raise ValueError(f"Unsupported data type: {data_type}") - if crop_start > 0: - last_hidden_state = last_hidden_state[:, crop_start:] - attention_mask = (attention_mask[:, crop_start:] if use_attention_mask else None) - - if output_hidden_states: - return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states) - return TextEncoderModelOutput(last_hidden_state, attention_mask) - - def forward( - self, - text, - use_attention_mask=None, - output_hidden_states=False, - do_sample=False, - hidden_state_skip_layer=None, - return_texts=False, - ): - batch_encoding = self.text2tokens(text) - return self.encode( - batch_encoding, - use_attention_mask=use_attention_mask, - output_hidden_states=output_hidden_states, - do_sample=do_sample, - hidden_state_skip_layer=hidden_state_skip_layer, - return_texts=return_texts, - ) diff --git a/fastvideo/v1/models/loader.py b/fastvideo/v1/models/loader.py index f74ef62d0..2de63095f 100644 --- a/fastvideo/v1/models/loader.py +++ b/fastvideo/v1/models/loader.py @@ -156,7 +156,7 @@ class VAELoader(ComponentLoader): def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): """Load the VAE based on the model path, architecture, and inference args.""" # TODO(will): move this to a constants file - from fastvideo.v1.pipelines.hunyuan.constants import PRECISION_TO_TYPE + from fastvideo.v1.utils import PRECISION_TO_TYPE config = get_diffusers_config(model=model_path) diff --git a/fastvideo/v1/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py index 114dc6a82..64f32dca0 100644 --- a/fastvideo/v1/pipelines/stages/decoding.py +++ b/fastvideo/v1/pipelines/stages/decoding.py @@ -6,7 +6,7 @@ import torch from fastvideo.v1.pipelines.stages.base import PipelineStage from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.inference_args import InferenceArgs -from fastvideo.v1.models.hunyuan.constants import PRECISION_TO_TYPE +from fastvideo.v1.utils import PRECISION_TO_TYPE from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py index e9285cb14..6f2997c60 100644 --- a/fastvideo/v1/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -11,9 +11,10 @@ from einops import rearrange from fastvideo.v1.pipelines.stages.base import PipelineStage from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.inference_args import InferenceArgs -from fastvideo.v1.models.hunyuan.constants import PRECISION_TO_TYPE -from fastvideo.v1.utils.parallel_states import nccl_info, get_sequence_parallel_state -from fastvideo.v1.utils.communications import all_gather +from fastvideo.v1.utils import PRECISION_TO_TYPE +# TODO(will-refactor): change this to fastvideo.distributed +from fastvideo.utils.parallel_states import nccl_info, get_sequence_parallel_state +from fastvideo.utils.communications import all_gather from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/v1/sample/v1_fastvideo_inference.py b/fastvideo/v1/sample/v1_fastvideo_inference.py index b0ae9842a..c61b71f6a 100644 --- a/fastvideo/v1/sample/v1_fastvideo_inference.py +++ b/fastvideo/v1/sample/v1_fastvideo_inference.py @@ -11,7 +11,7 @@ from einops import rearrange # Fix the import path from fastvideo.v1.inference_engine import InferenceEngine from fastvideo.v1.inference_args import InferenceArgs -from fastvideo.v1.utils.utils import FlexibleArgumentParser +from fastvideo.v1.utils import FlexibleArgumentParser from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/v1/utils/utils.py b/fastvideo/v1/utils.py similarity index 95% rename from fastvideo/v1/utils/utils.py rename to fastvideo/v1/utils.py index 2b5c5ee68..8f0741ed7 100644 --- a/fastvideo/v1/utils/utils.py +++ b/fastvideo/v1/utils.py @@ -3,6 +3,7 @@ import fastvideo.v1.envs as envs import inspect from fastvideo.v1.logger import init_logger import argparse +import math import sys from typing import List, Dict, Union, Type, Any, TypeVar import yaml @@ -13,6 +14,14 @@ logger = init_logger(__name__) T = TypeVar("T") +# TODO(will): used to convert inference_args.precision to torch.dtype. Find a +# cleaner way to do this. +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + def find_nccl_library() -> str: """ @@ -288,3 +297,16 @@ def warn_for_unimplemented_methods(cls: Type[T]) -> Type[T]: type.__setattr__(cls, '__init__', wrapped_init) return cls + + +def align_to(value, alignment): + """align height, width according to alignment + + Args: + value (int): height or width + alignment (int): target alignment factor + + Returns: + int: the aligned value + """ + return int(math.ceil(value / alignment) * alignment) From fb44fbaa1ce7a29b12ab474d70029d5bd33cbbf0 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 15 Mar 2025 19:30:12 +0000 Subject: [PATCH 08/12] debugging encoders --- fastvideo/v1/attention/__init__.py | 4 +- fastvideo/v1/attention/flash_attn.py | 97 +-------- fastvideo/v1/loader/loader.py | 55 +---- fastvideo/v1/models/encoders/clip.py | 43 ++-- fastvideo/v1/models/encoders/llama.py | 40 ++-- fastvideo/v1/models/loader.py | 2 +- fastvideo/v1/models/text_encoder.py | 2 +- fastvideo/v1/pipelines/__init__.py | 1 - .../v1/pipelines/stages/prompt_encoding.py | 5 +- fastvideo/v1/tests/test_clip_encoder.py | 62 ++++-- fastvideo/v1/tests/test_hunyuan_vae.py | 4 +- fastvideo/v1/tests/test_llama_encoder.py | 195 ++++++++++++++++++ 12 files changed, 308 insertions(+), 202 deletions(-) diff --git a/fastvideo/v1/attention/__init__.py b/fastvideo/v1/attention/__init__.py index 208b46685..7dc81b89b 100644 --- a/fastvideo/v1/attention/__init__.py +++ b/fastvideo/v1/attention/__init__.py @@ -1,3 +1,3 @@ -from .flash_attn import (DistributedAttention, TextEncoderFlashAttention, LocalAttention) +from .flash_attn import (DistributedAttention, LocalAttention) -__all__ = ["DistributedAttention", "TextEncoderFlashAttention", "LocalAttention"] \ No newline at end of file +__all__ = ["DistributedAttention", "LocalAttention"] \ No newline at end of file diff --git a/fastvideo/v1/attention/flash_attn.py b/fastvideo/v1/attention/flash_attn.py index d8c7711b1..5ba432389 100644 --- a/fastvideo/v1/attention/flash_attn.py +++ b/fastvideo/v1/attention/flash_attn.py @@ -135,99 +135,4 @@ class LocalAttention(nn.Module): causal=self.causal ) - return output - -class TextEncoderFlashAttention: - """Simplified Flash Attention implementation for text encoders in diffusion models. - - This implementation removes KV caching and other generation-specific features, - focusing only on the encoder attention pattern needed for text encoders in - diffusion models like DiT. - """ - - def __init__( - self, - softmax_scale: Optional[float] = None, - ) -> None: - """Initialize the text encoder flash attention implementation. - - Args: - num_heads: Number of attention heads - head_size: Size of each attention head - scale: Attention scale factor. If None, defaults to 1/sqrt(head_size) - """ - self.softmax_scale = softmax_scale - - # Set defaults for optional parameters - self.alibi_slopes = None - self.sliding_window = (-1, -1) # No sliding window by default - self.logits_soft_cap = 0 # No soft cap by default - - - - def forward( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - seq_lens: Optional[List[int]] = None, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Forward pass with FlashAttention for text encoder. - - Args: - query: shape = [num_tokens, num_heads, head_size] - key: shape = [num_tokens, num_kv_heads, head_size] - value: shape = [num_tokens, num_kv_heads, head_size] - seq_lens: List of sequence lengths for batched input - output: shape = [num_tokens, num_heads, head_size] - Pre-allocated output tensor (optional) - - Returns: - torch.Tensor: Output tensor with shape [num_tokens, num_heads, head_size] - """ - device = query.device - batch_size = len(seq_lens) if seq_lens is not None else 1 - - if output is None: - output = torch.empty_like(query) - - if seq_lens is None: - # Single sequence case - use standard flash attention - return flash_attn_varlen_func( - q=query, - k=key, - v=value, - softmax_scale=self.softmax_scale, - causal=False, # Non-causal for encoder - # window_size=self.sliding_window, - # alibi_slopes=self.alibi_slopes, - # softcap=self.logits_soft_cap, - out=output, - # fa_version=self.vllm_flash_attn_version, - ) - else: - assert False, "Not implemented" - # Multiple sequences case - use varlen flash attention - max_seq_len = max(seq_lens) - seq_start_loc = list(accumulate(seq_lens, initial=0)) - seq_start_loc_tensor = torch.tensor(seq_start_loc, - dtype=torch.int32, - device=device) - - return flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=seq_start_loc_tensor, - cu_seqlens_k=seq_start_loc_tensor, - max_seqlen_q=max_seq_len, - max_seqlen_k=max_seq_len, - softmax_scale=self.scale, - causal=False, # Non-causal for encoder - window_size=self.sliding_window, - alibi_slopes=self.alibi_slopes, - softcap=self.logits_soft_cap, - out=output, - fa_version=self.vllm_flash_attn_version, - ) + return output \ No newline at end of file diff --git a/fastvideo/v1/loader/loader.py b/fastvideo/v1/loader/loader.py index 09a8c670e..b0ae7a189 100644 --- a/fastvideo/v1/loader/loader.py +++ b/fastvideo/v1/loader/loader.py @@ -87,14 +87,12 @@ logger = init_logger(__name__) def _initialize_model( model_path: str, model_config, - inference_args: InferenceArgs, ) -> nn.Module: """Initialize a model with the given configurations. Args: model_name: The name of the model architecture model_path: Path to the model directory - inference_args: Inference arguments component_name: Optional component name (e.g., "text_encoder", "vae") If provided, will load the config from the component subdirectory @@ -115,41 +113,9 @@ def _initialize_model( return model_cls(model_config) - # Store the config for later use - if not hasattr(inference_args, 'hf_configs'): - inference_args.hf_configs = {} - - if component_name: - inference_args.hf_configs[component_name] = hf_config - else: - inference_args.hf_config = hf_config - - # Initialize the model based on the model_name - if model_name == "CLIPVisionModel": - from fastvideo.v1.models.encoders.clip import CLIPVisionModel - from transformers import CLIPVisionConfig - - # If the config is not a CLIPVisionConfig, convert it - if not isinstance(hf_config, CLIPVisionConfig): - # This might need adjustment based on your specific requirements - clip_config = CLIPVisionConfig.from_dict(hf_config.to_dict()) - else: - clip_config = hf_config - - # Initialize the model with the config - model = CLIPVisionModel(config=clip_config) - else: - # Add more model types as needed - raise ValueError(f"Unsupported model type: {model_name}") - - return model - class BaseModelLoader(ABC): """Base class for model loaders.""" - def __init__(self, inference_args: InferenceArgs): - self.inference_args = inference_args - @abstractmethod def download_model(self, inference_args: InferenceArgs) -> None: """Download a model so that it can be immediately loaded.""" @@ -186,9 +152,6 @@ class DefaultModelLoader(BaseModelLoader): counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 - def __init__(self, inference_args: InferenceArgs): - super().__init__(inference_args) - def _prepare_weights( self, model_name_or_path: str, @@ -344,7 +307,7 @@ class DefaultModelLoader(BaseModelLoader): # TODO(will): add support for other dtypes with set_default_torch_dtype(torch.bfloat16): with target_device: - model = _initialize_model(model_path, model_config, inference_args) + model = _initialize_model(model_path, model_config) weights_to_load = {name for name, _ in model.named_parameters()} model_config.model = model_path @@ -357,12 +320,12 @@ class DefaultModelLoader(BaseModelLoader): self.counter_before_loading_weights) # We only enable strict check for non-quantized models # that have loaded weights tracking currently. - if loaded_weights is not None: - weights_not_loaded = weights_to_load - loaded_weights - if weights_not_loaded: - raise ValueError( - "Following weights were not initialized from " - f"checkpoint: {weights_not_loaded}") + # if loaded_weights is not None: + weights_not_loaded = weights_to_load - loaded_weights + if weights_not_loaded: + raise ValueError( + "Following weights were not initialized from " + f"checkpoint: {weights_not_loaded}") # TODO(will): add support for training/finetune return model.eval() @@ -371,9 +334,9 @@ class DefaultModelLoader(BaseModelLoader): -def get_model_loader(inference_args: InferenceArgs) -> BaseModelLoader: +def get_model_loader() -> BaseModelLoader: """Get a model loader based on the load format.""" # if isinstance(load_config.load_format, type): # return load_config.load_format(load_config) - return DefaultModelLoader(inference_args) \ No newline at end of file + return DefaultModelLoader() \ No newline at end of file diff --git a/fastvideo/v1/models/encoders/clip.py b/fastvideo/v1/models/encoders/clip.py index 0960c3821..9be9332da 100644 --- a/fastvideo/v1/models/encoders/clip.py +++ b/fastvideo/v1/models/encoders/clip.py @@ -7,8 +7,10 @@ import torch import torch.nn as nn from transformers import CLIPVisionConfig, CLIPTextConfig from transformers.modeling_outputs import BaseModelOutputWithPooling +# from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask from vllm.attention.layer import MultiHeadAttention +# from fastvideo.v1.attention.flash_attn import LocalAttention from fastvideo.v1.distributed import divide, get_tensor_model_parallel_world_size from fastvideo.v1.layers.activation import get_act_fn from fastvideo.v1.layers.linear import (ColumnParallelLinear, @@ -21,6 +23,10 @@ from vllm.model_executor.models.interfaces import SupportsQuant from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs +from fastvideo.v1.logger import init_logger + +logger = init_logger(__name__) + class QuantizationConfig: pass @@ -229,7 +235,7 @@ class CLIPEncoderLayer(nn.Module): def __init__( self, - config: CLIPVisionConfig, + config: CLIPTextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: @@ -368,26 +374,26 @@ class CLIPTextTransformer(nn.Module): # CLIP's text model uses causal mask, prepare it here. # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 - causal_attention_mask = _create_4d_causal_attention_mask( - input_shape, hidden_states.dtype, device=hidden_states.device - ) + # causal_attention_mask = _create_4d_causal_attention_mask( + # input_shape, hidden_states.dtype, device=hidden_states.device + # ) - # expand attention_mask - if attention_mask is not None and not self._use_flash_attention_2: - raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer") - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + # # expand attention_mask + # if attention_mask is not None and not self._use_flash_attention_2: + # raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer") + # # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + # attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) encoder_outputs = self.encoder( inputs_embeds=hidden_states, - attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + # attention_mask=attention_mask, + # causal_attention_mask=causal_attention_mask, + # output_attentions=output_attentions, + return_all_hidden_states=output_hidden_states, + # return_dict=return_dict, ) - last_hidden_state = encoder_outputs[0] + last_hidden_state = encoder_outputs[-1] last_hidden_state = self.final_layer_norm(last_hidden_state) if self.eos_token_id == 2: @@ -415,11 +421,12 @@ class CLIPTextTransformer(nn.Module): if not return_dict: return (last_hidden_state, pooled_output) + encoder_outputs[1:] + # return last_hidden_state return BaseModelOutputWithPooling( last_hidden_state=last_hidden_state, pooler_output=pooled_output, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, + hidden_states=encoder_outputs, + # attentions=encoder_outputs.attentions, ) @@ -456,7 +463,7 @@ class CLIPTextModel(nn.Module): position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=None, ) def load_weights(self, weights: Iterable[Tuple[str, diff --git a/fastvideo/v1/models/encoders/llama.py b/fastvideo/v1/models/encoders/llama.py index e02226e47..21750b293 100644 --- a/fastvideo/v1/models/encoders/llama.py +++ b/fastvideo/v1/models/encoders/llama.py @@ -29,7 +29,7 @@ from torch import nn from transformers import LlamaConfig from transformers.modeling_outputs import BaseModelOutputWithPast -from fastvideo.v1.attention import TextEncoderFlashAttention +from vllm.attention.layer import MultiHeadAttention from fastvideo.v1.distributed import get_tensor_model_parallel_world_size from fastvideo.v1.layers.activation import SiluAndMul @@ -65,6 +65,7 @@ class LlamaMLP(nn.Module): self.gate_up_proj = MergedColumnParallelLinear( input_size=hidden_size, output_sizes=[intermediate_size] * 2, + # output_size=intermediate_size, bias=bias, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", @@ -130,6 +131,9 @@ class LlamaAttention(nn.Module): self.scaling = self.head_dim**-0.5 self.rope_theta = rope_theta self.max_position_embeddings = max_position_embeddings + print("head_dim", self.head_dim) + print("total_num_heads", self.total_num_heads) + print("total_num_kv_heads", self.total_num_kv_heads) self.qkv_proj = QKVParallelLinear( hidden_size=hidden_size, @@ -163,26 +167,10 @@ class LlamaAttention(nn.Module): is_neox_style=is_neox_style, ) - if hasattr(config, "interleaved_sliding_window"): - interleaved_sliding_window = config.interleaved_sliding_window - if isinstance(interleaved_sliding_window, int): - sliding_window = interleaved_sliding_window - elif isinstance(interleaved_sliding_window, list): - sw_idx = layer_idx % len(interleaved_sliding_window) - sliding_window = interleaved_sliding_window[sw_idx] - else: - raise ValueError( - f"{type(interleaved_sliding_window)} is not supported.") - else: - sliding_window = None - - self.attn = TextEncoderFlashAttention( - self.scaling, - # num_kv_heads=self.num_kv_heads, - # quant_config=quant_config, - # per_layer_sliding_window=sliding_window, - # prefix=f"{prefix}.attn", - ) + self.attn = MultiHeadAttention(self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads) def forward( self, @@ -192,6 +180,7 @@ class LlamaAttention(nn.Module): qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output @@ -315,7 +304,8 @@ class LlamaModel(nn.Module): def forward( self, input_ids: Optional[torch.Tensor], - positions: torch.Tensor, + positions: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, output_hidden_states: Optional[bool] = None, ) -> torch.Tensor: @@ -327,6 +317,12 @@ class LlamaModel(nn.Module): else: hidden_states = self.get_input_embeddings(input_ids) residual = None + print("after embed hidden_states", hidden_states.shape) + + if positions is None: + positions = torch.arange( + 0, hidden_states.shape[1], device=hidden_states.device + ).unsqueeze(0) all_hidden_states = () if output_hidden_states else None for layer in self.layers: diff --git a/fastvideo/v1/models/loader.py b/fastvideo/v1/models/loader.py index 2de63095f..4a5511c9f 100644 --- a/fastvideo/v1/models/loader.py +++ b/fastvideo/v1/models/loader.py @@ -95,7 +95,7 @@ class TextEncoderLoader(ComponentLoader): inference_args=inference_args, ) logger.info(f"HF Model config: {model_config}") - model_loader = get_model_loader(inference_args) + model_loader = get_model_loader() model = model_loader.load_model(model_path, model_config, inference_args) return model diff --git a/fastvideo/v1/models/text_encoder.py b/fastvideo/v1/models/text_encoder.py index 4449cfd9b..8268d224b 100644 --- a/fastvideo/v1/models/text_encoder.py +++ b/fastvideo/v1/models/text_encoder.py @@ -203,7 +203,7 @@ class TextEncoder(nn.Module): hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer) do_sample = use_default(do_sample, not self.reproduce) attention_mask = (batch_encoding["attention_mask"].to(device) if use_attention_mask else None) - print(f"batch_encoding: {batch_encoding}") + # print(f"batch_encoding: {batch_encoding}") # logger.info(f"attention_mask: {attention_mask}") outputs = self.model( input_ids=batch_encoding["input_ids"].to(device), diff --git a/fastvideo/v1/pipelines/__init__.py b/fastvideo/v1/pipelines/__init__.py index 2aefab8fd..e100e85d1 100644 --- a/fastvideo/v1/pipelines/__init__.py +++ b/fastvideo/v1/pipelines/__init__.py @@ -16,7 +16,6 @@ from huggingface_hub import snapshot_download from transformers import PretrainedConfig from fastvideo.v1.models.hf_transformer_utils import get_hf_config, get_diffusers_config from fastvideo.v1.models import get_scheduler -from fastvideo.v1.loader.loader import get_model_loader import glob from fastvideo.v1.loader.fsdp_load import load_fsdp_model from fastvideo.v1.models.loader import PipelineComponentLoader diff --git a/fastvideo/v1/pipelines/stages/prompt_encoding.py b/fastvideo/v1/pipelines/stages/prompt_encoding.py index 3234e9d88..24ff64dbf 100644 --- a/fastvideo/v1/pipelines/stages/prompt_encoding.py +++ b/fastvideo/v1/pipelines/stages/prompt_encoding.py @@ -87,13 +87,16 @@ class PromptEncodingStage(PipelineStage): attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len) if text_encoder is not None: - prompt_embeds_dtype = text_encoder.dtype + # TODO(will-refactor): use text_encoder.dtype + prompt_embeds_dtype = torch.float16 elif self.transformer is not None: prompt_embeds_dtype = self.transformer.dtype else: prompt_embeds_dtype = prompt_embeds.dtype prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) + # print("prompt_embeds", type(prompt_embeds)) + # logger.info(f"prompt_embeds shape: {prompt_embeds.shape}") if prompt_embeds.ndim == 2: bs_embed, _ = prompt_embeds.shape diff --git a/fastvideo/v1/tests/test_clip_encoder.py b/fastvideo/v1/tests/test_clip_encoder.py index 26e71432d..a67e1e6f9 100644 --- a/fastvideo/v1/tests/test_clip_encoder.py +++ b/fastvideo/v1/tests/test_clip_encoder.py @@ -1,13 +1,14 @@ from fastvideo.v1.models.encoders.clip import CLIPTextModel -from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.v0_reference_src.models.hunyuan.text_encoder import TextEncoder, load_text_encoder, load_tokenizer import os import torch import torch.nn as nn import argparse import numpy as np -from transformers import CLIPTextModel as HFCLIPTextModel from fastvideo.v1.logger import init_logger -from fastvideo.v1.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer +from transformers import AutoConfig +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel +# from fastvideo.v1.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer logger = init_logger(__name__) @@ -52,11 +53,13 @@ def setup_args(): parser = argparse.ArgumentParser(description='CLIP Text Encoder Test') parser.add_argument('--model-path', type=str, default="openai/clip-vit-large-patch14", help='Path to the CLIP model') - parser.add_argument('--precision', type=str, default="bfloat16", + parser.add_argument('--precision', type=str, default="float16", help='Precision to use for the model (float32, float16, bfloat16)') return parser.parse_args() def test_clip_encoder(): + init_distributed_environment(world_size=1, rank=0, distributed_init_method="env://", local_rank=0, backend="nccl") + initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1, backend="nccl") args = setup_args() # Set fixed random seed for reproducibility @@ -67,25 +70,56 @@ def test_clip_encoder(): # Initialize the two model implementations logger.info(f"Loading models from {args.model_path}") + model_path = "/home/ubuntu/src/FastVideo/data/hunyuanvideo-community/HunyuanVideo/text_encoder_2" + + # config = json.load(open(os.path.join(model_path, "config.json"))) + + hf_config = AutoConfig.from_pretrained(model_path) + print(hf_config) + print(hf_config.use_return_dict) # Load our implementation using the loader from text_encoder/__init__.py model1, _ = load_text_encoder( text_encoder_type="clipL", - text_encoder_precision=args.precision, - text_encoder_path=args.model_path, + text_encoder_precision='fp16', + text_encoder_path=model_path, logger=logger, device=device ) + + from fastvideo.v1.models.loader import TextEncoderLoader + from fastvideo.v1.loader.loader import DefaultModelLoader + loader = DefaultModelLoader() + args.device_str = "cuda:0" + model2 = loader.load_model(model_path, hf_config, args) # Load the HuggingFace implementation directly - model2 = HFCLIPTextModel.from_pretrained(args.model_path) - if args.precision == "bfloat16": - model2 = model2.to(torch.bfloat16) - elif args.precision == "float16": - model2 = model2.to(torch.float16) + # model2 = CLIPTextModel(hf_config) + model2 = model2.to(torch.float16) model2 = model2.to(device) 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(f"Model1 has {len(params1)} parameters") + logger.info(f"Model2 has {len(params2)} parameters") + + # Compare a few key parameters + weight_diffs = [] + for (name1, param1), (name2, param2) in zip( + sorted(params1.items()), sorted(params2.items()) + ): + # if len(weight_diffs) < 5: # Just check a few parameters + max_diff = torch.max(torch.abs(param1 - param2)).item() + mean_diff = torch.mean(torch.abs(param1 - param2)).item() + weight_diffs.append((name1, name2, max_diff, mean_diff)) + logger.info(f"Parameter: {name1} vs {name2}") + logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") + # Load tokenizer tokenizer, _ = load_tokenizer( tokenizer_type="clipL", @@ -121,17 +155,21 @@ def test_clip_encoder(): attention_mask=tokens.attention_mask, output_hidden_states=True ) + + logger.info(f"Testing model2") # Get embeddings from HuggingFace implementation outputs2 = model2( input_ids=tokens.input_ids, - attention_mask=tokens.attention_mask, + # attention_mask=tokens.attention_mask, output_hidden_states=True ) # Compare last hidden states last_hidden_state1 = outputs1.last_hidden_state last_hidden_state2 = outputs2.last_hidden_state + # print("last_hidden_state1", last_hidden_state1) + # print("last_hidden_state2", last_hidden_state2) 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}" diff --git a/fastvideo/v1/tests/test_hunyuan_vae.py b/fastvideo/v1/tests/test_hunyuan_vae.py index cb3f0e88b..39e0ce4d2 100644 --- a/fastvideo/v1/tests/test_hunyuan_vae.py +++ b/fastvideo/v1/tests/test_hunyuan_vae.py @@ -1,4 +1,4 @@ -from fastvideo.v1.models.vaes import AutoencoderKLHunyuanVideo as MyHunyuanVAE +from fastvideo.v1.models.vaes.hunyuanvae import AutoencoderKLHunyuanVideo as MyHunyuanVAE from diffusers import AutoencoderKLHunyuanVideo as DiffusersHunyuanVAE import os @@ -73,7 +73,7 @@ def test_hunyuan_vae(): device = torch.device("cuda:0") print # Initialize the two model implementations - path = "/mbz/users/hao.zhang/peiyuan/FastVideo/data/hunyuanvideo_community/vae" + path = "/home/ubuntu/src/FastVideo/data/hunyuanvideo-community/HunyuanVideo/vae" config_path = os.path.join(path, "config.json") config = json.load(open(config_path)) config.pop("_class_name") diff --git a/fastvideo/v1/tests/test_llama_encoder.py b/fastvideo/v1/tests/test_llama_encoder.py index e69de29bb..9f6444d96 100644 --- a/fastvideo/v1/tests/test_llama_encoder.py +++ b/fastvideo/v1/tests/test_llama_encoder.py @@ -0,0 +1,195 @@ +from fastvideo.v1.models.encoders.llama import LlamaModel +from fastvideo.v1.v0_reference_src.models.hunyuan.text_encoder import TextEncoder, load_text_encoder, load_tokenizer +import os +import torch +import torch.nn as nn +import argparse +import numpy as np +from fastvideo.v1.logger import init_logger +from transformers import AutoConfig +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel + +logger = init_logger(__name__) + +def initialize_identical_weights(model1, model2, seed=42): + """Initialize both models with identical weights using a fixed seed for reproducibility.""" + # Get all parameters from both models + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Initialize each layer with identical values + with torch.no_grad(): + # Initialize weights + for name1, param1 in params1.items(): + if 'weight' in name1: + # Set seed before each weight initialization + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + + for name2, param2 in params2.items(): + if 'weight' in name2: + # Reset seed to get same initialization + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + + # Initialize biases + for name1, param1 in params1.items(): + if 'bias' in name1: + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + param1.data = param1.data.to(torch.bfloat16) + + for name2, param2 in params2.items(): + if 'bias' in name2: + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + param2.data = param2.data.to(torch.bfloat16) + + logger.info("Both models initialized with identical weights in bfloat16") + return model1, model2 + +def setup_args(): + parser = argparse.ArgumentParser(description='LLaMA Encoder Test') + parser.add_argument('--model-path', type=str, default="meta-llama/Llama-2-7b-hf", + help='Path to the LLaMA model') + parser.add_argument('--precision', type=str, default="float16", + help='Precision to use for the model (float32, float16, bfloat16)') + return parser.parse_args() + +def test_llama_encoder(): + init_distributed_environment(world_size=1, rank=0, distributed_init_method="env://", local_rank=0, backend="nccl") + initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1, backend="nccl") + args = setup_args() + + # Set fixed random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + # Initialize the two model implementations + logger.info(f"Loading models from {args.model_path}") + model_path = "/home/ubuntu/src/FastVideo/data/hunyuanvideo-community/HunyuanVideo/text_encoder" + + hf_config = AutoConfig.from_pretrained(model_path) + print(hf_config) + + # Load our implementation using the loader from text_encoder/__init__.py + model1, _ = load_text_encoder( + text_encoder_type="llm", + text_encoder_precision='fp16', + text_encoder_path=model_path, + logger=logger, + device=device + ) + + from fastvideo.v1.models.loader import TextEncoderLoader + from fastvideo.v1.loader.loader import DefaultModelLoader + loader = DefaultModelLoader() + args.device_str = "cuda:0" + model2 = loader.load_model(model_path, hf_config, args) + + # Convert to float16 and move to device + model2 = model2.to(torch.float16) + model2 = model2.to(device) + 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(f"Model1 has {len(params1)} parameters") + logger.info(f"Model2 has {len(params2)} parameters") + + # Compare a few key parameters + weight_diffs = [] + weights = ["layers.{}.input_layernorm.weight", "layers.{}.post_attention_layernorm.weight"] + # for (name1, param1), (name2, param2) in zip( + # sorted(params1.items()), sorted(params2.items()) + # ): + for l in range(hf_config.num_hidden_layers): + for w in weights: + name1 = w.format(l) + name2 = w.format(l) + p1 = params1[name1] + p2 = params2[name2] + print(type(p2)) + if "gate_up" in name2: + print("skipping gate_up") + continue + try: + logger.info(f"Parameter: {name1} vs {name2}") + max_diff = torch.max(torch.abs(p1 - p2)).item() + mean_diff = torch.mean(torch.abs(p1 - p2)).item() + weight_diffs.append((name1, name2, max_diff, mean_diff)) + logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") + except Exception as e: + logger.info(f"Error comparing {name1} and {name2}: {e}") + + tokenizer_path = "/home/ubuntu/src/FastVideo/data/hunyuanvideo-community/HunyuanVideo/tokenizer" + # Load tokenizer + tokenizer, _ = load_tokenizer( + tokenizer_type="llm", + tokenizer_path=tokenizer_path, + logger=logger + ) + + # 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 LLaMA encoder with sample prompts") + + with torch.no_grad(): + for prompt in prompts: + logger.info(f"Testing prompt: '{prompt}'") + + # Tokenize the prompt + tokens = tokenizer( + prompt, + padding="max_length", + max_length=128, + truncation=True, + return_tensors="pt" + ).to(device) + + # Get outputs from our implementation + outputs1 = model1( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True + ) + + logger.info(f"Testing model2") + + # Get outputs from HuggingFace implementation + outputs2 = model2( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True + ) + + # Compare last hidden states + last_hidden_state1 = outputs1.last_hidden_state + last_hidden_state2 = outputs2.last_hidden_state + + 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(f"Maximum difference in last hidden states: {max_diff_hidden.item()}") + logger.info(f"Mean difference in last hidden states: {mean_diff_hidden.item()}") + + + logger.info("Test passed! Both LLaMA encoder implementations produce similar outputs.") + logger.info("Test completed successfully") + +if __name__ == "__main__": + test_llama_encoder() From 759f243cf30dd15108691aed50a423f057ea6f00 Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Sat, 15 Mar 2025 21:13:38 +0000 Subject: [PATCH 09/12] fix attn --- fastvideo/v1/loader/loader.py | 2 +- fastvideo/v1/models/encoders/clip.py | 10 +++++++-- fastvideo/v1/models/encoders/llama.py | 21 ++++++++++++------ fastvideo/v1/tests/test_clip_encoder.py | 27 ++++++++++++------------ fastvideo/v1/tests/test_hunyuan_vae.py | 2 +- fastvideo/v1/tests/test_llama_encoder.py | 20 ++++++++++++------ 6 files changed, 52 insertions(+), 30 deletions(-) diff --git a/fastvideo/v1/loader/loader.py b/fastvideo/v1/loader/loader.py index b0ae7a189..1980aefaf 100644 --- a/fastvideo/v1/loader/loader.py +++ b/fastvideo/v1/loader/loader.py @@ -305,7 +305,7 @@ class DefaultModelLoader(BaseModelLoader): logger.info(f"Loading model on device: {inference_args.device_str}") target_device = torch.device(inference_args.device_str) # TODO(will): add support for other dtypes - with set_default_torch_dtype(torch.bfloat16): + with set_default_torch_dtype(torch.float16): with target_device: model = _initialize_model(model_path, model_config) diff --git a/fastvideo/v1/models/encoders/clip.py b/fastvideo/v1/models/encoders/clip.py index 9be9332da..c1bdc1a11 100644 --- a/fastvideo/v1/models/encoders/clip.py +++ b/fastvideo/v1/models/encoders/clip.py @@ -195,8 +195,14 @@ class CLIPAttention(nn.Module): qkv_states, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) - out = self.attn(query_states, key_states, value_states) - attn_output, _ = self.out_proj(out) + # use flash_attn_func + from flash_attn import flash_attn_func + query_states = query_states.reshape(query_states.shape[0], query_states.shape[1], self.num_heads, self.head_dim) + key_states = key_states.reshape(key_states.shape[0], key_states.shape[1], self.num_heads, self.head_dim) + value_states = value_states.reshape(value_states.shape[0], value_states.shape[1], self.num_heads, self.head_dim) + attn_output = flash_attn_func(query_states, key_states, value_states,softmax_scale=self.scale, causal=True) + attn_output = attn_output.reshape(attn_output.shape[0], attn_output.shape[1], self.embed_dim) + attn_output, _ = self.out_proj(attn_output) return attn_output, None diff --git a/fastvideo/v1/models/encoders/llama.py b/fastvideo/v1/models/encoders/llama.py index 21750b293..5bcc6b59a 100644 --- a/fastvideo/v1/models/encoders/llama.py +++ b/fastvideo/v1/models/encoders/llama.py @@ -131,9 +131,6 @@ class LlamaAttention(nn.Module): self.scaling = self.head_dim**-0.5 self.rope_theta = rope_theta self.max_position_embeddings = max_position_embeddings - print("head_dim", self.head_dim) - print("total_num_heads", self.total_num_heads) - print("total_num_kv_heads", self.total_num_kv_heads) self.qkv_proj = QKVParallelLinear( hidden_size=hidden_size, @@ -179,9 +176,20 @@ class LlamaAttention(nn.Module): ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) + # print(f"query_states sum before apply_rotary_pos_emb: {q.float().sum()}") + # q, k = self.rotary_emb(positions, q, k) + # attn_output = self.attn(q, k, v) + # use flash_attn_func + from flash_attn import flash_attn_func + # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim) + batch_size = q.shape[0] + seq_len = q.shape[1] + q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True) + attn_output = attn_output.reshape(batch_size, seq_len, self.total_num_heads * self.head_dim) - attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output @@ -253,8 +261,10 @@ class LlamaDecoderLayer(nn.Module): else: hidden_states, residual = self.input_layernorm( hidden_states, residual) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) + # Fully Connected hidden_states, residual = self.post_attention_layernorm( @@ -317,7 +327,6 @@ class LlamaModel(nn.Module): else: hidden_states = self.get_input_embeddings(input_ids) residual = None - print("after embed hidden_states", hidden_states.shape) if positions is None: positions = torch.arange( diff --git a/fastvideo/v1/tests/test_clip_encoder.py b/fastvideo/v1/tests/test_clip_encoder.py index a67e1e6f9..c7a8297af 100644 --- a/fastvideo/v1/tests/test_clip_encoder.py +++ b/fastvideo/v1/tests/test_clip_encoder.py @@ -70,7 +70,7 @@ def test_clip_encoder(): # Initialize the two model implementations logger.info(f"Loading models from {args.model_path}") - model_path = "/home/ubuntu/src/FastVideo/data/hunyuanvideo-community/HunyuanVideo/text_encoder_2" + model_path = "data/hunyuanvideo-community/HunyuanVideo/text_encoder_2" # config = json.load(open(os.path.join(model_path, "config.json"))) @@ -109,16 +109,17 @@ def test_clip_encoder(): logger.info(f"Model2 has {len(params2)} parameters") # Compare a few key parameters - weight_diffs = [] - for (name1, param1), (name2, param2) in zip( - sorted(params1.items()), sorted(params2.items()) - ): - # if len(weight_diffs) < 5: # Just check a few parameters - max_diff = torch.max(torch.abs(param1 - param2)).item() - mean_diff = torch.mean(torch.abs(param1 - param2)).item() - weight_diffs.append((name1, name2, max_diff, mean_diff)) - logger.info(f"Parameter: {name1} vs {name2}") - logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") + + # weight_diffs = [] + # for (name1, param1), (name2, param2) in zip( + # sorted(params1.items()), sorted(params2.items()) + # ): + # # if len(weight_diffs) < 5: # Just check a few parameters + # max_diff = torch.max(torch.abs(param1 - param2)).item() + # mean_diff = torch.mean(torch.abs(param1 - param2)).item() + # weight_diffs.append((name1, name2, max_diff, mean_diff)) + # logger.info(f"Parameter: {name1} vs {name2}") + # logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") # Load tokenizer tokenizer, _ = load_tokenizer( @@ -166,8 +167,8 @@ def test_clip_encoder(): ) # Compare last hidden states - last_hidden_state1 = outputs1.last_hidden_state - last_hidden_state2 = outputs2.last_hidden_state + last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask==1] + last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask==1] # print("last_hidden_state1", last_hidden_state1) # print("last_hidden_state2", last_hidden_state2) diff --git a/fastvideo/v1/tests/test_hunyuan_vae.py b/fastvideo/v1/tests/test_hunyuan_vae.py index 39e0ce4d2..e73624379 100644 --- a/fastvideo/v1/tests/test_hunyuan_vae.py +++ b/fastvideo/v1/tests/test_hunyuan_vae.py @@ -73,7 +73,7 @@ def test_hunyuan_vae(): device = torch.device("cuda:0") print # Initialize the two model implementations - path = "/home/ubuntu/src/FastVideo/data/hunyuanvideo-community/HunyuanVideo/vae" + path = "data/hunyuanvideo-community/HunyuanVideo/vae" config_path = os.path.join(path, "config.json") config = json.load(open(config_path)) config.pop("_class_name") diff --git a/fastvideo/v1/tests/test_llama_encoder.py b/fastvideo/v1/tests/test_llama_encoder.py index 9f6444d96..11153d195 100644 --- a/fastvideo/v1/tests/test_llama_encoder.py +++ b/fastvideo/v1/tests/test_llama_encoder.py @@ -69,7 +69,7 @@ def test_llama_encoder(): # Initialize the two model implementations logger.info(f"Loading models from {args.model_path}") - model_path = "/home/ubuntu/src/FastVideo/data/hunyuanvideo-community/HunyuanVideo/text_encoder" + model_path = "data/hunyuanvideo-community/HunyuanVideo/text_encoder" hf_config = AutoConfig.from_pretrained(model_path) print(hf_config) @@ -105,6 +105,9 @@ def test_llama_encoder(): # Compare a few key parameters weight_diffs = [] + # check if embed_tokens are the same + print(model1.embed_tokens.weight.shape, model2.embed_tokens.weight.shape) + assert torch.allclose(model1.embed_tokens.weight, model2.embed_tokens.weight) weights = ["layers.{}.input_layernorm.weight", "layers.{}.post_attention_layernorm.weight"] # for (name1, param1), (name2, param2) in zip( # sorted(params1.items()), sorted(params2.items()) @@ -128,7 +131,7 @@ def test_llama_encoder(): except Exception as e: logger.info(f"Error comparing {name1} and {name2}: {e}") - tokenizer_path = "/home/ubuntu/src/FastVideo/data/hunyuanvideo-community/HunyuanVideo/tokenizer" + tokenizer_path = "data/hunyuanvideo-community/HunyuanVideo/tokenizer" # Load tokenizer tokenizer, _ = load_tokenizer( tokenizer_type="llm", @@ -139,8 +142,8 @@ def test_llama_encoder(): # Test with some sample prompts prompts = [ "Once upon a time", - "The quick brown fox jumps over", - "In a galaxy far, far away" + # "The quick brown fox jumps over", + # "In a galaxy far, far away" ] logger.info("Testing LLaMA encoder with sample prompts") @@ -159,12 +162,15 @@ def test_llama_encoder(): ).to(device) # Get outputs from our 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, output_hidden_states=True ) - + print("--------------------------------") logger.info(f"Testing model2") # Get outputs from HuggingFace implementation @@ -175,8 +181,8 @@ def test_llama_encoder(): ) # Compare last hidden states - last_hidden_state1 = outputs1.last_hidden_state - last_hidden_state2 = outputs2.last_hidden_state + last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask==1] + last_hidden_state2 = outputs2.last_hidden_state[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}" From f51e9d486b78b09b89965f93321187de5a0b41b3 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 15 Mar 2025 23:38:37 +0000 Subject: [PATCH 10/12] running, correctness isues --- fastvideo/v1/models/__init__.py | 2 +- fastvideo/v1/models/encoders/clip.py | 8 +-- fastvideo/v1/models/encoders/llama.py | 9 +++- fastvideo/v1/models/loader.py | 2 +- .../hunyuan/hunyuan_pipeline.py | 54 ++++++++++++++++++- .../v1/pipelines/stages/prompt_encoding.py | 3 ++ fastvideo/v1/sample/v1_fastvideo_inference.py | 2 +- scripts/inference/v1_inference_hunyuan.sh | 4 ++ 8 files changed, 74 insertions(+), 10 deletions(-) diff --git a/fastvideo/v1/models/__init__.py b/fastvideo/v1/models/__init__.py index e6667e618..76d2ac2dc 100644 --- a/fastvideo/v1/models/__init__.py +++ b/fastvideo/v1/models/__init__.py @@ -24,7 +24,7 @@ def get_transformers_model(model_config: PretrainedConfig, model_path: str, infe Returns: The loaded model """ - model_loader = get_model_loader(inference_args) + model_loader = get_model_loader() return model_loader.load_model(model_config, model_path, inference_args) def get_diffusers_model(model_config: Dict[str, Any], model_path: str, inference_args: InferenceArgs) -> nn.Module: diff --git a/fastvideo/v1/models/encoders/clip.py b/fastvideo/v1/models/encoders/clip.py index c1bdc1a11..b141db13f 100644 --- a/fastvideo/v1/models/encoders/clip.py +++ b/fastvideo/v1/models/encoders/clip.py @@ -197,11 +197,11 @@ class CLIPAttention(nn.Module): query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) # use flash_attn_func from flash_attn import flash_attn_func - query_states = query_states.reshape(query_states.shape[0], query_states.shape[1], self.num_heads, self.head_dim) - key_states = key_states.reshape(key_states.shape[0], key_states.shape[1], self.num_heads, self.head_dim) - value_states = value_states.reshape(value_states.shape[0], value_states.shape[1], self.num_heads, self.head_dim) + query_states = query_states.reshape(query_states.shape[0], query_states.shape[1], self.num_heads_per_partition, self.head_dim) + key_states = key_states.reshape(key_states.shape[0], key_states.shape[1], self.num_heads_per_partition, self.head_dim) + value_states = value_states.reshape(value_states.shape[0], value_states.shape[1], self.num_heads_per_partition, self.head_dim) attn_output = flash_attn_func(query_states, key_states, value_states,softmax_scale=self.scale, causal=True) - attn_output = attn_output.reshape(attn_output.shape[0], attn_output.shape[1], self.embed_dim) + attn_output = attn_output.reshape(attn_output.shape[0], attn_output.shape[1], self.num_heads_per_partition * self.head_dim) attn_output, _ = self.out_proj(attn_output) return attn_output, None diff --git a/fastvideo/v1/models/encoders/llama.py b/fastvideo/v1/models/encoders/llama.py index 5bcc6b59a..d5b8169c1 100644 --- a/fastvideo/v1/models/encoders/llama.py +++ b/fastvideo/v1/models/encoders/llama.py @@ -177,7 +177,7 @@ class LlamaAttention(nn.Module): qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # print(f"query_states sum before apply_rotary_pos_emb: {q.float().sum()}") - # q, k = self.rotary_emb(positions, q, k) + q, k = self.rotary_emb(positions, q, k) # attn_output = self.attn(q, k, v) # use flash_attn_func from flash_attn import flash_attn_func @@ -187,8 +187,13 @@ class LlamaAttention(nn.Module): q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + # import pdb; pdb.set_trace() attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True) - attn_output = attn_output.reshape(batch_size, seq_len, self.total_num_heads * self.head_dim) + print(f"shape of attn_output: {attn_output.shape}") + print(f"num_kv_heads: {self.num_kv_heads}") + print(f"num_heads: {self.num_heads}") + print(f"head_dim: {self.head_dim}") + attn_output = attn_output.reshape(batch_size, seq_len, self.num_heads * self.head_dim) output, _ = self.o_proj(attn_output) return output diff --git a/fastvideo/v1/models/loader.py b/fastvideo/v1/models/loader.py index 4a5511c9f..375e774a4 100644 --- a/fastvideo/v1/models/loader.py +++ b/fastvideo/v1/models/loader.py @@ -79,7 +79,7 @@ class TextEncoderLoader(ComponentLoader): def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): """Load the text encoders based on the model path, architecture, and inference args.""" - use_v1 = True + use_v1 = False if not use_v1: return self.load_v0(model_path, architecture, inference_args) else: diff --git a/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py index 4b1c39dfc..2804b80e1 100644 --- a/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py +++ b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py @@ -25,7 +25,6 @@ from fastvideo.v1.pipelines.stages.prompt_encoding import PromptEncodingStage from fastvideo.v1.inference_args import InferenceArgs # from fastvideo.v1.pipelines.composed.composed_pipeline_base import DiffusionPipelineOutput from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.v1.models.text_encoder import TextEncoder # TODO(will): move PRECISION_TO_TYPE to better place from .constants import PROMPT_TEMPLATE, PRECISION_TO_TYPE @@ -54,10 +53,63 @@ class DiffusionPipelineOutput(BaseOutput): class HunyuanVideoPipeline(ComposedPipelineBase): def initialize_encoders(self, modules: Dict[str, Any], inference_args: InferenceArgs): + use_v1 = False + if use_v1: + self.initialize_encoders_v1(modules, inference_args) + else: + self.initialize_encoders_v0(modules, inference_args) + + def initialize_encoders_v0(self, modules: Dict[str, Any], inference_args: InferenceArgs): + from fastvideo.v1.v0_reference_src.models.hunyuan.text_encoder import TextEncoder + # Text encoder + if inference_args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0) + elif inference_args.prompt_template is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template].get("crop_start", 0) + else: + crop_start = 0 + max_length = inference_args.text_len + crop_start + + # prompt_template + prompt_template = (PROMPT_TEMPLATE[inference_args.prompt_template] if inference_args.prompt_template is not None else None) + + # prompt_template_video + prompt_template_video = (PROMPT_TEMPLATE[inference_args.prompt_template_video] + if inference_args.prompt_template_video is not None else None) + + text_encoder = TextEncoder( + text_encoder_type=inference_args.text_encoder, + max_length=max_length, + text_encoder_precision=inference_args.text_encoder_precision, + tokenizer_type=inference_args.tokenizer, + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=inference_args.hidden_state_skip_layer, + apply_final_norm=inference_args.apply_final_norm, + reproduce=inference_args.reproduce, + logger=logger, + device=self.device if not inference_args.use_cpu_offload else "cpu", + ) + text_encoder_2 = None + if inference_args.text_encoder_2 is not None: + text_encoder_2 = TextEncoder( + text_encoder_type=inference_args.text_encoder_2, + max_length=inference_args.text_len_2, + text_encoder_precision=inference_args.text_encoder_precision_2, + tokenizer_type=inference_args.tokenizer_2, + reproduce=inference_args.reproduce, + logger=logger, + device=self.device if not inference_args.use_cpu_offload else "cpu", + ) + modules["text_encoder"] = text_encoder + modules["text_encoder_2"] = text_encoder_2 + + def initialize_encoders_v1(self, modules: Dict[str, Any], inference_args: InferenceArgs): """ Initialize the encoders. Will remove the encoders/tokenizers modules from the modules. Will add the TextEncoder or ImageEncoder to the modules. """ + from fastvideo.v1.models.text_encoder import TextEncoder if inference_args.prompt_template_video is not None: crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0) elif inference_args.prompt_template is not None: diff --git a/fastvideo/v1/pipelines/stages/prompt_encoding.py b/fastvideo/v1/pipelines/stages/prompt_encoding.py index 24ff64dbf..9117e8ae0 100644 --- a/fastvideo/v1/pipelines/stages/prompt_encoding.py +++ b/fastvideo/v1/pipelines/stages/prompt_encoding.py @@ -97,6 +97,9 @@ class PromptEncodingStage(PipelineStage): prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) # print("prompt_embeds", type(prompt_embeds)) # logger.info(f"prompt_embeds shape: {prompt_embeds.shape}") + if prompt_embeds.ndim == 1: + prompt_embeds = prompt_embeds.unsqueeze(0) + # attention_mask = attention_mask.unsqueeze(0) if prompt_embeds.ndim == 2: bs_embed, _ = prompt_embeds.shape diff --git a/fastvideo/v1/sample/v1_fastvideo_inference.py b/fastvideo/v1/sample/v1_fastvideo_inference.py index c61b71f6a..bfd9560fb 100644 --- a/fastvideo/v1/sample/v1_fastvideo_inference.py +++ b/fastvideo/v1/sample/v1_fastvideo_inference.py @@ -34,7 +34,7 @@ def main(inference_args: InferenceArgs): # Initialize tensor model parallel groups initialize_model_parallel( sequence_model_parallel_size=inference_args.sp_size, - tensor_model_parallel_size=inference_args.sp_size + tensor_model_parallel_size=inference_args.tp_size, ) # initialize_sequence_parallel_state(world_size) # initialize_distributed() diff --git a/scripts/inference/v1_inference_hunyuan.sh b/scripts/inference/v1_inference_hunyuan.sh index 4d650b432..9cedb5242 100644 --- a/scripts/inference/v1_inference_hunyuan.sh +++ b/scripts/inference/v1_inference_hunyuan.sh @@ -3,9 +3,13 @@ num_gpus=4 export MODEL_BASE=FastVideo/FastHunyuan-diffusers # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo +# Note that the tp_size and sp_size should be the same and equal to the number +# of GPUs. They are used for different parallel groups. sp_size is used for +# dit model and tp_size is used for encoder models. torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ fastvideo/v1/sample/v1_fastvideo_inference.py \ --sp_size 4 \ + --tp_size 4 \ --height 720 \ --width 1280 \ --num_frames 125 \ From 796eaf809f529f0ce551e28a850781244c43c466 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 16 Mar 2025 00:51:52 +0000 Subject: [PATCH 11/12] add toggle flags for v0 pipeline components --- fastvideo/v1/inference_args.py | 18 +++ fastvideo/v1/models/loader.py | 105 +++++++++++------- .../hunyuan/hunyuan_pipeline.py | 2 +- fastvideo/v1/pipelines/stages/decoding.py | 2 + .../models/flash_attn_no_pad.py | 28 +++++ .../models/hunyuan/constants.py | 8 +- .../models/hunyuan/modules/__init__.py | 42 +++---- .../models/hunyuan/modules/attenion.py | 6 +- .../models/hunyuan/modules/models.py | 4 +- .../models/hunyuan/text_encoder/__init__.py | 4 +- .../v1/v0_reference_src/pipeline_base.py | 14 +-- .../v0_reference_src/utils/communications.py | 2 +- scripts/inference/v1_inference_hunyuan.sh | 3 +- 13 files changed, 156 insertions(+), 82 deletions(-) create mode 100644 fastvideo/v1/v0_reference_src/models/flash_attn_no_pad.py diff --git a/fastvideo/v1/inference_args.py b/fastvideo/v1/inference_args.py index 371a41d00..5ba176b30 100644 --- a/fastvideo/v1/inference_args.py +++ b/fastvideo/v1/inference_args.py @@ -29,6 +29,9 @@ class InferenceArgs: # Model and path configuration model_path: str model: str = "HYVideo-T/2-cfgdistill" + use_v1_text_encoder: bool = False + use_v1_vae: bool = False + use_v1_transformer: bool = False dit_weight: Optional[str] = None model_dir: Optional[str] = None @@ -138,6 +141,21 @@ class InferenceArgs: @staticmethod def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument( + "--use-v1-text-encoder", + action="store_true", + help="Use the v1 text encoder", + ) + parser.add_argument( + "--use-v1-vae", + action="store_true", + help="Use the v1 vae", + ) + parser.add_argument( + "--use-v1-transformer", + action="store_true", + help="Use the v1 transformer", + ) # Model and path configuration parser.add_argument( "--model-path", diff --git a/fastvideo/v1/models/loader.py b/fastvideo/v1/models/loader.py index 375e774a4..2d2257a98 100644 --- a/fastvideo/v1/models/loader.py +++ b/fastvideo/v1/models/loader.py @@ -79,11 +79,10 @@ class TextEncoderLoader(ComponentLoader): def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): """Load the text encoders based on the model path, architecture, and inference args.""" - use_v1 = False - if not use_v1: - return self.load_v0(model_path, architecture, inference_args) - else: - return self.load_v1(model_path, architecture, inference_args) + # should always use v1 here. If inference_args.use_v1_text_encoder is False, + # the pipeline will overwrite this text encoder with a v0 text encoder + # during initialize_encoders() + return self.load_v1(model_path, architecture, inference_args) def load_v1(self, model_path: str, architecture: str, inference_args: InferenceArgs): """Load the text encoders based on the model path, architecture, and inference args.""" @@ -99,40 +98,6 @@ class TextEncoderLoader(ComponentLoader): model = model_loader.load_model(model_path, model_config, inference_args) return model - def load_v0(self, model_path: str, architecture: str, inference_args: InferenceArgs): - from fastvideo.v1.pipelines.hunyuan.constants import PROMPT_TEMPLATE - from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder - # Text encoder - if inference_args.prompt_template_video is not None: - crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0) - elif inference_args.prompt_template is not None: - crop_start = PROMPT_TEMPLATE[inference_args.prompt_template].get("crop_start", 0) - else: - crop_start = 0 - max_length = inference_args.text_len + crop_start - - # prompt_template - prompt_template = (PROMPT_TEMPLATE[inference_args.prompt_template] if inference_args.prompt_template is not None else None) - - # prompt_template_video - prompt_template_video = (PROMPT_TEMPLATE[inference_args.prompt_template_video] - if inference_args.prompt_template_video is not None else None) - - text_encoder = TextEncoder( - text_encoder_type=inference_args.text_encoder, - text_encoder_path=model_path, - max_length=max_length, - text_encoder_precision=inference_args.text_encoder_precision, - tokenizer_type=inference_args.tokenizer, - prompt_template=prompt_template, - prompt_template_video=prompt_template_video, - hidden_state_skip_layer=inference_args.hidden_state_skip_layer, - apply_final_norm=inference_args.apply_final_norm, - reproduce=inference_args.reproduce, - logger=logger, - device=self.device if not inference_args.use_cpu_offload else "cpu", - ) - return text_encoder class TokenizerLoader(ComponentLoader): """Loader for tokenizers.""" @@ -150,10 +115,32 @@ class TokenizerLoader(ComponentLoader): logger.info(f"Loaded tokenizer: {tokenizer.__class__.__name__}") return tokenizer + class VAELoader(ComponentLoader): """Loader for VAE.""" - def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the VAE based on the model path, architecture, and inference args.""" + use_v1 = inference_args.use_v1_vae + if not use_v1: + return self.load_v0(model_path, architecture, inference_args) + else: + return self.load_v1(model_path, architecture, inference_args) + + def load_v0(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Custom VAE loading for Hunyuan""" + # TODO(will): replace this with abstracted model + from fastvideo.v1.v0_reference_src.models.hunyuan.vae import load_vae + vae, _, s_ratio, t_ratio = load_vae( + inference_args.vae, + inference_args.vae_precision, + logger=logger, + device=self.device + ) + vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio} + vae.kwargs = vae_kwargs + return vae + + def load_v1(self, model_path: str, architecture: str, inference_args: InferenceArgs): """Load the VAE based on the model path, architecture, and inference args.""" # TODO(will): move this to a constants file from fastvideo.v1.utils import PRECISION_TO_TYPE @@ -189,10 +176,44 @@ class VAELoader(ComponentLoader): class TransformerLoader(ComponentLoader): """Loader for transformer.""" - + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): """Load the transformer based on the model path, architecture, and inference args.""" - + use_v1 = inference_args.use_v1_transformer + if not use_v1: + return self.load_v0(model_path, architecture, inference_args) + else: + return self.load_v1(model_path, architecture, inference_args) + + def load_v0(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Custom transformer loading for Hunyuan""" + # TODO(will): replace this with abstracted model + from fastvideo.v1.v0_reference_src.models.hunyuan.modules import load_model + from fastvideo.v1.utils import PRECISION_TO_TYPE + # Disable gradient + torch.set_grad_enabled(False) + + # =========================== Build main model =========================== + logger.info("Building model...") + factor_kwargs = {"device": self.device, "dtype": PRECISION_TO_TYPE[inference_args.precision]} + in_channels = inference_args.latent_channels + out_channels = inference_args.latent_channels + + model = load_model( + inference_args, + in_channels=in_channels, + out_channels=out_channels, + factor_kwargs=factor_kwargs, + ) + model = model.to(self.device) + model = self._load_transformer_state_dict(inference_args, model, inference_args.model_path) + if inference_args.enable_torch_compile: + model = torch.compile(model) + model.eval() + return model + + def load_v1(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the transformer based on the model path, architecture, and inference args.""" model_config = get_diffusers_config(model=model_path) cls_name = model_config.pop("_class_name") if cls_name is None: diff --git a/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py index 2804b80e1..b78b5691e 100644 --- a/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py +++ b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py @@ -53,7 +53,7 @@ class DiffusionPipelineOutput(BaseOutput): class HunyuanVideoPipeline(ComposedPipelineBase): def initialize_encoders(self, modules: Dict[str, Any], inference_args: InferenceArgs): - use_v1 = False + use_v1 = inference_args.use_v1_text_encoder if use_v1: self.initialize_encoders_v1(modules, inference_args) else: diff --git a/fastvideo/v1/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py index 64f32dca0..bd8d6bd5c 100644 --- a/fastvideo/v1/pipelines/stages/decoding.py +++ b/fastvideo/v1/pipelines/stages/decoding.py @@ -55,6 +55,8 @@ class DecodingStage(PipelineStage): with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): if inference_args.vae_tiling: self.vae.enable_tiling() + # if inference_args.vae_sp: + # self.vae.enable_parallel() image = self.vae.decode(latents) diff --git a/fastvideo/v1/v0_reference_src/models/flash_attn_no_pad.py b/fastvideo/v1/v0_reference_src/models/flash_attn_no_pad.py new file mode 100644 index 000000000..c0334a9e6 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/flash_attn_no_pad.py @@ -0,0 +1,28 @@ +from einops import rearrange +from flash_attn import flash_attn_varlen_qkvpacked_func +from flash_attn.bert_padding import pad_input, unpad_input + + +def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None): + # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27 + batch_size = qkv.shape[0] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask) + + x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) + output_unpad = flash_attn_varlen_qkvpacked_func( + x_unpad, + cu_seqlens, + max_s, + dropout_p, + softmax_scale=softmax_scale, + causal=causal, + ) + 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 \ No newline at end of file diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py index 3dd67b248..8baee0883 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py @@ -78,9 +78,13 @@ TEXT_ENCODER_PATH = { } # Tokenizer +# TOKENIZER_PATH = { +# "clipL": f"{MODEL_BASE}/text_encoder_2", +# "llm": f"{MODEL_BASE}/text_encoder", +# } TOKENIZER_PATH = { - "clipL": f"{MODEL_BASE}/text_encoder_2", - "llm": f"{MODEL_BASE}/text_encoder", + "clipL": f"{MODEL_BASE}/tokenizer_2", + "llm": f"{MODEL_BASE}/tokenizer", } TEXT_PROJECTION = { diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py index 3bb947c4e..c85b51a4c 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py @@ -1,25 +1,25 @@ -# from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer +from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer -# def load_model(args, in_channels, out_channels, factor_kwargs): -# """load hunyuan video model +def load_model(args, in_channels, out_channels, factor_kwargs): + """load hunyuan video model -# Args: -# args (dict): model args -# in_channels (int): input channels number -# out_channels (int): output channels number -# factor_kwargs (dict): factor kwargs + Args: + args (dict): model args + in_channels (int): input channels number + out_channels (int): output channels number + factor_kwargs (dict): factor kwargs -# Returns: -# model (nn.Module): The hunyuan video model -# """ -# if args.model in HUNYUAN_VIDEO_CONFIG.keys(): -# model = HYVideoDiffusionTransformer( -# in_channels=in_channels, -# out_channels=out_channels, -# **HUNYUAN_VIDEO_CONFIG[args.model], -# **factor_kwargs, -# ) -# return model -# else: -# raise NotImplementedError() + Returns: + model (nn.Module): The hunyuan video model + """ + if args.model in HUNYUAN_VIDEO_CONFIG.keys(): + model = HYVideoDiffusionTransformer( + in_channels=in_channels, + out_channels=out_channels, + **HUNYUAN_VIDEO_CONFIG[args.model], + **factor_kwargs, + ) + return model + else: + raise NotImplementedError() diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py index b275523a9..0567b08aa 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py @@ -8,10 +8,10 @@ except ImportError: print("Could not load Sliding Tile Attention.") sliding_tile_attention = None -from fastvideo.v1.models.flash_attn_no_pad import flash_attn_no_pad +from fastvideo.v1.v0_reference_src.models.flash_attn_no_pad import flash_attn_no_pad from flash_attn import flash_attn_func -from fastvideo.v1.utils.communications import all_gather, all_to_all_4D -from fastvideo.v1.utils.parallel_states import get_sequence_parallel_state, nccl_info +from fastvideo.v1.v0_reference_src.utils.communications import all_gather, all_to_all_4D +from fastvideo.v1.v0_reference_src.utils.parallel_states import get_sequence_parallel_state, nccl_info def attention( diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py index 42b2b9c85..7f4e0d2ad 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py @@ -6,8 +6,8 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models import ModelMixin from einops import rearrange -from fastvideo.v1.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_embed -from fastvideo.v1.utils.parallel_states import nccl_info +from fastvideo.v1.v0_reference_src.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_embed +from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info from .activation_layers import get_activation_layer from .attenion import parallel_attention diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py index aadfaa87f..bf6f8023c 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py @@ -50,8 +50,8 @@ def load_text_encoder( def load_tokenizer(tokenizer_type, tokenizer_path=None, padding_side="right", logger=None): - if tokenizer_path is None: - tokenizer_path = TOKENIZER_PATH[tokenizer_type] + # if tokenizer_path is None: + tokenizer_path = TOKENIZER_PATH[tokenizer_type] if logger is not None: logger.info(f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}") diff --git a/fastvideo/v1/v0_reference_src/pipeline_base.py b/fastvideo/v1/v0_reference_src/pipeline_base.py index 0349c3b76..5229ff1a8 100644 --- a/fastvideo/v1/v0_reference_src/pipeline_base.py +++ b/fastvideo/v1/v0_reference_src/pipeline_base.py @@ -13,10 +13,10 @@ from diffusers.utils.torch_utils import randn_tensor from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.inference_args import InferenceArgs -from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder -from fastvideo.v1.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D -from fastvideo.v1.models.hunyuan.constants import PRECISION_TO_TYPE -from fastvideo.v1.distributed.parallel_state import get_sp_group +from fastvideo.v1.v0_reference_src.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.v0_reference_src.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from fastvideo.v1.v0_reference_src.models.hunyuan.constants import PRECISION_TO_TYPE +from fastvideo.v1.v0_reference_src.distributed.parallel_state import get_sp_group from fastvideo.v1.logger import init_logger from einops import rearrange from diffusers.utils import BaseOutput @@ -420,7 +420,7 @@ class DiffusionPipelineBase(ABC): # 5. Prepare latent variables batch = self.prepare_latents(batch) - from fastvideo.v1.utils.parallel_states import nccl_info + from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group # world_size = 1 # rank = 0 @@ -541,8 +541,8 @@ class DiffusionPipelineBase(ABC): # if sp_group: # latents = sp_group.all_gather(latents, dim=2) - from fastvideo.v1.utils.communications import all_gather - from fastvideo.v1.utils.parallel_states import get_sequence_parallel_state + from fastvideo.v1.v0_reference_src.utils.communications import all_gather + from fastvideo.v1.v0_reference_src.utils.parallel_states import get_sequence_parallel_state if get_sequence_parallel_state(): latents = all_gather(latents, dim=2) diff --git a/fastvideo/v1/v0_reference_src/utils/communications.py b/fastvideo/v1/v0_reference_src/utils/communications.py index 04fdc8194..38056738a 100644 --- a/fastvideo/v1/v0_reference_src/utils/communications.py +++ b/fastvideo/v1/v0_reference_src/utils/communications.py @@ -9,7 +9,7 @@ import torch import torch.distributed as dist from torch import Tensor -from fastvideo.v1.utils.parallel_states import nccl_info +from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info def broadcast(input_: torch.Tensor): diff --git a/scripts/inference/v1_inference_hunyuan.sh b/scripts/inference/v1_inference_hunyuan.sh index 9cedb5242..bab145435 100644 --- a/scripts/inference/v1_inference_hunyuan.sh +++ b/scripts/inference/v1_inference_hunyuan.sh @@ -1,13 +1,14 @@ #!/bin/bash num_gpus=4 -export MODEL_BASE=FastVideo/FastHunyuan-diffusers +export MODEL_BASE=data/FastHunyuan-diffusers # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo # Note that the tp_size and sp_size should be the same and equal to the number # of GPUs. They are used for different parallel groups. sp_size is used for # dit model and tp_size is used for encoder models. torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ fastvideo/v1/sample/v1_fastvideo_inference.py \ + --use-v1-transformer \ --sp_size 4 \ --tp_size 4 \ --height 720 \ From 691f9d1064dcbcf09f9cc58fb783f2dad698ac57 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 16 Mar 2025 01:37:58 +0000 Subject: [PATCH 12/12] vae update --- fastvideo/v1/layers/activation.py | 3 +- .../models/hunyuan/constants.py | 3 +- .../models/hunyuan/vae/__init__.py | 2 +- .../hunyuan/vae/autoencoder_kl_causal_3d.py | 2 +- .../hunyuan/vae/unet_causal_3d_blocks.py | 58 +++++++++---------- 5 files changed, 34 insertions(+), 34 deletions(-) diff --git a/fastvideo/v1/layers/activation.py b/fastvideo/v1/layers/activation.py index dc7db7d22..6217ff1a5 100644 --- a/fastvideo/v1/layers/activation.py +++ b/fastvideo/v1/layers/activation.py @@ -11,8 +11,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size) from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform -from vllm.utils import LazyDict +from fastvideo.v1.platforms import current_platform @CustomOp.register("fatrelu_and_mul") diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py index 8baee0883..53c9b0628 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py @@ -69,7 +69,8 @@ MODEL_BASE = os.getenv("MODEL_BASE", "./data/hunyuan") DATA_TYPE = {"image", "video", "image_video"} # 3D VAE -VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"} +# VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"} +VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/vae"} # Text Encoder TEXT_ENCODER_PATH = { diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py index 761c7cc1f..aa3de3c84 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py @@ -35,7 +35,7 @@ def load_vae( else: vae = AutoencoderKLCausal3D.from_config(config) - vae_ckpt = Path(vae_path) / "pytorch_model.pt" + vae_ckpt = Path("/home/ubuntu/src/FastVideo/data/hunyuan/hunyuan-video-t2v-720p/vae/pytorch_model.pt") assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}" ckpt = torch.load(vae_ckpt, map_location=vae.device) diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py index 5a63783cd..b0973853b 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py @@ -25,7 +25,7 @@ import torch.nn as nn import torch.distributed as dist from diffusers.configuration_utils import ConfigMixin, register_to_config -from fastvideo.v1.utils.parallel_states import nccl_info +from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info from fastvideo.v1.distributed.parallel_state import get_sp_group from fastvideo.v1.logger import init_logger diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py index 3c466f55f..6b16b26bd 100644 --- a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py @@ -448,20 +448,20 @@ def get_down_block3d( attention_head_dim = num_attention_heads down_block_type = (down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type) - if down_block_type == "DownEncoderBlockCausal3D": - return DownEncoderBlockCausal3D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - dropout=dropout, - add_downsample=add_downsample, - downsample_stride=downsample_stride, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - resnet_time_scale_shift=resnet_time_scale_shift, - ) + # if down_block_type == "DownEncoderBlockCausal3D": + return DownEncoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + downsample_stride=downsample_stride, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) raise ValueError(f"{down_block_type} does not exist.") @@ -502,21 +502,21 @@ def get_up_block3d( attention_head_dim = num_attention_heads up_block_type = (up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type) - if up_block_type == "UpDecoderBlockCausal3D": - return UpDecoderBlockCausal3D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - upsample_scale_factor=upsample_scale_factor, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - temb_channels=temb_channels, - ) + # if up_block_type == "UpDecoderBlockCausal3D": + return UpDecoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + resolution_idx=resolution_idx, + dropout=dropout, + add_upsample=add_upsample, + upsample_scale_factor=upsample_scale_factor, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + ) raise ValueError(f"{up_block_type} does not exist.")