# Copyright 2023 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, Union import diffusers import pkg_resources import torch import torch.nn.functional as F import torch.nn.init as init from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.attention import Attention, FeedForward from diffusers.models.attention_processor import (Attention, AttentionProcessor, AttnProcessor2_0, HunyuanAttnProcessor2_0) from diffusers.models.embeddings import (SinusoidalPositionalEmbedding, TimestepEmbedding, Timesteps, get_3d_sincos_pos_embed) from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.modeling_utils import ModelMixin from diffusers.models.normalization import (AdaLayerNorm, AdaLayerNormZero, CogVideoXLayerNormZero) from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging from diffusers.utils.import_utils import is_xformers_available from diffusers.utils.torch_utils import maybe_allow_in_graph from einops import rearrange, repeat from torch import nn from .motion_module import PositionalEncoding, get_motion_module from .norm import AdaLayerNormShift, FP32LayerNorm, EasyAnimateLayerNormZero from .processor import (EasyAnimateAttnProcessor2_0, LazyKVCompressionProcessor2_0) if is_xformers_available(): import xformers import xformers.ops else: xformers = None def zero_module(module): # Zero out the parameters of a module and return it. for p in module.parameters(): p.detach().zero_() return module @maybe_allow_in_graph class GatedSelfAttentionDense(nn.Module): r""" A gated self-attention dense layer that combines visual features and object features. Parameters: query_dim (`int`): The number of channels in the query. context_dim (`int`): The number of channels in the context. n_heads (`int`): The number of heads to use for attention. d_head (`int`): The number of channels in each head. """ def __init__(self, query_dim: int, context_dim: int, n_heads: int, d_head: int): super().__init__() # we need a linear projection since we need cat visual feature and obj feature self.linear = nn.Linear(context_dim, query_dim) self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head) self.ff = FeedForward(query_dim, activation_fn="geglu") self.norm1 = FP32LayerNorm(query_dim) self.norm2 = FP32LayerNorm(query_dim) self.register_parameter("alpha_attn", nn.Parameter(torch.tensor(0.0))) self.register_parameter("alpha_dense", nn.Parameter(torch.tensor(0.0))) self.enabled = True def forward(self, x: torch.Tensor, objs: torch.Tensor) -> torch.Tensor: if not self.enabled: return x n_visual = x.shape[1] objs = self.linear(objs) x = x + self.alpha_attn.tanh() * self.attn(self.norm1(torch.cat([x, objs], dim=1)))[:, :n_visual, :] x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x)) return x class LazyKVCompressionAttention(Attention): def __init__( self, sr_ratio=2, *args, **kwargs ): super().__init__(*args, **kwargs) self.sr_ratio = sr_ratio self.k_compression = nn.Conv2d( kwargs["query_dim"], kwargs["query_dim"], groups=kwargs["query_dim"], kernel_size=sr_ratio, stride=sr_ratio, bias=True ) self.v_compression = nn.Conv2d( kwargs["query_dim"], kwargs["query_dim"], groups=kwargs["query_dim"], kernel_size=sr_ratio, stride=sr_ratio, bias=True ) init.constant_(self.k_compression.weight, 1 / (sr_ratio * sr_ratio)) init.constant_(self.v_compression.weight, 1 / (sr_ratio * sr_ratio)) init.constant_(self.k_compression.bias, 0) init.constant_(self.v_compression.bias, 0) @maybe_allow_in_graph class TemporalTransformerBlock(nn.Module): r""" A Temporal 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, num_attention_heads: int, attention_head_dim: int, dropout=0.0, cross_attention_dim: Optional[int] = None, activation_fn: str = "geglu", num_embeds_ada_norm: Optional[int] = None, attention_bias: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, upcast_attention: bool = False, norm_elementwise_affine: bool = True, norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single' norm_eps: float = 1e-5, final_dropout: bool = False, attention_type: str = "default", positional_embeddings: Optional[str] = None, num_positional_embeddings: Optional[int] = None, # motion module kwargs motion_module_type = "VanillaGrid", motion_module_kwargs = None, qk_norm = False, after_norm = False, ): super().__init__() self.only_cross_attention = only_cross_attention self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" self.use_ada_layer_norm_single = norm_type == "ada_norm_single" self.use_layer_norm = norm_type == "layer_norm" if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: raise ValueError( f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." ) if positional_embeddings and (num_positional_embeddings is None): raise ValueError( "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." ) if positional_embeddings == "sinusoidal": self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) else: self.pos_embed = None # Define 3 blocks. Each block has its own normalization layer. # 1. Self-Attn if self.use_ada_layer_norm: self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) elif self.use_ada_layer_norm_zero: self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) else: self.norm1 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) self.attn1 = Attention( query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, dropout=dropout, bias=attention_bias, cross_attention_dim=cross_attention_dim if only_cross_attention else None, upcast_attention=upcast_attention, qk_norm="layer_norm" if qk_norm else None, processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), ) self.attn_temporal = get_motion_module( in_channels = dim, motion_module_type = motion_module_type, motion_module_kwargs = motion_module_kwargs, ) # 2. Cross-Attn if cross_attention_dim is not None or double_self_attention: # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during # the second cross attention block. self.norm2 = ( AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) ) self.attn2 = Attention( query_dim=dim, cross_attention_dim=cross_attention_dim if not double_self_attention else None, heads=num_attention_heads, dim_head=attention_head_dim, dropout=dropout, bias=attention_bias, upcast_attention=upcast_attention, qk_norm="layer_norm" if qk_norm else None, processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), ) # is self-attn if encoder_hidden_states is none else: self.norm2 = None self.attn2 = None # 3. Feed-forward if not self.use_ada_layer_norm_single: self.norm3 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) if after_norm: self.norm4 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) else: self.norm4 = None # 4. Fuser if attention_type == "gated" or attention_type == "gated-text-image": self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) # 5. Scale-shift for PixArt-Alpha. if self.use_ada_layer_norm_single: self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) # let chunk size default to None self._chunk_size = None self._chunk_dim = 0 def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): # Sets chunk feed-forward self._chunk_size = chunk_size self._chunk_dim = dim def forward( self, hidden_states: torch.FloatTensor, attention_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, timestep: Optional[torch.LongTensor] = None, cross_attention_kwargs: Dict[str, Any] = None, class_labels: Optional[torch.LongTensor] = None, num_frames: int = 16, height: int = 32, width: int = 32, ) -> torch.FloatTensor: # Notice that normalization is always applied before the real computation in the following blocks. # 0. Self-Attention batch_size = hidden_states.shape[0] if self.use_ada_layer_norm: norm_hidden_states = self.norm1(hidden_states, timestep) elif self.use_ada_layer_norm_zero: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype ) elif self.use_layer_norm: norm_hidden_states = self.norm1(hidden_states) elif self.use_ada_layer_norm_single: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) ).chunk(6, dim=1) norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa norm_hidden_states = norm_hidden_states.squeeze(1) else: raise ValueError("Incorrect norm used") if self.pos_embed is not None: norm_hidden_states = self.pos_embed(norm_hidden_states) # 1. Retrieve lora scale. lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 # 2. Prepare GLIGEN inputs cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} gligen_kwargs = cross_attention_kwargs.pop("gligen", None) norm_hidden_states = rearrange(norm_hidden_states, "b (f d) c -> (b f) d c", f=num_frames) attn_output = self.attn1( norm_hidden_states, encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, attention_mask=attention_mask, **cross_attention_kwargs, ) attn_output = rearrange(attn_output, "(b f) d c -> b (f d) c", f=num_frames) if self.use_ada_layer_norm_zero: attn_output = gate_msa.unsqueeze(1) * attn_output elif self.use_ada_layer_norm_single: attn_output = gate_msa * attn_output hidden_states = attn_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) # 2.5 GLIGEN Control if gligen_kwargs is not None: hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) # 2.75. Temp-Attention if self.attn_temporal is not None: attn_output = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=num_frames, h=height, w=width) attn_output = self.attn_temporal(attn_output) hidden_states = rearrange(attn_output, "b c f h w -> b (f h w) c") # 3. Cross-Attention if self.attn2 is not None: if self.use_ada_layer_norm: norm_hidden_states = self.norm2(hidden_states, timestep) elif self.use_ada_layer_norm_zero or self.use_layer_norm: norm_hidden_states = self.norm2(hidden_states) elif self.use_ada_layer_norm_single: # For PixArt norm2 isn't applied here: # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 norm_hidden_states = hidden_states else: raise ValueError("Incorrect norm") if self.pos_embed is not None and self.use_ada_layer_norm_single is None: norm_hidden_states = self.pos_embed(norm_hidden_states) if norm_hidden_states.dtype != encoder_hidden_states.dtype or norm_hidden_states.dtype != encoder_attention_mask.dtype: norm_hidden_states = norm_hidden_states.to(encoder_hidden_states.dtype) attn_output = self.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=encoder_attention_mask, **cross_attention_kwargs, ) hidden_states = attn_output + hidden_states # 4. Feed-forward if not self.use_ada_layer_norm_single: norm_hidden_states = self.norm3(hidden_states) if self.use_ada_layer_norm_zero: norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] if self.use_ada_layer_norm_single: norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp if self._chunk_size is not None: # "feed_forward_chunk_size" can be used to save memory if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: raise ValueError( f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." ) num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size ff_output = torch.cat( [ self.ff(hid_slice, scale=lora_scale) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim) ], dim=self._chunk_dim, ) else: ff_output = self.ff(norm_hidden_states, scale=lora_scale) if self.norm4 is not None: ff_output = self.norm4(ff_output) if self.use_ada_layer_norm_zero: ff_output = gate_mlp.unsqueeze(1) * ff_output elif self.use_ada_layer_norm_single: ff_output = gate_mlp * ff_output hidden_states = ff_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) return hidden_states @maybe_allow_in_graph class SelfAttentionTemporalTransformerBlock(nn.Module): r""" A Temporal 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, num_attention_heads: int, attention_head_dim: int, dropout=0.0, cross_attention_dim: Optional[int] = None, activation_fn: str = "geglu", num_embeds_ada_norm: Optional[int] = None, attention_bias: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, upcast_attention: bool = False, norm_elementwise_affine: bool = True, norm_type: str = "layer_norm", norm_eps: float = 1e-5, final_dropout: bool = False, attention_type: str = "default", positional_embeddings: Optional[str] = None, num_positional_embeddings: Optional[int] = None, qk_norm = False, after_norm = False, ): super().__init__() self.only_cross_attention = only_cross_attention self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" self.use_ada_layer_norm_single = norm_type == "ada_norm_single" self.use_layer_norm = norm_type == "layer_norm" if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: raise ValueError( f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." ) if positional_embeddings and (num_positional_embeddings is None): raise ValueError( "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." ) if positional_embeddings == "sinusoidal": self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) else: self.pos_embed = None # Define 3 blocks. Each block has its own normalization layer. # 1. Self-Attn if self.use_ada_layer_norm: self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) elif self.use_ada_layer_norm_zero: self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) else: self.norm1 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) self.attn1 = Attention( query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, dropout=dropout, bias=attention_bias, cross_attention_dim=cross_attention_dim if only_cross_attention else None, upcast_attention=upcast_attention, qk_norm="layer_norm" if qk_norm else None, processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), ) # 2. Cross-Attn if cross_attention_dim is not None or double_self_attention: # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during # the second cross attention block. self.norm2 = ( AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) ) self.attn2 = Attention( query_dim=dim, cross_attention_dim=cross_attention_dim if not double_self_attention else None, heads=num_attention_heads, dim_head=attention_head_dim, dropout=dropout, bias=attention_bias, upcast_attention=upcast_attention, qk_norm="layer_norm" if qk_norm else None, processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), ) # is self-attn if encoder_hidden_states is none else: self.norm2 = None self.attn2 = None # 3. Feed-forward if not self.use_ada_layer_norm_single: self.norm3 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) if after_norm: self.norm4 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) else: self.norm4 = None # 4. Fuser if attention_type == "gated" or attention_type == "gated-text-image": self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) # 5. Scale-shift for PixArt-Alpha. if self.use_ada_layer_norm_single: self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) # let chunk size default to None self._chunk_size = None self._chunk_dim = 0 def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): # Sets chunk feed-forward self._chunk_size = chunk_size self._chunk_dim = dim def forward( self, hidden_states: torch.FloatTensor, attention_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, timestep: Optional[torch.LongTensor] = None, cross_attention_kwargs: Dict[str, Any] = None, class_labels: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: # Notice that normalization is always applied before the real computation in the following blocks. # 0. Self-Attention batch_size = hidden_states.shape[0] if self.use_ada_layer_norm: norm_hidden_states = self.norm1(hidden_states, timestep) elif self.use_ada_layer_norm_zero: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype ) elif self.use_layer_norm: norm_hidden_states = self.norm1(hidden_states) elif self.use_ada_layer_norm_single: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) ).chunk(6, dim=1) norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa norm_hidden_states = norm_hidden_states.squeeze(1) else: raise ValueError("Incorrect norm used") if self.pos_embed is not None: norm_hidden_states = self.pos_embed(norm_hidden_states) # 1. Retrieve lora scale. lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 # 2. Prepare GLIGEN inputs cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} gligen_kwargs = cross_attention_kwargs.pop("gligen", None) attn_output = self.attn1( norm_hidden_states, encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, attention_mask=attention_mask, **cross_attention_kwargs, ) if self.use_ada_layer_norm_zero: attn_output = gate_msa.unsqueeze(1) * attn_output elif self.use_ada_layer_norm_single: attn_output = gate_msa * attn_output hidden_states = attn_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) # 2.5 GLIGEN Control if gligen_kwargs is not None: hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) # 3. Cross-Attention if self.attn2 is not None: if self.use_ada_layer_norm: norm_hidden_states = self.norm2(hidden_states, timestep) elif self.use_ada_layer_norm_zero or self.use_layer_norm: norm_hidden_states = self.norm2(hidden_states) elif self.use_ada_layer_norm_single: # For PixArt norm2 isn't applied here: # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 norm_hidden_states = hidden_states else: raise ValueError("Incorrect norm") if self.pos_embed is not None and self.use_ada_layer_norm_single is None: norm_hidden_states = self.pos_embed(norm_hidden_states) attn_output = self.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=encoder_attention_mask, **cross_attention_kwargs, ) hidden_states = attn_output + hidden_states # 4. Feed-forward if not self.use_ada_layer_norm_single: norm_hidden_states = self.norm3(hidden_states) if self.use_ada_layer_norm_zero: norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] if self.use_ada_layer_norm_single: norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp if self._chunk_size is not None: # "feed_forward_chunk_size" can be used to save memory if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: raise ValueError( f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." ) num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size ff_output = torch.cat( [ self.ff(hid_slice, scale=lora_scale) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim) ], dim=self._chunk_dim, ) else: ff_output = self.ff(norm_hidden_states, scale=lora_scale) if self.norm4 is not None: ff_output = self.norm4(ff_output) if self.use_ada_layer_norm_zero: ff_output = gate_mlp.unsqueeze(1) * ff_output elif self.use_ada_layer_norm_single: ff_output = gate_mlp * ff_output hidden_states = ff_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) return hidden_states class GEGLU(nn.Module): def __init__(self, dim_in, dim_out, norm_elementwise_affine): super().__init__() self.norm = FP32LayerNorm(dim_in, dim_in, norm_elementwise_affine) self.proj = nn.Linear(dim_in, dim_out * 2) def forward(self, x): x, gate = self.proj(self.norm(x)).chunk(2, dim=-1) return x * F.gelu(gate) @maybe_allow_in_graph class HunyuanDiTBlock(nn.Module): r""" Transformer block used in Hunyuan-DiT model (https://github.com/Tencent/HunyuanDiT). Allow skip connection and QKNorm Parameters: dim (`int`): The number of channels in the input and output. num_attention_heads (`int`): The number of headsto use for multi-head attention. cross_attention_dim (`int`,*optional*): The size of the encoder_hidden_states vector for cross attention. dropout(`float`, *optional*, defaults to 0.0): The dropout probability to use. activation_fn (`str`,*optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. . norm_elementwise_affine (`bool`, *optional*, defaults to `True`): Whether to use learnable elementwise affine parameters for normalization. norm_eps (`float`, *optional*, defaults to 1e-6): A small constant added to the denominator in normalization layers to prevent division by zero. final_dropout (`bool` *optional*, defaults to False): Whether to apply a final dropout after the last feed-forward layer. ff_inner_dim (`int`, *optional*): The size of the hidden layer in the feed-forward block. Defaults to `None`. ff_bias (`bool`, *optional*, defaults to `True`): Whether to use bias in the feed-forward block. skip (`bool`, *optional*, defaults to `False`): Whether to use skip connection. Defaults to `False` for down-blocks and mid-blocks. qk_norm (`bool`, *optional*, defaults to `True`): Whether to use normalization in QK calculation. Defaults to `True`. """ def __init__( self, dim: int, num_attention_heads: int, cross_attention_dim: int = 1024, dropout=0.0, activation_fn: str = "geglu", norm_elementwise_affine: bool = True, norm_eps: float = 1e-6, final_dropout: bool = False, ff_inner_dim: Optional[int] = None, ff_bias: bool = True, skip: bool = False, qk_norm: bool = True, time_position_encoding: bool = False, after_norm: bool = False, is_local_attention: bool = False, local_attention_frames: int = 2, enable_inpaint: bool = False, kvcompression = False, ): super().__init__() # Define 3 blocks. Each block has its own normalization layer. # NOTE: when new version comes, check norm2 and norm 3 # 1. Self-Attn self.norm1 = AdaLayerNormShift(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) self.t_embed = PositionalEncoding(dim, dropout=0., max_len=512) \ if time_position_encoding else nn.Identity() self.is_local_attention = is_local_attention self.local_attention_frames = local_attention_frames self.kvcompression = kvcompression if kvcompression: self.attn1 = LazyKVCompressionAttention( query_dim=dim, cross_attention_dim=None, dim_head=dim // num_attention_heads, heads=num_attention_heads, qk_norm="layer_norm" if qk_norm else None, eps=1e-6, bias=True, processor=LazyKVCompressionProcessor2_0(), ) else: self.attn1 = Attention( query_dim=dim, cross_attention_dim=None, dim_head=dim // num_attention_heads, heads=num_attention_heads, qk_norm="layer_norm" if qk_norm else None, eps=1e-6, bias=True, processor=HunyuanAttnProcessor2_0(), ) # 2. Cross-Attn self.norm2 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) if self.is_local_attention: from mamba_ssm import Mamba2 self.mamba_norm_in = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) self.in_linear = nn.Linear(dim, 1536) self.mamba_norm_1 = FP32LayerNorm(1536, norm_eps, norm_elementwise_affine) self.mamba_norm_2 = FP32LayerNorm(1536, norm_eps, norm_elementwise_affine) self.mamba_block_1 = Mamba2( d_model=1536, d_state=64, d_conv=4, expand=2, ) self.mamba_block_2 = Mamba2( d_model=1536, d_state=64, d_conv=4, expand=2, ) self.mamba_norm_after_mamba_block = FP32LayerNorm(1536, norm_eps, norm_elementwise_affine) self.out_linear = nn.Linear(1536, dim) self.out_linear = zero_module(self.out_linear) self.mamba_norm_out = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) self.attn2 = Attention( query_dim=dim, cross_attention_dim=cross_attention_dim, dim_head=dim // num_attention_heads, heads=num_attention_heads, qk_norm="layer_norm" if qk_norm else None, eps=1e-6, bias=True, processor=HunyuanAttnProcessor2_0(), ) if enable_inpaint: self.norm_clip = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) self.attn_clip = Attention( query_dim=dim, cross_attention_dim=cross_attention_dim, dim_head=dim // num_attention_heads, heads=num_attention_heads, qk_norm="layer_norm" if qk_norm else None, eps=1e-6, bias=True, processor=HunyuanAttnProcessor2_0(), ) self.gate_clip = GEGLU(dim, dim, norm_elementwise_affine) self.norm_clip_out = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) else: self.attn_clip = None self.norm_clip = None self.gate_clip = None self.norm_clip_out = None # 3. Feed-forward self.norm3 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) self.ff = FeedForward( dim, dropout=dropout, ### 0.0 activation_fn=activation_fn, ### approx GeLU final_dropout=final_dropout, ### 0.0 inner_dim=ff_inner_dim, ### int(dim * mlp_ratio) bias=ff_bias, ) # 4. Skip Connection if skip: self.skip_norm = FP32LayerNorm(2 * dim, norm_eps, elementwise_affine=True) self.skip_linear = nn.Linear(2 * dim, dim) else: self.skip_linear = None if after_norm: self.norm4 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) else: self.norm4 = None # let chunk size default to None self._chunk_size = None self._chunk_dim = 0 def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): # Sets chunk feed-forward self._chunk_size = chunk_size self._chunk_dim = dim def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None, image_rotary_emb=None, skip=None, num_frames: int = 1, height: int = 32, width: int = 32, clip_encoder_hidden_states: Optional[torch.Tensor] = None, disable_image_rotary_emb_in_attn1=False, ) -> torch.Tensor: # Notice that normalization is always applied before the real computation in the following blocks. # 0. Long Skip Connection if self.skip_linear is not None: cat = torch.cat([hidden_states, skip], dim=-1) cat = self.skip_norm(cat) hidden_states = self.skip_linear(cat) if image_rotary_emb is not None: image_rotary_emb = (torch.cat([image_rotary_emb[0] for i in range(num_frames)], dim=0), torch.cat([image_rotary_emb[1] for i in range(num_frames)], dim=0)) if num_frames != 1: # add time embedding hidden_states = rearrange(hidden_states, "b (f d) c -> (b d) f c", f=num_frames) if self.t_embed is not None: hidden_states = self.t_embed(hidden_states) hidden_states = rearrange(hidden_states, "(b d) f c -> b (f d) c", d=height * width) # 1. Self-Attention norm_hidden_states = self.norm1(hidden_states, temb) ### checked: self.norm1 is correct if num_frames > 2 and self.is_local_attention: if image_rotary_emb is not None: attn1_image_rotary_emb = (image_rotary_emb[0][:int(height * width * 2)], image_rotary_emb[1][:int(height * width * 2)]) else: attn1_image_rotary_emb = image_rotary_emb norm_hidden_states_1 = rearrange(norm_hidden_states, "b (f d) c -> b f d c", d=height * width) norm_hidden_states_1 = rearrange(norm_hidden_states_1, "b (f p) d c -> (b f) (p d) c", p = 2) attn_output = self.attn1( norm_hidden_states_1, image_rotary_emb=attn1_image_rotary_emb if not disable_image_rotary_emb_in_attn1 else None, ) attn_output = rearrange(attn_output, "(b f) (p d) c -> b (f p) d c", p = 2, f = num_frames // 2) norm_hidden_states_2 = rearrange(norm_hidden_states, "b (f d) c -> b f d c", d = height * width)[:, 1:-1] local_attention_frames_num = norm_hidden_states_2.size()[1] // 2 norm_hidden_states_2 = rearrange(norm_hidden_states_2, "b (f p) d c -> (b f) (p d) c", p = 2) attn_output_2 = self.attn1( norm_hidden_states_2, image_rotary_emb=attn1_image_rotary_emb if not disable_image_rotary_emb_in_attn1 else None, ) attn_output_2 = rearrange(attn_output_2, "(b f) (p d) c -> b (f p) d c", p = 2, f = local_attention_frames_num) attn_output[:, 1:-1] = (attn_output[:, 1:-1] + attn_output_2) / 2 attn_output = rearrange(attn_output, "b f d c -> b (f d) c") else: if self.kvcompression: norm_hidden_states = rearrange(norm_hidden_states, "b (f h w) c -> b c f h w", f = num_frames, h = height, w = width) attn_output = self.attn1( norm_hidden_states, image_rotary_emb=image_rotary_emb if not disable_image_rotary_emb_in_attn1 else None, ) else: attn_output = self.attn1( norm_hidden_states, image_rotary_emb=image_rotary_emb if not disable_image_rotary_emb_in_attn1 else None, ) hidden_states = hidden_states + attn_output if num_frames > 2 and self.is_local_attention: hidden_states_in = self.in_linear(self.mamba_norm_in(hidden_states)) hidden_states = hidden_states + self.mamba_norm_out( self.out_linear( self.mamba_norm_after_mamba_block( self.mamba_block_1( self.mamba_norm_1(hidden_states_in) ) + self.mamba_block_2( self.mamba_norm_2(hidden_states_in.flip(1)) ).flip(1) ) ) ) # 2. Cross-Attention hidden_states = hidden_states + self.attn2( self.norm2(hidden_states), encoder_hidden_states=encoder_hidden_states, image_rotary_emb=image_rotary_emb, ) if self.attn_clip is not None: hidden_states = hidden_states + self.norm_clip_out( self.gate_clip( self.attn_clip( self.norm_clip(hidden_states), encoder_hidden_states=clip_encoder_hidden_states, image_rotary_emb=image_rotary_emb, ) ) ) # FFN Layer ### TODO: switch norm2 and norm3 in the state dict mlp_inputs = self.norm3(hidden_states) if self.norm4 is not None: hidden_states = hidden_states + self.norm4(self.ff(mlp_inputs)) else: hidden_states = hidden_states + self.ff(mlp_inputs) return hidden_states @maybe_allow_in_graph class EasyAnimateDiTBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, time_embed_dim: int, dropout: float = 0.0, activation_fn: str = "gelu-approximate", norm_elementwise_affine: bool = True, norm_eps: float = 1e-6, final_dropout: bool = True, ff_inner_dim: Optional[int] = None, ff_bias: bool = True, qk_norm: bool = True, after_norm: bool = False, norm_type: str="fp32_layer_norm" ): super().__init__() # Attention Part self.norm1 = EasyAnimateLayerNormZero( time_embed_dim, dim, norm_elementwise_affine, norm_eps, norm_type=norm_type, bias=True ) self.attn1 = Attention( query_dim=dim, dim_head=attention_head_dim, heads=num_attention_heads, qk_norm="layer_norm" if qk_norm else None, eps=1e-6, bias=True, processor=EasyAnimateAttnProcessor2_0(), ) self.attn2 = Attention( query_dim=dim, dim_head=attention_head_dim, heads=num_attention_heads, qk_norm="layer_norm" if qk_norm else None, eps=1e-6, bias=True, processor=EasyAnimateAttnProcessor2_0(), ) # FFN Part self.norm2 = EasyAnimateLayerNormZero( time_embed_dim, dim, norm_elementwise_affine, norm_eps, norm_type=norm_type, bias=True ) self.ff = FeedForward( dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout, inner_dim=ff_inner_dim, bias=ff_bias, ) self.txt_ff = FeedForward( dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout, inner_dim=ff_inner_dim, bias=ff_bias, ) if after_norm: self.norm3 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) else: self.norm3 = None def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: # Norm norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1( hidden_states, encoder_hidden_states, temb ) # Attn attn_hidden_states, attn_encoder_hidden_states = self.attn1( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, image_rotary_emb=image_rotary_emb, attn2=self.attn2, ) hidden_states = hidden_states + gate_msa * attn_hidden_states encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states # Norm norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2( hidden_states, encoder_hidden_states, temb ) # FFN if self.norm3 is not None: norm_hidden_states = self.norm3(self.ff(norm_hidden_states)) norm_encoder_hidden_states = self.norm3(self.txt_ff(norm_encoder_hidden_states)) else: norm_hidden_states = self.ff(norm_hidden_states) norm_encoder_hidden_states = self.txt_ff(norm_encoder_hidden_states) hidden_states = hidden_states + gate_ff * norm_hidden_states encoder_hidden_states = encoder_hidden_states + enc_gate_ff * norm_encoder_hidden_states return hidden_states, encoder_hidden_states