AnchoredAI/llm-grounded-diffusion
0
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import warnings15from typing import Callable, Optional, Union16 17import torch18import torch.nn.functional as F19from torch import nn20 21from diffusers.utils import deprecate, logging, maybe_allow_in_graph22 23logger = logging.get_logger(__name__) # pylint: disable=invalid-name24 25@maybe_allow_in_graph26class Attention(nn.Module):27 r"""28 A cross attention layer.29 30 Parameters:31 query_dim (`int`): The number of channels in the query.32 cross_attention_dim (`int`, *optional*):33 The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.34 heads (`int`, *optional*, defaults to 8): The number of heads to use for multi-head attention.35 dim_head (`int`, *optional*, defaults to 64): The number of channels in each head.36 dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.37 bias (`bool`, *optional*, defaults to False):38 Set to `True` for the query, key, and value linear layers to contain a bias parameter.39 """40 41 def __init__(42 self,43 query_dim: int,44 cross_attention_dim: Optional[int] = None,45 heads: int = 8,46 dim_head: int = 64,47 dropout: float = 0.0,48 bias=False,49 upcast_attention: bool = False,50 upcast_softmax: bool = False,51 cross_attention_norm: Optional[str] = None,52 cross_attention_norm_num_groups: int = 32,53 added_kv_proj_dim: Optional[int] = None,54 norm_num_groups: Optional[int] = None,55 spatial_norm_dim: Optional[int] = None,56 out_bias: bool = True,57 scale_qk: bool = True,58 only_cross_attention: bool = False,59 eps: float = 1e-5,60 rescale_output_factor: float = 1.0,61 residual_connection: bool = False,62 _from_deprecated_attn_block=False,63 processor: Optional["AttnProcessor"] = None,64 ):65 super().__init__()66 inner_dim = dim_head * heads67 cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim68 self.upcast_attention = upcast_attention69 self.upcast_softmax = upcast_softmax70 self.rescale_output_factor = rescale_output_factor71 self.residual_connection = residual_connection72 73 # we make use of this private variable to know whether this class is loaded74 # with an deprecated state dict so that we can convert it on the fly75 self._from_deprecated_attn_block = _from_deprecated_attn_block76 77 self.scale_qk = scale_qk78 self.scale = dim_head**-0.5 if self.scale_qk else 1.079 80 self.heads = heads81 # for slice_size > 0 the attention score computation82 # is split across the batch axis to save memory83 # You can set slice_size with `set_attention_slice`84 self.sliceable_head_dim = heads85 86 self.added_kv_proj_dim = added_kv_proj_dim87 self.only_cross_attention = only_cross_attention88 89 if self.added_kv_proj_dim is None and self.only_cross_attention:90 raise ValueError(91 "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."92 )93 94 if norm_num_groups is not None:95 self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True)96 else:97 self.group_norm = None98 99 if spatial_norm_dim is not None:100 self.spatial_norm = SpatialNorm(f_channels=query_dim, zq_channels=spatial_norm_dim)101 else:102 self.spatial_norm = None103 104 if cross_attention_norm is None:105 self.norm_cross = None106 elif cross_attention_norm == "layer_norm":107 self.norm_cross = nn.LayerNorm(cross_attention_dim)108 elif cross_attention_norm == "group_norm":109 if self.added_kv_proj_dim is not None:110 # The given `encoder_hidden_states` are initially of shape111 # (batch_size, seq_len, added_kv_proj_dim) before being projected112 # to (batch_size, seq_len, cross_attention_dim). The norm is applied113 # before the projection, so we need to use `added_kv_proj_dim` as114 # the number of channels for the group norm.115 norm_cross_num_channels = added_kv_proj_dim116 else:117 norm_cross_num_channels = cross_attention_dim118 119 self.norm_cross = nn.GroupNorm(120 num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True121 )122 else:123 raise ValueError(124 f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"125 )126 127 self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)128 129 if not self.only_cross_attention:130 # only relevant for the `AddedKVProcessor` classes131 self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias)132 self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias)133 else:134 self.to_k = None135 self.to_v = None136 137 if self.added_kv_proj_dim is not None:138 self.add_k_proj = nn.Linear(added_kv_proj_dim, inner_dim)139 self.add_v_proj = nn.Linear(added_kv_proj_dim, inner_dim)140 141 self.to_out = nn.ModuleList([])142 self.to_out.append(nn.Linear(inner_dim, query_dim, bias=out_bias))143 self.to_out.append(nn.Dropout(dropout))144 145 # set attention processor146 # We use the AttnProcessor2_0 by default when torch 2.x is used which uses147 # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention148 # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1149 if processor is None:150 # processor = (151 # AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()152 # )153 # Note: efficient attention is not used. We can use efficient attention to speed up.154 processor = AttnProcessor()155 self.set_processor(processor)156 157 def set_processor(self, processor: "AttnProcessor"):158 # if current processor is in `self._modules` and if passed `processor` is not, we need to159 # pop `processor` from `self._modules`160 if (161 hasattr(self, "processor")162 and isinstance(self.processor, torch.nn.Module)163 and not isinstance(processor, torch.nn.Module)164 ):165 logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}")166 self._modules.pop("processor")167 168 self.processor = processor169 170 def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, return_attntion_probs=False, **cross_attention_kwargs):171 # The `Attention` class can call different attention processors / attention functions172 # here we simply pass along all tensors to the selected processor class173 # For standard processors that are defined here, `**cross_attention_kwargs` is empty174 return self.processor(175 self,176 hidden_states,177 encoder_hidden_states=encoder_hidden_states,178 attention_mask=attention_mask,179 return_attntion_probs=return_attntion_probs,180 **cross_attention_kwargs,181 )182 183 def batch_to_head_dim(self, tensor):184 head_size = self.heads185 batch_size, seq_len, dim = tensor.shape186 tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)187 tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)188 return tensor189 190 def head_to_batch_dim(self, tensor, out_dim=3):191 head_size = self.heads192 batch_size, seq_len, dim = tensor.shape193 tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)194 tensor = tensor.permute(0, 2, 1, 3)195 196 if out_dim == 3:197 tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)198 199 return tensor200 201 def get_attention_scores(self, query, key, attention_mask=None):202 dtype = query.dtype203 if self.upcast_attention:204 query = query.float()205 key = key.float()206 207 if attention_mask is None:208 baddbmm_input = torch.empty(209 query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device210 )211 beta = 0212 else:213 baddbmm_input = attention_mask214 beta = 1215 216 attention_scores = torch.baddbmm(217 baddbmm_input,218 query,219 key.transpose(-1, -2),220 beta=beta,221 alpha=self.scale,222 )223 del baddbmm_input224 225 if self.upcast_softmax:226 attention_scores = attention_scores.float()227 228 attention_probs = attention_scores.softmax(dim=-1)229 del attention_scores230 231 attention_probs = attention_probs.to(dtype)232 233 return attention_probs234 235 def prepare_attention_mask(self, attention_mask, target_length, batch_size=None, out_dim=3):236 if batch_size is None:237 deprecate(238 "batch_size=None",239 "0.0.15",240 (241 "Not passing the `batch_size` parameter to `prepare_attention_mask` can lead to incorrect"242 " attention mask preparation and is deprecated behavior. Please make sure to pass `batch_size` to"243 " `prepare_attention_mask` when preparing the attention_mask."244 ),245 )246 batch_size = 1247 248 head_size = self.heads249 if attention_mask is None:250 return attention_mask251 252 current_length: int = attention_mask.shape[-1]253 if current_length != target_length:254 if attention_mask.device.type == "mps":255 # HACK: MPS: Does not support padding by greater than dimension of input tensor.256 # Instead, we can manually construct the padding tensor.257 padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)258 padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)259 attention_mask = torch.cat([attention_mask, padding], dim=2)260 else:261 # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:262 # we want to instead pad by (0, remaining_length), where remaining_length is:263 # remaining_length: int = target_length - current_length264 # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding265 attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)266 267 if out_dim == 3:268 if attention_mask.shape[0] < batch_size * head_size:269 attention_mask = attention_mask.repeat_interleave(head_size, dim=0)270 elif out_dim == 4:271 attention_mask = attention_mask.unsqueeze(1)272 attention_mask = attention_mask.repeat_interleave(head_size, dim=1)273 274 return attention_mask275 276 def norm_encoder_hidden_states(self, encoder_hidden_states):277 assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states"278 279 if isinstance(self.norm_cross, nn.LayerNorm):280 encoder_hidden_states = self.norm_cross(encoder_hidden_states)281 elif isinstance(self.norm_cross, nn.GroupNorm):282 # Group norm norms along the channels dimension and expects283 # input to be in the shape of (N, C, *). In this case, we want284 # to norm along the hidden dimension, so we need to move285 # (batch_size, sequence_length, hidden_size) ->286 # (batch_size, hidden_size, sequence_length)287 encoder_hidden_states = encoder_hidden_states.transpose(1, 2)288 encoder_hidden_states = self.norm_cross(encoder_hidden_states)289 encoder_hidden_states = encoder_hidden_states.transpose(1, 2)290 else:291 assert False292 293 return encoder_hidden_states294 295 296class AttnProcessor:297 r"""298 Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).299 """300 301 def __init__(self):302 if not hasattr(F, "scaled_dot_product_attention"):303 raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")304 305 def __call_fast__(306 self,307 attn: Attention,308 hidden_states,309 encoder_hidden_states=None,310 attention_mask=None,311 temb=None,312 ):313 residual = hidden_states314 315 if attn.spatial_norm is not None:316 hidden_states = attn.spatial_norm(hidden_states, temb)317 318 input_ndim = hidden_states.ndim319 320 if input_ndim == 4:321 batch_size, channel, height, width = hidden_states.shape322 hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)323 324 batch_size, sequence_length, _ = (325 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape326 )327 inner_dim = hidden_states.shape[-1]328 329 if attention_mask is not None:330 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)331 # scaled_dot_product_attention expects attention_mask shape to be332 # (batch, heads, source_length, target_length)333 attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])334 335 if attn.group_norm is not None:336 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)337 338 query = attn.to_q(hidden_states)339 340 if encoder_hidden_states is None:341 encoder_hidden_states = hidden_states342 elif attn.norm_cross:343 encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)344 345 key = attn.to_k(encoder_hidden_states)346 value = attn.to_v(encoder_hidden_states)347 348 head_dim = inner_dim // attn.heads349 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)350 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)351 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)352 353 # the output of sdp = (batch, num_heads, seq_len, head_dim)354 # TODO: add support for attn.scale when we move to Torch 2.1355 hidden_states = F.scaled_dot_product_attention(356 query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False357 )358 359 hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)360 hidden_states = hidden_states.to(query.dtype)361 362 # linear proj363 hidden_states = attn.to_out[0](hidden_states)364 # dropout365 hidden_states = attn.to_out[1](hidden_states)366 367 if input_ndim == 4:368 hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)369 370 if attn.residual_connection:371 hidden_states = hidden_states + residual372 373 hidden_states = hidden_states / attn.rescale_output_factor374 375 return hidden_states376 377 def __call__(378 self,379 attn: Attention,380 hidden_states,381 encoder_hidden_states=None,382 attention_mask=None,383 temb=None,384 return_attntion_probs=False,385 attn_key=None,386 attn_process_fn=None,387 return_cond_ca_only=False,388 return_token_ca_only=None,389 offload_cross_attn_to_cpu=False,390 save_attn_to_dict=None,391 save_keys=None,392 enable_flash_attn=True,393 ):394 """395 attn_key: current key (a tuple of hierarchy index (up/mid/down, stage id, block id, sub-block id), sub block id should always be 0 in SD UNet)396 save_attn_to_dict: pass in a dict to save to dict397 """398 cross_attn = encoder_hidden_states is not None399 400 if (not cross_attn) or (401 (attn_process_fn is None) 402 and not (save_attn_to_dict is not None and (save_keys is None or (tuple(attn_key) in save_keys))) 403 and not return_attntion_probs):404 with torch.backends.cuda.sdp_kernel(enable_flash=enable_flash_attn, enable_math=True, enable_mem_efficient=enable_flash_attn):405 return self.__call_fast__(attn, hidden_states, encoder_hidden_states, attention_mask, temb)406 407 residual = hidden_states408 409 if attn.spatial_norm is not None:410 hidden_states = attn.spatial_norm(hidden_states, temb)411 412 input_ndim = hidden_states.ndim413 414 if input_ndim == 4:415 batch_size, channel, height, width = hidden_states.shape416 hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)417 418 batch_size, sequence_length, _ = (419 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape420 )421 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)422 423 if attn.group_norm is not None:424 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)425 426 query = attn.to_q(hidden_states)427 428 if encoder_hidden_states is None:429 encoder_hidden_states = hidden_states430 elif attn.norm_cross:431 encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)432 433 key = attn.to_k(encoder_hidden_states)434 value = attn.to_v(encoder_hidden_states)435 436 query = attn.head_to_batch_dim(query)437 key = attn.head_to_batch_dim(key)438 value = attn.head_to_batch_dim(value)439 440 attention_probs = attn.get_attention_scores(query, key, attention_mask)441 # Currently only process cross-attention442 if attn_process_fn is not None and cross_attn:443 attention_probs_before_process = attention_probs.clone()444 attention_probs = attn_process_fn(attention_probs, query, key, value, attn_key=attn_key, cross_attn=cross_attn, batch_size=batch_size, heads=attn.heads)445 else:446 attention_probs_before_process = attention_probs447 hidden_states = torch.bmm(attention_probs, value)448 hidden_states = attn.batch_to_head_dim(hidden_states)449 450 # linear proj451 hidden_states = attn.to_out[0](hidden_states)452 # dropout453 hidden_states = attn.to_out[1](hidden_states)454 455 if input_ndim == 4:456 hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)457 458 if attn.residual_connection:459 hidden_states = hidden_states + residual460 461 hidden_states = hidden_states / attn.rescale_output_factor462 463 if return_attntion_probs or save_attn_to_dict is not None:464 # Recover batch dimension: (batch_size, heads, flattened_2d, text_tokens)465 attention_probs_unflattened = attention_probs_before_process.unflatten(dim=0, sizes=(batch_size, attn.heads))466 if return_token_ca_only is not None:467 # (batch size, n heads, 2d dimension, num text tokens)468 if isinstance(return_token_ca_only, int):469 # return_token_ca_only: an integer470 attention_probs_unflattened = attention_probs_unflattened[:, :, :, return_token_ca_only:return_token_ca_only+1]471 else:472 # return_token_ca_only: A 1d index tensor473 attention_probs_unflattened = attention_probs_unflattened[:, :, :, return_token_ca_only]474 if return_cond_ca_only:475 assert batch_size % 2 == 0, f"Samples are not in pairs: {batch_size} samples"476 attention_probs_unflattened = attention_probs_unflattened[batch_size // 2:]477 if offload_cross_attn_to_cpu:478 attention_probs_unflattened = attention_probs_unflattened.cpu()479 if save_attn_to_dict is not None and (save_keys is None or (tuple(attn_key) in save_keys)):480 save_attn_to_dict[tuple(attn_key)] = attention_probs_unflattened481 if return_attntion_probs:482 return hidden_states, attention_probs_unflattened483 return hidden_states484 485# For typing486AttentionProcessor = AttnProcessor487 488class SpatialNorm(nn.Module):489 """490 Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002491 """492 493 def __init__(494 self,495 f_channels,496 zq_channels,497 ):498 super().__init__()499 self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True)500 self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)501 self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)502 503 def forward(self, f, zq):504 f_size = f.shape[-2:]505 zq = F.interpolate(zq, size=f_size, mode="nearest")506 norm_f = self.norm_layer(f)507 new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)508 return new_f509