Thread DeMemWM frame memory pose metadata

algorithms/dememwm/models/diffusion.py

@@ -171,9 +171,13 @@ class Diffusion(nn.Module):
             except:
                 pass
         t = t.permute(1,0)
+        if frame_memory_segments is not None and frame_memory_pose is not None:
+            # Algorithm preprocessing keeps packed metadata as T_all x B x 5;
+            # DiT receives the same batch-first layout as x, t, and actions.
+            frame_memory_pose = frame_memory_pose.permute(1,0,2)
         model_kwargs = {
             "current_frame": current_frame,
-            "pose_cond": pose_cond,
+            "pose_cond": None if frame_memory_segments is not None else pose_cond,
             "mode": mode,
             "reference_length": reference_length,
             "frame_idx": frame_idx,

tests/test_dememwm_diffusion.py

@@ -106,6 +106,54 @@ class DeMemWMDiffusionTargetOnlyTests(unittest.TestCase):
         self.assertEqual(tuple(x_pred.shape), tuple(x.shape))
         self.assertEqual(tuple(loss.shape), tuple(x.shape))
 
+    def test_frame_memory_pose_is_batch_first_and_separate_from_pose_cond(self):
+        torch.manual_seed(0)
+        diffusion = _make_diffusion(output_frames=2)
+        x = torch.zeros((5, 2, 1, 1, 1), dtype=torch.float32)
+        action_cond = torch.zeros((5, 2, 3), dtype=torch.float32)
+        noise_levels = torch.zeros((5, 2), dtype=torch.long)
+        pose_cond = torch.full((5, 2, 5), 99.0, dtype=torch.float32)
+        frame_memory_pose = torch.arange(5 * 2 * 5, dtype=torch.float32).view(5, 2, 5)
+        image_hw = torch.tensor([[360, 640], [720, 1280]], dtype=torch.long)
+        segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1}
+
+        diffusion(
+            x,
+            action_cond,
+            pose_cond,
+            noise_levels=noise_levels,
+            reference_length=0,
+            frame_memory_segments=segments,
+            frame_memory_pose=frame_memory_pose,
+            image_hw=image_hw,
+        )
+
+        kwargs = diffusion.model.calls[0]["kwargs"]
+        self.assertIsNone(kwargs["pose_cond"])
+        self.assertEqual(tuple(kwargs["frame_memory_pose"].shape), (2, 5, 5))
+        self.assertTrue(torch.equal(kwargs["frame_memory_pose"], frame_memory_pose.permute(1, 0, 2)))
+        self.assertIs(kwargs["image_hw"], image_hw)
+
+    def test_baseline_pose_cond_still_reaches_dit_batch_first(self):
+        torch.manual_seed(0)
+        diffusion = _make_diffusion()
+        x = torch.zeros((5, 2, 1, 1, 1), dtype=torch.float32)
+        action_cond = torch.zeros((5, 2, 3), dtype=torch.float32)
+        noise_levels = torch.zeros((5, 2), dtype=torch.long)
+        pose_cond = torch.arange(5 * 2 * 5, dtype=torch.float32).view(5, 2, 5)
+
+        diffusion(
+            x,
+            action_cond,
+            pose_cond,
+            noise_levels=noise_levels,
+            reference_length=0,
+        )
+
+        kwargs = diffusion.model.calls[0]["kwargs"]
+        self.assertTrue(torch.equal(kwargs["pose_cond"], pose_cond.permute(1, 0, 2)))
+        self.assertNotIn("frame_memory_pose", kwargs)
+
     def test_posterior_and_sample_steps_return_target_frames_for_frame_memory(self):
         torch.manual_seed(0)
         x, action_cond, _, segments = _packed_inputs()