Add DeMemWM model acceptance checks

algorithms/dememwm/models/dit.py

@@ -400,6 +400,7 @@ class SpatioTemporalDiTBlock(nn.Module):
                 frame_memory_geometry,
             )
         elif self.use_memory_attention:
+            reference_length = self.reference_length if reference_length is None else reference_length
             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)
 
             if pose_cond is not None:
@@ -443,12 +444,12 @@ class SpatioTemporalDiTBlock(nn.Module):
                 else:
                     # pose_cond *= 0
                     x = x + gate(self.r_attn(modulate(self.r_norm1(x+pose_cond[:,:,None, None]), r_shift_msa, r_scale_msa), 
-                                    current_frame=current_frame, timestep=timestep,
-                                    is_last_block=is_last_block,
                                     reference_length=reference_length), r_gate_msa)
             else:
-                x = x + gate(self.r_attn(modulate(self.r_norm1(x), r_shift_msa, r_scale_msa), current_frame=current_frame, timestep=timestep,
-                                is_last_block=is_last_block), r_gate_msa)
+                x = x + gate(self.r_attn(
+                    modulate(self.r_norm1(x), r_shift_msa, r_scale_msa),
+                    reference_length=reference_length,
+                ), r_gate_msa)
 
             x = x + gate(self.r_mlp(modulate(self.r_norm2(x), r_shift_mlp, r_scale_mlp)), r_gate_mlp)
 

tests/test_dememwm_diffusion.py

@@ -106,6 +106,31 @@ class DeMemWMDiffusionTargetOnlyTests(unittest.TestCase):
         self.assertEqual(tuple(x_pred.shape), tuple(x.shape))
         self.assertEqual(tuple(loss.shape), tuple(x.shape))
 
+    def test_padded_frame_memory_masks_keep_target_prediction_and_loss_shapes(self):
+        torch.manual_seed(0)
+        diffusion = _make_diffusion(output_frames=2)
+        x, action_cond, noise_levels, segments = _packed_inputs()
+        masks = {
+            "target": torch.ones((1, 2), dtype=torch.bool),
+            "anchor": torch.ones((1, 1), dtype=torch.bool),
+            "dynamic": torch.zeros((1, 1), dtype=torch.bool),
+            "revisit": torch.ones((1, 1), dtype=torch.bool),
+        }
+
+        x_pred, loss = diffusion(
+            x,
+            action_cond,
+            None,
+            noise_levels=noise_levels,
+            reference_length=0,
+            frame_memory_segments=segments,
+            frame_memory_masks=masks,
+        )
+
+        self.assertIs(diffusion.model.calls[0]["kwargs"]["frame_memory_masks"], masks)
+        self.assertEqual(tuple(x_pred.shape), (2, 1, 1, 1, 1))
+        self.assertEqual(tuple(loss.shape), (2, 1, 1, 1, 1))
+
     def test_frame_memory_pose_is_batch_first_and_separate_from_pose_cond(self):
         torch.manual_seed(0)
         diffusion = _make_diffusion(output_frames=2)
@@ -198,6 +223,19 @@ class DeMemWMDiffusionTargetOnlyTests(unittest.TestCase):
         self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1))
         self.assertEqual(tuple(ddim.shape), (2, 1, 1, 1, 1))
 
+        diffusion = _make_diffusion(output_frames=2)
+        sample = diffusion.sample_step(
+            x,
+            action_cond,
+            None,
+            curr_noise_level=torch.tensor([[3], [2], [0], [0], [0]], dtype=torch.long),
+            next_noise_level=torch.tensor([[2], [1], [0], [0], [0]], dtype=torch.long),
+            reference_length=0,
+            frame_memory_segments=segments,
+        )
+        self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1))
+        self.assertEqual(tuple(sample.shape), (2, 1, 1, 1, 1))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_dememwm_temporal_attention.py

@@ -61,6 +61,19 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         expected = torch.tensor([10.0, 15.0, 20.0, 100.0, 1.0, 3.0, 200.0])
         self.assertTrue(torch.allclose(out.flatten(), expected))
 
+    def test_frame_memory_temporal_targets_do_not_depend_on_appended_memory(self):
+        attn = _averaging_temporal_attention()
+        segments = {"target": 3, "anchor": 1, "dynamic": 2, "revisit": 1}
+        x = torch.tensor([10.0, 20.0, 30.0, 100.0, 1.0, 3.0, 200.0]).view(1, 7, 1, 1, 1)
+        changed_memory = torch.tensor([10.0, 20.0, 30.0, -500.0, 700.0, 900.0, -300.0]).view(1, 7, 1, 1, 1)
+
+        out = attn(x, frame_memory_segments=segments)
+        changed_out = attn(changed_memory, frame_memory_segments=segments)
+
+        expected_target = torch.tensor([10.0, 15.0, 20.0]).view(1, 3, 1, 1, 1)
+        self.assertTrue(torch.allclose(out[:, :3], expected_target))
+        self.assertTrue(torch.allclose(out[:, :3], changed_out[:, :3]))
+
     def test_block_threads_frame_memory_metadata_to_temporal_attention(self):
         class ZeroAttention(nn.Module):
             def forward(self, x):
@@ -135,6 +148,52 @@ 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_baseline_and_packed_frame_memory_output_lengths_with_real_blocks(self):
+        torch.manual_seed(0)
+        model = DiT(
+            input_h=2,
+            input_w=2,
+            patch_size=1,
+            in_channels=1,
+            hidden_size=8,
+            depth=1,
+            num_heads=1,
+            mlp_ratio=1.0,
+            action_cond_dim=3,
+            pose_cond_dim=0,
+            reference_length=1,
+            use_memory_attention=True,
+        )
+        model.eval()
+
+        x = torch.randn((1, 3, 1, 2, 2))
+        t = torch.zeros((1, 3), dtype=torch.long)
+        action_cond = torch.zeros((1, 3, 3))
+        with torch.no_grad():
+            baseline = model(x, t, action_cond, reference_length=1)
+        self.assertEqual(tuple(baseline.shape), (1, 3, 1, 2, 2))
+
+        packed = torch.randn((1, 5, 1, 2, 2))
+        packed_t = torch.zeros((1, 5), dtype=torch.long)
+        packed_action = torch.zeros((1, 5, 3))
+        segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1}
+        masks = {
+            "target": torch.ones((1, 2), dtype=torch.bool),
+            "anchor": torch.ones((1, 1), dtype=torch.bool),
+            "dynamic": torch.zeros((1, 1), dtype=torch.bool),
+            "revisit": torch.ones((1, 1), dtype=torch.bool),
+        }
+        with torch.no_grad():
+            packed_out = model(
+                packed,
+                packed_t,
+                packed_action,
+                reference_length=1,
+                frame_memory_segments=segments,
+                frame_memory_masks=masks,
+            )
+        self.assertEqual(tuple(packed_out.shape), (1, 2, 1, 2, 2))
+
     def test_dit_builds_geometry_cache_once_and_threads_to_reference_attention(self):
         class ZeroAttention(nn.Module):
             def forward(self, x):