Class AttnProcessor

src/models/attention_processor.py:11–156

class AttnProcessor(nn.Module):
    def __init__(
        self,
        hidden_size=None,
        cross_attention_dim=None,
        camera=False,
    ):
        super().__init__()
        self.start_point = 60
        if camera == "True":
            self.start_point = 20

    def __call__(
        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,
    ):
        start_point = self.start_point
        batch_size, sequence_length, _ = hidden_states.shape

        encoder_hidden_states = encoder_hidden_states

        if encoder_hidden_states is not None:
            is_self_attention = False
        else:
            is_self_attention = True

        if attn.group_norm is not None:
            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)

        query = attn.to_q(hidden_states)
        dim = query.shape[-1]

        query = attn.head_to_batch_dim(query)

        if attn.added_kv_proj_dim is not None:
            raise NotImplementedError

        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states

        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)

        if attn.updown == 'up' and iter_cur >= start_point and is_self_attention and not save_kv and (long_context is not None):
            key_ref = torch.cat([attn.buffer_key[iter_cur][c] for c in long_context[0]], dim=0).to('cuda', dtype=query.dtype)
            value_ref = torch.cat([attn.buffer_value[iter_cur][c] for c in long_context[0]], dim=0).to('cuda', dtype=query.dtype)
            if camera_movement:
                target_width = math.sqrt(value_ref.size()[1])
                target_height = target_width