Thread DeMemWM frame memory diffusion metadata

algorithms/dememwm/models/diffusion.py

@@ -155,8 +155,14 @@ class Diffusion(nn.Module):
     def add_shape_channels(self, x):
         return rearrange(x, f"... -> ...{' 1' * len(self.x_shape)}")
 
+    @staticmethod
+    def _target_frames(x, frame_memory_segments):
+        return frame_memory_segments["target"] if frame_memory_segments else x.shape[0]
+
     def model_predictions(self, x, t, action_cond=None, current_frame=None, 
-        pose_cond=None, mode="training", reference_length=None, frame_idx=None):
+        pose_cond=None, mode="training", reference_length=None, frame_idx=None,
+        frame_memory_segments=None, frame_memory_masks=None, frame_memory_pose=None,
+        image_hw=None):
         x = x.permute(1,0,2,3,4)
         action_cond = action_cond.permute(1,0,2)
         if pose_cond is not None and pose_cond[0] is not None:
@@ -165,8 +171,21 @@ class Diffusion(nn.Module):
             except:
                 pass
         t = t.permute(1,0)
-        model_output = self.model(x, t, action_cond, current_frame=current_frame, pose_cond=pose_cond, 
-            mode=mode, reference_length=reference_length, frame_idx=frame_idx)
+        model_kwargs = {
+            "current_frame": current_frame,
+            "pose_cond": pose_cond,
+            "mode": mode,
+            "reference_length": reference_length,
+            "frame_idx": frame_idx,
+        }
+        if frame_memory_segments is not None:
+            model_kwargs.update({
+                "frame_memory_segments": frame_memory_segments,
+                "frame_memory_masks": frame_memory_masks,
+                "frame_memory_pose": frame_memory_pose,
+                "image_hw": image_hw,
+            })
+        model_output = self.model(x, t, action_cond, **model_kwargs)
         model_output = model_output.permute(1,0,2,3,4)
         x = x.permute(1,0,2,3,4)
         t = t.permute(1,0)        
@@ -234,10 +253,14 @@ class Diffusion(nn.Module):
             + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
         )
 
-    def p_mean_variance(self, x, t, action_cond=None, pose_cond=None, reference_length=None):
+    def p_mean_variance(self, x, t, action_cond=None, pose_cond=None, reference_length=None,
+        frame_idx=None, frame_memory_segments=None, frame_memory_masks=None,
+        frame_memory_pose=None, image_hw=None):
         model_pred = self.model_predictions(x=x, t=t, action_cond=action_cond, 
             pose_cond=pose_cond, reference_length=reference_length,
-            frame_idx=frame_idx)
+            frame_idx=frame_idx, frame_memory_segments=frame_memory_segments,
+            frame_memory_masks=frame_memory_masks,
+            frame_memory_pose=frame_memory_pose, image_hw=image_hw)
         x_start = model_pred.pred_x_start
         return self.q_posterior(x_start=x_start, x_t=x, t=t)
 
@@ -286,7 +309,11 @@ class Diffusion(nn.Module):
         pose_cond,
         noise_levels: torch.Tensor,
         reference_length,
-        frame_idx=None
+        frame_idx=None,
+        frame_memory_segments=None,
+        frame_memory_masks=None,
+        frame_memory_pose=None,
+        image_hw=None
     ):
         noise = torch.randn_like(x)
         noise = torch.clamp(noise, -self.clip_noise, self.clip_noise)
@@ -294,7 +321,10 @@ class Diffusion(nn.Module):
         noised_x = self.q_sample(x_start=x, t=noise_levels, noise=noise)
 
         model_pred = self.model_predictions(x=noised_x, t=noise_levels, action_cond=action_cond, 
-                                    pose_cond=pose_cond,reference_length=reference_length, frame_idx=frame_idx)
+                                    pose_cond=pose_cond,reference_length=reference_length, frame_idx=frame_idx,
+                                    frame_memory_segments=frame_memory_segments,
+                                    frame_memory_masks=frame_memory_masks,
+                                    frame_memory_pose=frame_memory_pose, image_hw=image_hw)
 
         pred = model_pred.model_out
         x_pred = model_pred.pred_x_start
@@ -329,7 +359,11 @@ class Diffusion(nn.Module):
         current_frame=None,
         mode="training",
         reference_length=None,
-        frame_idx=None
+        frame_idx=None,
+        frame_memory_segments=None,
+        frame_memory_masks=None,
+        frame_memory_pose=None,
+        image_hw=None
     ):
         real_steps = torch.linspace(-1, self.timesteps - 1, steps=self.sampling_timesteps + 1, device=x.device).long()
 
@@ -348,7 +382,11 @@ class Diffusion(nn.Module):
                 current_frame=current_frame,
                 mode=mode,
                 reference_length=reference_length,
-                frame_idx=frame_idx
+                frame_idx=frame_idx,
+                frame_memory_segments=frame_memory_segments,
+                frame_memory_masks=frame_memory_masks,
+                frame_memory_pose=frame_memory_pose,
+                image_hw=image_hw
             )
 
         # FIXME: temporary code for checking ddpm sampling
@@ -367,7 +405,11 @@ class Diffusion(nn.Module):
             curr_noise_level=curr_noise_level,
             guidance_fn=guidance_fn,
             reference_length=reference_length,
-            frame_idx=frame_idx
+            frame_idx=frame_idx,
+            frame_memory_segments=frame_memory_segments,
+            frame_memory_masks=frame_memory_masks,
+            frame_memory_pose=frame_memory_pose,
+            image_hw=image_hw
         )
 
     def ddpm_sample_step(
@@ -379,6 +421,10 @@ class Diffusion(nn.Module):
         guidance_fn: Optional[Callable] = None,
         reference_length=None,
         frame_idx=None,
+        frame_memory_segments=None,
+        frame_memory_masks=None,
+        frame_memory_pose=None,
+        image_hw=None,
     ):
         clipped_curr_noise_level = torch.where(
             curr_noise_level < 0,
@@ -405,7 +451,11 @@ class Diffusion(nn.Module):
                 action_cond=action_cond,
                 pose_cond=pose_cond,
                 reference_length=reference_length,
-                frame_idx=frame_idx
+                frame_idx=frame_idx,
+                frame_memory_segments=frame_memory_segments,
+                frame_memory_masks=frame_memory_masks,
+                frame_memory_pose=frame_memory_pose,
+                image_hw=image_hw
             )
 
         noise = torch.where(
@@ -430,7 +480,11 @@ class Diffusion(nn.Module):
         current_frame=None,
         mode="training",
         reference_length=None,
-        frame_idx=None
+        frame_idx=None,
+        frame_memory_segments=None,
+        frame_memory_masks=None,
+        frame_memory_pose=None,
+        image_hw=None
     ):
         # convert noise level -1 to self.stabilization_level - 1
         clipped_curr_noise_level = torch.where(
@@ -477,7 +531,11 @@ class Diffusion(nn.Module):
                     current_frame=current_frame,
                     mode=mode,
                     reference_length=reference_length,
-                    frame_idx=frame_idx
+                    frame_idx=frame_idx,
+                    frame_memory_segments=frame_memory_segments,
+                    frame_memory_masks=frame_memory_masks,
+                    frame_memory_pose=frame_memory_pose,
+                    image_hw=image_hw
                 )
 
                 guidance_loss = guidance_fn(model_pred.pred_x_start)
@@ -499,7 +557,11 @@ class Diffusion(nn.Module):
                 current_frame=current_frame,
                 mode=mode,
                 reference_length=reference_length,
-                frame_idx=frame_idx
+                frame_idx=frame_idx,
+                frame_memory_segments=frame_memory_segments,
+                frame_memory_masks=frame_memory_masks,
+                frame_memory_pose=frame_memory_pose,
+                image_hw=image_hw
             )
             x_start = model_pred.pred_x_start
             pred_noise = model_pred.pred_noise

algorithms/dememwm/models/dit.py

@@ -480,7 +480,8 @@ class DiT(nn.Module):
         return imgs
 
     def forward(self, x, t, action_cond=None, pose_cond=None, current_frame=None, mode=None, 
-                    reference_length=None, frame_idx=None):
+                    reference_length=None, frame_idx=None, frame_memory_segments=None,
+                    frame_memory_masks=None, frame_memory_pose=None, image_hw=None):
         """
         Forward pass of DiT.
         x: (B, T, C, H, W) tensor of spatial inputs (images or latent representations of images)