Create src/linfusion/attention.py

src/linfusion/attention.py

@@ -0,0 +1,94 @@
+import torch
+from diffusers.models.attention_processor import Attention
+import torch.nn.functional as F
+
+try:
+    from fla.ops.linear_attn import chunk_linear_attn
+    FLA_ENABLE = True
+except ImportError:
+    print("Warning: FLA is not installed, falling back to default attention.")
+    FLA_ENABLE = False
+
+
+def get_none_linear_projection(query_dim, mid_dim=None):
+    # If mid_dim is None, then the mid_dim is the same as query_dim
+    # If mid_dim is -1, then no non-linear projection is used, and the identity is returned
+    return (
+        torch.nn.Sequential(
+            torch.nn.Linear(query_dim, mid_dim or query_dim),
+            torch.nn.LayerNorm(mid_dim or query_dim),
+            torch.nn.LeakyReLU(inplace=True),
+            torch.nn.Linear(mid_dim or query_dim, query_dim),
+        )
+        if mid_dim != -1
+        else torch.nn.Identity()
+    )
+
+
+class GeneralizedLinearAttention(Attention):
+    def __init__(self, *args, projection_mid_dim=None, **kwargs):
+        """
+        Args:
+            query_dim: the dimension of the query.
+            out_dim: the dimension of the output.
+            dim_head: the dimension of the head. (dim_head * num_heads = query_dim)
+            projection_mid_dim: the dimension of the intermediate layer in the non-linear projection. 
+              If `None`, then the dimension is the same as the query dimension.
+              If `-1`, then no non-linear projection is used, and the identity is returned.
+        """
+        super().__init__(*args, **kwargs)
+        self.add_non_linear_model(projection_mid_dim)
+
+    def from_attention_instance(attention_instance, projection_mid_dim=None):
+        assert isinstance(attention_instance, Attention)
+        new_instance = GeneralizedLinearAttention(128)
+        new_instance.__dict__ = attention_instance.__dict__
+        new_instance.add_non_linear_model(mid_dim = projection_mid_dim)
+        return new_instance
+
+    def add_non_linear_model(self, mid_dim=None, **kwargs):
+        query_dim = self.to_q.weight.shape[0]
+        self.to_q_ = get_none_linear_projection(query_dim, mid_dim, **kwargs)
+        self.to_k_ = get_none_linear_projection(query_dim, mid_dim, **kwargs)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        **kwargs,
+    ):
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        _, sequence_length, _ = hidden_states.shape
+
+        query = self.to_q(hidden_states + self.to_q_(hidden_states))
+        key = self.to_k(encoder_hidden_states + self.to_k_(encoder_hidden_states))
+        value = self.to_v(encoder_hidden_states)
+
+        query = self.head_to_batch_dim(query)
+        key = self.head_to_batch_dim(key)
+        value = self.head_to_batch_dim(value)
+
+        query = F.elu(query) + 1.0
+        key = F.elu(key) + 1.0
+
+        if FLA_ENABLE and False:
+            # TODO: there is a bug in the FLA implementation
+            raise NotImplementedError
+        else:
+            z = query @ key.mean(dim=-2, keepdim=True).transpose(-2, -1) + 1e-4
+            kv = (key.transpose(-2, -1) * (sequence_length**-0.5)) @ (
+                value * (sequence_length**-0.5)
+            )
+            hidden_states = query @ kv / z
+
+        hidden_states = self.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = self.to_out[0](hidden_states)
+        # dropout
+        hidden_states = self.to_out[1](hidden_states)
+
+        return hidden_states