r""" Processor for implementing sliced attention. Args: slice_size (`int`, *optional*): The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and `attention_head_dim` must be a multiple of the `slice_size`.
| 1494 | |
| 1495 | |
| 1496 | class SlicedAttnProcessor: |
| 1497 | r""" |
| 1498 | Processor for implementing sliced attention. |
| 1499 | |
| 1500 | Args: |
| 1501 | slice_size (`int`, *optional*): |
| 1502 | The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and |
| 1503 | `attention_head_dim` must be a multiple of the `slice_size`. |
| 1504 | """ |
| 1505 | |
| 1506 | def __init__(self, slice_size: int): |
| 1507 | self.slice_size = slice_size |
| 1508 | |
| 1509 | def __call__( |
| 1510 | self, |
| 1511 | attn: Attention, |
| 1512 | hidden_states: torch.FloatTensor, |
| 1513 | encoder_hidden_states: Optional[torch.FloatTensor] = None, |
| 1514 | attention_mask: Optional[torch.FloatTensor] = None, |
| 1515 | ) -> torch.FloatTensor: |
| 1516 | residual = hidden_states |
| 1517 | |
| 1518 | input_ndim = hidden_states.ndim |
| 1519 | |
| 1520 | if input_ndim == 4: |
| 1521 | batch_size, channel, height, width = hidden_states.shape |
| 1522 | hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) |
| 1523 | |
| 1524 | batch_size, sequence_length, _ = ( |
| 1525 | hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape |
| 1526 | ) |
| 1527 | attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) |
| 1528 | |
| 1529 | if attn.group_norm is not None: |
| 1530 | hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) |
| 1531 | |
| 1532 | query = attn.to_q(hidden_states) |
| 1533 | dim = query.shape[-1] |
| 1534 | query = attn.head_to_batch_dim(query) |
| 1535 | |
| 1536 | if encoder_hidden_states is None: |
| 1537 | encoder_hidden_states = hidden_states |
| 1538 | elif attn.norm_cross: |
| 1539 | encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) |
| 1540 | |
| 1541 | key = attn.to_k(encoder_hidden_states) |
| 1542 | value = attn.to_v(encoder_hidden_states) |
| 1543 | key = attn.head_to_batch_dim(key) |
| 1544 | value = attn.head_to_batch_dim(value) |
| 1545 | |
| 1546 | batch_size_attention, query_tokens, _ = query.shape |
| 1547 | hidden_states = torch.zeros( |
| 1548 | (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype |
| 1549 | ) |
| 1550 | |
| 1551 | for i in range(batch_size_attention // self.slice_size): |
| 1552 | start_idx = i * self.slice_size |
| 1553 | end_idx = (i + 1) * self.slice_size |