Add DeMemWM geometry-aware memory attention

algorithms/dememwm/models/dit.py

@@ -150,11 +150,39 @@ class FrameMemoryReferenceAttention(nn.Module):
         self.to_q = nn.Linear(hidden_size, hidden_size, bias=False)
         self.to_k = nn.Linear(hidden_size, hidden_size, bias=False)
         self.to_v = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.query_pose_proj = nn.Linear(6, hidden_size, bias=True)
+        self.key_pose_proj = nn.Linear(6, hidden_size, bias=True)
         self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True)
 
     def _split_heads(self, x):
         return rearrange(x, "b n (h d) -> b h n d", h=self.num_heads)
 
+    def _geometry_rays(
+        self,
+        geometry_cache,
+        B,
+        T_target,
+        T_memory,
+        H,
+        W,
+        device,
+        dtype,
+    ):
+        if not isinstance(geometry_cache, dict):
+            return None, None
+        query_rays = geometry_cache.get("query_rays")
+        relative_rays = geometry_cache.get("relative_rays")
+        if not torch.is_tensor(query_rays) or not torch.is_tensor(relative_rays):
+            return None, None
+        if tuple(query_rays.shape) != (B, T_target, H, W, 6):
+            return None, None
+        if tuple(relative_rays.shape) != (B, T_target, T_memory, H, W, 6):
+            return None, None
+        return (
+            query_rays.to(device=device, dtype=dtype),
+            relative_rays.to(device=device, dtype=dtype),
+        )
+
     def forward(self, target_hidden, memory_hidden, memory_mask=None, geometry_cache=None):
         B, T_target, H, W, D = target_hidden.shape
         if memory_hidden is None or int(memory_hidden.shape[1]) == 0:
@@ -162,18 +190,34 @@ class FrameMemoryReferenceAttention(nn.Module):
 
         T_memory = memory_hidden.shape[1]
         P = H * W
-        q = rearrange(target_hidden, "b t h w d -> (b t) (h w) d")
+        query_rays, relative_rays = self._geometry_rays(
+            geometry_cache,
+            B,
+            T_target,
+            T_memory,
+            H,
+            W,
+            target_hidden.device,
+            target_hidden.dtype,
+        )
+        q_input = target_hidden
+        if query_rays is not None:
+            # Values stay visual-only; Plucker features only steer Q/K matching.
+            q_input = q_input + self.query_pose_proj(query_rays)
+        q = rearrange(q_input, "b t h w d -> (b t) (h w) d")
         q = self._split_heads(self.to_q(q))
 
-        memory_tokens = rearrange(memory_hidden, "b m h w d -> b (m h w) d")
-        k = self._split_heads(self.to_k(memory_tokens))
-        v = self._split_heads(self.to_v(memory_tokens))
-
-        # A zero dummy key keeps all-padded samples finite; the projected delta
-        # is masked back to zero below so padded memory cannot contribute.
-        dummy = k.new_zeros((B, self.num_heads, 1, self.head_dim))
-        k = torch.cat([k, dummy], dim=2)
-        v = torch.cat([v, dummy], dim=2)
+        memory_tokens = rearrange(memory_hidden, "b m h w d -> b m (h w) d")
+        memory_flat = rearrange(memory_tokens, "b m p d -> b (m p) d")
+        if relative_rays is None:
+            k = self._split_heads(self.to_k(memory_flat))
+        else:
+            key_pose = self.key_pose_proj(
+                rearrange(relative_rays, "b t m h w r -> b t m (h w) r")
+            )
+            k = self.to_k(memory_tokens[:, None] + key_pose)
+            k = self._split_heads(rearrange(k, "b t m p d -> (b t) (m p) d"))
+        v = self._split_heads(self.to_v(memory_flat))
 
         if memory_mask is None:
             frame_valid = torch.ones((B, T_memory), device=target_hidden.device, dtype=torch.bool)
@@ -184,13 +228,22 @@ class FrameMemoryReferenceAttention(nn.Module):
         key_valid = torch.cat([key_valid, ~active[:, None]], dim=1)
 
         row_batch = torch.arange(B, device=target_hidden.device).repeat_interleave(T_target)
+        if relative_rays is None:
+            k = k.index_select(0, row_batch)
+        # A zero dummy key keeps all-padded samples finite; the projected delta
+        # is masked back to zero below so padded memory cannot contribute.
+        dummy_k = k.new_zeros((k.shape[0], self.num_heads, 1, self.head_dim))
+        k = torch.cat([k, dummy_k], dim=2)
+        dummy_v = v.new_zeros((B, self.num_heads, 1, self.head_dim))
+        v = torch.cat([v, dummy_v], dim=2).index_select(0, row_batch)
+
         attn_bias = torch.zeros((B * T_target, 1, 1, key_valid.shape[1]), device=q.device, dtype=q.dtype)
         attn_bias = attn_bias.masked_fill(~key_valid[row_batch][:, None, None], float("-inf"))
 
         x = F.scaled_dot_product_attention(
             query=q.contiguous(),
-            key=k.index_select(0, row_batch).contiguous(),
-            value=v.index_select(0, row_batch).contiguous(),
+            key=k.contiguous(),
+            value=v.contiguous(),
             attn_mask=attn_bias,
         )
         x = rearrange(x, "n h p d -> n p (h d)")
@@ -658,16 +711,18 @@ class DiT(nn.Module):
         else:
             pc = None
         
-        frame_memory_geometry = self._build_frame_memory_geometry(
-            frame_memory_segments,
-            frame_memory_masks,
-            frame_memory_pose,
-            image_hw,
-            x.shape[2],
-            x.shape[3],
-            x.device,
-            x.dtype,
-        )
+        frame_memory_geometry = None
+        if self.use_memory_attention and self.use_plucker:
+            frame_memory_geometry = self._build_frame_memory_geometry(
+                frame_memory_segments,
+                frame_memory_masks,
+                frame_memory_pose,
+                image_hw,
+                x.shape[2],
+                x.shape[3],
+                x.device,
+                x.dtype,
+            )
 
         for i, block in enumerate(self.blocks):
             x = block(x, c, current_frame=current_frame, timestep=t, is_last_block= (i+1 == len(self.blocks)), 

tests/test_dememwm_temporal_attention.py

@@ -264,6 +264,7 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
             action_cond_dim=3,
             pose_cond_dim=0,
             reference_length=0,
+            use_plucker=True,
             use_memory_attention=True,
         )
         spies = []
@@ -328,5 +329,48 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         self.assertTrue(torch.allclose(second_only, torch.full_like(second_only, 10.0)))
         self.assertTrue(torch.equal(empty, torch.zeros_like(empty)))
 
+    def test_frame_memory_reference_attention_uses_geometry_qk(self):
+        attn = FrameMemoryReferenceAttention(hidden_size=2, num_heads=1)
+        with torch.no_grad():
+            attn.to_q.weight.zero_()
+            attn.to_q.weight[0, 0] = 1.0
+            attn.to_k.weight.zero_()
+            attn.to_k.weight[0, 0] = 1.0
+            attn.to_v.weight.zero_()
+            attn.to_v.weight[0, 1] = 1.0
+            attn.query_pose_proj.weight.zero_()
+            attn.query_pose_proj.bias.zero_()
+            attn.query_pose_proj.weight[0, 5] = 1.0
+            attn.key_pose_proj.weight.zero_()
+            attn.key_pose_proj.bias.zero_()
+            attn.key_pose_proj.weight[0, 5] = 1.0
+            attn.out_proj.weight.zero_()
+            attn.out_proj.weight[0, 0] = 1.0
+            attn.out_proj.bias.zero_()
+
+        target = torch.zeros((1, 1, 1, 1, 2))
+        memory = torch.zeros((1, 2, 1, 1, 2))
+        memory[0, :, 0, 0, 1] = torch.tensor([1.0, 10.0])
+        query_rays = torch.zeros((1, 1, 1, 1, 6))
+        query_rays[..., 5] = 1.0
+        relative_a = torch.zeros((1, 1, 2, 1, 1, 6))
+        relative_b = torch.zeros_like(relative_a)
+        relative_a[0, 0, 0, 0, 0, 5] = 2.0
+        relative_b[0, 0, 1, 0, 0, 5] = 2.0
+
+        out_a = attn(
+            target,
+            memory,
+            geometry_cache={"query_rays": query_rays, "relative_rays": relative_a},
+        )
+        out_b = attn(
+            target,
+            memory,
+            geometry_cache={"query_rays": query_rays, "relative_rays": relative_b},
+        )
+
+        self.assertLess(out_a[0, 0, 0, 0, 0].item(), out_b[0, 0, 0, 0, 0].item())
+        self.assertFalse(torch.allclose(out_a, out_b))
+
 if __name__ == "__main__":
     unittest.main()