Use configured focal length in DeMemWM geometry

algorithms/dememwm/df_video.py

@@ -690,7 +690,8 @@ class DeMemWMMinecraft(DiffusionForcingBase):
             state_embed_only_on_qk=self.state_embed_only_on_qk,
             use_memory_attention=self.use_memory_attention,
             add_timestamp_embedding=self.add_timestamp_embedding,
-            ref_mode=self.ref_mode
+            ref_mode=self.ref_mode,
+            focal_length=self.focal_length,
         )
 
         self.validation_lpips_model = LearnedPerceptualImagePatchSimilarity(sync_on_compute=False)

algorithms/dememwm/models/diffusion.py

@@ -29,7 +29,8 @@ class Diffusion(nn.Module):
         state_embed_only_on_qk=False,
         use_memory_attention=False,
         add_timestamp_embedding=False,
-        ref_mode='sequential'
+        ref_mode='sequential',
+        focal_length=0.35,
     ):
         super().__init__()
         self.cfg = cfg
@@ -58,6 +59,7 @@ class Diffusion(nn.Module):
         self.use_memory_attention = use_memory_attention
         self.add_timestamp_embedding = add_timestamp_embedding
         self.ref_mode = ref_mode
+        self.focal_length = focal_length
 
         self._build_model()
         self._build_buffer()
@@ -72,7 +74,8 @@ class Diffusion(nn.Module):
                                             state_embed_only_on_qk=self.state_embed_only_on_qk,
                                             use_memory_attention=self.use_memory_attention,
                                             add_timestamp_embedding=self.add_timestamp_embedding,
-                                            ref_mode=self.ref_mode)
+                                            ref_mode=self.ref_mode,
+                                            focal_length=self.focal_length)
         else:
             raise NotImplementedError
 

algorithms/dememwm/models/dit.py

@@ -459,7 +459,8 @@ class DiT(nn.Module):
         state_embed_only_on_qk=False,
         use_memory_attention=False,
         add_timestamp_embedding=False,
-        ref_mode='sequential'
+        ref_mode='sequential',
+        focal_length=0.35,
     ):
         super().__init__()
         self.in_channels = in_channels
@@ -467,6 +468,7 @@ class DiT(nn.Module):
         self.patch_size = patch_size
         self.num_heads = num_heads
         self.max_frames = max_frames
+        self.focal_length = focal_length
 
         self.x_embedder = PatchEmbed(input_h, input_w, patch_size, in_channels, hidden_size, flatten=False)
         self.t_embedder = TimestepEmbedder(hidden_size)
@@ -615,7 +617,7 @@ class DiT(nn.Module):
             # Token grid sets position count; image_hw keeps camera scaling in original image coordinates.
             y = (torch.arange(grid_h, device=device, dtype=torch.float32).view(1, grid_h, 1) + 0.5) * (image_h / grid_h)
             x = (torch.arange(grid_w, device=device, dtype=torch.float32).view(1, 1, grid_w) + 0.5) * (image_w / grid_w)
-            fx, fy = 0.35 * image_w, 0.35 * image_h
+            fx, fy = self.focal_length * image_w, self.focal_length * image_h
             cx, cy = 0.5 * image_w, 0.5 * image_h
             directions = torch.stack(
                 (
@@ -758,7 +760,7 @@ class DiT(nn.Module):
 def DiT_S_2(action_cond_dim, pose_cond_dim, reference_length, 
 use_plucker, relative_embedding, 
 state_embed_only_on_qk, use_memory_attention, add_timestamp_embedding,
-ref_mode):
+ref_mode, focal_length=0.35):
     return DiT(
         patch_size=2,
         hidden_size=1024,
@@ -772,7 +774,8 @@ ref_mode):
         state_embed_only_on_qk=state_embed_only_on_qk,
         use_memory_attention=use_memory_attention,
         add_timestamp_embedding=add_timestamp_embedding,
-        ref_mode=ref_mode
+        ref_mode=ref_mode,
+        focal_length=focal_length,
     )