751 lines
28 KiB
Python
751 lines
28 KiB
Python
# Copyright 2024 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.
|
|
|
|
"""该模型是自定义的attn_processor，实现特殊功能的 Attn功能。
|
|
相对而言，开源代码经常会重新定义Attention 类，
|
|
|
|
This module implements special AttnProcessor function with custom attn_processor class.
|
|
While other open source code always modify Attention class.
|
|
"""
|
|
# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from typing import Any, Callable, Optional
|
|
import logging
|
|
|
|
from einops import rearrange, repeat
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import xformers
|
|
from MuseVdiffusers.models.lora import LoRACompatibleLinear
|
|
|
|
from MuseVdiffusers.utils.torch_utils import maybe_allow_in_graph
|
|
from MuseVdiffusers.models.attention_processor import (
|
|
Attention as DiffusersAttention,
|
|
AttnProcessor,
|
|
AttnProcessor2_0,
|
|
)
|
|
from ..data.data_util import (
|
|
batch_concat_two_tensor_with_index,
|
|
batch_index_select,
|
|
align_repeat_tensor_single_dim,
|
|
batch_adain_conditioned_tensor,
|
|
)
|
|
|
|
from . import Model_Register
|
|
|
|
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
@maybe_allow_in_graph
|
|
class IPAttention(DiffusersAttention):
|
|
r"""
|
|
Modified Attention class which has special layer, like ip_apadapter_to_k, ip_apadapter_to_v,
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
query_dim: int,
|
|
cross_attention_dim: int | None = None,
|
|
heads: int = 8,
|
|
dim_head: int = 64,
|
|
dropout: float = 0,
|
|
bias=False,
|
|
upcast_attention: bool = False,
|
|
upcast_softmax: bool = False,
|
|
cross_attention_norm: str | None = None,
|
|
cross_attention_norm_num_groups: int = 32,
|
|
added_kv_proj_dim: int | None = None,
|
|
norm_num_groups: int | None = None,
|
|
spatial_norm_dim: int | None = None,
|
|
out_bias: bool = True,
|
|
scale_qk: bool = True,
|
|
only_cross_attention: bool = False,
|
|
eps: float = 0.00001,
|
|
rescale_output_factor: float = 1,
|
|
residual_connection: bool = False,
|
|
_from_deprecated_attn_block=False,
|
|
processor: AttnProcessor | None = None,
|
|
cross_attn_temporal_cond: bool = False,
|
|
image_scale: float = 1.0,
|
|
ip_adapter_dim: int = None,
|
|
need_t2i_facein: bool = False,
|
|
facein_dim: int = None,
|
|
need_t2i_ip_adapter_face: bool = False,
|
|
ip_adapter_face_dim: int = None,
|
|
):
|
|
super().__init__(
|
|
query_dim,
|
|
cross_attention_dim,
|
|
heads,
|
|
dim_head,
|
|
dropout,
|
|
bias,
|
|
upcast_attention,
|
|
upcast_softmax,
|
|
cross_attention_norm,
|
|
cross_attention_norm_num_groups,
|
|
added_kv_proj_dim,
|
|
norm_num_groups,
|
|
spatial_norm_dim,
|
|
out_bias,
|
|
scale_qk,
|
|
only_cross_attention,
|
|
eps,
|
|
rescale_output_factor,
|
|
residual_connection,
|
|
_from_deprecated_attn_block,
|
|
processor,
|
|
)
|
|
self.cross_attn_temporal_cond = cross_attn_temporal_cond
|
|
self.image_scale = image_scale
|
|
# 面向首帧的 ip_adapter
|
|
# ip_apdater
|
|
if cross_attn_temporal_cond:
|
|
self.to_k_ip = LoRACompatibleLinear(ip_adapter_dim, query_dim, bias=False)
|
|
self.to_v_ip = LoRACompatibleLinear(ip_adapter_dim, query_dim, bias=False)
|
|
# facein
|
|
self.need_t2i_facein = need_t2i_facein
|
|
self.facein_dim = facein_dim
|
|
if need_t2i_facein:
|
|
raise NotImplementedError("facein")
|
|
|
|
# ip_adapter_face
|
|
self.need_t2i_ip_adapter_face = need_t2i_ip_adapter_face
|
|
self.ip_adapter_face_dim = ip_adapter_face_dim
|
|
if need_t2i_ip_adapter_face:
|
|
self.ip_adapter_face_to_k_ip = LoRACompatibleLinear(
|
|
ip_adapter_face_dim, query_dim, bias=False
|
|
)
|
|
self.ip_adapter_face_to_v_ip = LoRACompatibleLinear(
|
|
ip_adapter_face_dim, query_dim, bias=False
|
|
)
|
|
|
|
def set_use_memory_efficient_attention_xformers(
|
|
self,
|
|
use_memory_efficient_attention_xformers: bool,
|
|
attention_op: Callable[..., Any] | None = None,
|
|
):
|
|
if (
|
|
"XFormers" in self.processor.__class__.__name__
|
|
or "IP" in self.processor.__class__.__name__
|
|
):
|
|
pass
|
|
else:
|
|
return super().set_use_memory_efficient_attention_xformers(
|
|
use_memory_efficient_attention_xformers, attention_op
|
|
)
|
|
|
|
|
|
@Model_Register.register
|
|
class BaseIPAttnProcessor(nn.Module):
|
|
print_idx = 0
|
|
|
|
def __init__(self, *args, **kwargs) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
|
@Model_Register.register
|
|
class T2IReferencenetIPAdapterXFormersAttnProcessor(BaseIPAttnProcessor):
|
|
r"""
|
|
面向 ref_image的 self_attn的 IPAdapter
|
|
"""
|
|
print_idx = 0
|
|
|
|
def __init__(
|
|
self,
|
|
attention_op: Optional[Callable] = None,
|
|
):
|
|
super().__init__()
|
|
|
|
self.attention_op = attention_op
|
|
|
|
def __call__(
|
|
self,
|
|
attn: IPAttention,
|
|
hidden_states: torch.FloatTensor,
|
|
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
|
attention_mask: Optional[torch.FloatTensor] = None,
|
|
temb: Optional[torch.FloatTensor] = None,
|
|
scale: float = 1.0,
|
|
num_frames: int = None,
|
|
sample_index: torch.LongTensor = None,
|
|
vision_conditon_frames_sample_index: torch.LongTensor = None,
|
|
refer_emb: torch.Tensor = None,
|
|
vision_clip_emb: torch.Tensor = None,
|
|
ip_adapter_scale: float = 1.0,
|
|
face_emb: torch.Tensor = None,
|
|
facein_scale: float = 1.0,
|
|
ip_adapter_face_emb: torch.Tensor = None,
|
|
ip_adapter_face_scale: float = 1.0,
|
|
do_classifier_free_guidance: bool = False,
|
|
):
|
|
residual = hidden_states
|
|
|
|
if attn.spatial_norm is not None:
|
|
hidden_states = attn.spatial_norm(hidden_states, temb)
|
|
|
|
input_ndim = hidden_states.ndim
|
|
|
|
if input_ndim == 4:
|
|
batch_size, channel, height, width = hidden_states.shape
|
|
hidden_states = hidden_states.view(
|
|
batch_size, channel, height * width
|
|
).transpose(1, 2)
|
|
|
|
batch_size, key_tokens, _ = (
|
|
hidden_states.shape
|
|
if encoder_hidden_states is None
|
|
else encoder_hidden_states.shape
|
|
)
|
|
|
|
attention_mask = attn.prepare_attention_mask(
|
|
attention_mask, key_tokens, batch_size
|
|
)
|
|
if attention_mask is not None:
|
|
# expand our mask's singleton query_tokens dimension:
|
|
# [batch*heads, 1, key_tokens] ->
|
|
# [batch*heads, query_tokens, key_tokens]
|
|
# so that it can be added as a bias onto the attention scores that xformers computes:
|
|
# [batch*heads, query_tokens, key_tokens]
|
|
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
|
_, query_tokens, _ = hidden_states.shape
|
|
attention_mask = attention_mask.expand(-1, query_tokens, -1)
|
|
|
|
if attn.group_norm is not None:
|
|
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
|
|
1, 2
|
|
)
|
|
|
|
query = attn.to_q(hidden_states)
|
|
|
|
if encoder_hidden_states is None:
|
|
encoder_hidden_states = hidden_states
|
|
elif attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(
|
|
encoder_hidden_states
|
|
)
|
|
encoder_hidden_states = align_repeat_tensor_single_dim(
|
|
encoder_hidden_states, target_length=hidden_states.shape[0], dim=0
|
|
)
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
# for facein
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"T2IReferencenetIPAdapterXFormersAttnProcessor,type(face_emb)={type(face_emb)}, facein_scale={facein_scale}"
|
|
)
|
|
if facein_scale > 0 and face_emb is not None:
|
|
raise NotImplementedError("facein")
|
|
|
|
query = attn.head_to_batch_dim(query).contiguous()
|
|
key = attn.head_to_batch_dim(key).contiguous()
|
|
value = attn.head_to_batch_dim(value).contiguous()
|
|
hidden_states = xformers.ops.memory_efficient_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
attn_bias=attention_mask,
|
|
op=self.attention_op,
|
|
scale=attn.scale,
|
|
)
|
|
|
|
# ip-adapter start
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"T2IReferencenetIPAdapterXFormersAttnProcessor,type(vision_clip_emb)={type(vision_clip_emb)}"
|
|
)
|
|
if ip_adapter_scale > 0 and vision_clip_emb is not None:
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"T2I cross_attn, ipadapter, vision_clip_emb={vision_clip_emb.shape}, hidden_states={hidden_states.shape}, batch_size={batch_size}"
|
|
)
|
|
ip_key = attn.to_k_ip(vision_clip_emb)
|
|
ip_value = attn.to_v_ip(vision_clip_emb)
|
|
ip_key = align_repeat_tensor_single_dim(
|
|
ip_key, target_length=batch_size, dim=0
|
|
)
|
|
ip_value = align_repeat_tensor_single_dim(
|
|
ip_value, target_length=batch_size, dim=0
|
|
)
|
|
ip_key = attn.head_to_batch_dim(ip_key).contiguous()
|
|
ip_value = attn.head_to_batch_dim(ip_value).contiguous()
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"query={query.shape}, ip_key={ip_key.shape}, ip_value={ip_value.shape}"
|
|
)
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
hidden_states_from_ip = xformers.ops.memory_efficient_attention(
|
|
query,
|
|
ip_key,
|
|
ip_value,
|
|
attn_bias=attention_mask,
|
|
op=self.attention_op,
|
|
scale=attn.scale,
|
|
)
|
|
hidden_states = hidden_states + ip_adapter_scale * hidden_states_from_ip
|
|
# ip-adapter end
|
|
|
|
# ip-adapter face start
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"T2IReferencenetIPAdapterXFormersAttnProcessor,type(ip_adapter_face_emb)={type(ip_adapter_face_emb)}"
|
|
)
|
|
if ip_adapter_face_scale > 0 and ip_adapter_face_emb is not None:
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"T2I cross_attn, ipadapter face, ip_adapter_face_emb={vision_clip_emb.shape}, hidden_states={hidden_states.shape}, batch_size={batch_size}"
|
|
)
|
|
ip_key = attn.ip_adapter_face_to_k_ip(ip_adapter_face_emb)
|
|
ip_value = attn.ip_adapter_face_to_v_ip(ip_adapter_face_emb)
|
|
ip_key = align_repeat_tensor_single_dim(
|
|
ip_key, target_length=batch_size, dim=0
|
|
)
|
|
ip_value = align_repeat_tensor_single_dim(
|
|
ip_value, target_length=batch_size, dim=0
|
|
)
|
|
ip_key = attn.head_to_batch_dim(ip_key).contiguous()
|
|
ip_value = attn.head_to_batch_dim(ip_value).contiguous()
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"query={query.shape}, ip_key={ip_key.shape}, ip_value={ip_value.shape}"
|
|
)
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
hidden_states_from_ip = xformers.ops.memory_efficient_attention(
|
|
query,
|
|
ip_key,
|
|
ip_value,
|
|
attn_bias=attention_mask,
|
|
op=self.attention_op,
|
|
scale=attn.scale,
|
|
)
|
|
hidden_states = (
|
|
hidden_states + ip_adapter_face_scale * hidden_states_from_ip
|
|
)
|
|
# ip-adapter face end
|
|
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
hidden_states = attn.batch_to_head_dim(hidden_states)
|
|
|
|
# linear proj
|
|
hidden_states = attn.to_out[0](hidden_states)
|
|
# dropout
|
|
hidden_states = attn.to_out[1](hidden_states)
|
|
|
|
if input_ndim == 4:
|
|
hidden_states = hidden_states.transpose(-1, -2).reshape(
|
|
batch_size, channel, height, width
|
|
)
|
|
|
|
if attn.residual_connection:
|
|
hidden_states = hidden_states + residual
|
|
|
|
hidden_states = hidden_states / attn.rescale_output_factor
|
|
self.print_idx += 1
|
|
return hidden_states
|
|
|
|
|
|
@Model_Register.register
|
|
class NonParamT2ISelfReferenceXFormersAttnProcessor(BaseIPAttnProcessor):
|
|
r"""
|
|
面向首帧的 referenceonly attn,适用于 T2I的 self_attn
|
|
referenceonly with vis_cond as key, value, in t2i self_attn.
|
|
"""
|
|
print_idx = 0
|
|
|
|
def __init__(
|
|
self,
|
|
attention_op: Optional[Callable] = None,
|
|
):
|
|
super().__init__()
|
|
|
|
self.attention_op = attention_op
|
|
|
|
def __call__(
|
|
self,
|
|
attn: IPAttention,
|
|
hidden_states: torch.FloatTensor,
|
|
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
|
attention_mask: Optional[torch.FloatTensor] = None,
|
|
temb: Optional[torch.FloatTensor] = None,
|
|
scale: float = 1.0,
|
|
num_frames: int = None,
|
|
sample_index: torch.LongTensor = None,
|
|
vision_conditon_frames_sample_index: torch.LongTensor = None,
|
|
refer_emb: torch.Tensor = None,
|
|
face_emb: torch.Tensor = None,
|
|
vision_clip_emb: torch.Tensor = None,
|
|
ip_adapter_scale: float = 1.0,
|
|
facein_scale: float = 1.0,
|
|
ip_adapter_face_emb: torch.Tensor = None,
|
|
ip_adapter_face_scale: float = 1.0,
|
|
do_classifier_free_guidance: bool = False,
|
|
):
|
|
residual = hidden_states
|
|
|
|
if attn.spatial_norm is not None:
|
|
hidden_states = attn.spatial_norm(hidden_states, temb)
|
|
|
|
input_ndim = hidden_states.ndim
|
|
|
|
if input_ndim == 4:
|
|
batch_size, channel, height, width = hidden_states.shape
|
|
hidden_states = hidden_states.view(
|
|
batch_size, channel, height * width
|
|
).transpose(1, 2)
|
|
|
|
batch_size, key_tokens, _ = (
|
|
hidden_states.shape
|
|
if encoder_hidden_states is None
|
|
else encoder_hidden_states.shape
|
|
)
|
|
|
|
attention_mask = attn.prepare_attention_mask(
|
|
attention_mask, key_tokens, batch_size
|
|
)
|
|
if attention_mask is not None:
|
|
# expand our mask's singleton query_tokens dimension:
|
|
# [batch*heads, 1, key_tokens] ->
|
|
# [batch*heads, query_tokens, key_tokens]
|
|
# so that it can be added as a bias onto the attention scores that xformers computes:
|
|
# [batch*heads, query_tokens, key_tokens]
|
|
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
|
_, query_tokens, _ = hidden_states.shape
|
|
attention_mask = attention_mask.expand(-1, query_tokens, -1)
|
|
|
|
# vision_cond in same unet attn start
|
|
if (
|
|
vision_conditon_frames_sample_index is not None and num_frames > 1
|
|
) or refer_emb is not None:
|
|
batchsize_timesize = hidden_states.shape[0]
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"NonParamT2ISelfReferenceXFormersAttnProcessor 0, hidden_states={hidden_states.shape}, vision_conditon_frames_sample_index={vision_conditon_frames_sample_index}"
|
|
)
|
|
encoder_hidden_states = rearrange(
|
|
hidden_states, "(b t) hw c -> b t hw c", t=num_frames
|
|
)
|
|
# if False:
|
|
if vision_conditon_frames_sample_index is not None and num_frames > 1:
|
|
ip_hidden_states = batch_index_select(
|
|
encoder_hidden_states,
|
|
dim=1,
|
|
index=vision_conditon_frames_sample_index,
|
|
).contiguous()
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"NonParamT2ISelfReferenceXFormersAttnProcessor 1, vis_cond referenceonly, encoder_hidden_states={encoder_hidden_states.shape}, ip_hidden_states={ip_hidden_states.shape}"
|
|
)
|
|
#
|
|
ip_hidden_states = rearrange(
|
|
ip_hidden_states, "b t hw c -> b 1 (t hw) c"
|
|
)
|
|
ip_hidden_states = align_repeat_tensor_single_dim(
|
|
ip_hidden_states,
|
|
dim=1,
|
|
target_length=num_frames,
|
|
)
|
|
# b t hw c -> b t hw + hw c
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"NonParamT2ISelfReferenceXFormersAttnProcessor 2, vis_cond referenceonly, encoder_hidden_states={encoder_hidden_states.shape}, ip_hidden_states={ip_hidden_states.shape}"
|
|
)
|
|
encoder_hidden_states = torch.concat(
|
|
[encoder_hidden_states, ip_hidden_states], dim=2
|
|
)
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"NonParamT2ISelfReferenceXFormersAttnProcessor 3, hidden_states={hidden_states.shape}, ip_hidden_states={ip_hidden_states.shape}"
|
|
)
|
|
# if False:
|
|
if refer_emb is not None: # and num_frames > 1:
|
|
refer_emb = rearrange(refer_emb, "b c t h w->b 1 (t h w) c")
|
|
refer_emb = align_repeat_tensor_single_dim(
|
|
refer_emb, target_length=num_frames, dim=1
|
|
)
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"NonParamT2ISelfReferenceXFormersAttnProcessor4, referencenet, encoder_hidden_states={encoder_hidden_states.shape}, refer_emb={refer_emb.shape}"
|
|
)
|
|
encoder_hidden_states = torch.concat(
|
|
[encoder_hidden_states, refer_emb], dim=2
|
|
)
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"NonParamT2ISelfReferenceXFormersAttnProcessor5, referencenet, encoder_hidden_states={encoder_hidden_states.shape}, refer_emb={refer_emb.shape}"
|
|
)
|
|
encoder_hidden_states = rearrange(
|
|
encoder_hidden_states, "b t hw c -> (b t) hw c"
|
|
)
|
|
# vision_cond in same unet attn end
|
|
|
|
if attn.group_norm is not None:
|
|
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
|
|
1, 2
|
|
)
|
|
|
|
query = attn.to_q(hidden_states)
|
|
|
|
if encoder_hidden_states is None:
|
|
encoder_hidden_states = hidden_states
|
|
elif attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(
|
|
encoder_hidden_states
|
|
)
|
|
encoder_hidden_states = align_repeat_tensor_single_dim(
|
|
encoder_hidden_states, target_length=hidden_states.shape[0], dim=0
|
|
)
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
query = attn.head_to_batch_dim(query).contiguous()
|
|
key = attn.head_to_batch_dim(key).contiguous()
|
|
value = attn.head_to_batch_dim(value).contiguous()
|
|
|
|
hidden_states = xformers.ops.memory_efficient_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
attn_bias=attention_mask,
|
|
op=self.attention_op,
|
|
scale=attn.scale,
|
|
)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
hidden_states = attn.batch_to_head_dim(hidden_states)
|
|
|
|
# linear proj
|
|
hidden_states = attn.to_out[0](hidden_states)
|
|
# dropout
|
|
hidden_states = attn.to_out[1](hidden_states)
|
|
|
|
if input_ndim == 4:
|
|
hidden_states = hidden_states.transpose(-1, -2).reshape(
|
|
batch_size, channel, height, width
|
|
)
|
|
|
|
if attn.residual_connection:
|
|
hidden_states = hidden_states + residual
|
|
|
|
hidden_states = hidden_states / attn.rescale_output_factor
|
|
self.print_idx += 1
|
|
|
|
return hidden_states
|
|
|
|
|
|
@Model_Register.register
|
|
class NonParamReferenceIPXFormersAttnProcessor(
|
|
NonParamT2ISelfReferenceXFormersAttnProcessor
|
|
):
|
|
def __init__(self, attention_op: Callable[..., Any] | None = None):
|
|
super().__init__(attention_op)
|
|
|
|
|
|
@maybe_allow_in_graph
|
|
class ReferEmbFuseAttention(IPAttention):
|
|
"""使用 attention 融合 refernet 中的 emb 到 unet 对应的 latens 中
|
|
# TODO: 目前只支持 bt hw c 的融合，后续考虑增加对 视频 bhw t c、b thw c的融合
|
|
residual_connection: bool = True, 默认， 从不产生影响开始学习
|
|
|
|
use attention to fuse referencenet emb into unet latents
|
|
# TODO: by now, only support bt hw c, later consider to support bhw t c, b thw c
|
|
residual_connection: bool = True, default, start from no effect
|
|
|
|
Args:
|
|
IPAttention (_type_): _description_
|
|
"""
|
|
|
|
print_idx = 0
|
|
|
|
def __init__(
|
|
self,
|
|
query_dim: int,
|
|
cross_attention_dim: int | None = None,
|
|
heads: int = 8,
|
|
dim_head: int = 64,
|
|
dropout: float = 0,
|
|
bias=False,
|
|
upcast_attention: bool = False,
|
|
upcast_softmax: bool = False,
|
|
cross_attention_norm: str | None = None,
|
|
cross_attention_norm_num_groups: int = 32,
|
|
added_kv_proj_dim: int | None = None,
|
|
norm_num_groups: int | None = None,
|
|
spatial_norm_dim: int | None = None,
|
|
out_bias: bool = True,
|
|
scale_qk: bool = True,
|
|
only_cross_attention: bool = False,
|
|
eps: float = 0.00001,
|
|
rescale_output_factor: float = 1,
|
|
residual_connection: bool = True,
|
|
_from_deprecated_attn_block=False,
|
|
processor: AttnProcessor | None = None,
|
|
cross_attn_temporal_cond: bool = False,
|
|
image_scale: float = 1,
|
|
):
|
|
super().__init__(
|
|
query_dim,
|
|
cross_attention_dim,
|
|
heads,
|
|
dim_head,
|
|
dropout,
|
|
bias,
|
|
upcast_attention,
|
|
upcast_softmax,
|
|
cross_attention_norm,
|
|
cross_attention_norm_num_groups,
|
|
added_kv_proj_dim,
|
|
norm_num_groups,
|
|
spatial_norm_dim,
|
|
out_bias,
|
|
scale_qk,
|
|
only_cross_attention,
|
|
eps,
|
|
rescale_output_factor,
|
|
residual_connection,
|
|
_from_deprecated_attn_block,
|
|
processor,
|
|
cross_attn_temporal_cond,
|
|
image_scale,
|
|
)
|
|
self.processor = None
|
|
# 配合residual,使一开始不影响之前结果
|
|
nn.init.zeros_(self.to_out[0].weight)
|
|
nn.init.zeros_(self.to_out[0].bias)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.FloatTensor,
|
|
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
|
attention_mask: Optional[torch.FloatTensor] = None,
|
|
temb: Optional[torch.FloatTensor] = None,
|
|
scale: float = 1.0,
|
|
num_frames: int = None,
|
|
) -> torch.Tensor:
|
|
"""fuse referencenet emb b c t2 h2 w2 into unet latents b c t1 h1 w1 with attn
|
|
refer to musev/models/attention_processor.py::NonParamT2ISelfReferenceXFormersAttnProcessor
|
|
|
|
Args:
|
|
hidden_states (torch.FloatTensor): unet latents, (b t1) c h1 w1
|
|
encoder_hidden_states (Optional[torch.FloatTensor], optional): referencenet emb b c2 t2 h2 w2. Defaults to None.
|
|
attention_mask (Optional[torch.FloatTensor], optional): _description_. Defaults to None.
|
|
temb (Optional[torch.FloatTensor], optional): _description_. Defaults to None.
|
|
scale (float, optional): _description_. Defaults to 1.0.
|
|
num_frames (int, optional): _description_. Defaults to None.
|
|
|
|
Returns:
|
|
torch.Tensor: _description_
|
|
"""
|
|
residual = hidden_states
|
|
# start
|
|
hidden_states = rearrange(
|
|
hidden_states, "(b t) c h w -> b c t h w", t=num_frames
|
|
)
|
|
batch_size, channel, t1, height, width = hidden_states.shape
|
|
if self.print_idx == 0:
|
|
logger.debug(
|
|
f"hidden_states={hidden_states.shape},encoder_hidden_states={encoder_hidden_states.shape}"
|
|
)
|
|
# concat with hidden_states b c t1 h1 w1 in hw channel into bt (t2 + 1)hw c
|
|
encoder_hidden_states = rearrange(
|
|
encoder_hidden_states, " b c t2 h w-> b (t2 h w) c"
|
|
)
|
|
encoder_hidden_states = repeat(
|
|
encoder_hidden_states, " b t2hw c -> (b t) t2hw c", t=t1
|
|
)
|
|
hidden_states = rearrange(hidden_states, " b c t h w-> (b t) (h w) c")
|
|
# bt (t2+1)hw d
|
|
encoder_hidden_states = torch.concat(
|
|
[encoder_hidden_states, hidden_states], dim=1
|
|
)
|
|
# encoder_hidden_states = align_repeat_tensor_single_dim(
|
|
# encoder_hidden_states, target_length=hidden_states.shape[0], dim=0
|
|
# )
|
|
# end
|
|
|
|
if self.spatial_norm is not None:
|
|
hidden_states = self.spatial_norm(hidden_states, temb)
|
|
|
|
_, key_tokens, _ = (
|
|
hidden_states.shape
|
|
if encoder_hidden_states is None
|
|
else encoder_hidden_states.shape
|
|
)
|
|
|
|
attention_mask = self.prepare_attention_mask(
|
|
attention_mask, key_tokens, batch_size
|
|
)
|
|
if attention_mask is not None:
|
|
# expand our mask's singleton query_tokens dimension:
|
|
# [batch*heads, 1, key_tokens] ->
|
|
# [batch*heads, query_tokens, key_tokens]
|
|
# so that it can be added as a bias onto the attention scores that xformers computes:
|
|
# [batch*heads, query_tokens, key_tokens]
|
|
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
|
_, query_tokens, _ = hidden_states.shape
|
|
attention_mask = attention_mask.expand(-1, query_tokens, -1)
|
|
|
|
if self.group_norm is not None:
|
|
hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(
|
|
1, 2
|
|
)
|
|
|
|
query = self.to_q(hidden_states)
|
|
|
|
if encoder_hidden_states is None:
|
|
encoder_hidden_states = hidden_states
|
|
elif self.norm_cross:
|
|
encoder_hidden_states = self.norm_encoder_hidden_states(
|
|
encoder_hidden_states
|
|
)
|
|
|
|
key = self.to_k(encoder_hidden_states)
|
|
value = self.to_v(encoder_hidden_states)
|
|
|
|
query = self.head_to_batch_dim(query).contiguous()
|
|
key = self.head_to_batch_dim(key).contiguous()
|
|
value = self.head_to_batch_dim(value).contiguous()
|
|
|
|
# query: b t hw d
|
|
# key/value: bt (t1+1)hw d
|
|
hidden_states = xformers.ops.memory_efficient_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
attn_bias=attention_mask,
|
|
scale=self.scale,
|
|
)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
hidden_states = self.batch_to_head_dim(hidden_states)
|
|
|
|
# linear proj
|
|
hidden_states = self.to_out[0](hidden_states)
|
|
# dropout
|
|
hidden_states = self.to_out[1](hidden_states)
|
|
|
|
hidden_states = rearrange(
|
|
hidden_states,
|
|
"bt (h w) c-> bt c h w",
|
|
h=height,
|
|
w=width,
|
|
)
|
|
if self.residual_connection:
|
|
hidden_states = hidden_states + residual
|
|
|
|
hidden_states = hidden_states / self.rescale_output_factor
|
|
self.print_idx += 1
|
|
return hidden_states
|