Build DeMemWM frame memory geometry cache

algorithms/dememwm/models/dit.py

@@ -155,7 +155,7 @@ class FrameMemoryReferenceAttention(nn.Module):
     def _split_heads(self, x):
         return rearrange(x, "b n (h d) -> b h n d", h=self.num_heads)
 
-    def forward(self, target_hidden, memory_hidden, memory_mask=None):
+    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:
             return target_hidden.new_zeros(target_hidden.shape)
@@ -321,7 +321,7 @@ class SpatioTemporalDiTBlock(nn.Module):
             return None
         return frame_memory_masks[stream_name].to(device=stream_hidden.device, dtype=torch.bool)
 
-    def _apply_frame_memory_reference_attention(self, x, c, frame_memory_segments, frame_memory_masks):
+    def _apply_frame_memory_reference_attention(self, x, c, frame_memory_segments, frame_memory_masks, frame_memory_geometry=None):
         x_target, x_anchor, x_dynamic, x_revisit = self._split_frame_memory(x, frame_memory_segments)
         if int(x_target.shape[1]) == 0:
             return x
@@ -340,7 +340,8 @@ class SpatioTemporalDiTBlock(nn.Module):
             if int(stream_hidden.shape[1]) == 0:
                 continue
             stream_mask = self._frame_memory_stream_mask(frame_memory_masks, stream_name, stream_hidden)
-            deltas.append(r_attn(attn_target, stream_attn, stream_mask))
+            stream_geometry = None if frame_memory_geometry is None else frame_memory_geometry.get(stream_name)
+            deltas.append(r_attn(attn_target, stream_attn, stream_mask, geometry_cache=stream_geometry))
             if stream_mask is None:
                 active_stream_count = active_stream_count + 1
             else:
@@ -362,7 +363,7 @@ class SpatioTemporalDiTBlock(nn.Module):
     def forward(self, x, c, current_frame=None, timestep=None, is_last_block=False, 
         pose_cond=None, mode="training", c_action_cond=None, reference_length=None,
         frame_memory_segments=None, frame_memory_masks=None, frame_memory_pose=None,
-        image_hw=None):
+        image_hw=None, frame_memory_geometry=None):
         B, T, H, W, D = x.shape
 
         # spatial block
@@ -396,6 +397,7 @@ class SpatioTemporalDiTBlock(nn.Module):
                 c,
                 frame_memory_segments,
                 frame_memory_masks,
+                frame_memory_geometry,
             )
         elif self.use_memory_attention:
             r_shift_msa, r_scale_msa, r_gate_msa, r_shift_mlp, r_scale_mlp, r_gate_mlp = self.r_adaLN_modulation(c).chunk(6, dim=-1)
@@ -611,6 +613,84 @@ class DiT(nn.Module):
         imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
         return imgs
 
+    def _frame_memory_pose_to_c2w(self, pose):
+        B, T = pose.shape[:2]
+        x, y, z, pitch, yaw = pose.unbind(dim=-1)
+        pitch, yaw = torch.deg2rad(pitch), torch.deg2rad(yaw)
+        cp, sp = torch.cos(pitch), torch.sin(pitch)
+        cy, sy = torch.cos(yaw), torch.sin(yaw)
+        one, zero = torch.ones_like(pitch), torch.zeros_like(pitch)
+        r_pitch = torch.stack((one, zero, zero, zero, cp, -sp, zero, sp, cp), dim=-1).reshape(B, T, 3, 3)
+        r_yaw = torch.stack((cy, zero, sy, zero, one, zero, -sy, zero, cy), dim=-1).reshape(B, T, 3, 3)
+        c2w = torch.eye(4, device=pose.device, dtype=pose.dtype).view(1, 1, 4, 4).repeat(B, T, 1, 1)
+        c2w[:, :, :3, :3] = torch.matmul(r_yaw, r_pitch)
+        c2w[:, :, :3, 3] = torch.stack((x, y, z), dim=-1)
+        return c2w
+
+    def _build_frame_memory_geometry(self, frame_memory_segments, frame_memory_masks, frame_memory_pose, image_hw, grid_h, grid_w, device, dtype):
+        if frame_memory_segments is None or frame_memory_pose is None or image_hw is None:
+            return None
+
+        pose = frame_memory_pose.to(device=device, dtype=dtype)
+        B = pose.shape[0]
+        target_frames = int(frame_memory_segments.get("target", 0))
+        image_hw = image_hw if torch.is_tensor(image_hw) else torch.as_tensor(image_hw)
+        image_hw = image_hw.to(device=device, dtype=dtype)
+        if image_hw.ndim == 1:
+            image_hw = image_hw.unsqueeze(0)
+        if image_hw.shape[0] == 1 and B != 1:
+            image_hw = image_hw.expand(B, -1)
+        image_hw = image_hw.contiguous()
+        image_h, image_w = image_hw[:, 0].view(B, 1, 1), image_hw[:, 1].view(B, 1, 1)
+
+        # Token grid sets position count; image_hw keeps camera scaling in original image coordinates.
+        y = (torch.arange(grid_h, device=device, dtype=dtype).view(1, grid_h, 1) + 0.5) * (image_h / grid_h)
+        x = (torch.arange(grid_w, device=device, dtype=dtype).view(1, 1, grid_w) + 0.5) * (image_w / grid_w)
+        fx, fy = 0.35 * image_w, 0.35 * image_h
+        cx, cy = 0.5 * image_w, 0.5 * image_h
+        directions = torch.stack(
+            (
+                (-(x - cx) / fx).expand(-1, grid_h, -1),
+                (-(y - cy) / fy).expand(-1, -1, grid_w),
+                torch.ones((B, grid_h, grid_w), device=device, dtype=dtype),
+            ),
+            dim=-1,
+        )
+        directions = F.normalize(directions, dim=-1)
+        query_rays = torch.cat([torch.zeros_like(directions), directions], dim=-1)[:, None].expand(-1, target_frames, -1, -1, -1)
+        intrinsics = torch.stack((fx[:, 0, 0], fy[:, 0, 0], cx[:, 0, 0], cy[:, 0, 0]), dim=-1)
+
+        c2w = self._frame_memory_pose_to_c2w(pose)
+        target_pose = pose[:, :target_frames]
+        target_c2w = c2w[:, :target_frames]
+        target_rot = target_c2w[:, :, :3, :3].transpose(-1, -2)
+        target_t = target_c2w[:, :, :3, 3]
+        target_w2c = torch.eye(4, device=device, dtype=dtype).view(1, 1, 4, 4).repeat(B, target_frames, 1, 1)
+        target_w2c[:, :, :3, :3] = target_rot
+        target_w2c[:, :, :3, 3] = -torch.matmul(target_rot, target_t.unsqueeze(-1)).squeeze(-1)
+
+        cache = {"query_rays": query_rays, "target_pose": target_pose, "image_hw": image_hw, "intrinsics": intrinsics}
+        cursor = target_frames
+        for stream_name in ("anchor", "dynamic", "revisit"):
+            stream_len = int(frame_memory_segments.get(stream_name, 0))
+            stream_pose = pose[:, cursor:cursor + stream_len]
+            relative_c2w = torch.matmul(target_w2c[:, :, None], c2w[:, None, cursor:cursor + stream_len])
+            rays_d = torch.einsum("bhwj,btmij->btmhwi", directions, relative_c2w[..., :3, :3])
+            rays_d = F.normalize(rays_d, dim=-1)
+            rays_o = relative_c2w[..., :3, 3][:, :, :, None, None, :].expand(-1, -1, -1, grid_h, grid_w, -1)
+            stream_mask = None if frame_memory_masks is None or frame_memory_masks.get(stream_name) is None else frame_memory_masks[stream_name].to(device=device, dtype=torch.bool)
+            cache[stream_name] = {
+                "query_rays": query_rays,
+                "relative_pose": stream_pose[:, None] - target_pose[:, :, None],
+                "relative_c2w": relative_c2w,
+                "relative_rays": torch.cat([torch.linalg.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1),
+                "image_hw": image_hw,
+                "intrinsics": intrinsics,
+                "mask": stream_mask,
+            }
+            cursor += stream_len
+        return cache
+
     def forward(self, x, t, action_cond=None, pose_cond=None, current_frame=None, mode=None, 
                     reference_length=None, frame_idx=None, frame_memory_segments=None,
                     frame_memory_masks=None, frame_memory_pose=None, image_hw=None):
@@ -670,13 +750,25 @@ 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,
+        )
+
         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)), 
                 pose_cond=pc, mode=mode, c_action_cond=c_action_cond, reference_length=reference_length,
                 frame_memory_segments=frame_memory_segments,
                 frame_memory_masks=frame_memory_masks,
                 frame_memory_pose=frame_memory_pose,
-                image_hw=image_hw)  # (N, T, H, W, D)
+                image_hw=image_hw,
+                frame_memory_geometry=frame_memory_geometry)  # (N, T, H, W, D)
         x = self.final_layer(x, c) # (N, T, H, W, patch_size ** 2 * out_channels)
         # unpatchify
         x = rearrange(x, "b t h w d -> (b t) h w d")

tests/test_dememwm_temporal_attention.py

@@ -134,6 +134,80 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         self.assertIs(spy.calls[0]["frame_memory_segments"], segments)
         self.assertIs(spy.calls[0]["frame_memory_masks"], masks)
 
+    def test_dit_builds_geometry_cache_once_and_threads_to_reference_attention(self):
+        class ZeroAttention(nn.Module):
+            def forward(self, x):
+                return torch.zeros_like(x)
+
+        class ZeroTemporalAttention(nn.Module):
+            def forward(self, x, frame_memory_segments=None, frame_memory_masks=None):
+                return torch.zeros_like(x)
+
+        class SpyReferenceAttention(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.geometry = []
+
+            def forward(self, target_hidden, memory_hidden, memory_mask=None, geometry_cache=None):
+                self.geometry.append(geometry_cache)
+                return torch.zeros_like(target_hidden)
+
+        model = DiT(
+            input_h=2,
+            input_w=2,
+            patch_size=1,
+            in_channels=1,
+            hidden_size=8,
+            depth=2,
+            num_heads=1,
+            mlp_ratio=1.0,
+            action_cond_dim=3,
+            pose_cond_dim=0,
+            reference_length=0,
+            use_memory_attention=True,
+        )
+        spies = []
+        for block in model.blocks:
+            block.s_attn = ZeroAttention()
+            block.t_attn = ZeroTemporalAttention()
+            anchor = SpyReferenceAttention()
+            dynamic = SpyReferenceAttention()
+            revisit = SpyReferenceAttention()
+            block.r_attn_anchor = anchor
+            block.r_attn_dynamic = dynamic
+            block.r_attn_revisit = revisit
+            spies.append((anchor, dynamic, revisit))
+        model.eval()
+
+        x = torch.zeros((1, 5, 1, 2, 2))
+        t = torch.zeros((1, 5), dtype=torch.long)
+        action_cond = torch.zeros((1, 5, 3))
+        frame_memory_pose = torch.zeros((1, 5, 5))
+        frame_memory_pose[0, :, 0] = torch.arange(5, dtype=torch.float32)
+        image_hw = torch.tensor([[4, 8]], dtype=torch.long)
+        segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1}
+        masks = {"dynamic": torch.tensor([[False]])}
+
+        with torch.no_grad():
+            model(
+                x,
+                t,
+                action_cond,
+                frame_memory_segments=segments,
+                frame_memory_masks=masks,
+                frame_memory_pose=frame_memory_pose,
+                image_hw=image_hw,
+            )
+
+        anchor_cache = spies[0][0].geometry[0]
+        self.assertIs(anchor_cache, spies[1][0].geometry[0])
+        self.assertEqual(tuple(anchor_cache["query_rays"].shape), (1, 2, 2, 2, 6))
+        self.assertEqual(tuple(anchor_cache["relative_rays"].shape), (1, 2, 1, 2, 2, 6))
+        self.assertTrue(torch.equal(anchor_cache["image_hw"], torch.tensor([[4.0, 8.0]])))
+        self.assertTrue(torch.allclose(anchor_cache["intrinsics"], torch.tensor([[2.8, 1.4, 4.0, 2.0]])))
+        self.assertTrue(torch.equal(anchor_cache["relative_pose"][0, :, 0, 0], torch.tensor([2.0, 1.0])))
+        self.assertTrue(torch.equal(spies[0][1].geometry[0]["mask"], masks["dynamic"]))
+
     def test_frame_memory_reference_attention_masks_padded_keys(self):
         attn = FrameMemoryReferenceAttention(hidden_size=1, num_heads=1)
         with torch.no_grad():