# 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.
# Adapted from https://github.com/huggingface/diffusers/blob/64bf5d33b7ef1b1deac256bed7bd99b55020c4e0/src/diffusers/models/attention.py
from __future__ import annotations
from copy import deepcopy
from typing import Any, Dict, List, Literal, Optional, Callable, Tuple
import logging
from einops import rearrange
import torch
import torch.nn.functional as F
from torch import nn
from diffusers.models.embeddings import CombinedTimestepLabelEmbeddings
from diffusers.utils.torch_utils import maybe_allow_in_graph
from diffusers.models.attention_processor import Attention as DiffusersAttention
from diffusers.models.attention import (
BasicTransformerBlock as DiffusersBasicTransformerBlock,
AdaLayerNormZero,
AdaLayerNorm,
FeedForward,
)
from diffusers.models.attention_processor import AttnProcessor
from .attention_processor import IPAttention, BaseIPAttnProcessor
logger = logging.getLogger(__name__)
def not_use_xformers_anyway(
use_memory_efficient_attention_xformers: bool,
attention_op: Optional[Callable] = None,
):
return None
@maybe_allow_in_graph
class BasicTransformerBlock(DiffusersBasicTransformerBlock):
print_idx = 0
def __init__(
self,
dim: int,
num_attention_heads: int,
attention_head_dim: int,
dropout=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",
final_dropout: bool = False,
attention_type: str = "default",
allow_xformers: bool = True,
cross_attn_temporal_cond: bool = False,
image_scale: float = 1.0,
processor: AttnProcessor | None = None,
ip_adapter_cross_attn: bool = False,
need_t2i_facein: bool = False,
need_t2i_ip_adapter_face: bool = False,
):
if not only_cross_attention and double_self_attention:
cross_attention_dim = None
super().__init__(
dim,
num_attention_heads,
attention_head_dim,
dropout,
cross_attention_dim,
activation_fn,
num_embeds_ada_norm,
attention_bias,
only_cross_attention,
double_self_attention,
upcast_attention,
norm_elementwise_affine,
norm_type,
final_dropout,
attention_type,
)
self.attn1 = IPAttention(
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,
cross_attn_temporal_cond=cross_attn_temporal_cond,
image_scale=image_scale,
ip_adapter_dim=cross_attention_dim
if only_cross_attention
else attention_head_dim,
facein_dim=cross_attention_dim
if only_cross_attention
else attention_head_dim,
processor=processor,
)
# 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 = IPAttention(
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,
cross_attn_temporal_cond=ip_adapter_cross_attn,
need_t2i_facein=need_t2i_facein,
need_t2i_ip_adapter_face=need_t2i_ip_adapter_face,
image_scale=image_scale,
ip_adapter_dim=cross_attention_dim
if not double_self_attention
else attention_head_dim,
facein_dim=cross_attention_dim
if not double_self_attention
else attention_head_dim,
ip_adapter_face_dim=cross_attention_dim
if not double_self_attention
else attention_head_dim,
processor=processor,
) # is self-attn if encoder_hidden_states is none
else:
self.norm2 = None
self.attn2 = None
if self.attn1 is not None:
if not allow_xformers:
self.attn1.set_use_memory_efficient_attention_xformers = (
not_use_xformers_anyway
)
if self.attn2 is not None:
if not allow_xformers:
self.attn2.set_use_memory_efficient_attention_xformers = (
not_use_xformers_anyway
)
self.double_self_attention = double_self_attention
self.only_cross_attention = only_cross_attention
self.cross_attn_temporal_cond = cross_attn_temporal_cond
self.image_scale = image_scale
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,
self_attn_block_embs: Optional[Tuple[List[torch.Tensor], List[None]]] = None,
self_attn_block_embs_mode: Literal["read", "write"] = "write",
) -> 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
)
if cross_attention_kwargs is None:
cross_attention_kwargs = {}
# 特殊AttnProcessor需要的入参 在 cross_attention_kwargs 准备
# special AttnProcessor needs input parameters in cross_attention_kwargs
original_cross_attention_kwargs = {
k: v
for k, v in cross_attention_kwargs.items()
if k
not in [
"num_frames",
"sample_index",
"vision_conditon_frames_sample_index",
"vision_cond",
"vision_clip_emb",
"ip_adapter_scale",
"face_emb",
"facein_scale",
"ip_adapter_face_emb",
"ip_adapter_face_scale",
"do_classifier_free_guidance",
]
}
if "do_classifier_free_guidance" in cross_attention_kwargs:
do_classifier_free_guidance = cross_attention_kwargs[
"do_classifier_free_guidance"
]
else:
do_classifier_free_guidance = False
# 2. Prepare GLIGEN inputs
original_cross_attention_kwargs = (
original_cross_attention_kwargs.copy()
if original_cross_attention_kwargs is not None
else {}
)
gligen_kwargs = original_cross_attention_kwargs.pop("gligen", None)
# 返回self_attn的结果，适用于referencenet的输出给其他Unet来使用
# return the result of self_attn, which is suitable for the output of referencenet to be used by other Unet
if (
self_attn_block_embs is not None
and self_attn_block_embs_mode.lower() == "write"
):
# self_attn_block_emb = self.attn1.head_to_batch_dim(attn_output, out_dim=4)
self_attn_block_emb = norm_hidden_states
if not hasattr(self, "spatial_self_attn_idx"):
raise ValueError(
"must call unet.insert_spatial_self_attn_idx to generate spatial attn index"
)
basick_transformer_idx = self.spatial_self_attn_idx
if self.print_idx == 0:
logger.debug(
f"self_attn_block_embs, self_attn_block_embs_mode={self_attn_block_embs_mode}, "
f"basick_transformer_idx={basick_transformer_idx}, length={len(self_attn_block_embs)}, shape={self_attn_block_emb.shape}, "
# f"attn1 processor, {type(self.attn1.processor)}"
)
self_attn_block_embs[basick_transformer_idx] = self_attn_block_emb
# read and put referencenet emb into cross_attention_kwargs, which would be fused into attn_processor
if (
self_attn_block_embs is not None
and self_attn_block_embs_mode.lower() == "read"
):
basick_transformer_idx = self.spatial_self_attn_idx
if not hasattr(self, "spatial_self_attn_idx"):
raise ValueError(
"must call unet.insert_spatial_self_attn_idx to generate spatial attn index"
)
if self.print_idx == 0:
logger.debug(
f"refer_self_attn_emb: , self_attn_block_embs_mode={self_attn_block_embs_mode}, "
f"length={len(self_attn_block_embs)}, idx={basick_transformer_idx}, "
# f"attn1 processor, {type(self.attn1.processor)}, "
)
ref_emb = self_attn_block_embs[basick_transformer_idx]
cross_attention_kwargs["refer_emb"] = ref_emb
if self.print_idx == 0:
logger.debug(
f"unet attention read, {self.spatial_self_attn_idx}",
)
# ------------------------------warning-----------------------
# 这两行由于使用了ref_emb会导致和checkpoint_train相关的训练错误，具体未知，留在这里作为警示
# bellow annoated code will cause training error, keep it here as a warning
# logger.debug(f"ref_emb shape,{ref_emb.shape}, {ref_emb.mean()}")
# logger.debug(
# f"norm_hidden_states shape, {norm_hidden_states.shape}, {norm_hidden_states.mean()}",
# )
if self.attn1 is None:
self.print_idx += 1
return norm_hidden_states
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 isinstance(self.attn1.processor, BaseIPAttnProcessor)
else original_cross_attention_kwargs
),
)
if self.use_ada_layer_norm_zero:
attn_output = gate_msa.unsqueeze(1) * attn_output
hidden_states = attn_output + hidden_states
# 推断的时候，对于uncondition_部分独立生成，排除掉 refer_emb，
# 首帧等的影响，避免生成参考了refer_emb、首帧等，又在uncond上去除了
# in inference stage, eliminate influence of refer_emb, vis_cond on unconditionpart
# to avoid use that, and then eliminate in pipeline
# refer to moore-animate anyone
# do_classifier_free_guidance = False
if self.print_idx == 0:
logger.debug(f"do_classifier_free_guidance={do_classifier_free_guidance},")
if do_classifier_free_guidance:
hidden_states_c = attn_output.clone()
_uc_mask = (
torch.Tensor(
[1] * (norm_hidden_states.shape[0] // 2)
+ [0] * (norm_hidden_states.shape[0] // 2)
)
.to(norm_hidden_states.device)
.bool()
)
hidden_states_c[_uc_mask] = self.attn1(
norm_hidden_states[_uc_mask],
encoder_hidden_states=norm_hidden_states[_uc_mask],
attention_mask=attention_mask,
)
attn_output = hidden_states_c.clone()
if "refer_emb" in cross_attention_kwargs:
del cross_attention_kwargs["refer_emb"]
# 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)
)
# 特殊AttnProcessor需要的入参 在 cross_attention_kwargs 准备
# special AttnProcessor needs input parameters in cross_attention_kwargs
attn_output = self.attn2(
norm_hidden_states,
encoder_hidden_states=encoder_hidden_states
if not self.double_self_attention
else None,
attention_mask=encoder_attention_mask,
**(
original_cross_attention_kwargs
if not isinstance(self.attn2.processor, BaseIPAttnProcessor)
else cross_attention_kwargs
),
)
if self.print_idx == 0:
logger.debug(
f"encoder_hidden_states, type={type(encoder_hidden_states)}"
)
if encoder_hidden_states is not None:
logger.debug(
f"encoder_hidden_states, ={encoder_hidden_states.shape}"
)
# encoder_hidden_states_tmp = (
# encoder_hidden_states
# if not self.double_self_attention
# else norm_hidden_states
# )
# if do_classifier_free_guidance:
# hidden_states_c = attn_output.clone()
# _uc_mask = (
# torch.Tensor(
# [1] * (norm_hidden_states.shape[0] // 2)
# + [0] * (norm_hidden_states.shape[0] // 2)
# )
# .to(norm_hidden_states.device)
# .bool()
# )
# hidden_states_c[_uc_mask] = self.attn2(
# norm_hidden_states[_uc_mask],
# encoder_hidden_states=encoder_hidden_states_tmp[_uc_mask],
# attention_mask=attention_mask,
# )
# attn_output = hidden_states_c.clone()
hidden_states = attn_output + hidden_states
# 4. Feed-forward
if self.norm3 is not None and self.ff is not None:
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
self.print_idx += 1
return hidden_states