1481 lines
60 KiB
Python
1481 lines
60 KiB
Python
# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
|
|
import math
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torch.fft as fft
|
|
import pdb
|
|
|
|
|
|
class AttnProcessor(nn.Module):
|
|
r"""
|
|
Default processor for performing attention-related computations.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size=None,
|
|
cross_attention_dim=None,
|
|
save_in_unet='down',
|
|
atten_control=None,
|
|
):
|
|
super().__init__()
|
|
self.atten_control = atten_control
|
|
self.save_in_unet = save_in_unet
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
|
|
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)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
query = attn.head_to_batch_dim(query)
|
|
key = attn.head_to_batch_dim(key)
|
|
value = attn.head_to_batch_dim(value)
|
|
|
|
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
|
hidden_states = torch.bmm(attention_probs, value)
|
|
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
|
|
|
|
return hidden_states
|
|
|
|
|
|
class IPAttnProcessor(nn.Module):
|
|
r"""
|
|
Attention processor for IP-Adapater.
|
|
Args:
|
|
hidden_size (`int`):
|
|
The hidden size of the attention layer.
|
|
cross_attention_dim (`int`):
|
|
The number of channels in the `encoder_hidden_states`.
|
|
scale (`float`, defaults to 1.0):
|
|
the weight scale of image prompt.
|
|
num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
|
|
The context length of the image features.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4, skip=False,save_in_unet='down', atten_control=None):
|
|
super().__init__()
|
|
|
|
self.hidden_size = hidden_size
|
|
self.cross_attention_dim = cross_attention_dim
|
|
self.scale = scale
|
|
self.num_tokens = num_tokens
|
|
self.skip = skip
|
|
|
|
self.atten_control = atten_control
|
|
self.save_in_unet = save_in_unet
|
|
|
|
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
|
|
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
|
|
else:
|
|
# get encoder_hidden_states, ip_hidden_states
|
|
end_pos = encoder_hidden_states.shape[1] - self.num_tokens
|
|
encoder_hidden_states, ip_hidden_states = (
|
|
encoder_hidden_states[:, :end_pos, :],
|
|
encoder_hidden_states[:, end_pos:, :],
|
|
)
|
|
if attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
query = attn.head_to_batch_dim(query)
|
|
key = attn.head_to_batch_dim(key)
|
|
value = attn.head_to_batch_dim(value)
|
|
|
|
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
|
hidden_states = torch.bmm(attention_probs, value)
|
|
hidden_states = attn.batch_to_head_dim(hidden_states)
|
|
|
|
if not self.skip:
|
|
# for ip-adapter
|
|
ip_key = self.to_k_ip(ip_hidden_states)
|
|
ip_value = self.to_v_ip(ip_hidden_states)
|
|
|
|
ip_key = attn.head_to_batch_dim(ip_key)
|
|
ip_value = attn.head_to_batch_dim(ip_value)
|
|
|
|
ip_attention_probs = attn.get_attention_scores(query, ip_key, None)
|
|
self.attn_map = ip_attention_probs
|
|
ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
|
|
ip_hidden_states = attn.batch_to_head_dim(ip_hidden_states)
|
|
|
|
hidden_states = hidden_states + self.scale * ip_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
|
|
|
|
return hidden_states
|
|
|
|
|
|
class AttnProcessor2_0(torch.nn.Module):
|
|
r"""
|
|
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size=None,
|
|
cross_attention_dim=None,
|
|
save_in_unet='down',
|
|
atten_control=None,
|
|
):
|
|
super().__init__()
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
self.atten_control = atten_control
|
|
self.save_in_unet = save_in_unet
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
# scaled_dot_product_attention expects attention_mask shape to be
|
|
# (batch, heads, source_length, target_length)
|
|
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-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)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
# 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
|
|
|
|
return hidden_states
|
|
|
|
|
|
class IPAttnProcessor2_0(torch.nn.Module):
|
|
r"""
|
|
Attention processor for IP-Adapater for PyTorch 2.0.
|
|
Args:
|
|
hidden_size (`int`):
|
|
The hidden size of the attention layer.
|
|
cross_attention_dim (`int`):
|
|
The number of channels in the `encoder_hidden_states`.
|
|
scale (`float`, defaults to 1.0):
|
|
the weight scale of image prompt.
|
|
num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
|
|
The context length of the image features.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4, skip=False,save_in_unet='down', atten_control=None):
|
|
super().__init__()
|
|
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
|
|
self.hidden_size = hidden_size
|
|
self.cross_attention_dim = cross_attention_dim
|
|
self.scale = scale
|
|
self.num_tokens = num_tokens
|
|
self.skip = skip
|
|
|
|
self.atten_control = atten_control
|
|
self.save_in_unet = save_in_unet
|
|
|
|
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
# scaled_dot_product_attention expects attention_mask shape to be
|
|
# (batch, heads, source_length, target_length)
|
|
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-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
|
|
else:
|
|
# get encoder_hidden_states, ip_hidden_states
|
|
end_pos = encoder_hidden_states.shape[1] - self.num_tokens
|
|
encoder_hidden_states, ip_hidden_states = (
|
|
encoder_hidden_states[:, :end_pos, :],
|
|
encoder_hidden_states[:, end_pos:, :],
|
|
)
|
|
if attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
if not self.skip:
|
|
# for ip-adapter
|
|
ip_key = self.to_k_ip(ip_hidden_states)
|
|
ip_value = self.to_v_ip(ip_hidden_states)
|
|
|
|
ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
ip_hidden_states = F.scaled_dot_product_attention(
|
|
query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
|
|
)
|
|
with torch.no_grad():
|
|
self.attn_map = query @ ip_key.transpose(-2, -1).softmax(dim=-1)
|
|
#print(self.attn_map.shape)
|
|
|
|
ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
ip_hidden_states = ip_hidden_states.to(query.dtype)
|
|
|
|
hidden_states = hidden_states + self.scale * ip_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
|
|
|
|
return hidden_states
|
|
|
|
|
|
class IP_CS_AttnProcessor2_0(torch.nn.Module):
|
|
r"""
|
|
Attention processor for IP-Adapater for PyTorch 2.0.
|
|
Args:
|
|
hidden_size (`int`):
|
|
The hidden size of the attention layer.
|
|
cross_attention_dim (`int`):
|
|
The number of channels in the `encoder_hidden_states`.
|
|
scale (`float`, defaults to 1.0):
|
|
the weight scale of image prompt.
|
|
num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
|
|
The context length of the image features.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, cross_attention_dim=None, content_scale=1.0,style_scale=1.0, num_content_tokens=4,num_style_tokens=4,
|
|
skip=False,content=False, style=False):
|
|
super().__init__()
|
|
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
|
|
self.hidden_size = hidden_size
|
|
self.cross_attention_dim = cross_attention_dim
|
|
self.content_scale = content_scale
|
|
self.style_scale = style_scale
|
|
self.num_content_tokens = num_content_tokens
|
|
self.num_style_tokens = num_style_tokens
|
|
self.skip = skip
|
|
|
|
self.content = content
|
|
self.style = style
|
|
|
|
if self.content or self.style:
|
|
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_k_ip_content =None
|
|
self.to_v_ip_content =None
|
|
|
|
def set_content_ipa(self,content_scale=1.0):
|
|
|
|
self.to_k_ip_content = nn.Linear(self.cross_attention_dim or self.hidden_size, self.hidden_size, bias=False)
|
|
self.to_v_ip_content = nn.Linear(self.cross_attention_dim or self.hidden_size, self.hidden_size, bias=False)
|
|
self.content_scale=content_scale
|
|
self.content =True
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
# scaled_dot_product_attention expects attention_mask shape to be
|
|
# (batch, heads, source_length, target_length)
|
|
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-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
|
|
else:
|
|
# get encoder_hidden_states, ip_hidden_states
|
|
end_pos = encoder_hidden_states.shape[1] - self.num_content_tokens-self.num_style_tokens
|
|
encoder_hidden_states, ip_content_hidden_states,ip_style_hidden_states = (
|
|
encoder_hidden_states[:, :end_pos, :],
|
|
encoder_hidden_states[:, end_pos:end_pos + self.num_content_tokens, :],
|
|
encoder_hidden_states[:, end_pos + self.num_content_tokens:, :],
|
|
)
|
|
if attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
if self.content is True:
|
|
exit()
|
|
if not self.skip and self.content is True:
|
|
# print('content#####################################################')
|
|
# for ip-content-adapter
|
|
if self.to_k_ip_content is None:
|
|
|
|
ip_content_key = self.to_k_ip(ip_content_hidden_states)
|
|
ip_content_value = self.to_v_ip(ip_content_hidden_states)
|
|
else:
|
|
ip_content_key = self.to_k_ip_content(ip_content_hidden_states)
|
|
ip_content_value = self.to_v_ip_content(ip_content_hidden_states)
|
|
|
|
ip_content_key = ip_content_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
ip_content_value = ip_content_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
ip_content_hidden_states = F.scaled_dot_product_attention(
|
|
query, ip_content_key, ip_content_value, attn_mask=None, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
|
|
ip_content_hidden_states = ip_content_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
ip_content_hidden_states = ip_content_hidden_states.to(query.dtype)
|
|
|
|
hidden_states = hidden_states + self.content_scale * ip_content_hidden_states
|
|
|
|
if not self.skip and self.style is True:
|
|
# for ip-style-adapter
|
|
ip_style_key = self.to_k_ip(ip_style_hidden_states)
|
|
ip_style_value = self.to_v_ip(ip_style_hidden_states)
|
|
|
|
ip_style_key = ip_style_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
ip_style_value = ip_style_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
ip_style_hidden_states = F.scaled_dot_product_attention(
|
|
query, ip_style_key, ip_style_value, attn_mask=None, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
ip_style_hidden_states = ip_style_hidden_states.transpose(1, 2).reshape(batch_size, -1,
|
|
attn.heads * head_dim)
|
|
ip_style_hidden_states = ip_style_hidden_states.to(query.dtype)
|
|
|
|
hidden_states = hidden_states + self.style_scale * ip_style_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
|
|
|
|
return hidden_states
|
|
|
|
## for controlnet
|
|
class CNAttnProcessor:
|
|
r"""
|
|
Default processor for performing attention-related computations.
|
|
"""
|
|
|
|
def __init__(self, num_tokens=4,save_in_unet='down',atten_control=None):
|
|
self.num_tokens = num_tokens
|
|
self.atten_control = atten_control
|
|
self.save_in_unet = save_in_unet
|
|
|
|
def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None):
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
|
|
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
|
|
else:
|
|
end_pos = encoder_hidden_states.shape[1] - self.num_tokens
|
|
encoder_hidden_states = encoder_hidden_states[:, :end_pos] # only use text
|
|
if attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
query = attn.head_to_batch_dim(query)
|
|
key = attn.head_to_batch_dim(key)
|
|
value = attn.head_to_batch_dim(value)
|
|
|
|
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
|
hidden_states = torch.bmm(attention_probs, value)
|
|
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
|
|
|
|
return hidden_states
|
|
|
|
class CNAttnProcessor2_0:
|
|
r"""
|
|
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
|
"""
|
|
|
|
def __init__(self, num_tokens=4, save_in_unet='down', atten_control=None):
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
self.num_tokens = num_tokens
|
|
self.atten_control = atten_control
|
|
self.save_in_unet = save_in_unet
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
# scaled_dot_product_attention expects attention_mask shape to be
|
|
# (batch, heads, source_length, target_length)
|
|
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-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
|
|
else:
|
|
end_pos = encoder_hidden_states.shape[1] - self.num_tokens
|
|
encoder_hidden_states = encoder_hidden_states[:, :end_pos] # only use text
|
|
if attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
# 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
|
|
|
|
return hidden_states
|
|
|
|
## for exp
|
|
class IP_FuAd_AttnProcessor2_0(torch.nn.Module):
|
|
r"""
|
|
Attention processor for IP-Adapater for PyTorch 2.0.
|
|
Args:
|
|
hidden_size (`int`):
|
|
The hidden size of the attention layer.
|
|
cross_attention_dim (`int`):
|
|
The number of channels in the `encoder_hidden_states`.
|
|
scale (`float`, defaults to 1.0):
|
|
the weight scale of image prompt.
|
|
num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
|
|
The context length of the image features.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, cross_attention_dim=None, content_scale=1.0,style_scale=1.0, num_content_tokens=4,num_style_tokens=4,
|
|
skip=False,content=False, style=False, fuAttn=False, fuIPAttn=False, adainIP=False,
|
|
fuScale=0, end_fusion=0, attn_name=None):
|
|
super().__init__()
|
|
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
|
|
self.hidden_size = hidden_size
|
|
self.cross_attention_dim = cross_attention_dim
|
|
self.content_scale = content_scale
|
|
self.style_scale = style_scale
|
|
self.num_style_tokens = num_style_tokens
|
|
self.skip = skip
|
|
|
|
self.content = content
|
|
self.style = style
|
|
|
|
self.fuAttn = fuAttn
|
|
self.fuIPAttn = fuIPAttn
|
|
self.adainIP = adainIP
|
|
self.fuScale = fuScale
|
|
self.denoise_step = 0
|
|
self.end_fusion = end_fusion
|
|
self.name = attn_name
|
|
|
|
if self.content or self.style:
|
|
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_k_ip_content =None
|
|
self.to_v_ip_content =None
|
|
|
|
# def set_content_ipa(self,content_scale=1.0):
|
|
|
|
# self.to_k_ip_content = nn.Linear(self.cross_attention_dim or self.hidden_size, self.hidden_size, bias=False)
|
|
# self.to_v_ip_content = nn.Linear(self.cross_attention_dim or self.hidden_size, self.hidden_size, bias=False)
|
|
# self.content_scale=content_scale
|
|
# self.content =True
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
self.denoise_step += 1
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
# scaled_dot_product_attention expects attention_mask shape to be
|
|
# (batch, heads, source_length, target_length)
|
|
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-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
|
|
else:
|
|
# get encoder_hidden_states, ip_hidden_states
|
|
end_pos = encoder_hidden_states.shape[1] -self.num_style_tokens
|
|
encoder_hidden_states, ip_style_hidden_states = (
|
|
encoder_hidden_states[:, :end_pos, :],
|
|
encoder_hidden_states[:, end_pos:, :],
|
|
)
|
|
if attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
# # modified the attnMap of the Stylization Image
|
|
|
|
if self.fuAttn and self.denoise_step <= self.end_fusion:
|
|
assert query.shape[0] == 4
|
|
scale_factor = 1 / math.sqrt(torch.tensor(head_dim, dtype=query.dtype))
|
|
text_attn_probs = (torch.matmul(query, key.transpose(-2, -1)) * scale_factor).softmax(dim=-1)
|
|
text_attn_probs[1] = self.fuScale*text_attn_probs[1] + (1-self.fuScale)*text_attn_probs[0]
|
|
text_attn_probs[3] = self.fuScale*text_attn_probs[3] + (1-self.fuScale)*text_attn_probs[2]
|
|
hidden_states = torch.matmul(text_attn_probs, value)
|
|
else:
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
raw_hidden_states = hidden_states
|
|
|
|
if not self.skip and self.style is True:
|
|
|
|
# for ip-style-adapter
|
|
ip_style_key = self.to_k_ip(ip_style_hidden_states)
|
|
ip_style_value = self.to_v_ip(ip_style_hidden_states)
|
|
|
|
ip_style_key = ip_style_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
ip_style_value = ip_style_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
if self.fuIPAttn and self.denoise_step <= self.end_fusion:
|
|
assert query.shape[0] == 4
|
|
# if "down" in self.name:
|
|
# print("wrong! coding")
|
|
# exit()
|
|
scale_factor = 1 / math.sqrt(torch.tensor(head_dim, dtype=query.dtype))
|
|
ip_attn_probs = torch.matmul(query, ip_style_key.transpose(-2, -1)) * scale_factor
|
|
ip_attn_probs = F.softmax(ip_attn_probs, dim=-1)
|
|
ip_attn_probs[1] = self.fuScale*ip_attn_probs[1] + (1-self.fuScale)*ip_attn_probs[0]
|
|
ip_attn_probs[3] = self.fuScale*ip_attn_probs[3] + (1-self.fuScale)*ip_attn_probs[2]
|
|
ip_style_hidden_states = torch.matmul(ip_attn_probs, ip_style_value)
|
|
else:
|
|
ip_style_hidden_states = F.scaled_dot_product_attention(
|
|
query, ip_style_key, ip_style_value, attn_mask=None, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
ip_style_hidden_states = ip_style_hidden_states.transpose(1, 2).reshape(batch_size, -1,
|
|
attn.heads * head_dim)
|
|
ip_style_hidden_states = ip_style_hidden_states.to(query.dtype)
|
|
|
|
if not self.adainIP:
|
|
hidden_states = hidden_states + self.style_scale * ip_style_hidden_states
|
|
else:
|
|
# print("adain")
|
|
def adain(content, style):
|
|
content_mean = content.mean(dim=1, keepdim=True)
|
|
content_std = content.std(dim=1, keepdim=True)
|
|
style_mean = style.mean(dim=1, keepdim=True)
|
|
style_std = style.std(dim=1, keepdim=True)
|
|
normalized_content = (content - content_mean) / content_std
|
|
stylized_content = normalized_content * style_std + style_mean
|
|
return stylized_content
|
|
hidden_states = adain(content=hidden_states, style=ip_style_hidden_states)
|
|
|
|
if hidden_states.shape[0] == 4:
|
|
hidden_states[0] = raw_hidden_states[0]
|
|
hidden_states[2] = raw_hidden_states[2]
|
|
# hidden_states = raw_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.reset_denoise_step()
|
|
return hidden_states
|
|
class IP_FuAd_AttnProcessor2_0_exp(torch.nn.Module):
|
|
r"""
|
|
Attention processor for IP-Adapater for PyTorch 2.0.
|
|
Args:
|
|
hidden_size (`int`):
|
|
The hidden size of the attention layer.
|
|
cross_attention_dim (`int`):
|
|
The number of channels in the `encoder_hidden_states`.
|
|
scale (`float`, defaults to 1.0):
|
|
the weight scale of image prompt.
|
|
num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
|
|
The context length of the image features.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, cross_attention_dim=None, content_scale=1.0,style_scale=1.0, num_content_tokens=4,num_style_tokens=4,
|
|
skip=False,content=False, style=False, fuAttn=False, fuIPAttn=False, adainIP=False,
|
|
fuScale=0, end_fusion=0, attn_name=None, save_attn_map=False):
|
|
super().__init__()
|
|
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
|
|
self.hidden_size = hidden_size
|
|
self.cross_attention_dim = cross_attention_dim
|
|
self.content_scale = content_scale
|
|
self.style_scale = style_scale
|
|
self.num_style_tokens = num_style_tokens
|
|
self.skip = skip
|
|
|
|
self.content = content
|
|
self.style = style
|
|
|
|
self.fuAttn = fuAttn
|
|
self.fuIPAttn = fuIPAttn
|
|
self.adainIP = adainIP
|
|
self.fuScale = fuScale
|
|
self.denoise_step = 0
|
|
self.end_fusion = end_fusion
|
|
self.name = attn_name
|
|
|
|
self.save_attn_map = save_attn_map
|
|
|
|
if self.content or self.style:
|
|
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_k_ip_content =None
|
|
self.to_v_ip_content =None
|
|
|
|
# def set_content_ipa(self,content_scale=1.0):
|
|
|
|
# self.to_k_ip_content = nn.Linear(self.cross_attention_dim or self.hidden_size, self.hidden_size, bias=False)
|
|
# self.to_v_ip_content = nn.Linear(self.cross_attention_dim or self.hidden_size, self.hidden_size, bias=False)
|
|
# self.content_scale=content_scale
|
|
# self.content =True
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
self.denoise_step += 1
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
# scaled_dot_product_attention expects attention_mask shape to be
|
|
# (batch, heads, source_length, target_length)
|
|
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-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
|
|
else:
|
|
# get encoder_hidden_states, ip_hidden_states
|
|
end_pos = encoder_hidden_states.shape[1] - self.num_content_tokens-self.num_style_tokens
|
|
encoder_hidden_states, ip_style_hidden_states = (
|
|
encoder_hidden_states[:, :end_pos, :],
|
|
encoder_hidden_states[:, end_pos:, :],
|
|
)
|
|
if attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
## attention map
|
|
if self.save_attn_map:
|
|
attention_probs = attn.get_attention_scores(attn.head_to_batch_dim(query), attn.head_to_batch_dim(value), attention_mask)
|
|
if attention_probs is not None:
|
|
if not hasattr(attn, "attn_map"):
|
|
setattr(attn, "attn_map", {})
|
|
setattr(attn, "inference_step", 0)
|
|
else:
|
|
attn.inference_step += 1
|
|
|
|
# # maybe we need to save all the timestep
|
|
# if attn.inference_step in self.attn_map_save_steps:
|
|
attn.attn_map[attn.inference_step] = attention_probs.clone().cpu().detach()
|
|
# attn.attn_map[attn.inference_step] = attention_probs.detach()
|
|
## end of attention map
|
|
else:
|
|
print(f"{attn} didn't get the attention probs")
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
# # modified the attnMap of the Stylization Image
|
|
|
|
if self.fuAttn and self.denoise_step <= self.end_fusion:
|
|
assert query.shape[0] == 4
|
|
scale_factor = 1 / math.sqrt(torch.tensor(head_dim, dtype=query.dtype))
|
|
text_attn_probs = (torch.matmul(query, key.transpose(-2, -1)) * scale_factor).softmax(dim=-1)
|
|
text_attn_probs[1] = self.fuScale*text_attn_probs[1] + (1-self.fuScale)*text_attn_probs[0]
|
|
text_attn_probs[3] = self.fuScale*text_attn_probs[3] + (1-self.fuScale)*text_attn_probs[2]
|
|
hidden_states = torch.matmul(text_attn_probs, value)
|
|
else:
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
raw_hidden_states = hidden_states
|
|
|
|
if not self.skip and self.style is True:
|
|
|
|
# for ip-style-adapter
|
|
ip_style_key = self.to_k_ip(ip_style_hidden_states)
|
|
ip_style_value = self.to_v_ip(ip_style_hidden_states)
|
|
|
|
ip_style_key = ip_style_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
ip_style_value = ip_style_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
if self.fuIPAttn and self.denoise_step <= self.end_fusion:
|
|
assert query.shape[0] == 4
|
|
# if "down" in self.name:
|
|
# print("wrong! coding")
|
|
# exit()
|
|
scale_factor = 1 / math.sqrt(torch.tensor(head_dim, dtype=query.dtype))
|
|
ip_attn_probs = torch.matmul(query, ip_style_key.transpose(-2, -1)) * scale_factor
|
|
ip_attn_probs = F.softmax(ip_attn_probs, dim=-1)
|
|
ip_attn_probs[1] = self.fuScale*ip_attn_probs[1] + (1-self.fuScale)*ip_attn_probs[0]
|
|
ip_attn_probs[3] = self.fuScale*ip_attn_probs[3] + (1-self.fuScale)*ip_attn_probs[2]
|
|
ip_style_hidden_states = torch.matmul(ip_attn_probs, ip_style_value)
|
|
else:
|
|
ip_style_hidden_states = F.scaled_dot_product_attention(
|
|
query, ip_style_key, ip_style_value, attn_mask=None, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
ip_style_hidden_states = ip_style_hidden_states.transpose(1, 2).reshape(batch_size, -1,
|
|
attn.heads * head_dim)
|
|
ip_style_hidden_states = ip_style_hidden_states.to(query.dtype)
|
|
|
|
# if self.adainIP and self.denoise_step >= self.start_adain:
|
|
if self.adainIP:
|
|
# print("adain")
|
|
def adain(content, style):
|
|
content_mean = content.mean(dim=1, keepdim=True)
|
|
content_std = content.std(dim=1, keepdim=True)
|
|
print("exp code")
|
|
pdb.set_trace()
|
|
style_mean = style.mean(dim=1, keepdim=True)
|
|
style_std = style.std(dim=1, keepdim=True)
|
|
normalized_content = (content - content_mean) / content_std
|
|
stylized_content = normalized_content * style_std + style_mean
|
|
return stylized_content
|
|
pdb.set_trace()
|
|
hidden_states = adain(content=hidden_states, style=ip_style_hidden_states)
|
|
else:
|
|
hidden_states = hidden_states + self.style_scale * ip_style_hidden_states
|
|
|
|
if hidden_states.shape[0] == 4:
|
|
hidden_states[0] = raw_hidden_states[0]
|
|
hidden_states[2] = raw_hidden_states[2]
|
|
# hidden_states = raw_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.reset_denoise_step()
|
|
return hidden_states
|
|
|
|
class AttnProcessor2_0_hijack(torch.nn.Module):
|
|
r"""
|
|
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size=None,
|
|
cross_attention_dim=None,
|
|
save_in_unet='down',
|
|
atten_control=None,
|
|
fuSAttn=False,
|
|
fuScale=0,
|
|
end_fusion=0,
|
|
attn_name=None,
|
|
num_inference_step=50,
|
|
):
|
|
super().__init__()
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
self.atten_control = atten_control
|
|
self.save_in_unet = save_in_unet
|
|
|
|
self.fuSAttn = fuSAttn
|
|
self.fuScale = fuScale
|
|
self.denoise_step = 0
|
|
self.end_fusion = end_fusion
|
|
self.name = attn_name
|
|
self.num_inference_step = num_inference_step
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
self.denoise_step += 1
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
# scaled_dot_product_attention expects attention_mask shape to be
|
|
# (batch, heads, source_length, target_length)
|
|
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-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)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
if self.fuSAttn and self.denoise_step <= self.end_fusion:
|
|
assert query.shape[0] == 4
|
|
scale_factor = 1 / math.sqrt(torch.tensor(head_dim, dtype=query.dtype))
|
|
attn_probs = (torch.matmul(query, key.transpose(-2, -1)) * scale_factor).softmax(dim=-1)
|
|
attn_probs[1] = attn_probs[0]
|
|
attn_probs[3] = attn_probs[2]
|
|
hidden_states = torch.matmul(attn_probs, value)
|
|
else:
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
# 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
|
|
|
|
if self.denoise_step == self.num_inference_step:
|
|
self.denoise_step = 0
|
|
return hidden_states
|
|
|
|
class IPAttnProcessor2_0_cross_modal(torch.nn.Module):
|
|
r"""
|
|
Attention processor for IP-Adapater for PyTorch 2.0.
|
|
Args:
|
|
hidden_size (`int`):
|
|
The hidden size of the attention layer.
|
|
cross_attention_dim (`int`):
|
|
The number of channels in the `encoder_hidden_states`.
|
|
scale (`float`, defaults to 1.0):
|
|
the weight scale of image prompt.
|
|
num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
|
|
The context length of the image features.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4, skip=False,
|
|
fuAttn=False, fuIPAttn=False, adainIP=False, end_fusion=0, attn_name=None,
|
|
num_inference_step=50):
|
|
super().__init__()
|
|
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
|
|
self.hidden_size = hidden_size
|
|
self.cross_attention_dim = cross_attention_dim
|
|
self.scale = scale
|
|
self.num_tokens = num_tokens
|
|
self.skip = skip
|
|
|
|
self.fuAttn = fuAttn
|
|
self.fuIPAttn = fuIPAttn
|
|
self.adainIP = adainIP
|
|
self.denoise_step = 0
|
|
self.end_fusion = end_fusion
|
|
self.name = attn_name
|
|
|
|
self.num_inference_step = num_inference_step
|
|
|
|
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
|
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
temb=None,
|
|
):
|
|
self.denoise_step += 1
|
|
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, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
# scaled_dot_product_attention expects attention_mask shape to be
|
|
# (batch, heads, source_length, target_length)
|
|
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-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
|
|
else:
|
|
# get encoder_hidden_states, ip_hidden_states
|
|
end_pos = encoder_hidden_states.shape[1] - self.num_tokens
|
|
encoder_hidden_states, ip_hidden_states = (
|
|
encoder_hidden_states[:, :end_pos, :],
|
|
encoder_hidden_states[:, end_pos:, :],
|
|
)
|
|
if attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
if self.fuAttn and self.denoise_step <= self.end_fusion:
|
|
assert query.shape[0] == 4
|
|
scale_factor = 1 / math.sqrt(torch.tensor(head_dim, dtype=query.dtype))
|
|
text_attn_probs = (torch.matmul(query, key.transpose(-2, -1)) * scale_factor).softmax(dim=-1)
|
|
text_attn_probs[1] = text_attn_probs[0]
|
|
text_attn_probs[3] = text_attn_probs[2]
|
|
hidden_states = torch.matmul(text_attn_probs, value)
|
|
else:
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
raw_hidden_states = hidden_states
|
|
|
|
if not self.skip:
|
|
# for ip-adapter
|
|
ip_key = self.to_k_ip(ip_hidden_states)
|
|
ip_value = self.to_v_ip(ip_hidden_states)
|
|
|
|
ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
|
# TODO: add support for attn.scale when we move to Torch 2.1
|
|
if self.fuIPAttn and self.denoise_step <= self.end_fusion:
|
|
assert query.shape[0] == 4
|
|
print("fuIPAttn")
|
|
scale_factor = 1 / math.sqrt(torch.tensor(head_dim, dtype=query.dtype))
|
|
ip_attn_probs = torch.matmul(query, ip_key.transpose(-2, -1)) * scale_factor
|
|
ip_attn_probs = F.softmax(ip_attn_probs, dim=-1)
|
|
ip_attn_probs[1] = ip_attn_probs[0]
|
|
ip_attn_probs[3] = ip_attn_probs[2]
|
|
ip_hidden_states = torch.matmul(ip_attn_probs, ip_value)
|
|
else:
|
|
ip_hidden_states = F.scaled_dot_product_attention(
|
|
query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
with torch.no_grad():
|
|
self.attn_map = query @ ip_key.transpose(-2, -1).softmax(dim=-1)
|
|
#print(self.attn_map.shape)
|
|
|
|
ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
ip_hidden_states = ip_hidden_states.to(query.dtype)
|
|
|
|
if self.adainIP:
|
|
def adain(content, style):
|
|
# 计算内容特征的均值和标准差
|
|
content_mean = content.mean(dim=1, keepdim=True)
|
|
content_std = content.std(dim=1, keepdim=True)
|
|
# 计算风格特征的均值和标准差
|
|
style_mean = style.mean(dim=1, keepdim=True)
|
|
style_std = style.std(dim=1, keepdim=True)
|
|
# 归一化内容特征并应用风格特征的均值和方差
|
|
normalized_content = (content - content_mean) / content_std
|
|
stylized_content = normalized_content * style_std + style_mean
|
|
return stylized_content
|
|
hidden_states = adain(content=hidden_states, style=ip_hidden_states)
|
|
else:
|
|
hidden_states = hidden_states + self.scale * ip_hidden_states
|
|
|
|
if hidden_states.shape[0] == 4:
|
|
hidden_states[0] = raw_hidden_states[0]
|
|
hidden_states[2] = raw_hidden_states[2]
|
|
|
|
# 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
|
|
|
|
if self.denoise_step == self.num_inference_step:
|
|
self.denoise_step = 0
|
|
|
|
return hidden_states
|