@property

def attn_processors(self) -> dict[str, AttentionProcessor]:

r"""

# 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

def set_attn_processor(self, processor: AttentionProcessor | dict[str, AttentionProcessor]):

r"""

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 fuse_qkv_projections(self):

"""

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.")

for module in self.modules():

if isinstance(module, AttentionModuleMixin) and module._supports_qkv_fusion:

module.fuse_projections()

def unfuse_qkv_projections(self):

"""Disables the fused QKV projection if enabled.

for module in self.modules():

if isinstance(module, AttentionModuleMixin) and module._supports_qkv_fusion:

_default_processor_cls = None

_available_processors = []

_supports_qkv_fusion = True

fused_projections = False

def set_processor(self, processor: AttentionProcessor) -> None:

"""

# if current processor is in `self._modules` and if passed `processor` is not, we need to

# pop `processor` from `self._modules`

if (

hasattr(self, "processor")

and isinstance(self.processor, torch.nn.Module)

and not isinstance(processor, torch.nn.Module)

):

logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}")

self._modules.pop("processor")

self.processor = processor

def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProcessor":

"""

if not return_deprecated_lora:

return self.processor

def set_attention_backend(self, backend: str):

from .attention_dispatch import AttentionBackendName

available_backends = {x.value for x in AttentionBackendName.__members__.values()}

if backend not in available_backends:

raise ValueError(f"`{backend=}` must be one of the following: " + ", ".join(available_backends))

backend = AttentionBackendName(backend.lower())

self.processor._attention_backend = backend

def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None:

"""

if use_npu_flash_attention:

if not is_torch_npu_available():

raise ImportError("torch_npu is not available")

self.set_attention_backend("_native_npu")

def set_use_xla_flash_attention(

self,

use_xla_flash_attention: bool,

partition_spec: tuple[str | None, ...] | None = None,

is_flux=False,

) -> None:

"""

if use_xla_flash_attention:

if not is_torch_xla_available():

raise ImportError("torch_xla is not available")

self.set_attention_backend("_native_xla")

def set_use_memory_efficient_attention_xformers(

self, use_memory_efficient_attention_xformers: bool, attention_op: Callable | None = None

) -> None:

"""

if use_memory_efficient_attention_xformers:

if not is_xformers_available():

raise ModuleNotFoundError(

"Refer to https://github.com/facebookresearch/xformers for more information on how to install xformers",

name="xformers",

)

elif not torch.cuda.is_available():

raise ValueError(

"torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"

" only available for GPU "

)

else:

try:

# Make sure we can run the memory efficient attention

if is_xformers_available():

dtype = None

if attention_op is not None:

op_fw, op_bw = attention_op

dtype, *_ = op_fw.SUPPORTED_DTYPES

q = torch.randn((1, 2, 40), device="cuda", dtype=dtype)

_ = xops.ops.memory_efficient_attention(q, q, q)

except Exception as e:

raise e

self.set_attention_backend("xformers")

@torch.no_grad()

def fuse_projections(self):

"""

# Skip if the AttentionModuleMixin subclass does not support fusion (for example, the QKV projections in Flux2

# single stream blocks are always fused)

if not self._supports_qkv_fusion:

logger.debug(

f"{self.__class__.__name__} does not support fusing QKV projections, so `fuse_projections` will no-op."

)

return

# Skip if already fused

if getattr(self, "fused_projections", False):

return

device = self.to_q.weight.data.device

dtype = self.to_q.weight.data.dtype

if hasattr(self, "is_cross_attention") and self.is_cross_attention:

# Fuse cross-attention key-value projections

concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data])

in_features = concatenated_weights.shape[1]

out_features = concatenated_weights.shape[0]

self.to_kv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)

self.to_kv.weight.copy_(concatenated_weights)

if hasattr(self, "use_bias") and self.use_bias:

concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data])

self.to_kv.bias.copy_(concatenated_bias)

else:

# Fuse self-attention projections

concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])

in_features = concatenated_weights.shape[1]

out_features = concatenated_weights.shape[0]

self.to_qkv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)

self.to_qkv.weight.copy_(concatenated_weights)

if hasattr(self, "use_bias") and self.use_bias:

concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])

self.to_qkv.bias.copy_(concatenated_bias)

# Handle added projections for models like SD3, Flux, etc.

if (

getattr(self, "add_q_proj", None) is not None

and getattr(self, "add_k_proj", None) is not None

and getattr(self, "add_v_proj", None) is not None

):

concatenated_weights = torch.cat(

[self.add_q_proj.weight.data, self.add_k_proj.weight.data, self.add_v_proj.weight.data]

)

in_features = concatenated_weights.shape[1]

out_features = concatenated_weights.shape[0]

self.to_added_qkv = nn.Linear(

in_features, out_features, bias=self.added_proj_bias, device=device, dtype=dtype

)

self.to_added_qkv.weight.copy_(concatenated_weights)

if self.added_proj_bias:

concatenated_bias = torch.cat(

[self.add_q_proj.bias.data, self.add_k_proj.bias.data, self.add_v_proj.bias.data]

)

self.to_added_qkv.bias.copy_(concatenated_bias)

self.fused_projections = True

@torch.no_grad()

def unfuse_projections(self):

"""

# Skip if the AttentionModuleMixin subclass does not support fusion (for example, the QKV projections in Flux2

# single stream blocks are always fused)

if not self._supports_qkv_fusion:

return

# Skip if not fused

if not getattr(self, "fused_projections", False):

return

# Remove fused projection layers

if hasattr(self, "to_qkv"):

delattr(self, "to_qkv")

if hasattr(self, "to_kv"):

delattr(self, "to_kv")

if hasattr(self, "to_added_qkv"):

delattr(self, "to_added_qkv")

self.fused_projections = False

def set_attention_slice(self, slice_size: int) -> None:

"""

if hasattr(self, "sliceable_head_dim") and slice_size is not None and slice_size > self.sliceable_head_dim:

raise ValueError(f"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.")

processor = None

# Try to get a compatible processor for sliced attention

if slice_size is not None:

processor = self._get_compatible_processor("sliced")

# If no processor was found or slice_size is None, use default processor

if processor is None:

processor = self.default_processor_cls()

self.set_processor(processor)

def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:

"""

head_size = self.heads

batch_size, seq_len, dim = tensor.shape

tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)

tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)

return tensor

def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:

"""

head_size = self.heads

if tensor.ndim == 3:

batch_size, seq_len, dim = tensor.shape

extra_dim = 1

else:

batch_size, extra_dim, seq_len, dim = tensor.shape

tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size)

tensor = tensor.permute(0, 2, 1, 3)

if out_dim == 3:

tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size)

return tensor

def get_attention_scores(

self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor | None = None

) -> torch.Tensor:

"""

dtype = query.dtype

if self.upcast_attention:

query = query.float()

key = key.float()

if attention_mask is None:

baddbmm_input = torch.empty(

query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device

)

beta = 0

else:

baddbmm_input = attention_mask

beta = 1

attention_scores = torch.baddbmm(

baddbmm_input,

query,

key.transpose(-1, -2),

beta=beta,

alpha=self.scale,

)

del baddbmm_input

if self.upcast_softmax:

attention_scores = attention_scores.float()

attention_probs = attention_scores.softmax(dim=-1)

del attention_scores

attention_probs = attention_probs.to(dtype)

return attention_probs

def prepare_attention_mask(

self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3

) -> torch.Tensor:

"""

head_size = self.heads

if attention_mask is None:

return attention_mask

current_length: int = attention_mask.shape[-1]

if current_length != target_length:

if attention_mask.device.type == "mps":

# HACK: MPS: Does not support padding by greater than dimension of input tensor.

# Instead, we can manually construct the padding tensor.

padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)

padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)

attention_mask = torch.cat([attention_mask, padding], dim=2)

else:

# TODO: for pipelines such as stable-diffusion, padding cross-attn mask:

# we want to instead pad by (0, remaining_length), where remaining_length is:

# remaining_length: int = target_length - current_length

# TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding

attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)

if out_dim == 3:

if attention_mask.shape[0] < batch_size * head_size:

attention_mask = attention_mask.repeat_interleave(head_size, dim=0)

elif out_dim == 4:

attention_mask = attention_mask.unsqueeze(1)

attention_mask = attention_mask.repeat_interleave(head_size, dim=1)

return attention_mask

def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:

"""

assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states"

if isinstance(self.norm_cross, nn.LayerNorm):

encoder_hidden_states = self.norm_cross(encoder_hidden_states)

elif isinstance(self.norm_cross, nn.GroupNorm):

# Group norm norms along the channels dimension and expects

# input to be in the shape of (N, C, *). In this case, we want

# to norm along the hidden dimension, so we need to move

# (batch_size, sequence_length, hidden_size) ->

# (batch_size, hidden_size, sequence_length)

encoder_hidden_states = encoder_hidden_states.transpose(1, 2)

encoder_hidden_states = self.norm_cross(encoder_hidden_states)

encoder_hidden_states = encoder_hidden_states.transpose(1, 2)

else:

assert False

# "feed_forward_chunk_size" can be used to save memory

if hidden_states.shape[chunk_dim] % chunk_size != 0:

raise ValueError(

f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."

)

num_chunks = hidden_states.shape[chunk_dim] // chunk_size

ff_output = torch.cat(

[ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)],

dim=chunk_dim,

)

r"""

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 = nn.LayerNorm(query_dim)

self.norm2 = nn.LayerNorm(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))

r"""

def __init__(

self,

dim: int,

num_attention_heads: int,

attention_head_dim: int,

context_pre_only: bool = False,

qk_norm: str | None = None,

use_dual_attention: bool = False,

):

super().__init__()

self.use_dual_attention = use_dual_attention

self.context_pre_only = context_pre_only

context_norm_type = "ada_norm_continous" if context_pre_only else "ada_norm_zero"

if use_dual_attention:

self.norm1 = SD35AdaLayerNormZeroX(dim)

else:

self.norm1 = AdaLayerNormZero(dim)

if context_norm_type == "ada_norm_continous":

self.norm1_context = AdaLayerNormContinuous(

dim, dim, elementwise_affine=False, eps=1e-6, bias=True, norm_type="layer_norm"

)

elif context_norm_type == "ada_norm_zero":

self.norm1_context = AdaLayerNormZero(dim)

else:

raise ValueError(

f"Unknown context_norm_type: {context_norm_type}, currently only support `ada_norm_continous`, `ada_norm_zero`"

)

if hasattr(F, "scaled_dot_product_attention"):

processor = JointAttnProcessor2_0()

else:

raise ValueError(

"The current PyTorch version does not support the `scaled_dot_product_attention` function."

)

self.attn = Attention(

query_dim=dim,

cross_attention_dim=None,

added_kv_proj_dim=dim,

dim_head=attention_head_dim,

heads=num_attention_heads,

out_dim=dim,

context_pre_only=context_pre_only,

bias=True,

processor=processor,

qk_norm=qk_norm,

eps=1e-6,

)

if use_dual_attention:

self.attn2 = Attention(

query_dim=dim,

cross_attention_dim=None,

dim_head=attention_head_dim,

heads=num_attention_heads,

out_dim=dim,

bias=True,

processor=processor,

qk_norm=qk_norm,

eps=1e-6,

)

else:

self.attn2 = None

self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)

self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")

if not context_pre_only:

self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)

self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")

else:

self.norm2_context = None

self.ff_context = None

# let chunk size default to None

self._chunk_size = None

self._chunk_dim = 0

# Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward

def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0):

# Sets chunk feed-forward

self._chunk_size = chunk_size

self._chunk_dim = dim

def forward(

self,

hidden_states: torch.FloatTensor,

encoder_hidden_states: torch.FloatTensor,

temb: torch.FloatTensor,

joint_attention_kwargs: dict[str, Any] | None = None,

) -> tuple[torch.Tensor, torch.Tensor]:

joint_attention_kwargs = joint_attention_kwargs or {}

if self.use_dual_attention:

norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1(

hidden_states, emb=temb

)

else:

norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)

if self.context_pre_only:

norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)

else:

norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(

encoder_hidden_states, emb=temb

)

# Attention.

attn_output, context_attn_output = self.attn(

hidden_states=norm_hidden_states,

encoder_hidden_states=norm_encoder_hidden_states,

**joint_attention_kwargs,

)

# Process attention outputs for the `hidden_states`.

attn_output = gate_msa.unsqueeze(1) * attn_output

hidden_states = hidden_states + attn_output

if self.use_dual_attention:

attn_output2 = self.attn2(hidden_states=norm_hidden_states2, **joint_attention_kwargs)

attn_output2 = gate_msa2.unsqueeze(1) * attn_output2

hidden_states = hidden_states + attn_output2

norm_hidden_states = self.norm2(hidden_states)

norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

if self._chunk_size is not None:

# "feed_forward_chunk_size" can be used to save memory

ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)

else:

ff_output = self.ff(norm_hidden_states)

ff_output = gate_mlp.unsqueeze(1) * ff_output

hidden_states = hidden_states + ff_output

# Process attention outputs for the `encoder_hidden_states`.

if self.context_pre_only:

encoder_hidden_states = None

else:

context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output

encoder_hidden_states = encoder_hidden_states + context_attn_output

norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)

norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]

if self._chunk_size is not None:

# "feed_forward_chunk_size" can be used to save memory

context_ff_output = _chunked_feed_forward(

self.ff_context, norm_encoder_hidden_states, self._chunk_dim, self._chunk_size

)

else:

context_ff_output = self.ff_context(norm_encoder_hidden_states)

encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output

r"""

def __init__(

self,

dim: int,

num_attention_heads: int,

attention_head_dim: int,

dropout=0.0,

cross_attention_dim: int | None = None,

activation_fn: str = "geglu",

num_embeds_ada_norm: int | None = 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', 'ada_norm_continuous', 'layer_norm_i2vgen'

norm_eps: float = 1e-5,

final_dropout: bool = False,

attention_type: str = "default",

positional_embeddings: str | None = None,

num_positional_embeddings: int | None = None,

ada_norm_continous_conditioning_embedding_dim: int | None = None,

ada_norm_bias: int | None = None,

ff_inner_dim: int | None = None,

ff_bias: bool = True,

attention_out_bias: bool = True,

):

super().__init__()

self.dim = dim

self.num_attention_heads = num_attention_heads

self.attention_head_dim = attention_head_dim

self.dropout = dropout

self.cross_attention_dim = cross_attention_dim

self.activation_fn = activation_fn

self.attention_bias = attention_bias

self.double_self_attention = double_self_attention

self.norm_elementwise_affine = norm_elementwise_affine

self.positional_embeddings = positional_embeddings

self.num_positional_embeddings = num_positional_embeddings

self.only_cross_attention = only_cross_attention

# We keep these boolean flags for backward-compatibility.

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"

self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous"

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}."

)

self.norm_type = norm_type

self.num_embeds_ada_norm = num_embeds_ada_norm

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 norm_type == "ada_norm":

self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)

elif norm_type == "ada_norm_zero":

self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)

elif norm_type == "ada_norm_continuous":

self.norm1 = AdaLayerNormContinuous(

dim,

ada_norm_continous_conditioning_embedding_dim,

norm_elementwise_affine,

norm_eps,

ada_norm_bias,

"rms_norm",

)

else:

self.norm1 = nn.LayerNorm(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,

out_bias=attention_out_bias,

)

# 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.

if norm_type == "ada_norm":

self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm)

elif norm_type == "ada_norm_continuous":

self.norm2 = AdaLayerNormContinuous(

dim,

ada_norm_continous_conditioning_embedding_dim,

norm_elementwise_affine,

norm_eps,

ada_norm_bias,

"rms_norm",

)

else:

self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)

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,

out_bias=attention_out_bias,

) # is self-attn if encoder_hidden_states is none

else:

if norm_type == "ada_norm_single": # For Latte

self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)

else:

self.norm2 = None

self.attn2 = None

# 3. Feed-forward

if norm_type == "ada_norm_continuous":

self.norm3 = AdaLayerNormContinuous(

dim,

ada_norm_continous_conditioning_embedding_dim,

norm_elementwise_affine,

norm_eps,

ada_norm_bias,

"layer_norm",

)

elif norm_type in ["ada_norm_zero", "ada_norm", "layer_norm"]:

self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)

elif norm_type == "layer_norm_i2vgen":

self.norm3 = None

self.ff = FeedForward(

dim,

dropout=dropout,

activation_fn=activation_fn,

final_dropout=final_dropout,

inner_dim=ff_inner_dim,

bias=ff_bias,

)

# 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 norm_type == "ada_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: int | None, dim: int = 0):

# Sets chunk feed-forward

self._chunk_size = chunk_size

self._chunk_dim = dim

def forward(

self,

hidden_states: torch.Tensor,

attention_mask: torch.Tensor | None = None,

encoder_hidden_states: torch.Tensor | None = None,

encoder_attention_mask: torch.Tensor | None = None,

timestep: torch.LongTensor | None = None,

cross_attention_kwargs: dict[str, Any] = None,

class_labels: torch.LongTensor | None = None,

added_cond_kwargs: dict[str, torch.Tensor] | None = None,

) -> torch.Tensor:

if cross_attention_kwargs is not None:

if cross_attention_kwargs.get("scale", None) is not None:

logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")

# 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.norm_type == "ada_norm":

norm_hidden_states = self.norm1(hidden_states, timestep)

elif self.norm_type == "ada_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.norm_type in ["layer_norm", "layer_norm_i2vgen"]:

norm_hidden_states = self.norm1(hidden_states)

elif self.norm_type == "ada_norm_continuous":

norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"])

elif self.norm_type == "ada_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

else:

raise ValueError("Incorrect norm used")

if self.pos_embed is not None:

norm_hidden_states = self.pos_embed(norm_hidden_states)