r""" Processor for performing attention-related computations with extra learnable key and value matrices for the text encoder.
| 1260 | |
| 1261 | |
| 1262 | class AttnAddedKVProcessor: |
| 1263 | r""" |
| 1264 | Processor for performing attention-related computations with extra learnable key and value matrices for the text |
| 1265 | encoder. |
| 1266 | """ |
| 1267 | |
| 1268 | def __call__( |
| 1269 | self, |
| 1270 | attn: Attention, |
| 1271 | hidden_states: torch.Tensor, |
| 1272 | encoder_hidden_states: Optional[torch.Tensor] = None, |
| 1273 | attention_mask: Optional[torch.Tensor] = None, |
| 1274 | *args, |
| 1275 | **kwargs, |
| 1276 | ) -> torch.Tensor: |
| 1277 | if len(args) > 0 or kwargs.get("scale", None) is not None: |
| 1278 | deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." |
| 1279 | deprecate("scale", "1.0.0", deprecation_message) |
| 1280 | |
| 1281 | residual = hidden_states |
| 1282 | |
| 1283 | hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2) |
| 1284 | batch_size, sequence_length, _ = hidden_states.shape |
| 1285 | |
| 1286 | attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) |
| 1287 | |
| 1288 | if encoder_hidden_states is None: |
| 1289 | encoder_hidden_states = hidden_states |
| 1290 | elif attn.norm_cross: |
| 1291 | encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) |
| 1292 | |
| 1293 | hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) |
| 1294 | |
| 1295 | query = attn.to_q(hidden_states) |
| 1296 | query = attn.head_to_batch_dim(query) |
| 1297 | |
| 1298 | encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) |
| 1299 | encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) |
| 1300 | encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj) |
| 1301 | encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj) |
| 1302 | |
| 1303 | if not attn.only_cross_attention: |
| 1304 | key = attn.to_k(hidden_states) |
| 1305 | value = attn.to_v(hidden_states) |
| 1306 | key = attn.head_to_batch_dim(key) |
| 1307 | value = attn.head_to_batch_dim(value) |
| 1308 | key = torch.cat([encoder_hidden_states_key_proj, key], dim=1) |
| 1309 | value = torch.cat([encoder_hidden_states_value_proj, value], dim=1) |
| 1310 | else: |
| 1311 | key = encoder_hidden_states_key_proj |
| 1312 | value = encoder_hidden_states_value_proj |
| 1313 | |
| 1314 | attention_probs = attn.get_attention_scores(query, key, attention_mask) |
| 1315 | hidden_states = torch.bmm(attention_probs, value) |
| 1316 | hidden_states = attn.batch_to_head_dim(hidden_states) |
| 1317 | |
| 1318 | # linear proj |
| 1319 | hidden_states = attn.to_out[0](hidden_states) |