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