code clean up

algorithms/dememwm/models/dit.py

@@ -158,23 +158,6 @@ class FrameMemoryReferenceAttention(nn.Module):
     def _split_heads(self, x):
         return rearrange(x, "b n (h d) -> b h n d", h=self.num_heads)
 
-    def _geometry_rays(
-        self,
-        geometry_cache,
-        B,
-        T_target,
-        T_memory,
-        H,
-        W,
-        device,
-        dtype,
-    ):
-        if geometry_cache is None:
-            return None, None
-        return (
-            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):
         B, T_target, H, W, D = target_hidden.shape
@@ -183,16 +166,12 @@ class FrameMemoryReferenceAttention(nn.Module):
 
         T_memory = memory_hidden.shape[1]
         P = H * W
-        query_rays, relative_rays = self._geometry_rays(
-            geometry_cache,
-            B,
-            T_target,
-            T_memory,
-            H,
-            W,
-            target_hidden.device,
-            target_hidden.dtype,
-        )
+        if geometry_cache is not None:
+            query_rays = geometry_cache["query_rays"].to(device=target_hidden.device, dtype=target_hidden.dtype)
+            relative_rays = geometry_cache["relative_rays"].to(device=target_hidden.device, dtype=target_hidden.dtype)
+        else:
+            query_rays = None
+        
         q_input = target_hidden
         if query_rays is not None:
             # Values stay visual-only; Plucker features only steer Q/K matching.