guysss/ACE-Step
0
1# Copyright 2024 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 Optional, Union, Tuple15 16import torch17import torch.nn.functional as F18from torch import nn19 20from diffusers.utils import logging21from diffusers.models.attention_processor import Attention22 23logger = logging.get_logger(__name__) # pylint: disable=invalid-name24 25 26class CustomLiteLAProcessor2_0:27 """Attention processor used typically in processing the SD3-like self-attention projections. add rms norm for query and key and apply RoPE"""28 29 def __init__(self):30 self.kernel_func = nn.ReLU(inplace=False)31 self.eps = 1e-1532 self.pad_val = 1.033 34 def apply_rotary_emb(35 self,36 x: torch.Tensor,37 freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],38 ) -> Tuple[torch.Tensor, torch.Tensor]:39 """40 Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings41 to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are42 reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting43 tensors contain rotary embeddings and are returned as real tensors.44 45 Args:46 x (`torch.Tensor`):47 Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply48 freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)49 50 Returns:51 Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.52 """53 cos, sin = freqs_cis # [S, D]54 cos = cos[None, None]55 sin = sin[None, None]56 cos, sin = cos.to(x.device), sin.to(x.device)57 58 x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]59 x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)60 out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)61 62 return out63 64 def __call__(65 self,66 attn: Attention,67 hidden_states: torch.FloatTensor,68 encoder_hidden_states: torch.FloatTensor = None,69 attention_mask: Optional[torch.FloatTensor] = None,70 encoder_attention_mask: Optional[torch.FloatTensor] = None,71 rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,72 rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,73 *args,74 **kwargs,75 ) -> torch.FloatTensor:76 hidden_states_len = hidden_states.shape[1]77 78 input_ndim = hidden_states.ndim79 if input_ndim == 4:80 batch_size, channel, height, width = hidden_states.shape81 hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)82 if encoder_hidden_states is not None:83 context_input_ndim = encoder_hidden_states.ndim84 if context_input_ndim == 4:85 batch_size, channel, height, width = encoder_hidden_states.shape86 encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)87 88 batch_size = hidden_states.shape[0]89 90 # `sample` projections.91 dtype = hidden_states.dtype92 query = attn.to_q(hidden_states)93 key = attn.to_k(hidden_states)94 value = attn.to_v(hidden_states)95 96 # `context` projections.97 has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")98 if encoder_hidden_states is not None and has_encoder_hidden_state_proj:99 encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)100 encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)101 encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)102 103 # attention104 if not attn.is_cross_attention:105 query = torch.cat([query, encoder_hidden_states_query_proj], dim=1)106 key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)107 value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)108 else:109 query = hidden_states110 key = encoder_hidden_states111 value = encoder_hidden_states112 113 inner_dim = key.shape[-1]114 head_dim = inner_dim // attn.heads115 116 query = query.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)117 key = key.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1).transpose(-1, -2)118 value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)119 120 # RoPE需要 [B, H, S, D] 输入121 # 此时 query是 [B, H, D, S], 需要转成 [B, H, S, D] 才能应用RoPE122 query = query.permute(0, 1, 3, 2) # [B, H, S, D] (从 [B, H, D, S])123 124 # Apply query and key normalization if needed125 if attn.norm_q is not None:126 query = attn.norm_q(query)127 if attn.norm_k is not None:128 key = attn.norm_k(key)129 130 # Apply RoPE if needed131 if rotary_freqs_cis is not None:132 query = self.apply_rotary_emb(query, rotary_freqs_cis)133 if not attn.is_cross_attention:134 key = self.apply_rotary_emb(key, rotary_freqs_cis)135 elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:136 key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)137 138 # 此时 query是 [B, H, S, D]，需要还原成 [B, H, D, S]139 query = query.permute(0, 1, 3, 2) # [B, H, D, S]140 141 if attention_mask is not None:142 # attention_mask: [B, S] -> [B, 1, S, 1]143 attention_mask = attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S, 1]144 query = query * attention_mask.permute(0, 1, 3, 2) # [B, H, S, D] * [B, 1, S, 1]145 if not attn.is_cross_attention:146 key = key * attention_mask # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘147 value = value * attention_mask.permute(0, 1, 3, 2) # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度148 149 if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:150 encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S_enc, 1]151 # 此时 key: [B, h, S_enc, D], value: [B, h, D, S_enc]152 key = key * encoder_attention_mask # [B, h, S_enc, D] * [B, 1, S_enc, 1]153 value = value * encoder_attention_mask.permute(0, 1, 3, 2) # [B, h, D, S_enc] * [B, 1, 1, S_enc]154 155 query = self.kernel_func(query)156 key = self.kernel_func(key)157 158 query, key, value = query.float(), key.float(), value.float()159 160 value = F.pad(value, (0, 0, 0, 1), mode="constant", value=self.pad_val)161 162 vk = torch.matmul(value, key)163 164 hidden_states = torch.matmul(vk, query)165 166 if hidden_states.dtype in [torch.float16, torch.bfloat16]:167 hidden_states = hidden_states.float()168 169 hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)170 171 hidden_states = hidden_states.view(batch_size, attn.heads * head_dim, -1).permute(0, 2, 1)172 173 hidden_states = hidden_states.to(dtype)174 if encoder_hidden_states is not None:175 encoder_hidden_states = encoder_hidden_states.to(dtype)176 177 # Split the attention outputs.178 if encoder_hidden_states is not None and not attn.is_cross_attention and has_encoder_hidden_state_proj:179 hidden_states, encoder_hidden_states = (180 hidden_states[:, : hidden_states_len],181 hidden_states[:, hidden_states_len:],182 )183 184 # linear proj185 hidden_states = attn.to_out[0](hidden_states)186 # dropout187 hidden_states = attn.to_out[1](hidden_states)188 if encoder_hidden_states is not None and not attn.context_pre_only and not attn.is_cross_attention and hasattr(attn, "to_add_out"):189 encoder_hidden_states = attn.to_add_out(encoder_hidden_states)190 191 if input_ndim == 4:192 hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)193 if encoder_hidden_states is not None and context_input_ndim == 4:194 encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)195 196 if torch.get_autocast_gpu_dtype() == torch.float16:197 hidden_states = hidden_states.clip(-65504, 65504)198 if encoder_hidden_states is not None:199 encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)200 201 return hidden_states, encoder_hidden_states202 203 204class CustomerAttnProcessor2_0:205 r"""206 Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).207 """208 209 def __init__(self):210 if not hasattr(F, "scaled_dot_product_attention"):211 raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")212 213 def apply_rotary_emb(214 self,215 x: torch.Tensor,216 freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],217 ) -> Tuple[torch.Tensor, torch.Tensor]:218 """219 Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings220 to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are221 reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting222 tensors contain rotary embeddings and are returned as real tensors.223 224 Args:225 x (`torch.Tensor`):226 Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply227 freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)228 229 Returns:230 Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.231 """232 cos, sin = freqs_cis # [S, D]233 cos = cos[None, None]234 sin = sin[None, None]235 cos, sin = cos.to(x.device), sin.to(x.device)236 237 x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]238 x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)239 out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)240 241 return out242 243 def __call__(244 self,245 attn: Attention,246 hidden_states: torch.FloatTensor,247 encoder_hidden_states: torch.FloatTensor = None,248 attention_mask: Optional[torch.FloatTensor] = None,249 encoder_attention_mask: Optional[torch.FloatTensor] = None,250 rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,251 rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,252 *args,253 **kwargs,254 ) -> torch.Tensor:255 256 residual = hidden_states257 input_ndim = hidden_states.ndim258 259 if input_ndim == 4:260 batch_size, channel, height, width = hidden_states.shape261 hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)262 263 batch_size, sequence_length, _ = (264 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape265 )266 267 has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")268 269 if attn.group_norm is not None:270 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)271 272 query = attn.to_q(hidden_states)273 274 if encoder_hidden_states is None:275 encoder_hidden_states = hidden_states276 elif attn.norm_cross:277 encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)278 279 key = attn.to_k(encoder_hidden_states)280 value = attn.to_v(encoder_hidden_states)281 282 inner_dim = key.shape[-1]283 head_dim = inner_dim // attn.heads284 285 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)286 287 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)288 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)289 290 if attn.norm_q is not None:291 query = attn.norm_q(query)292 if attn.norm_k is not None:293 key = attn.norm_k(key)294 295 # Apply RoPE if needed296 if rotary_freqs_cis is not None:297 query = self.apply_rotary_emb(query, rotary_freqs_cis)298 if not attn.is_cross_attention:299 key = self.apply_rotary_emb(key, rotary_freqs_cis)300 elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:301 key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)302 303 if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:304 # attention_mask: N x S1305 # encoder_attention_mask: N x S2306 # cross attention 整合attention_mask和encoder_attention_mask307 combined_mask = attention_mask[:, :, None] * encoder_attention_mask[:, None, :]308 attention_mask = torch.where(combined_mask == 1, 0.0, -torch.inf)309 attention_mask = attention_mask[:, None, :, :].expand(-1, attn.heads, -1, -1).to(query.dtype)310 311 elif not attn.is_cross_attention and attention_mask is not None:312 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)313 # scaled_dot_product_attention expects attention_mask shape to be314 # (batch, heads, source_length, target_length)315 attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])316 317 # the output of sdp = (batch, num_heads, seq_len, head_dim)318 # TODO: add support for attn.scale when we move to Torch 2.1319 hidden_states = F.scaled_dot_product_attention(320 query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False321 )322 323 hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)324 hidden_states = hidden_states.to(query.dtype)325 326 # linear proj327 hidden_states = attn.to_out[0](hidden_states)328 # dropout329 hidden_states = attn.to_out[1](hidden_states)330 331 if input_ndim == 4:332 hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)333 334 if attn.residual_connection:335 hidden_states = hidden_states + residual336 337 hidden_states = hidden_states / attn.rescale_output_factor338 339 return hidden_states340