Warn on missing DeMemWM Plucker geometry

algorithms/dememwm/models/dit.py

@@ -6,6 +6,7 @@ References:
 """
 
 from typing import Optional, Literal
+import logging
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -19,6 +20,16 @@ from collections import namedtuple
 from typing import Optional, Callable
 from .cameractrl_module import SimpleCameraPoseEncoder
 
+logger = logging.getLogger(__name__)
+_FRAME_MEMORY_GEOMETRY_WARNINGS = set()
+
+
+def _warn_frame_memory_geometry_once(key, message):
+    if key not in _FRAME_MEMORY_GEOMETRY_WARNINGS:
+        _FRAME_MEMORY_GEOMETRY_WARNINGS.add(key)
+        logger.warning(message)
+
+
 def modulate(x, shift, scale):
     fixed_dims = [1] * len(shift.shape[1:])
     shift = shift.repeat(x.shape[0] // shift.shape[0], *fixed_dims)
@@ -173,10 +184,18 @@ class FrameMemoryReferenceAttention(nn.Module):
         query_rays = geometry_cache.get("query_rays")
         relative_rays = geometry_cache.get("relative_rays")
         if not torch.is_tensor(query_rays) or not torch.is_tensor(relative_rays):
+            _warn_frame_memory_geometry_once(
+                "missing_ray_tensors",
+                "DeMemWM Plucker memory geometry is missing query/relative ray tensors; using content-only memory attention.",
+            )
             return None, None
-        if tuple(query_rays.shape) != (B, T_target, H, W, 6):
-            return None, None
-        if tuple(relative_rays.shape) != (B, T_target, T_memory, H, W, 6):
+        expected_query_shape = (B, T_target, H, W, 6)
+        expected_relative_shape = (B, T_target, T_memory, H, W, 6)
+        if tuple(query_rays.shape) != expected_query_shape or tuple(relative_rays.shape) != expected_relative_shape:
+            _warn_frame_memory_geometry_once(
+                "malformed_ray_tensors",
+                "DeMemWM Plucker memory geometry has unexpected ray tensor shapes; using content-only memory attention.",
+            )
             return None, None
         return (
             query_rays.to(device=device, dtype=dtype),
@@ -590,17 +609,49 @@ class DiT(nn.Module):
 
     def _build_frame_memory_geometry(self, frame_memory_segments, frame_memory_masks, frame_memory_pose, image_hw, grid_h, grid_w, device, dtype):
         if frame_memory_segments is None or frame_memory_pose is None or image_hw is None:
+            _warn_frame_memory_geometry_once(
+                "missing_geometry_metadata",
+                "DeMemWM Plucker memory geometry metadata is missing; using content-only memory attention.",
+            )
             return None
 
         pose = frame_memory_pose.to(device=device, dtype=dtype)
+        if pose.ndim != 3 or pose.shape[-1] != 5:
+            _warn_frame_memory_geometry_once(
+                "malformed_pose_metadata",
+                "DeMemWM Plucker memory geometry pose metadata has an unexpected shape; using content-only memory attention.",
+            )
+            return None
         B = pose.shape[0]
         target_frames = int(frame_memory_segments.get("target", 0))
+        expected_pose_frames = target_frames + sum(
+            int(frame_memory_segments.get(stream_name, 0))
+            for stream_name in ("anchor", "dynamic", "revisit")
+        )
+        if pose.shape[1] < expected_pose_frames:
+            _warn_frame_memory_geometry_once(
+                "short_pose_metadata",
+                "DeMemWM Plucker memory geometry pose metadata does not cover the packed memory segments; using content-only memory attention.",
+            )
+            return None
         image_hw = image_hw if torch.is_tensor(image_hw) else torch.as_tensor(image_hw)
         image_hw = image_hw.to(device=device, dtype=dtype)
         if image_hw.ndim == 1:
             image_hw = image_hw.unsqueeze(0)
+        if image_hw.ndim != 2 or image_hw.shape[-1] < 2:
+            _warn_frame_memory_geometry_once(
+                "malformed_image_hw_metadata",
+                "DeMemWM Plucker memory geometry image size metadata has an unexpected shape; using content-only memory attention.",
+            )
+            return None
         if image_hw.shape[0] == 1 and B != 1:
             image_hw = image_hw.expand(B, -1)
+        if image_hw.shape[0] != B:
+            _warn_frame_memory_geometry_once(
+                "mismatched_image_hw_metadata",
+                "DeMemWM Plucker memory geometry image size metadata does not match the batch; using content-only memory attention.",
+            )
+            return None
         image_hw = image_hw.contiguous()
         image_h, image_w = image_hw[:, 0].view(B, 1, 1), image_hw[:, 1].view(B, 1, 1)