# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
import torch
import torch.nn as nn
import torch.nn.functional as F


class IPAttnProcessorDecoupled(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 (`List[float]`, defaults to 1.0):
            the weight scale of image prompt.
        num_tokens (`List[int]`, defaults to 4 when do ip_adapter_plus it should be 16):
            The context length of the image features.
    Notes:
        IP Attn 1: for face id embed
        IP Attn 2: for global image embed
    """

    def __init__(self, hidden_size, cross_attention_dim=None, scale=[0.7, 0.3], num_tokens=[16, 4]):
        super().__init__()

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

        self.to_k_ip_1 = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
        self.to_v_ip_1 = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
        self.to_k_ip_2 = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
        self.to_v_ip_2 = 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,
        *args,
        **kwargs,
    ):
        # 1. prepare 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
        )
        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)

        # 2. prepare encoder hidden states
        # 2.1 splitting data
        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        else:
            # get encoder_hidden_states, ip_hidden_states
            end_pos_2 = encoder_hidden_states.shape[1] - self.num_tokens[1]
            end_pos_1 = end_pos_2 - self.num_tokens[0]
            encoder_hidden_states, ip_hidden_states_1, ip_hidden_states_2 = (
                encoder_hidden_states[:, :end_pos_1, :],
                encoder_hidden_states[:, end_pos_1:end_pos_2, :],
                encoder_hidden_states[:, end_pos_2:, :],
            )
            if attn.norm_cross:
                encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)

        # 2.2 for text hidden states (key, value)
        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)

        # 2.3.1 for ip-adapter part 1  (key, value)
        ip_key_1 = self.to_k_ip_1(ip_hidden_states_1)
        ip_value_1 = self.to_v_ip_1(ip_hidden_states_1)

        ip_key_1 = attn.head_to_batch_dim(ip_key_1)
        ip_value_1 = attn.head_to_batch_dim(ip_value_1)

        ip_attention_probs_1 = attn.get_attention_scores(query, ip_key_1, None)
        self.attn_map_1 = ip_attention_probs_1
        ip_hidden_states_1 = torch.bmm(ip_attention_probs_1, ip_value_1)
        ip_hidden_states_1 = attn.batch_to_head_dim(ip_hidden_states_1)

        # 2.3.2 for ip-adapter part 2  (key, value)
        ip_key_2 = self.to_k_ip_2(ip_hidden_states_2)
        ip_value_2 = self.to_v_ip_2(ip_hidden_states_2)

        ip_key_2 = attn.head_to_batch_dim(ip_key_2)
        ip_value_2 = attn.head_to_batch_dim(ip_value_2)

        ip_attention_probs_2 = attn.get_attention_scores(query, ip_key_2, None)
        self.attn_map_2 = ip_attention_probs_2
        ip_hidden_states_2 = torch.bmm(ip_attention_probs_2, ip_value_2)
        ip_hidden_states_2 = attn.batch_to_head_dim(ip_hidden_states_2)

        # 3. Merge Hidden States
        hidden_states = hidden_states + self.scale[0] * ip_hidden_states_1 + self.scale[1] * ip_hidden_states_2

        # 4.1 linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # 4.2 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 IPAttnProcessorDecoupled2_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 (`List[float]`, defaults to 1.0):
            the weight scale of image prompt.
        num_tokens (`List[int]`, defaults to 4 when do ip_adapter_plus it should be 16):
            The context length of the image features.
    Notes:
        IP Attn 1: for face id embed
        IP Attn 2: for global image embed
    """

    def __init__(self, hidden_size, cross_attention_dim=None, scale=[0.7, 0.3], num_tokens=[16, 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_1 = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
        self.to_v_ip_1 = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
        self.to_k_ip_2 = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
        self.to_v_ip_2 = 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,
        *args,
        **kwargs,
    ):
        # 1. prepare 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)

        query = attn.to_q(hidden_states)

        # 2. prepare encoder hidden states
        # 2.1 splitting data
        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        else:
            # get encoder_hidden_states, ip_hidden_states
            end_pos_2 = encoder_hidden_states.shape[1] - self.num_tokens[1]
            end_pos_1 = end_pos_2 - self.num_tokens[0]
            encoder_hidden_states, ip_hidden_states_1, ip_hidden_states_2 = (
                encoder_hidden_states[:, :end_pos_1, :],
                encoder_hidden_states[:, end_pos_1:end_pos_2, :],
                encoder_hidden_states[:, end_pos_2:, :],
            )
            if attn.norm_cross:
                encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)

        # 2.2 for text hidden states (key, value)
        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)

        # the output of sdp = (batch, num_heads, seq_len, head_dim)
        # TODO: add support for attn.scale when we move to Torch 2.1
        # 2.3.1 for ip-adapter part 1  (key, value)
        ip_key_1 = self.to_k_ip_1(ip_hidden_states_1)
        ip_value_1 = self.to_v_ip_1(ip_hidden_states_1)

        ip_key_1 = ip_key_1.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        ip_value_1 = ip_value_1.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        ip_hidden_states_1 = F.scaled_dot_product_attention(
            query, ip_key_1, ip_value_1, attn_mask=None, dropout_p=0.0, is_causal=False
        )
        with torch.no_grad():
            self.attn_map = query @ ip_key_1.transpose(-2, -1).softmax(dim=-1)

        ip_hidden_states_1 = ip_hidden_states_1.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        ip_hidden_states_1 = ip_hidden_states_1.to(query.dtype)

        # 2.3.2 for ip-adapter part 2  (key, value)
        ip_key_2 = self.to_k_ip_2(ip_hidden_states_2)
        ip_value_2 = self.to_v_ip_2(ip_hidden_states_2)

        ip_key_2 = ip_key_2.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        ip_value_2 = ip_value_2.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        ip_hidden_states_2 = F.scaled_dot_product_attention(
            query, ip_key_2, ip_value_2, attn_mask=None, dropout_p=0.0, is_causal=False
        )
        with torch.no_grad():
            self.attn_map_2 = query @ ip_key_2.transpose(-2, -1).softmax(dim=-1)
            # print(self.attn_map.shape)

        ip_hidden_states_2 = ip_hidden_states_2.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        ip_hidden_states_2 = ip_hidden_states_2.to(query.dtype)


        # 3. Merge Hidden States
        hidden_states = hidden_states + self.scale[0] * ip_hidden_states_1 + self.scale[1] * ip_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

        return hidden_states
