ready cross attention for attending to and from latents

2

Author: lucidrains

rin_pytorch/rin_pytorch.py

@@ -137,18 +137,26 @@ def __init__(
         self.scale = dim_head ** -0.5
         self.heads = heads
         hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
+        self.to_q = nn.Conv2d(dim, hidden_dim, 1, bias = False)
+        self.to_kv = nn.Conv2d(dim, hidden_dim * 2, 1, bias = False)
         self.to_out = nn.Conv2d(hidden_dim, dim, 1)
 
-    def forward(self, x):
+    def forward(
+        self,
+        x,
+        context = None
+    ):
         b, c, h, w = x.shape
-        qkv = self.to_qkv(x).chunk(3, dim = 1)
+        context = default(context, x)
+
+        qkv = (self.to_q(x), *self.to_kv(context).chunk(2, dim = 1))
         q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv)
 
         q = q * self.scale
 
         sim = einsum('b h d i, b h d j -> b h i j', q, k)
         attn = sim.softmax(dim = -1)
+
         out = einsum('b h i j, b h d j -> b h i d', attn, v)
         out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = h, y = w)
         return self.to_out(out)