import cv2 import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.attention_processor import Attention from typing import Optional from .layers import MIFusion, MIFusionPrototype from .utils import get_masks, get_sigmoid class AttnProcessor2_0(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): 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__() def __call__( self, attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None, # Useless bboxes=[], obboxes=[], embeds_pooler=None, height=512, width=512, prototypes=None, ref_features=None, guidance_masks=None, supplement_mask=None, sigmoid_values=None, in_box=None, do_classifier_free_guidance=False, # End Useless *args, **kwargs, ) -> torch.Tensor: 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) if attn.norm_q is not None: query = attn.norm_q(query) if attn.norm_k is not None: key = attn.norm_k(key) # 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 MaskedProcessor2_0(nn.Module): def __init__(self, hidden_size, cross_attention_dim=None, use_ea_attn=False, **kwargs): super().__init__() self.hidden_size = hidden_size self.cross_attention_dim = cross_attention_dim self.use_ea_attn = use_ea_attn self.prototype_mode = use_ea_attn and (kwargs['phase'] == 'novel') and (kwargs.get('attn_prototype_switch') is True) if self.prototype_mode: self.fusion = MIFusionPrototype(hidden_size, context_dim=cross_attention_dim) elif use_ea_attn: self.fusion = MIFusion(hidden_size, context_dim=cross_attention_dim) # Train ; Test (do_classifier_free_guidance) # Train // Train2(self.use_ea_attn) ; Test // Test2(self.use_ea_attn) def __call__( self, attn: Attention, # Output the same size as hidden_states, encoder_hidden_states as Key and Value to inject information to hidden_states # shape[-2] 4096 as 64x64, 64 as 8x8; shape[-1] as hidden_size hidden_states, # [1, 4096, 320] // [1, 64, 1280] ; [2, 4096, 320] // [2, 64, 1280] && torch.all(hidden_states[0] == hidden_states[1]) is Ture encoder_hidden_states=None, # [16, 77, 768] ; [17, 77, 768] attention_mask=None, bboxes=[], obboxes=[], embeds_pooler=None, height=512, width=512, prototypes=None, ref_features=None, guidance_masks=None, supplement_mask=None, sigmoid_values=None, in_box=None, do_classifier_free_guidance=False, ): instance_num = len(obboxes[0]) # 15 if not self.use_ea_attn: # [1, 77, 768]; [2, 77, 768] encoder_hidden_states = encoder_hidden_states[:2, ...] if do_classifier_free_guidance else encoder_hidden_states[:1, ...] if self.use_ea_attn: if do_classifier_free_guidance: hidden_states = torch.cat([hidden_states[0:1], hidden_states[1:2].repeat(instance_num + 1, 1, 1)]) # ;//[17, 64, 1280] image_token = hidden_states[1:] else: hidden_states = hidden_states.repeat(instance_num + 1, 1, 1) # //[16, 64, 1280] image_token = hidden_states batch_size, sequence_length, _ = hidden_states.shape # _ is hidden_size 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]) query = attn.to_q(hidden_states) # [1, 4096, 320] // [16, 64, 1280] ; [2, 4096, 320]->[2, 4096, 320] // [17, 64, 1280]->[17, 64, 1280] key = attn.to_k(encoder_hidden_states) # [1, 77, 320] // [16, 77, 1280] ; [2, 77, 768]->[2, 77, 320] // [17, 77, 768 (cross_attention_dim)] -> [17, 77, 1280 (self.inner_kv_dim)] value = attn.to_v(encoder_hidden_states) # Same with key inner_dim = key.shape[-1] # 320 // 1280 head_dim = inner_dim // attn.heads # 40 // 160 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # [1, 8, 4096, 40] // [16, 8, 64, 160] key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # [1, 8, 77, 40] // [16, 8, 77, 160] value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # [1, 8, 77, 40] // [16, 8, 77, 160] if attn.norm_q is not None: query = attn.norm_q(query) if attn.norm_k is not None: key = attn.norm_k(key) hidden_states = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) # [1, 8, 4096, 40] // [16, 8, 64, 1280] hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) # [1, 4096, 320] // [16, 64, 1280] hidden_states = hidden_states.to(query.dtype) hidden_states = attn.to_out[0](hidden_states) # [1, 4096, 320] // [16, 64, 1280] ; [2, 4096, 320] // [17, 64, 1280] # Linear hidden_states = attn.to_out[1](hidden_states) # [1, 4096, 320] // [16, 64, 1280] ; [2, 4096, 320] // [17, 64, 1280] # Dropout if not self.use_ea_attn: return hidden_states assert self.use_ea_attn if do_classifier_free_guidance: hidden_states_uncond, hidden_states = hidden_states[0:1], hidden_states[1:] # torch.Size([1, HW, C]) other_info = {} other_info['image_token'] = image_token.unsqueeze(0) # [1, 16, 64, 1280] other_info['context'] = encoder_hidden_states[1:, ...] if do_classifier_free_guidance else encoder_hidden_states # [16, 77, 768] other_info['box'] = in_box # [1, 15, 8] other_info['context_pooler'] = embeds_pooler # [16, 1, 768] other_info['supplement_mask'] = supplement_mask # [1, 1, 64, 64] other_info['height'] = height # 512 other_info['width'] = width # 512 other_info['ref_features'] = ref_features # [15, 16, 768], [1, 16, 768] other_info['sigmoid_values'] = sigmoid_values # [1, 16, 768] other_info['guidance_masks'] = guidance_masks # [1, 15, 64, 64] other_info['instance_num'] = instance_num other_info['prototypes'] = prototypes hidden_states = self.fusion(hidden_states.unsqueeze(0), # [1, 16, 64, 1280] other_info=other_info) # hidden_states_cond.shape [1, 64, 1280] if do_classifier_free_guidance: hidden_states = torch.cat([hidden_states_uncond, hidden_states]) return hidden_states def set_processors(unet, **kwargs): attn_processors = {} for name, _ in unet.attn_processors.items(): use_ea_attn = False kwargs['attn_prototype_switch'] = False cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim if name.startswith("mid_block"): hidden_size = unet.config.block_out_channels[-1] # unet.config.block_out_channels [320, 640, 1280, 1280] use_ea_attn = True elif name.startswith("up_blocks"): block_id = int(name[len("up_blocks.")]) attention_id = int(name[len("up_blocks.2.attentions.")]) hidden_size = list(reversed(unet.config.block_out_channels))[block_id] if block_id == 1: use_ea_attn = True elif (block_id != 1) and (kwargs['phase'] == 'novel'): # run-4 use_ea_attn = True kwargs['attn_prototype_switch'] = True elif name.startswith("down_blocks"): block_id = int(name[len("down_blocks.")]) hidden_size = unet.config.block_out_channels[block_id] if cross_attention_dim is not None: attn_processors[name] = MaskedProcessor2_0(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, use_ea_attn=use_ea_attn, **kwargs) else: attn_processors[name] = AttnProcessor2_0() unet.set_attn_processor(attn_processors)