Trust DeMemWM Plucker geometry metadata

algorithms/dememwm/models/dit.py

@@ -6,7 +6,6 @@ References:
 """
 
 from typing import Optional, Literal
-import logging
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -20,15 +19,6 @@ 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:])
@@ -179,27 +169,11 @@ class FrameMemoryReferenceAttention(nn.Module):
         device,
         dtype,
     ):
-        if not isinstance(geometry_cache, dict):
-            return None, None
-        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
-        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.",
-            )
+        if geometry_cache is None:
             return None, None
         return (
-            query_rays.to(device=device, dtype=dtype),
-            relative_rays.to(device=device, dtype=dtype),
+            geometry_cache["query_rays"].to(device=device, dtype=dtype),
+            geometry_cache["relative_rays"].to(device=device, dtype=dtype),
         )
 
     def forward(self, target_hidden, memory_hidden, memory_mask=None, geometry_cache=None):
@@ -638,50 +612,18 @@ class DiT(nn.Module):
         return c2w
 
     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.",
-            )
+        if frame_memory_segments is None:
             return None
 
         pose = frame_memory_pose.to(device=device)
-        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)
         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()
         # Mixed precision can round camera offsets before relative poses/rays are
         # formed, so keep this local geometry block in fp32 and only cast caches.