fix: remove flash attention for now

a_unet/blocks.py

@@ -3,8 +3,6 @@
 
 import torch
 import torch.nn.functional as F
-import xformers
-import xformers.ops
 from einops import pack, rearrange, reduce, repeat, unpack
 from torch import Tensor, einsum, nn
 from typing_extensions import TypeGuard
@@ -230,18 +228,28 @@ def ConvNextV2Block(dim: int, channels: int) -> nn.Module:
 
 
 def AttentionBase(features: int, head_features: int, num_heads: int) -> nn.Module:
+ scale = head_features**-0.5
  mid_features = head_features * num_heads
  to_out = nn.Linear(in_features=mid_features, out_features=features, bias=False)
 
  def forward(
  q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None
  ) -> Tensor:
- # Use memory efficient attention
- out = xformers.ops.memory_efficient_attention(q, k, v)
+ h = num_heads
+ # Split heads
+ q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
+ # Compute similarity matrix and add eventual mask
+ sim = einsum("... n d, ... m d -> ... n m", q, k) * scale
+ # Get attention matrix with softmax
+ attn = sim.softmax(dim=-1)
+ # Compute values
+ out = einsum("... n m, ... m d -> ... n d", attn, v)
+ out = rearrange(out, "b h n d -> b n (h d)")
  return to_out(out)
 
  return Module([to_out], forward)
 
+
 def LinearAttentionBase(features: int, head_features: int, num_heads: int) -> nn.Module:
  scale = head_features**-0.5
  mid_features = head_features * num_heads
@@ -262,6 +270,7 @@ def forward(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
 
  return Module([to_out], forward)
 
+
 def FixedEmbedding(max_length: int, features: int):
  embedding = nn.Embedding(max_length, features)
 

setup.py

@@ -3,15 +3,20 @@
 setup(
  name="a-unet",
  packages=find_packages(exclude=[]),
- version="0.0.13",
+ version="0.0.14",
  license="MIT",
  description="A-UNet",
  long_description_content_type="text/markdown",
  author="Flavio Schneider",
  author_email="archinetai@protonmail.com",
  url="https://github.com/archinetai/a-unet",
  keywords=["artificial intelligence", "deep learning"],
- install_requires=["torch>=1.6", "data-science-types>=0.2", "einops>=0.6.0", "xformers>=0.0.13"],
+ install_requires=[
+ "torch>=1.6",
+ "data-science-types>=0.2",
+ "einops>=0.6.0",
+ "xformers>=0.0.13",
+ ],
  classifiers=[
  "Development Status :: 4 - Beta",
  "Intended Audience :: Developers",