Add DeMemWM memory sheet diagnostics

algorithms/dememwm/df_video.py

@@ -10,7 +10,7 @@ import torch.distributed as dist
 import torch.nn.functional as F
 import torchvision.transforms.functional as TF
 from torchvision.transforms import InterpolationMode
-from PIL import Image
+from PIL import Image, ImageDraw
 from einops import rearrange
 from tqdm import tqdm
 from omegaconf import DictConfig, open_dict
@@ -675,6 +675,10 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         self.require_pose_prediction = getattr(cfg, "require_pose_prediction", False)
         self.log_per_frame_metrics = bool(getattr(cfg, "log_per_frame_metrics", False))
         self.log_revisit_metrics = bool(getattr(cfg, "log_revisit_metrics", False))
+        self.log_memory_selection_sheet = bool(getattr(cfg, "log_memory_selection_sheet", False))
+        self.memory_sheet_generated_frames = tuple(
+            int(x) for x in getattr(cfg, "memory_sheet_generated_frames", (1, 25, 50, 100, 200, 350, 500))
+        )
         self.revisit_metrics_min_gap = int(getattr(cfg, "revisit_metrics_min_gap", 8))
         self.revisit_metrics_gap_bands = tuple(int(x) for x in getattr(cfg, "revisit_metrics_gap_bands", (32, 128)))
         self.revisit_metrics_self_crop_fraction = float(getattr(cfg, "revisit_metrics_self_crop_fraction", 0.5))
@@ -891,6 +895,126 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         base = Path(self.local_save_dir) if self.local_save_dir is not None else Path.cwd() / "eval_artifacts"
         return base / "metrics"
 
+    def _memory_sheet_artifact_dir(self) -> Path:
+        base = Path(self.local_save_dir) if self.local_save_dir is not None else Path.cwd() / "eval_artifacts"
+        return base / "memory_selection_sheets"
+
+    def _memory_sheet_target_indices(self, eval_start: int, target_length: int) -> set[int]:
+        return {
+            int(eval_start) + int(frame) - 1
+            for frame in self.memory_sheet_generated_frames
+            if int(frame) > 0 and int(eval_start) + int(frame) - 1 < int(target_length)
+        }
+
+    def _save_memory_selection_sheets(
+        self,
+        xs_pred: torch.Tensor,
+        xs_gt: torch.Tensor,
+        frame_indices: torch.Tensor,
+        records: list[list[dict]],
+        eval_start: int,
+        batch_idx: int,
+    ) -> None:
+        if not self.log_memory_selection_sheet or not self._is_global_rank_zero():
+            return
+        output_dir = self._memory_sheet_artifact_dir()
+        output_dir.mkdir(parents=True, exist_ok=True)
+        max_slots = {"anchor": 2, "dynamic": 4, "revisit": 2}
+        tile_w, tile_h = 192, 108
+        label_h = 34
+        gap = 6
+
+        def decode_indices(source: torch.Tensor, batch_i: int, indices: list[int]) -> dict[int, Image.Image]:
+            if not indices:
+                return {}
+            unique = sorted(set(int(i) for i in indices))
+            latents = torch.stack([source[i, batch_i] for i in unique], dim=0).unsqueeze(1)
+            with torch.no_grad():
+                decoded = self.decode(latents.to(source.device))[:, 0].detach().cpu().float().clamp(0.0, 1.0)
+            return {idx: TF.to_pil_image(decoded[row]).resize((tile_w, tile_h), Image.BILINEAR) for row, idx in enumerate(unique)}
+
+        def make_tile(image: Image.Image | None, label: str, fill=(236, 236, 236)) -> Image.Image:
+            tile = Image.new("RGB", (tile_w, tile_h + label_h), fill)
+            if image is not None:
+                tile.paste(image.convert("RGB"), (0, label_h))
+            draw = ImageDraw.Draw(tile)
+            draw.rectangle((0, 0, tile_w, label_h), fill=(18, 18, 18))
+            for row, line in enumerate(label.split("\n")[:2]):
+                draw.text((5, 4 + row * 14), line, fill=(255, 255, 255))
+            return tile
+
+        summary_rows = []
+        for batch_i, batch_records in enumerate(records):
+            if not batch_records:
+                continue
+            target_indices = [int(row["target_index"]) for row in batch_records]
+            memory_gt = []
+            memory_pred = []
+            for row in batch_records:
+                for stream in _DEMEMWM_STREAM_KEYS:
+                    for source_index in row["streams"][stream]:
+                        (memory_gt if source_index < eval_start else memory_pred).append(int(source_index))
+            target_gt_images = decode_indices(xs_gt, batch_i, target_indices)
+            target_pred_images = decode_indices(xs_pred, batch_i, target_indices)
+            memory_images = {}
+            memory_images.update(decode_indices(xs_gt, batch_i, memory_gt))
+            memory_images.update(decode_indices(xs_pred, batch_i, memory_pred))
+
+            columns = [("target_gt", -1), ("target_pred", -1)]
+            for stream in _DEMEMWM_STREAM_KEYS:
+                columns.extend((stream, slot) for slot in range(max_slots[stream]))
+            sheet = Image.new(
+                "RGB",
+                (len(columns) * tile_w + (len(columns) - 1) * gap, len(batch_records) * (tile_h + label_h) + (len(batch_records) - 1) * gap),
+                (245, 245, 245),
+            )
+            for record_row, row in enumerate(batch_records):
+                y = record_row * (tile_h + label_h + gap)
+                target_index = int(row["target_index"])
+                generated_frame = target_index - int(eval_start) + 1
+                raw_target = int(row["target_frame"])
+                for col, (stream, slot) in enumerate(columns):
+                    x = col * (tile_w + gap)
+                    if stream == "target_gt":
+                        image = target_gt_images.get(target_index)
+                        label = f"gen{generated_frame:03d} raw{raw_target}\ntarget GT"
+                    elif stream == "target_pred":
+                        image = target_pred_images.get(target_index)
+                        label = f"gen{generated_frame:03d} raw{raw_target}\ntarget pred"
+                    else:
+                        stream_indices = row["streams"][stream]
+                        source_index = int(stream_indices[slot]) if slot < len(stream_indices) else -1
+                        image = memory_images.get(source_index)
+                        if source_index >= 0:
+                            raw_source = int(frame_indices[source_index, batch_i].detach().cpu().item())
+                            source_type = "ctx" if source_index < eval_start else "pred"
+                            label = f"{stream}{slot} seq{source_index} raw{raw_source}\n{source_type} for gen{generated_frame:03d}"
+                            summary_rows.append(
+                                {
+                                    "batch": batch_idx,
+                                    "sample": batch_i,
+                                    "generated_frame": generated_frame,
+                                    "target_index": target_index,
+                                    "target_raw_frame": raw_target,
+                                    "stream": stream,
+                                    "slot": slot,
+                                    "source_index": source_index,
+                                    "source_raw_frame": raw_source,
+                                    "source_type": source_type,
+                                }
+                            )
+                        else:
+                            label = f"{stream}{slot}\nempty"
+                    sheet.paste(make_tile(image, label), (x, y))
+            sheet.save(output_dir / f"batch{batch_idx:06d}_sample{batch_i:02d}_memory_streams.jpg", quality=95)
+
+        if summary_rows:
+            csv_path = output_dir / f"batch{batch_idx:06d}_memory_streams.csv"
+            with csv_path.open("w", newline="", encoding="utf-8") as handle:
+                writer = csv.DictWriter(handle, fieldnames=list(summary_rows[0].keys()))
+                writer.writeheader()
+                writer.writerows(summary_rows)
+
     def _is_global_rank_zero(self) -> bool:
         return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0
 
@@ -1374,6 +1498,8 @@ class DeMemWMMinecraft(DiffusionForcingBase):
             if dynamic_policy not in {"recent", "event_triggered"}:
                 raise ValueError(f"memory_selection.dynamic.selection_policy must be recent or event_triggered; got {dynamic_policy!r}")
             event_dynamic_caches = [_new_online_event_cache(target_length) for _ in range(batch_size)] if dynamic_policy == "event_triggered" else None
+            memory_sheet_targets = self._memory_sheet_target_indices(n_context_frames, target_length) if self.log_memory_selection_sheet else set()
+            memory_sheet_records = [[] for _ in range(batch_size)] if memory_sheet_targets else None
 
             curr_frame = n_context_frames
             generated_frames = 0
@@ -1456,6 +1582,17 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                                 else None
                             ),
                         )
+                    if memory_sheet_records is not None and int(target_positions[0]) in memory_sheet_targets:
+                        memory_sheet_records[batch_i].append(
+                            {
+                                "target_index": int(target_positions[0]),
+                                "target_frame": int(target_frame_indices[int(target_positions[0]), batch_i].detach().cpu().item()),
+                                "streams": {
+                                    key: [int(v) for v in np.asarray(selected.get(key, []), dtype=np.int64) if 0 <= int(v) < curr_frame]
+                                    for key in _DEMEMWM_STREAM_KEYS
+                                },
+                            }
+                        )
                     selected_masks = {key: np.ones(len(value), dtype=bool) for key, value in selected.items()}
                     for key, indices in stream_indices.items():
                         count = indices.shape[1]
@@ -1562,6 +1699,8 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 latent_loss = xs_pred.new_tensor(0.0)
 
             self.log(f"{namespace}/latent_mse", latent_loss.detach())
+            if memory_sheet_records is not None:
+                self._save_memory_selection_sheets(xs_pred, xs, target_frame_indices, memory_sheet_records, eval_start, batch_idx)
             if eval_start < target_length:
                 xs_pred_decode = self.decode(xs_pred[eval_start:target_length].to(target_poses.device))
                 xs_decode = self.decode(xs[eval_start:target_length].to(target_poses.device))

datasets/video/memory_selection.py

@@ -329,18 +329,18 @@ def _select_anchor(candidates: np.ndarray, count: int, cfg, poses=None) -> np.nd
     poses_t = _as_pose_tensor(poses)
     candidates_t = torch.as_tensor(candidates.astype(np.int64), device=poses_t.device, dtype=torch.long)
     candidate_poses = poses_t.index_select(0, candidates_t).float()
-    pairwise = torch.cdist(candidate_poses, candidate_poses)
+    spatial = torch.cdist(candidate_poses[:, :3], candidate_poses[:, :3])
+    angular = torch.linalg.vector_norm(
+        _wrap_degrees(candidate_poses[:, None, 3:] - candidate_poses[None, :, 3:]),
+        dim=-1,
+    )
+    pairwise = torch.sqrt(spatial.square() + angular.square())
     if not bool((pairwise > 0).any().item()):
         return candidates[:count].astype(np.int64)
 
     available = torch.ones((int(candidates_t.numel()),), device=poses_t.device, dtype=torch.bool)
-    if count == 1:
-        selected = [int(pairwise.mean(dim=1).argmax().item())]
-    else:
-        first, second = divmod(int(pairwise.argmax().item()), int(pairwise.shape[1]))
-        selected = [int(first), int(second)]
-    for idx in selected:
-        available[idx] = False
+    selected = [0]
+    available[0] = False
 
     dists = pairwise[selected].min(dim=0).values.masked_fill(~available, float("-inf"))
     for _ in range(count - len(selected)):