update

algorithms/worldmem/df_video.py

@@ -354,10 +354,10 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
         self.is_interactive = cfg.get("is_interactive", False)
         if self.is_interactive:
-            self.frames = None
-            self.poses = None
-            self.memory_c2w = None
-            self.frame_idx = None
+            self_frames = None
+            self_poses = None
+            self_memory_c2w = None
+            self_frame_idx = None
 
         super().__init__(cfg)
             
@@ -791,21 +791,23 @@ class WorldMemMinecraft(DiffusionForcingBase):
         return
 
     @torch.no_grad()
-    def interactive(self, first_frame, curr_actions, first_pose, context_frames_idx, device):
+    def interactive(self, first_frame, curr_actions, first_pose, context_frames_idx, device,
+                    self_frames, self_poses, self_memory_c2w, self_frame_idx):
+    
         condition_similar_length = self.condition_similar_length
 
-        if self.frames is None:
+        if self_frames is None:
             first_frame_encode = self.encode(first_frame[None, None].to(device))
-            self.frames = first_frame_encode.cpu()
+            self_frames = first_frame_encode.cpu()
             self.actions = curr_actions[None, None].to(device)
-            self.poses = first_pose[None, None].to(device)
+            self_poses = first_pose[None, None].to(device)
             new_c2w_mat = euler_to_camera_to_world_matrix(first_pose)
-            self.memory_c2w = new_c2w_mat[None, None].to(device)
-            self.frame_idx = torch.tensor([[context_frames_idx]]).to(device)
-            return first_frame
+            self_memory_c2w = new_c2w_mat[None, None].to(device)
+            self_frame_idx = torch.tensor([[context_frames_idx]]).to(device)
+            return first_frame, self_frames, self_poses, self_memory_c2w, self_frame_idx
         else:
-            last_frame = self.frames[-1].clone()
-            last_pose_condition = self.poses[-1].clone()
+            last_frame = self_frames[-1].clone()
+            last_pose_condition = self_poses[-1].clone()
             last_pose_condition[:,3:] = last_pose_condition[:,3:] // 15
             new_pose_condition_offset = self.pose_prediction_model(last_frame.to(device), curr_actions[None].to(device), last_pose_condition)
             
@@ -814,15 +816,15 @@ class WorldMemMinecraft(DiffusionForcingBase):
             new_pose_condition[:,3:] = new_pose_condition[:,3:] * 15
             new_pose_condition[:,3:] %= 360
             self.actions = torch.cat([self.actions, curr_actions[None, None].to(device)])
-            self.poses = torch.cat([self.poses, new_pose_condition[None].to(device)])
+            self_poses = torch.cat([self_poses, new_pose_condition[None].to(device)])
             new_c2w_mat = euler_to_camera_to_world_matrix(new_pose_condition)
-            self.memory_c2w = torch.cat([self.memory_c2w, new_c2w_mat[None].to(device)])
-            self.frame_idx = torch.cat([self.frame_idx, torch.tensor([[context_frames_idx]]).to(device)])
+            self_memory_c2w = torch.cat([self_memory_c2w, new_c2w_mat[None].to(device)])
+            self_frame_idx = torch.cat([self_frame_idx, torch.tensor([[context_frames_idx]]).to(device)])
 
         conditions = self.actions.clone()
-        pose_conditions = self.poses.clone()
-        c2w_mat = self.memory_c2w .clone()
-        frame_idx = self.frame_idx.clone()
+        pose_conditions = self_poses.clone()
+        c2w_mat = self_memory_c2w .clone()
+        frame_idx = self_frame_idx.clone()
 
 
         curr_frame = 0
@@ -831,7 +833,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
         n_frames = curr_frame + horizon
         # context
         n_context_frames = context_frames_idx // self.frame_stack
-        xs_pred = self.frames[:n_context_frames].clone()
+        xs_pred = self_frames[:n_context_frames].clone()
         curr_frame += n_context_frames
 
         pbar = tqdm(total=n_frames, initial=curr_frame, desc="Sampling")
@@ -894,14 +896,15 @@ class WorldMemMinecraft(DiffusionForcingBase):
         curr_frame += horizon
         pbar.update(horizon)
 
-        self.frames = torch.cat([self.frames, xs_pred[n_context_frames:]])
+        self_frames = torch.cat([self_frames, xs_pred[n_context_frames:]])
 
         xs_pred = self.decode(xs_pred[n_context_frames:].to(device)).cpu()
-        return xs_pred[-1,0]
+
+        return xs_pred[-1,0], self_frames, self_poses, self_memory_c2w, self_frame_idx
 
 
     def reset(self):
-        self.frames = None
-        self.poses = None
-        self.memory_c2w = None
-        self.frame_idx = None
+        self_frames = None
+        self_poses = None
+        self_memory_c2w = None
+        self_frame_idx = None

app.py

@@ -182,15 +182,28 @@ poses = torch.zeros((1, 5))
 
 memory_frames.append(load_image_as_tensor(DEFAULT_IMAGE))
 
+self_frames = None
+self_poses = None
+self_memory_c2w = None
+self_frame_idx = None
+
+
 @spaces.GPU()
 def run_interactive(first_frame, action, first_pose, curr_frame, device):
-    global worldmem
-    new_frame = worldmem.interactive(first_frame,
+    global self_frames
+    global self_poses
+    global self_memory_c2w
+    global self_frame_idx
+
+    new_frame, self_frames, self_poses, self_memory_c2w, self_frame_idx = worldmem.interactive(first_frame,
                                     action,
                                     first_pose, 
                                     curr_frame, 
-                                    device=device)
-    print("algo frame:", len(worldmem.frames))
+                                    device=device,
+                                    self_frames=self_frames,
+                                    self_poses=self_poses,
+                                    self_memory_c2w=self_memory_c2w,
+                                    self_frame_idx=self_frame_idx)
     return new_frame
 
 def set_denoising_steps(denoising_steps, sampling_timesteps_state):
@@ -201,7 +214,7 @@ def set_denoising_steps(denoising_steps, sampling_timesteps_state):
     return sampling_timesteps_state
 
 def generate(keys):
-    print("algo frame:", len(worldmem.frames))
+    # print("algo frame:", len(worldmem.frames))
     actions = parse_input_to_tensor(keys)
     global input_history
     global memory_curr_frame
@@ -236,7 +249,11 @@ def reset():
     global input_history
     global memory_frames
 
-    worldmem.reset()
+    # worldmem.reset()
+    self_frames = None
+    self_poses = None
+    self_memory_c2w = None
+    self_frame_idx = None
     memory_frames = []
     memory_frames.append(load_image_as_tensor(DEFAULT_IMAGE))
     memory_curr_frame = 0