(
self,
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
attention_mask: Optional[torch.FloatTensor] = None,
temb: Optional[torch.FloatTensor] = None,
scale: float = 1.0,
)
| 1181 | raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") |
| 1182 | |
| 1183 | def __call__( |
| 1184 | self, |
| 1185 | attn: Attention, |
| 1186 | hidden_states: torch.FloatTensor, |
| 1187 | encoder_hidden_states: Optional[torch.FloatTensor] = None, |
| 1188 | attention_mask: Optional[torch.FloatTensor] = None, |
| 1189 | temb: Optional[torch.FloatTensor] = None, |
| 1190 | scale: float = 1.0, |
| 1191 | ) -> torch.FloatTensor: |
| 1192 | residual = hidden_states |
| 1193 | |
| 1194 | args = () if USE_PEFT_BACKEND else (scale,) |
| 1195 | |
| 1196 | if attn.spatial_norm is not None: |
| 1197 | hidden_states = attn.spatial_norm(hidden_states, temb) |
| 1198 | |
| 1199 | input_ndim = hidden_states.ndim |
| 1200 | |
| 1201 | if input_ndim == 4: |
| 1202 | batch_size, channel, height, width = hidden_states.shape |
| 1203 | hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) |
| 1204 | |
| 1205 | batch_size, sequence_length, _ = ( |
| 1206 | hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape |
| 1207 | ) |
| 1208 | |
| 1209 | if attention_mask is not None: |
| 1210 | attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) |
| 1211 | # scaled_dot_product_attention expects attention_mask shape to be |
| 1212 | # (batch, heads, source_length, target_length) |
| 1213 | attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) |
| 1214 | |
| 1215 | if attn.group_norm is not None: |
| 1216 | hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) |
| 1217 | |
| 1218 | args = () if USE_PEFT_BACKEND else (scale,) |
| 1219 | query = attn.to_q(hidden_states, *args) |
| 1220 | |
| 1221 | if encoder_hidden_states is None: |
| 1222 | encoder_hidden_states = hidden_states |
| 1223 | elif attn.norm_cross: |
| 1224 | encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) |
| 1225 | |
| 1226 | key = attn.to_k(encoder_hidden_states, *args) |
| 1227 | value = attn.to_v(encoder_hidden_states, *args) |
| 1228 | |
| 1229 | inner_dim = key.shape[-1] |
| 1230 | head_dim = inner_dim // attn.heads |
| 1231 | |
| 1232 | query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 1233 | |
| 1234 | key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 1235 | value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 1236 | |
| 1237 | # the output of sdp = (batch, num_heads, seq_len, head_dim) |
| 1238 | # TODO: add support for attn.scale when we move to Torch 2.1 |
| 1239 | hidden_states = F.scaled_dot_product_attention( |
| 1240 | query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False |
nothing calls this directly
no test coverage detected