Bubbliiiingandhkunzhe
9f34f7c8c5
* Update Flow * Update Flow * Update Flow * add image recaptioning * Fix bug in t2v * update train_reward_lora.py * update reward training * Update V5.1 and mix multi text_encoders to one pipeline * Update V5.1 training Code * Update ComfyUI * Update Comment * Delete files * update reward training * Update Readme * fix extract frames in compute_semantic_consistency * Update Readme && Remove to in prediction * Update Demo * Update Readme * Update ui * support vae gradient checkpointing in reward training * Update Training Readme --------- Co-authored-by: hkunzhe <huangkunzhe.hkz@alibaba-inc.com>
1163 lines
50 KiB
Python
1163 lines
50 KiB
Python
# 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, EasyAnimateLayerNormZero, FP32LayerNorm
|
|
from .processor import (EasyAnimateAttnProcessor2_0,
|
|
EasyAnimateSWAttnProcessor2_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",
|
|
is_mmdit_block: bool = True,
|
|
is_swa: bool = False,
|
|
):
|
|
super().__init__()
|
|
|
|
# Attention Part
|
|
self.norm1 = EasyAnimateLayerNormZero(
|
|
time_embed_dim, dim, norm_elementwise_affine, norm_eps, norm_type=norm_type, bias=True
|
|
)
|
|
|
|
self.is_swa = is_swa
|
|
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() if not is_swa else EasyAnimateSWAttnProcessor2_0(),
|
|
)
|
|
if is_mmdit_block:
|
|
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() if not is_swa else EasyAnimateSWAttnProcessor2_0(),
|
|
)
|
|
else:
|
|
self.attn2 = None
|
|
|
|
# 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,
|
|
)
|
|
if is_mmdit_block:
|
|
self.txt_ff = FeedForward(
|
|
dim,
|
|
dropout=dropout,
|
|
activation_fn=activation_fn,
|
|
final_dropout=final_dropout,
|
|
inner_dim=ff_inner_dim,
|
|
bias=ff_bias,
|
|
)
|
|
else:
|
|
self.txt_ff = None
|
|
|
|
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,
|
|
num_frames = None,
|
|
height = None,
|
|
width = 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
|
|
if self.is_swa:
|
|
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,
|
|
num_frames=num_frames,
|
|
height=height,
|
|
width=width,
|
|
)
|
|
else:
|
|
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))
|
|
if self.txt_ff is not None:
|
|
norm_encoder_hidden_states = self.norm3(self.txt_ff(norm_encoder_hidden_states))
|
|
else:
|
|
norm_encoder_hidden_states = self.norm3(self.ff(norm_encoder_hidden_states))
|
|
else:
|
|
norm_hidden_states = self.ff(norm_hidden_states)
|
|
if self.txt_ff is not None:
|
|
norm_encoder_hidden_states = self.txt_ff(norm_encoder_hidden_states)
|
|
else:
|
|
norm_encoder_hidden_states = self.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 |