declare-lab/tango2
92
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.14from typing import Callable, Optional, Union15 16import torch17import torch.nn.functional as F18from torch import nn19 20from ..utils import deprecate, logging21from ..utils.import_utils import is_xformers_available22 23 24logger = logging.get_logger(__name__) # pylint: disable=invalid-name25 26 27if is_xformers_available():28 import xformers29 import xformers.ops30else:31 xformers = None32 33 34class Attention(nn.Module):35 r"""36 A cross attention layer.37 38 Parameters:39 query_dim (`int`): The number of channels in the query.40 cross_attention_dim (`int`, *optional*):41 The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.42 heads (`int`, *optional*, defaults to 8): The number of heads to use for multi-head attention.43 dim_head (`int`, *optional*, defaults to 64): The number of channels in each head.44 dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.45 bias (`bool`, *optional*, defaults to False):46 Set to `True` for the query, key, and value linear layers to contain a bias parameter.47 """48 49 def __init__(50 self,51 query_dim: int,52 cross_attention_dim: Optional[int] = None,53 heads: int = 8,54 dim_head: int = 64,55 dropout: float = 0.0,56 bias=False,57 upcast_attention: bool = False,58 upcast_softmax: bool = False,59 cross_attention_norm: bool = False,60 added_kv_proj_dim: Optional[int] = None,61 norm_num_groups: Optional[int] = None,62 out_bias: bool = True,63 scale_qk: bool = True,64 processor: Optional["AttnProcessor"] = None,65 ):66 super().__init__()67 inner_dim = dim_head * heads68 cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim69 self.upcast_attention = upcast_attention70 self.upcast_softmax = upcast_softmax71 self.cross_attention_norm = cross_attention_norm72 73 self.scale = dim_head**-0.5 if scale_qk else 1.074 75 self.heads = heads76 # for slice_size > 0 the attention score computation77 # is split across the batch axis to save memory78 # You can set slice_size with `set_attention_slice`79 self.sliceable_head_dim = heads80 81 self.added_kv_proj_dim = added_kv_proj_dim82 83 if norm_num_groups is not None:84 self.group_norm = nn.GroupNorm(num_channels=inner_dim, num_groups=norm_num_groups, eps=1e-5, affine=True)85 else:86 self.group_norm = None87 88 if cross_attention_norm:89 self.norm_cross = nn.LayerNorm(cross_attention_dim)90 91 self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)92 self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias)93 self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias)94 95 if self.added_kv_proj_dim is not None:96 self.add_k_proj = nn.Linear(added_kv_proj_dim, cross_attention_dim)97 self.add_v_proj = nn.Linear(added_kv_proj_dim, cross_attention_dim)98 99 self.to_out = nn.ModuleList([])100 self.to_out.append(nn.Linear(inner_dim, query_dim, bias=out_bias))101 self.to_out.append(nn.Dropout(dropout))102 103 # set attention processor104 # We use the AttnProcessor2_0 by default when torch 2.x is used which uses105 # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention106 # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1107 if processor is None:108 processor = (109 AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and scale_qk else AttnProcessor()110 )111 self.set_processor(processor)112 113 def set_use_memory_efficient_attention_xformers(114 self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None115 ):116 is_lora = hasattr(self, "processor") and isinstance(117 self.processor, (LoRAAttnProcessor, LoRAXFormersAttnProcessor)118 )119 120 if use_memory_efficient_attention_xformers:121 if self.added_kv_proj_dim is not None:122 # TODO(Anton, Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP123 # which uses this type of cross attention ONLY because the attention mask of format124 # [0, ..., -10.000, ..., 0, ...,] is not supported125 raise NotImplementedError(126 "Memory efficient attention with `xformers` is currently not supported when"127 " `self.added_kv_proj_dim` is defined."128 )129 elif not is_xformers_available():130 raise ModuleNotFoundError(131 (132 "Refer to https://github.com/facebookresearch/xformers for more information on how to install"133 " xformers"134 ),135 name="xformers",136 )137 elif not torch.cuda.is_available():138 raise ValueError(139 "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"140 " only available for GPU "141 )142 else:143 try:144 # Make sure we can run the memory efficient attention145 _ = xformers.ops.memory_efficient_attention(146 torch.randn((1, 2, 40), device="cuda"),147 torch.randn((1, 2, 40), device="cuda"),148 torch.randn((1, 2, 40), device="cuda"),149 )150 except Exception as e:151 raise e152 153 if is_lora:154 processor = LoRAXFormersAttnProcessor(155 hidden_size=self.processor.hidden_size,156 cross_attention_dim=self.processor.cross_attention_dim,157 rank=self.processor.rank,158 attention_op=attention_op,159 )160 processor.load_state_dict(self.processor.state_dict())161 processor.to(self.processor.to_q_lora.up.weight.device)162 else:163 processor = XFormersAttnProcessor(attention_op=attention_op)164 else:165 if is_lora:166 processor = LoRAAttnProcessor(167 hidden_size=self.processor.hidden_size,168 cross_attention_dim=self.processor.cross_attention_dim,169 rank=self.processor.rank,170 )171 processor.load_state_dict(self.processor.state_dict())172 processor.to(self.processor.to_q_lora.up.weight.device)173 else:174 processor = AttnProcessor()175 176 self.set_processor(processor)177 178 def set_attention_slice(self, slice_size):179 if slice_size is not None and slice_size > self.sliceable_head_dim:180 raise ValueError(f"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.")181 182 if slice_size is not None and self.added_kv_proj_dim is not None:183 processor = SlicedAttnAddedKVProcessor(slice_size)184 elif slice_size is not None:185 processor = SlicedAttnProcessor(slice_size)186 elif self.added_kv_proj_dim is not None:187 processor = AttnAddedKVProcessor()188 else:189 processor = AttnProcessor()190 191 self.set_processor(processor)192 193 def set_processor(self, processor: "AttnProcessor"):194 # if current processor is in `self._modules` and if passed `processor` is not, we need to195 # pop `processor` from `self._modules`196 if (197 hasattr(self, "processor")198 and isinstance(self.processor, torch.nn.Module)199 and not isinstance(processor, torch.nn.Module)200 ):201 logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}")202 self._modules.pop("processor")203 204 self.processor = processor205 206 def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, **cross_attention_kwargs):207 # The `Attention` class can call different attention processors / attention functions208 # here we simply pass along all tensors to the selected processor class209 # For standard processors that are defined here, `**cross_attention_kwargs` is empty210 return self.processor(211 self,212 hidden_states,213 encoder_hidden_states=encoder_hidden_states,214 attention_mask=attention_mask,215 **cross_attention_kwargs,216 )217 218 def batch_to_head_dim(self, tensor):219 head_size = self.heads220 batch_size, seq_len, dim = tensor.shape221 tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)222 tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)223 return tensor224 225 def head_to_batch_dim(self, tensor):226 head_size = self.heads227 batch_size, seq_len, dim = tensor.shape228 tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)229 tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size, seq_len, dim // head_size)230 return tensor231 232 def get_attention_scores(self, query, key, attention_mask=None):233 dtype = query.dtype234 if self.upcast_attention:235 query = query.float()236 key = key.float()237 238 if attention_mask is None:239 baddbmm_input = torch.empty(240 query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device241 )242 beta = 0243 else:244 baddbmm_input = attention_mask245 beta = 1246 247 attention_scores = torch.baddbmm(248 baddbmm_input,249 query,250 key.transpose(-1, -2),251 beta=beta,252 alpha=self.scale,253 )254 255 if self.upcast_softmax:256 attention_scores = attention_scores.float()257 258 attention_probs = attention_scores.softmax(dim=-1)259 attention_probs = attention_probs.to(dtype)260 261 return attention_probs262 263 def prepare_attention_mask(self, attention_mask, target_length, batch_size=None):264 if batch_size is None:265 deprecate(266 "batch_size=None",267 "0.0.15",268 (269 "Not passing the `batch_size` parameter to `prepare_attention_mask` can lead to incorrect"270 " attention mask preparation and is deprecated behavior. Please make sure to pass `batch_size` to"271 " `prepare_attention_mask` when preparing the attention_mask."272 ),273 )274 batch_size = 1275 276 head_size = self.heads277 if attention_mask is None:278 return attention_mask279 280 current_length: int = attention_mask.shape[-1]281 if current_length > target_length:282 # we *could* trim the mask with:283 # attention_mask = attention_mask[:,:target_length]284 # but this is weird enough that it's more likely to be a mistake than a shortcut285 raise ValueError(f"mask's length ({current_length}) exceeds the sequence length ({target_length}).")286 elif current_length < target_length:287 if attention_mask.device.type == "mps":288 # HACK: MPS: Does not support padding by greater than dimension of input tensor.289 # Instead, we can manually construct the padding tensor.290 padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)291 padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)292 attention_mask = torch.cat([attention_mask, padding], dim=2)293 else:294 remaining_length: int = target_length - current_length295 attention_mask = F.pad(attention_mask, (0, remaining_length), value=0.0)296 297 if attention_mask.shape[0] < batch_size * head_size:298 attention_mask = attention_mask.repeat_interleave(head_size, dim=0)299 return attention_mask300 301 302class AttnProcessor:303 def __call__(304 self,305 attn: Attention,306 hidden_states,307 encoder_hidden_states=None,308 attention_mask=None,309 ):310 batch_size, sequence_length, _ = (311 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape312 )313 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)314 query = attn.to_q(hidden_states)315 316 if encoder_hidden_states is None:317 encoder_hidden_states = hidden_states318 elif attn.cross_attention_norm:319 encoder_hidden_states = attn.norm_cross(encoder_hidden_states)320 321 key = attn.to_k(encoder_hidden_states)322 value = attn.to_v(encoder_hidden_states)323 324 query = attn.head_to_batch_dim(query)325 key = attn.head_to_batch_dim(key)326 value = attn.head_to_batch_dim(value)327 328 attention_probs = attn.get_attention_scores(query, key, attention_mask)329 hidden_states = torch.bmm(attention_probs, value)330 hidden_states = attn.batch_to_head_dim(hidden_states)331 332 # linear proj333 hidden_states = attn.to_out[0](hidden_states)334 # dropout335 hidden_states = attn.to_out[1](hidden_states)336 337 return hidden_states338 339 340class LoRALinearLayer(nn.Module):341 def __init__(self, in_features, out_features, rank=4):342 super().__init__()343 344 if rank > min(in_features, out_features):345 raise ValueError(f"LoRA rank {rank} must be less or equal than {min(in_features, out_features)}")346 347 self.down = nn.Linear(in_features, rank, bias=False)348 self.up = nn.Linear(rank, out_features, bias=False)349 350 nn.init.normal_(self.down.weight, std=1 / rank)351 nn.init.zeros_(self.up.weight)352 353 def forward(self, hidden_states):354 orig_dtype = hidden_states.dtype355 dtype = self.down.weight.dtype356 357 down_hidden_states = self.down(hidden_states.to(dtype))358 up_hidden_states = self.up(down_hidden_states)359 360 return up_hidden_states.to(orig_dtype)361 362 363class LoRAAttnProcessor(nn.Module):364 def __init__(self, hidden_size, cross_attention_dim=None, rank=4):365 super().__init__()366 367 self.hidden_size = hidden_size368 self.cross_attention_dim = cross_attention_dim369 self.rank = rank370 371 self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank)372 self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank)373 self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank)374 self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank)375 376 def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0):377 batch_size, sequence_length, _ = (378 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape379 )380 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)381 382 query = attn.to_q(hidden_states) + scale * self.to_q_lora(hidden_states)383 query = attn.head_to_batch_dim(query)384 385 encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states386 387 key = attn.to_k(encoder_hidden_states) + scale * self.to_k_lora(encoder_hidden_states)388 value = attn.to_v(encoder_hidden_states) + scale * self.to_v_lora(encoder_hidden_states)389 390 key = attn.head_to_batch_dim(key)391 value = attn.head_to_batch_dim(value)392 393 attention_probs = attn.get_attention_scores(query, key, attention_mask)394 hidden_states = torch.bmm(attention_probs, value)395 hidden_states = attn.batch_to_head_dim(hidden_states)396 397 # linear proj398 hidden_states = attn.to_out[0](hidden_states) + scale * self.to_out_lora(hidden_states)399 # dropout400 hidden_states = attn.to_out[1](hidden_states)401 402 return hidden_states403 404 405class AttnAddedKVProcessor:406 def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None):407 residual = hidden_states408 hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)409 batch_size, sequence_length, _ = hidden_states.shape410 encoder_hidden_states = encoder_hidden_states.transpose(1, 2)411 412 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)413 414 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)415 416 query = attn.to_q(hidden_states)417 query = attn.head_to_batch_dim(query)418 419 key = attn.to_k(hidden_states)420 value = attn.to_v(hidden_states)421 key = attn.head_to_batch_dim(key)422 value = attn.head_to_batch_dim(value)423 424 encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)425 encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)426 encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)427 encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)428 429 key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)430 value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)431 432 attention_probs = attn.get_attention_scores(query, key, attention_mask)433 hidden_states = torch.bmm(attention_probs, value)434 hidden_states = attn.batch_to_head_dim(hidden_states)435 436 # linear proj437 hidden_states = attn.to_out[0](hidden_states)438 # dropout439 hidden_states = attn.to_out[1](hidden_states)440 441 hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)442 hidden_states = hidden_states + residual443 444 return hidden_states445 446 447class XFormersAttnProcessor:448 def __init__(self, attention_op: Optional[Callable] = None):449 self.attention_op = attention_op450 451 def __call__(452 self,453 attn: Attention,454 hidden_states: torch.FloatTensor,455 encoder_hidden_states: Optional[torch.FloatTensor] = None,456 attention_mask: Optional[torch.FloatTensor] = None,457 ):458 batch_size, key_tokens, _ = (459 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape460 )461 462 attention_mask = attn.prepare_attention_mask(attention_mask, key_tokens, batch_size)463 if attention_mask is not None:464 # xformers doesn't broadcast for us, so we expand our singleton dimension manually465 _, query_tokens, _ = hidden_states.shape466 attention_mask = attention_mask.expand(-1, query_tokens, -1)467 468 query = attn.to_q(hidden_states)469 470 if encoder_hidden_states is None:471 encoder_hidden_states = hidden_states472 elif attn.cross_attention_norm:473 encoder_hidden_states = attn.norm_cross(encoder_hidden_states)474 475 key = attn.to_k(encoder_hidden_states)476 value = attn.to_v(encoder_hidden_states)477 478 query = attn.head_to_batch_dim(query).contiguous()479 key = attn.head_to_batch_dim(key).contiguous()480 value = attn.head_to_batch_dim(value).contiguous()481 482 hidden_states = xformers.ops.memory_efficient_attention(483 query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale484 )485 hidden_states = hidden_states.to(query.dtype)486 hidden_states = attn.batch_to_head_dim(hidden_states)487 488 # linear proj489 hidden_states = attn.to_out[0](hidden_states)490 # dropout491 hidden_states = attn.to_out[1](hidden_states)492 return hidden_states493 494 495class AttnProcessor2_0:496 def __init__(self):497 if not hasattr(F, "scaled_dot_product_attention"):498 raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")499 500 def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None):501 batch_size, sequence_length, _ = (502 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape503 )504 inner_dim = hidden_states.shape[-1]505 506 if attention_mask is not None:507 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)508 # scaled_dot_product_attention expects attention_mask shape to be509 # (batch, heads, source_length, target_length)510 attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])511 512 query = attn.to_q(hidden_states)513 514 if encoder_hidden_states is None:515 encoder_hidden_states = hidden_states516 elif attn.cross_attention_norm:517 encoder_hidden_states = attn.norm_cross(encoder_hidden_states)518 519 key = attn.to_k(encoder_hidden_states)520 value = attn.to_v(encoder_hidden_states)521 522 head_dim = inner_dim // attn.heads523 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)524 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)525 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)526 527 # the output of sdp = (batch, num_heads, seq_len, head_dim)528 # TODO: add support for attn.scale when we move to Torch 2.1529 hidden_states = F.scaled_dot_product_attention(530 query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False531 )532 533 hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)534 hidden_states = hidden_states.to(query.dtype)535 536 # linear proj537 hidden_states = attn.to_out[0](hidden_states)538 # dropout539 hidden_states = attn.to_out[1](hidden_states)540 return hidden_states541 542 543class LoRAXFormersAttnProcessor(nn.Module):544 def __init__(self, hidden_size, cross_attention_dim, rank=4, attention_op: Optional[Callable] = None):545 super().__init__()546 547 self.hidden_size = hidden_size548 self.cross_attention_dim = cross_attention_dim549 self.rank = rank550 self.attention_op = attention_op551 552 self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank)553 self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank)554 self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank)555 self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank)556 557 def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0):558 batch_size, sequence_length, _ = (559 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape560 )561 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)562 563 query = attn.to_q(hidden_states) + scale * self.to_q_lora(hidden_states)564 query = attn.head_to_batch_dim(query).contiguous()565 566 encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states567 568 key = attn.to_k(encoder_hidden_states) + scale * self.to_k_lora(encoder_hidden_states)569 value = attn.to_v(encoder_hidden_states) + scale * self.to_v_lora(encoder_hidden_states)570 571 key = attn.head_to_batch_dim(key).contiguous()572 value = attn.head_to_batch_dim(value).contiguous()573 574 hidden_states = xformers.ops.memory_efficient_attention(575 query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale576 )577 hidden_states = attn.batch_to_head_dim(hidden_states)578 579 # linear proj580 hidden_states = attn.to_out[0](hidden_states) + scale * self.to_out_lora(hidden_states)581 # dropout582 hidden_states = attn.to_out[1](hidden_states)583 584 return hidden_states585 586 587class SlicedAttnProcessor:588 def __init__(self, slice_size):589 self.slice_size = slice_size590 591 def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None):592 batch_size, sequence_length, _ = (593 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape594 )595 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)596 597 query = attn.to_q(hidden_states)598 dim = query.shape[-1]599 query = attn.head_to_batch_dim(query)600 601 if encoder_hidden_states is None:602 encoder_hidden_states = hidden_states603 elif attn.cross_attention_norm:604 encoder_hidden_states = attn.norm_cross(encoder_hidden_states)605 606 key = attn.to_k(encoder_hidden_states)607 value = attn.to_v(encoder_hidden_states)608 key = attn.head_to_batch_dim(key)609 value = attn.head_to_batch_dim(value)610 611 batch_size_attention, query_tokens, _ = query.shape612 hidden_states = torch.zeros(613 (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype614 )615 616 for i in range(batch_size_attention // self.slice_size):617 start_idx = i * self.slice_size618 end_idx = (i + 1) * self.slice_size619 620 query_slice = query[start_idx:end_idx]621 key_slice = key[start_idx:end_idx]622 attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None623 624 attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)625 626 attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])627 628 hidden_states[start_idx:end_idx] = attn_slice629 630 hidden_states = attn.batch_to_head_dim(hidden_states)631 632 # linear proj633 hidden_states = attn.to_out[0](hidden_states)634 # dropout635 hidden_states = attn.to_out[1](hidden_states)636 637 return hidden_states638 639 640class SlicedAttnAddedKVProcessor:641 def __init__(self, slice_size):642 self.slice_size = slice_size643 644 def __call__(self, attn: "Attention", hidden_states, encoder_hidden_states=None, attention_mask=None):645 residual = hidden_states646 hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)647 encoder_hidden_states = encoder_hidden_states.transpose(1, 2)648 649 batch_size, sequence_length, _ = hidden_states.shape650 651 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)652 653 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)654 655 query = attn.to_q(hidden_states)656 dim = query.shape[-1]657 query = attn.head_to_batch_dim(query)658 659 key = attn.to_k(hidden_states)660 value = attn.to_v(hidden_states)661 encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)662 encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)663 664 key = attn.head_to_batch_dim(key)665 value = attn.head_to_batch_dim(value)666 encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)667 encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)668 669 key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)670 value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)671 672 batch_size_attention, query_tokens, _ = query.shape673 hidden_states = torch.zeros(674 (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype675 )676 677 for i in range(batch_size_attention // self.slice_size):678 start_idx = i * self.slice_size679 end_idx = (i + 1) * self.slice_size680 681 query_slice = query[start_idx:end_idx]682 key_slice = key[start_idx:end_idx]683 attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None684 685 attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)686 687 attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])688 689 hidden_states[start_idx:end_idx] = attn_slice690 691 hidden_states = attn.batch_to_head_dim(hidden_states)692 693 # linear proj694 hidden_states = attn.to_out[0](hidden_states)695 # dropout696 hidden_states = attn.to_out[1](hidden_states)697 698 hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)699 hidden_states = hidden_states + residual700 701 return hidden_states702 703 704AttentionProcessor = Union[705 AttnProcessor,706 XFormersAttnProcessor,707 SlicedAttnProcessor,708 AttnAddedKVProcessor,709 SlicedAttnAddedKVProcessor,710 LoRAAttnProcessor,711 LoRAXFormersAttnProcessor,712]713