adapter/attention_processor.py

from typing import Callable, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn
from diffusers.utils import USE_PEFT_BACKEND
from diffusers.models.lora import LoRALinearLayer





class CacheAttnProcessor2_0:
 r"""
 Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
 """

 def __init__(self):
 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.cache = {} # cache hidden states

 def __call__(
 self,
 attn,
 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,
 ) -> torch.FloatTensor:

 self.cache["hidden_states"] = hidden_states # cache hidden states

 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)

 args = () if USE_PEFT_BACKEND else (scale,)
 query = attn.to_q(hidden_states, *args)

 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, *args)
 value = attn.to_v(encoder_hidden_states, *args)

 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, *args)
 # 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 SAttnProcessor2_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, name, hidden_size, cross_attention_dim=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.")

 super().__init__()

 self.name = name
 self.hidden_size = hidden_size
 self.cross_attention_dim = cross_attention_dim

 def __call__(
 self,
 attn,
 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,
 cond_hidden_states=None,
 sa_hidden_states=None,
 ) -> torch.FloatTensor:
 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)

 args = () if USE_PEFT_BACKEND else (scale,)
 query = attn.to_q(hidden_states, *args)

 if encoder_hidden_states is None:
 # for reference adapter
 if sa_hidden_states is not None:
 ref_hidden_states = sa_hidden_states[self.name]
 encoder_hidden_states = torch.cat([hidden_states, ref_hidden_states], dim=1)
 else:
 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, *args)
 value = attn.to_v(encoder_hidden_states, *args)

 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, *args)
 # 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 CAttnProcessor2_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, name, hidden_size, cross_attention_dim=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.")

 super().__init__()

 self.name = name
 self.hidden_size = hidden_size
 self.cross_attention_dim = cross_attention_dim

 # 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: torch.FloatTensor,
 encoder_hidden_states: Optional[torch.FloatTensor] = None,
 attention_mask: Optional[torch.FloatTensor] = None,
 temb: Optional[torch.FloatTensor] = None,
 scale: float = 1.0,
 cond_hidden_states=None,
 sa_hidden_states=None,
 ) -> torch.FloatTensor:
 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)

 args = () if USE_PEFT_BACKEND else (scale,)
 query = attn.to_q(hidden_states, *args)

 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, *args)
 value = attn.to_v(encoder_hidden_states, *args)

 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)

 # for ip
 # if cond_hidden_states:
 # ip_hidden_states = cond_hidden_states
 # 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=attention_mask, dropout_p=0.0, is_causal=False
 # )
 # 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 + ip_hidden_states

 # linear proj
 hidden_states = attn.to_out[0](hidden_states, *args)
 # 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 RefLoraSAttnProcessor2_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, name, hidden_size, cross_attention_dim=None, scale=1.0, rank=128, network_alpha=None, lora_scale=1.0,):
 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.")

 super().__init__()

 self.name = name
 self.hidden_size = hidden_size
 self.cross_attention_dim = cross_attention_dim
 self.to_k_ref = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
 self.to_v_ref = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
 self.scale = scale

 self.rank = rank
 self.lora_scale = lora_scale
 self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
 self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
 self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
 self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
 
 def __call__(
 self,
 attn,
 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_images_per_prompt=1,
 cond_hidden_states=None,
 sa_hidden_states=None,

 ) -> torch.FloatTensor:
 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)

 args = () if USE_PEFT_BACKEND else (scale,)
 query = attn.to_q(hidden_states, *args) + self.lora_scale * self.to_q_lora(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, *args) + self.lora_scale * self.to_k_lora(encoder_hidden_states)
 value = attn.to_v(encoder_hidden_states, *args) + self.lora_scale * self.to_v_lora(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)

 # for ref adapter
 if sa_hidden_states is not None:
 ref_hidden_states = sa_hidden_states[self.name]
 # for ref
 ref_key = self.to_k_ref(ref_hidden_states)
 ref_value = self.to_v_ref(ref_hidden_states)
 ref_key = ref_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
 ref_value = ref_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
 ref_hidden_states = F.scaled_dot_product_attention(
 query, ref_key, ref_value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
 )
 ref_hidden_states = ref_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
 ref_hidden_states = ref_hidden_states.to(query.dtype)
 hidden_states = hidden_states + ref_hidden_states * self.scale

 # linear proj
 hidden_states = attn.to_out[0](hidden_states, *args) + self.lora_scale * self.to_out_lora(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 RefSAttnProcessor2_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, name, hidden_size, cross_attention_dim=None, scale=1.0):
 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.")

 super().__init__()

 self.name = name
 self.hidden_size = hidden_size
 self.cross_attention_dim = cross_attention_dim
 self.to_k_ref = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
 self.to_v_ref = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
 self.scale = scale

 def __call__(
 self,
 attn,
 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_images_per_prompt=1,
 cond_hidden_states=None,
 sa_hidden_states=None,

 ) -> torch.FloatTensor:
 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)

 args = () if USE_PEFT_BACKEND else (scale,)
 query = attn.to_q(hidden_states, *args)

 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, *args)
 value = attn.to_v(encoder_hidden_states, *args)

 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)

 # for ref adapter
 if sa_hidden_states is not None:
 ref_hidden_states = sa_hidden_states[self.name]
 # for ref
 ref_key = self.to_k_ref(ref_hidden_states)
 ref_value = self.to_v_ref(ref_hidden_states)
 ref_key = ref_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
 ref_value = ref_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
 ref_hidden_states = F.scaled_dot_product_attention(
 query, ref_key, ref_value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
 )
 ref_hidden_states = ref_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
 ref_hidden_states = ref_hidden_states.to(query.dtype)
 hidden_states = hidden_states + ref_hidden_states * self.scale

 # linear proj
 hidden_states = attn.to_out[0](hidden_states, *args)
 # 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):
 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.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,
 sa_hidden_states=None,
 scale: float = 1.0,
 ):
 # attn原始的attn模块
 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:
 if sa_hidden_states is not None:
 ref_hidden_states = sa_hidden_states[self.name]
 # print(ref_hidden_states.shape, hidden_states.shape)
 encoder_hidden_states = torch.cat([hidden_states, ref_hidden_states], dim=1)
 else:
 encoder_hidden_states = hidden_states
 else:
 # get encoder_hidden_states, ip_hidden_states
 end_pos = encoder_hidden_states.shape[1] - self.num_tokens
 if end_pos != 89:
 encoder_hidden_states = encoder_hidden_states
 ip_hidden_states = None
 else:
 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)

 # make sure the ipa is in the inference stage
 if ip_hidden_states is not None:
 # 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 LoRAIPAttnProcessor2_0(nn.Module):
 r"""
 Processor for implementing the LoRA attention mechanism.

 Args:
 hidden_size (`int`, *optional*):
 The hidden size of the attention layer.
 cross_attention_dim (`int`, *optional*):
 The number of channels in the `encoder_hidden_states`.
 rank (`int`, defaults to 4):
 The dimension of the LoRA update matrices.
 network_alpha (`int`, *optional*):
 Equivalent to `alpha` but it's usage is specific to Kohya (A1111) style LoRAs.
 """

 def __init__(self, hidden_size, cross_attention_dim=None, rank=128, network_alpha=None, lora_scale=1.0, scale=1.0,
 num_tokens=4):
 super().__init__()

 self.rank = rank
 self.lora_scale = lora_scale
 self.num_tokens = num_tokens

 self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
 self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
 self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
 self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)

 self.hidden_size = hidden_size
 self.cross_attention_dim = cross_attention_dim
 self.scale = scale

 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, scale=1.0, temb=None, *args,
 **kwargs,
 ):
 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) + self.lora_scale * self.to_q_lora(hidden_states)
 # query = attn.head_to_batch_dim(query)

 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)

 # for text
 key = attn.to_k(encoder_hidden_states) + self.lora_scale * self.to_k_lora(encoder_hidden_states)
 value = attn.to_v(encoder_hidden_states) + self.lora_scale * self.to_v_lora(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)

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

 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) + self.lora_scale * self.to_out_lora(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