# 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

import torch

from torch import nn

from ..utils import USE_PEFT_BACKEND

from ..utils.torch_utils import maybe_allow_in_graph

from .activations import GEGLU, GELU, ApproximateGELU

from .attention_processor import Attention

from .lora import LoRACompatibleLinear

from .normalization import AdaLayerNorm, AdaLayerNormZero

@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 = 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))

return x

@maybe_allow_in_graph

class BasicTransformerBlock(nn.Module):

r"""

A basic Transformer block.

Parameters:

dim (`int`): The number of channels in the input and output.

num_attention_heads (`int`): The number of heads to use for multi-head attention.

attention_head_dim (`int`): The number of channels in each head.

dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.

cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.

activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.

num_embeds_ada_norm (:

obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.

attention_bias (:

obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.

only_cross_attention (`bool`, *optional*):

Whether to use only cross-attention layers. In this case two cross attention layers are used.

double_self_attention (`bool`, *optional*):

Whether to use two self-attention layers. In this case no cross attention layers are used.

upcast_attention (`bool`, *optional*):

Whether to upcast the attention computation to float32. This is useful for mixed precision training.

norm_elementwise_affine (`bool`, *optional*, defaults to `True`):

Whether to use learnable elementwise affine parameters for normalization.

norm_type (`str`, *optional*, defaults to `"layer_norm"`):

The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.

final_dropout (`bool` *optional*, defaults to False):

Whether to apply a final dropout after the last feed-forward layer.

attention_type (`str`, *optional*, defaults to `"default"`):

The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.

"""

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",

final_dropout: bool = False,

attention_type: str = "default",

):

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"

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

)

# 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 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)

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,

)

# 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 nn.LayerNorm(dim, elementwise_affine=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,

) # is self-attn if encoder_hidden_states is none

else:

self.norm2 = None

self.attn2 = None

# 3. Feed-forward

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

self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout)

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

# 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

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

)

else:

norm_hidden_states = self.norm1(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

hidden_states = attn_output + hidden_states

# 2.5 GLIGEN Control

if gligen_kwargs is not None:

hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])

# 2.5 ends

# 3. Cross-Attention

if self.attn2 is not None:

norm_hidden_states = (

self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(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

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._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.use_ada_layer_norm_zero:

ff_output = gate_mlp.unsqueeze(1) * ff_output

hidden_states = ff_output + hidden_states

return hidden_states

class FeedForward(nn.Module):

r"""

A feed-forward layer.

Parameters:

dim (`int`): The number of channels in the input.

dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.

mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.

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.

final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.

"""

def __init__(

self,

dim: int,

dim_out: Optional[int] = None,

mult: int = 4,

dropout: float = 0.0,

activation_fn: str = "geglu",

final_dropout: bool = False,

):

super().__init__()

inner_dim = int(dim * mult)

dim_out = dim_out if dim_out is not None else dim

linear_cls = LoRACompatibleLinear if not USE_PEFT_BACKEND else nn.Linear

if activation_fn == "gelu":

act_fn = GELU(dim, inner_dim)

if activation_fn == "gelu-approximate":

act_fn = GELU(dim, inner_dim, approximate="tanh")

elif activation_fn == "geglu":

act_fn = GEGLU(dim, inner_dim)

elif activation_fn == "geglu-approximate":

act_fn = ApproximateGELU(dim, inner_dim)

self.net = nn.ModuleList([])

# project in

self.net.append(act_fn)

# project dropout

self.net.append(nn.Dropout(dropout))

# project out

self.net.append(linear_cls(inner_dim, dim_out))

# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout

if final_dropout:

self.net.append(nn.Dropout(dropout))

def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor:

compatible_cls = (GEGLU,) if USE_PEFT_BACKEND else (GEGLU, LoRACompatibleLinear)

for module in self.net:

if isinstance(module, compatible_cls):

hidden_states = module(hidden_states, scale)

else:

hidden_states = module(hidden_states)

return hidden_states