Refactor DeMemWM inference memory streams by key

algorithms/dememwm/df_video.py

@@ -1523,20 +1523,20 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 target_length_step = local_slice.stop - local_slice.start
                 query_offset = curr_frame - start_frame
 
-                stream_indices = {
+                stream_indices_by_key = {
                     key: torch.full((batch_size, int(stream_lengths[key])), -1, device=xs.device, dtype=torch.long)
                     for key in _DEMEMWM_STREAM_KEYS
                 }
-                stream_masks = {key: torch.zeros_like(indices, dtype=torch.bool) for key, indices in stream_indices.items()}
+                stream_masks_by_key = {key: torch.zeros_like(indices, dtype=torch.bool) for key, indices in stream_indices_by_key.items()}
                 target_positions = np.arange(target_slice.start, target_slice.stop, dtype=np.int64)
                 source_latents = xs_pred[:curr_frame]
                 source_actions = target_actions[:curr_frame]
                 source_poses = target_poses[:curr_frame]
                 source_frame_indices = target_frame_indices[:curr_frame]
                 for batch_i in range(batch_size):
-                    anchor_count = int(stream_indices["anchor"].shape[1])
-                    dynamic_count = int(stream_indices["dynamic"].shape[1])
-                    revisit_count = int(stream_indices["revisit"].shape[1])
+                    anchor_count = int(stream_indices_by_key["anchor"].shape[1])
+                    dynamic_count = int(stream_indices_by_key["dynamic"].shape[1])
+                    revisit_count = int(stream_indices_by_key["revisit"].shape[1])
                     selected = {"anchor": online_anchor_indices[batch_i]}
                     if dynamic_policy == "recent":
                         selected["dynamic"] = _select_online_dynamic_indices(
@@ -1583,35 +1583,53 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                             }
                         )
                     selected_masks = {key: np.ones(len(value), dtype=bool) for key, value in selected.items()}
-                    for key, indices in stream_indices.items():
+                    for key, indices in stream_indices_by_key.items():
                         count = indices.shape[1]
                         selected_key = np.asarray(selected.get(key, []), dtype=np.int64)[:count]
                         mask_key = np.asarray(selected_masks.get(key, []), dtype=bool)[:count]
                         indices[batch_i, : len(selected_key)] = torch.as_tensor(selected_key, device=xs.device, dtype=torch.long)
-                        stream_masks[key][batch_i, : len(mask_key)] = torch.as_tensor(mask_key, device=xs.device, dtype=torch.bool)
+                        stream_masks_by_key[key][batch_i, : len(mask_key)] = torch.as_tensor(mask_key, device=xs.device, dtype=torch.bool)
 
                 frame_memory_masks = {
                     "target": torch.ones((batch_size, target_length_step), device=xs.device, dtype=torch.bool)
                     if target_mask is None
                     else target_mask[:, local_slice],
-                    **stream_masks,
+                    **stream_masks_by_key,
+                }
+                stream_latents_by_key = {
+                    key: _gather_online_memory_tensor(source_latents, stream_indices_by_key[key], stream_masks_by_key[key])
+                    for key in _DEMEMWM_STREAM_KEYS
+                }
+                stream_poses_by_key = {
+                    key: _gather_online_memory_tensor(source_poses, stream_indices_by_key[key], stream_masks_by_key[key])
+                    for key in _DEMEMWM_STREAM_KEYS
+                }
+                stream_frame_indices_by_key = {
+                    key: _gather_online_memory_tensor(source_frame_indices, stream_indices_by_key[key], stream_masks_by_key[key])
+                    for key in _DEMEMWM_STREAM_KEYS
                 }
-                stream_latents = [_gather_online_memory_tensor(source_latents, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
-                stream_poses = [_gather_online_memory_tensor(source_poses, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
-                stream_frame_indices = [_gather_online_memory_tensor(source_frame_indices, stream_indices[key], stream_masks[key]) for key in _DEMEMWM_STREAM_KEYS]
                 target_conditions = target_tensors["action_conditions"][local_slice].to(device=xs.device)
-                memory_length = sum(mask.shape[1] for mask in stream_masks.values())
+                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], dim=0)
+                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], dim=0)
-                frame_indices = torch.cat([target_frame_indices[local_slice], *stream_frame_indices], 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": int(target_length_step),
-                    **{key: int(stream_masks[key].shape[1]) for key in _DEMEMWM_STREAM_KEYS},
+                    **{key: int(stream_masks_by_key[key].shape[1]) for key in _DEMEMWM_STREAM_KEYS},
                 }
                 scheduling_matrix = self._generate_scheduling_matrix(horizon)