Method attention_fn

sat/dit_video_concat.py:354–435

(
        self,
        query_layer,
        key_layer,
        value_layer,
        attention_mask,
        attention_dropout=None,
        log_attention_weights=None,
        scaling_attention_score=True,
        **kwargs,
    )

            return None

    def attention_fn(
        self,
        query_layer,
        key_layer,
        value_layer,
        attention_mask,
        attention_dropout=None,
        log_attention_weights=None,
        scaling_attention_score=True,
        **kwargs,
    ):
        attention_fn_default = HOOKS_DEFAULT["attention_fn"]

        if self.pnp:
            query_layer = self.rotary(query_layer, **kwargs)
            key_layer = self.rotary(key_layer, **kwargs)
            if self.rot_v:
                value_layer = self.rotary(value_layer)
        else:
            query_layer = torch.cat(
                (
                    query_layer[
                        :,
                        :,
                        : kwargs["text_length"],
                    ],
                    self.rotary(
                        query_layer[
                            :,
                            :,
                            kwargs["text_length"] :,
                        ]
                    ),
                ),
                dim=2,
            )
            key_layer = torch.cat(
                (
                    key_layer[
                        :,
                        :,
                        : kwargs["text_length"],
                    ],
                    self.rotary(
                        key_layer[
                            :,
                            :,
                            kwargs["text_length"] :,
                        ]
                    ),
                ),
                dim=2,
            )
            if self.rot_v:
                value_layer = torch.cat(
                    (
                        value_layer[
                            :,