(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
iter_cur=None,
save_kv=None,
source_masks=None,
target_masks=None,
camera_movement=True,
long_context=None,
inference_num=50,
)
| 21 | self.start_point = 20 |
| 22 | |
| 23 | def __call__( |
| 24 | self, |
| 25 | attn, |
| 26 | hidden_states, |
| 27 | encoder_hidden_states=None, |
| 28 | attention_mask=None, |
| 29 | iter_cur=None, |
| 30 | save_kv=None, |
| 31 | source_masks=None, |
| 32 | target_masks=None, |
| 33 | camera_movement=True, |
| 34 | long_context=None, |
| 35 | inference_num=50, |
| 36 | ): |
| 37 | start_point = self.start_point |
| 38 | batch_size, sequence_length, _ = hidden_states.shape |
| 39 | |
| 40 | encoder_hidden_states = encoder_hidden_states |
| 41 | |
| 42 | if encoder_hidden_states is not None: |
| 43 | is_self_attention = False |
| 44 | else: |
| 45 | is_self_attention = True |
| 46 | |
| 47 | if attn.group_norm is not None: |
| 48 | hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) |
| 49 | |
| 50 | query = attn.to_q(hidden_states) |
| 51 | dim = query.shape[-1] |
| 52 | |
| 53 | query = attn.head_to_batch_dim(query) |
| 54 | |
| 55 | if attn.added_kv_proj_dim is not None: |
| 56 | raise NotImplementedError |
| 57 | |
| 58 | encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states |
| 59 | |
| 60 | key = attn.to_k(encoder_hidden_states) |
| 61 | value = attn.to_v(encoder_hidden_states) |
| 62 | |
| 63 | if attn.updown == 'up' and iter_cur >= start_point and is_self_attention and not save_kv and (long_context is not None): |
| 64 | key_ref = torch.cat([attn.buffer_key[iter_cur][c] for c in long_context[0]], dim=0).to('cuda', dtype=query.dtype) |
| 65 | value_ref = torch.cat([attn.buffer_value[iter_cur][c] for c in long_context[0]], dim=0).to('cuda', dtype=query.dtype) |
| 66 | if camera_movement: |
| 67 | target_width = math.sqrt(value_ref.size()[1]) |
| 68 | target_height = target_width |
| 69 | source_masks = F.interpolate(source_masks, size=(int(target_height), int(target_width)), mode="nearest") |
| 70 | background_source_masks = 1-source_masks |
| 71 | background_source_masks = rearrange(background_source_masks, "f c h w -> f (h w) c") |
| 72 | key_ref = key_ref * background_source_masks |
| 73 | value_ref = value_ref * background_source_masks |
| 74 | key_ref = key_ref.repeat(2, 1, 1) |
| 75 | value_ref = value_ref.repeat(2, 1, 1) |
| 76 | target_masks = F.interpolate(target_masks, size=(int(target_height), int(target_width)), mode="nearest") |
| 77 | |
| 78 | foreground_target_masks = rearrange(target_masks, "f c h w -> f (h w) c") |
| 79 | foreground_target_masks = torch.cat([foreground_target_masks] * 2, dim=0) |
| 80 | key = key * foreground_target_masks |
nothing calls this directly
no test coverage detected