data

app.py

@@ -5,6 +5,33 @@ import torch.nn as nn
 from einops import rearrange
 
 
+class Attn(nn.Module):
+    def __init__(self, dim, dim_text, heads = 16, dim_head = 64):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_q = nn.Conv2d(dim, hidden_dim, 4, bias = False, stride=4)
+        self.to_kv = nn.Linear(dim_text, hidden_dim * 2, bias=False)
+        #self.norm = nn.LayerNorm(dim)
+        self.to_out = nn.Linear(hidden_dim, dim)
+
+    def forward(self, x, text):
+        b, c, h, w = x.shape
+        kv = self.to_kv(text).chunk(2, dim = -1)
+        k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), kv)
+        q = self.to_q(x)
+        q = rearrange(q, 'b (h c) x y -> b h (x y) c', h=self.heads)
+
+        #attn = torch.einsum('bhnd,bhed->bhne',q,k) * self.scale
+        attn = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+        attn = attn.softmax(dim=-1)
+        #print(attn.shape)
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        #print(out.shape)
+        return self.to_out(out)
+
 class RangeModel(nn.Module):
     def __init__(self):
         super(RangeModel, self).__init__()