Prune inactive DeMemWM inference memory streams

algorithms/dememwm/df_video.py

@@ -577,6 +577,72 @@ def _apply_memory_route_masks(frame_memory_masks, query_noise_levels, cfg, diffu
     return routed
 
 
+def _pack_active_inference_memory_streams(
+    target_latents,
+    target_conditions,
+    target_poses,
+    target_frame_indices,
+    target_mask,
+    stream_latents_by_key,
+    stream_poses_by_key,
+    stream_frame_indices_by_key,
+    routed_stream_masks,
+    memory_noise_levels_by_key,
+):
+    active_streams = []
+    pruned_streams = []
+    active_latents = [target_latents]
+    active_poses = [target_poses]
+    active_frame_indices = [target_frame_indices]
+    active_frame_memory_segments = {"target": int(target_latents.shape[0])}
+    active_frame_memory_masks = {"target": target_mask}
+    active_memory_noise_levels = []
+    active_memory_length = 0
+
+    for key in _DEMEMWM_STREAM_KEYS:
+        routed_mask = routed_stream_masks[key]
+        if bool(routed_mask.any().item()):
+            active_streams.append(key)
+            active_latents.append(stream_latents_by_key[key])
+            active_poses.append(stream_poses_by_key[key])
+            active_frame_indices.append(stream_frame_indices_by_key[key])
+            active_frame_memory_segments[key] = int(stream_latents_by_key[key].shape[0])
+            active_frame_memory_masks[key] = routed_mask
+            active_memory_noise_levels.append(memory_noise_levels_by_key[key])
+            active_memory_length += int(stream_latents_by_key[key].shape[0])
+        else:
+            pruned_streams.append(key)
+            active_frame_memory_segments[key] = 0
+            active_frame_memory_masks[key] = routed_mask[:, :0]
+
+    active_packed_latents = torch.cat(active_latents, dim=0)
+    active_packed_conditions = torch.cat(
+        [
+            target_conditions,
+            target_conditions.new_zeros((active_memory_length, target_conditions.shape[1], target_conditions.shape[-1])),
+        ],
+        dim=0,
+    )
+    active_frame_memory_pose = torch.cat(active_poses, dim=0)
+    active_frame_indices = torch.cat(active_frame_indices, dim=0)
+    if active_memory_noise_levels:
+        active_memory_noise_levels = torch.cat(active_memory_noise_levels, dim=0)
+    else:
+        active_memory_noise_levels = target_frame_indices.new_zeros((0, target_frame_indices.shape[1]))
+
+    return (
+        active_packed_latents,
+        active_packed_conditions,
+        active_frame_memory_pose,
+        active_frame_indices,
+        active_frame_memory_segments,
+        active_frame_memory_masks,
+        active_memory_noise_levels,
+        active_streams,
+        pruned_streams,
+    )
+
+
 def random_transform(tensor):
     """
     Apply the same random translation, rotation, and scaling to all frames in the batch.
@@ -1609,25 +1675,8 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                     for key in _DEMEMWM_STREAM_KEYS
                 }
                 target_conditions = target_tensors["action_conditions"][local_slice].to(device=xs.device)
-                memory_length = sum(stream_masks_by_key[key].shape[1] for key in _DEMEMWM_STREAM_KEYS)
-
-                packed_latents = torch.cat(
-                    [xs_pred[local_slice], *[stream_latents_by_key[key] for key in _DEMEMWM_STREAM_KEYS]],
-                    dim=0,
-                )
-                packed_conditions = torch.cat(
-                    [target_conditions, target_conditions.new_zeros((memory_length, batch_size, target_conditions.shape[-1]))],
-                    dim=0,
-                )
-                frame_memory_pose = torch.cat(
-                    [target_poses[local_slice], *[stream_poses_by_key[key] for key in _DEMEMWM_STREAM_KEYS]],
-                    dim=0,
-                )
-                frame_indices = torch.cat(
-                    [target_frame_indices[local_slice], *[stream_frame_indices_by_key[key] for key in _DEMEMWM_STREAM_KEYS]],
-                    dim=0,
-                )
-                frame_memory_segments = {
+                target_latents_step = xs_pred[local_slice]
+                full_frame_memory_segments = {
                     "target": int(target_length_step),
                     **{key: int(stream_masks_by_key[key].shape[1]) for key in _DEMEMWM_STREAM_KEYS},
                 }
@@ -1639,14 +1688,13 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                     context_levels = from_query.new_zeros((query_offset, batch_size))
                     from_target = torch.cat([context_levels, from_query], dim=0)
                     to_target = torch.cat([context_levels, to_query], dim=0)
-                    memory_noise_levels = _memory_noise_levels_for_streams(
+                    memory_noise_levels_by_key = _memory_noise_levels_for_streams(
                         getattr(self, "cfg", None),
                         self.diffusion_model,
                         from_query,
-                        frame_memory_segments,
+                        full_frame_memory_segments,
                         mode=namespace,
                     )
-                    memory_noise_levels = torch.cat([memory_noise_levels[key] for key in _DEMEMWM_STREAM_KEYS], dim=0)
                     routed_frame_memory_masks = _apply_memory_route_masks(
                         frame_memory_masks,
                         from_query,
@@ -1654,24 +1702,46 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                         self.diffusion_model,
                         mode=namespace,
                     )
+                    (
+                        active_packed_latents,
+                        active_packed_conditions,
+                        active_frame_memory_pose,
+                        active_frame_indices,
+                        active_frame_memory_segments,
+                        active_frame_memory_masks,
+                        active_memory_noise_levels,
+                        _active_streams,
+                        _pruned_streams,
+                    ) = _pack_active_inference_memory_streams(
+                        target_latents_step,
+                        target_conditions,
+                        target_poses[local_slice],
+                        target_frame_indices[local_slice],
+                        frame_memory_masks["target"],
+                        stream_latents_by_key,
+                        stream_poses_by_key,
+                        stream_frame_indices_by_key,
+                        routed_frame_memory_masks,
+                        memory_noise_levels_by_key,
+                    )
                     sampled_target = self.diffusion_model.sample_step(
-                        packed_latents,
-                        packed_conditions,
+                        active_packed_latents,
+                        active_packed_conditions,
                         None,
-                        torch.cat([from_target, memory_noise_levels], dim=0),
-                        torch.cat([to_target, memory_noise_levels], dim=0),
+                        torch.cat([from_target, active_memory_noise_levels], dim=0),
+                        torch.cat([to_target, active_memory_noise_levels], dim=0),
                         current_frame=curr_frame,
                         mode=namespace,
                         reference_length=0,
-                        frame_idx=frame_indices,
-                        frame_memory_segments=frame_memory_segments,
-                        frame_memory_masks=routed_frame_memory_masks,
-                        frame_memory_pose=frame_memory_pose,
+                        frame_idx=active_frame_indices,
+                        frame_memory_segments=active_frame_memory_segments,
+                        frame_memory_masks=active_frame_memory_masks,
+                        frame_memory_pose=active_frame_memory_pose,
                         image_hw=image_hw,
                     )
-                    packed_latents[:target_length_step] = sampled_target[:target_length_step]
+                    target_latents_step = sampled_target[:target_length_step]
 
-                xs_pred[local_slice] = packed_latents[:target_length_step]
+                xs_pred[local_slice] = target_latents_step
                 chunk_seconds = max(time.perf_counter() - chunk_start_time, 1e-9)
                 end_frame = curr_frame + horizon
                 curr_frame = end_frame