MuseV/musev/models/attention.py

# 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