| 427 | |
| 428 | def register_attention_control(model, controller): |
| 429 | def ca_forward(self, place_in_unet): |
| 430 | def forward( |
| 431 | hidden_states: torch.FloatTensor, |
| 432 | encoder_hidden_states: Optional[torch.FloatTensor] = None, |
| 433 | attention_mask: Optional[torch.FloatTensor] = None, |
| 434 | temb: Optional[torch.FloatTensor] = None, |
| 435 | # scale: float = 1.0, |
| 436 | ): |
| 437 | if self.spatial_norm is not None: |
| 438 | hidden_states = self.spatial_norm(hidden_states, temb) |
| 439 | |
| 440 | batch_size, sequence_length, _ = ( |
| 441 | hidden_states.shape |
| 442 | if encoder_hidden_states is None |
| 443 | else encoder_hidden_states.shape |
| 444 | ) |
| 445 | |
| 446 | if attention_mask is not None: |
| 447 | attention_mask = self.prepare_attention_mask( |
| 448 | attention_mask, sequence_length, batch_size |
| 449 | ) |
| 450 | # scaled_dot_product_attention expects attention_mask shape to be |
| 451 | # (batch, heads, source_length, target_length) |
| 452 | attention_mask = attention_mask.view( |
| 453 | batch_size, self.heads, -1, attention_mask.shape[-1] |
| 454 | ) # type: ignore |
| 455 | |
| 456 | if self.group_norm is not None: |
| 457 | hidden_states = self.group_norm( |
| 458 | hidden_states.transpose(1, 2) |
| 459 | ).transpose(1, 2) |
| 460 | |
| 461 | query = self.to_q(hidden_states) |
| 462 | |
| 463 | is_cross = encoder_hidden_states is not None |
| 464 | if encoder_hidden_states is None: |
| 465 | encoder_hidden_states = hidden_states |
| 466 | elif self.norm_cross: |
| 467 | encoder_hidden_states = self.norm_encoder_hidden_states( |
| 468 | encoder_hidden_states |
| 469 | ) |
| 470 | key = self.to_k(encoder_hidden_states) |
| 471 | value = self.to_v(encoder_hidden_states) |
| 472 | |
| 473 | def reshape_heads_to_batch_dim(tensor): |
| 474 | batch_size, seq_len, dim = tensor.shape |
| 475 | head_size = self.heads |
| 476 | tensor = tensor.reshape( |
| 477 | batch_size, seq_len, head_size, dim // head_size |
| 478 | ) |
| 479 | tensor = tensor.permute(0, 2, 1, 3).reshape( |
| 480 | batch_size * head_size, seq_len, dim // head_size |
| 481 | ) |
| 482 | return tensor |
| 483 | |
| 484 | query = reshape_heads_to_batch_dim(query) |
| 485 | key = reshape_heads_to_batch_dim(key) |
| 486 | value = reshape_heads_to_batch_dim(value) |