Spaces:
Running
on
Zero
Running
on
Zero
| # 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 | |
| def my_scaled_dot_product_attention( | |
| query, | |
| key, | |
| value, | |
| attn_mask=None, | |
| dropout_p=0.0, | |
| is_causal=False, | |
| scale=None, | |
| special_token_weight=1.0, | |
| special_token_indices=None, | |
| ) -> torch.Tensor: | |
| """ | |
| Computes the scaled dot-product attention with additional control over specific tokens. | |
| This function is a re-implementation of the scaled dot-product attention mechanism, | |
| designed to return both the attention map and the output of the attention operation. | |
| It also provides additional control via a scalar that modifies the attention map | |
| for specific tokens. | |
| """ | |
| L, S = query.size(-2), key.size(-2) | |
| scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale | |
| attn_bias = torch.zeros(L, S, dtype=query.dtype).cuda() | |
| if is_causal: | |
| assert attn_mask is None | |
| temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) | |
| attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) | |
| attn_bias.to(query.dtype) | |
| if attn_mask is not None: | |
| if attn_mask.dtype == torch.bool: | |
| attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) | |
| else: | |
| attn_bias += attn_mask | |
| attn_weight = query @ key.transpose(-2, -1) * scale_factor | |
| attn_weight += attn_bias | |
| if special_token_indices is not None and special_token_weight != 1.0: | |
| bs = attn_weight.shape[0] | |
| attn_weight[torch.arange(bs), :, :, special_token_indices] = torch.max( | |
| attn_weight[torch.arange(bs), :, :, special_token_indices], | |
| attn_weight[torch.arange(bs), :, :, special_token_indices] | |
| * special_token_weight, | |
| ) | |
| attn_weight = torch.softmax(attn_weight, dim=-1) | |
| attn_weight = torch.dropout(attn_weight, dropout_p, train=True) | |
| return attn_weight @ value, attn_weight | |
| class AttnProcessor(torch.nn.Module): | |
| r""" | |
| Processor for implementing scaled dot-product attention. | |
| """ | |
| def __init__( | |
| self, | |
| hidden_size=None, | |
| cross_attention_dim=None, | |
| ): | |
| super().__init__() | |
| if not hasattr(F, "scaled_dot_product_attention"): | |
| raise ImportError( | |
| "AttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." | |
| ) | |
| def __call__( | |
| self, | |
| attn, | |
| hidden_states, | |
| qformer_tokens_out=None, | |
| special_token_indices=None, | |
| inference_mode=None, | |
| encoder_hidden_states=None, | |
| attention_mask=None, | |
| temb=None, | |
| special_token_weight=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) | |
| 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 NestedAttnProcessor(torch.nn.Module): | |
| r""" | |
| Nested Attention processor for IP-Adapater for PyTorch 2.0. | |
| """ | |
| def __init__(self, hidden_size, cross_attention_dim=None, normalize_factor=1.0): | |
| super().__init__() | |
| if not hasattr(F, "scaled_dot_product_attention"): | |
| raise ImportError( | |
| "NestedAttnProcessor 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.normalize_factor = normalize_factor | |
| self.nested_to_k = nn.Linear( | |
| cross_attention_dim or hidden_size, hidden_size, bias=False | |
| ) | |
| self.nested_to_v = nn.Linear( | |
| cross_attention_dim or hidden_size, hidden_size, bias=False | |
| ) | |
| def __call__( | |
| self, | |
| attn, | |
| hidden_states, | |
| qformer_tokens_out, | |
| special_token_indices, | |
| inference_mode=False, | |
| encoder_hidden_states=None, | |
| attention_mask=None, | |
| temb=None, | |
| special_token_weight=1.0, | |
| ): | |
| assert ( | |
| special_token_indices.shape[0] > 0 | |
| ), "special_token_indices should not be empty" | |
| # if inference mode is set to True, the code assumes that CFG is used and the first half | |
| # of the batch is used for the null prompt and the second half is used for the prompt | |
| residual = hidden_states | |
| if attn.spatial_norm is not None: | |
| hidden_states = attn.spatial_norm(hidden_states, temb) | |
| input_ndim = hidden_states.ndim | |
| bs = hidden_states.shape[0] | |
| if input_ndim == 4: | |
| bs, channel, height, width = hidden_states.shape | |
| hidden_states = hidden_states.view(bs, channel, height * width).transpose( | |
| 1, 2 | |
| ) | |
| bs, 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, bs | |
| ) | |
| # scaled_dot_product_attention expects attention_mask shape to be | |
| # (batch, heads, source_length, target_length) | |
| attention_mask = attention_mask.view( | |
| bs, 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: | |
| 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(bs, -1, attn.heads, head_dim).transpose(1, 2) | |
| key = key.view(bs, -1, attn.heads, head_dim).transpose(1, 2) | |
| value = value.view(bs, -1, attn.heads, head_dim).transpose(1, 2) | |
| # nested attention | |
| nested_key = self.nested_to_k(qformer_tokens_out) | |
| nested_value = self.nested_to_v(qformer_tokens_out) | |
| nested_key = nested_key.view(bs, -1, attn.heads, head_dim).transpose(1, 2) | |
| nested_value = nested_value.view(bs, -1, attn.heads, head_dim).transpose(1, 2) | |
| nested_hidden_states = F.scaled_dot_product_attention( | |
| query, | |
| nested_key, | |
| nested_value, | |
| attn_mask=None, | |
| dropout_p=0.0, | |
| is_causal=False, | |
| ) | |
| # normalize V_q | |
| textual_values_norms = torch.norm( | |
| value[torch.arange(bs), :, special_token_indices], dim=-1 | |
| ) | |
| nested_hidden_states = ( | |
| torch.nn.functional.normalize(nested_hidden_states, p=2, dim=-1) | |
| * self.normalize_factor | |
| ) | |
| nested_hidden_states = ( | |
| textual_values_norms.view(bs, -1, 1, 1) * nested_hidden_states | |
| ) | |
| # outer attention | |
| value_without_special_tokens = value.clone() | |
| if inference_mode: | |
| value_without_special_tokens[bs // 2 : bs, :, special_token_indices, :] = ( | |
| 0.0 | |
| ) | |
| else: | |
| value_without_special_tokens[ | |
| torch.arange(bs), :, special_token_indices, : | |
| ] = 0.0 | |
| hidden_states_without_special_tokens, attn_weight = ( | |
| my_scaled_dot_product_attention( | |
| query, | |
| key, | |
| value_without_special_tokens, | |
| attn_mask=None, | |
| dropout_p=0.0, | |
| is_causal=False, | |
| special_token_weight=special_token_weight, | |
| special_token_indices=special_token_indices, | |
| ) | |
| ) | |
| # add the special token values | |
| if inference_mode: | |
| special_token_attn_weight = attn_weight[ | |
| bs // 2 : bs, :, :, special_token_indices | |
| ] | |
| else: | |
| special_token_attn_weight = attn_weight[ | |
| torch.arange(bs), :, :, special_token_indices | |
| ] | |
| if inference_mode: | |
| special_token_weighted_values = ( | |
| special_token_attn_weight * nested_hidden_states[bs // 2 : bs] | |
| ) | |
| else: | |
| special_token_weighted_values = ( | |
| special_token_attn_weight.unsqueeze(-1) * nested_hidden_states | |
| ) | |
| if inference_mode: | |
| hidden_states = hidden_states_without_special_tokens | |
| hidden_states[bs // 2 : bs] += special_token_weighted_values | |
| else: | |
| hidden_states = ( | |
| hidden_states_without_special_tokens + special_token_weighted_values | |
| ) | |
| # arrange hidden states | |
| hidden_states = hidden_states.transpose(1, 2).reshape( | |
| bs, -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( | |
| bs, channel, height, width | |
| ) | |
| if attn.residual_connection: | |
| hidden_states = hidden_states + residual | |
| hidden_states = hidden_states / attn.rescale_output_factor | |
| return hidden_states | |