Return target-only DeMemWM DiT outputs

algorithms/dememwm/models/dit.py

@@ -769,6 +769,11 @@ class DiT(nn.Module):
                 frame_memory_pose=frame_memory_pose,
                 image_hw=image_hw,
                 frame_memory_geometry=frame_memory_geometry)  # (N, T, H, W, D)
+        if frame_memory_segments is not None:
+            target_frames = int(frame_memory_segments.get("target", 0))
+            x = x[:, :target_frames]
+            c = c[:, :target_frames]
+            T = target_frames
         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

@@ -129,8 +129,9 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         masks = {"dynamic": torch.ones((1, 1), dtype=torch.bool)}
 
         with torch.no_grad():
-            model(x, t, action_cond, frame_memory_segments=segments, frame_memory_masks=masks)
+            out = model(x, t, action_cond, frame_memory_segments=segments, frame_memory_masks=masks)
 
+        self.assertEqual(tuple(out.shape), (1, 2, 1, 2, 2))
         self.assertIs(spy.calls[0]["frame_memory_segments"], segments)
         self.assertIs(spy.calls[0]["frame_memory_masks"], masks)