update

algorithms/worldmem/df_video.py

@@ -792,39 +792,46 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
     @torch.no_grad()
     def interactive(self, first_frame, curr_actions, first_pose, context_frames_idx, device,
-                    self_frames, self_poses, self_memory_c2w, self_frame_idx):
+                    self_frames, self_actions, self_poses, self_memory_c2w, self_frame_idx):
     
         condition_similar_length = self.condition_similar_length
 
         if self_frames is None:
+            first_frame = torch.from_numpy(first_frame)
+            curr_actions = torch.from_numpy(curr_actions)
+            first_pose = torch.from_numpy(first_pose)
             first_frame_encode = self.encode(first_frame[None, None].to(device))
             self_frames = first_frame_encode.cpu()
-            self.actions = curr_actions[None, None].to(device)
+            self_actions = curr_actions[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.cpu(), self_frames.cpu().numpy(), self_poses.cpu(), self_memory_c2w.cpu(), self_frame_idx.cpu()
+            return first_frame.cpu(), self_frames.cpu().numpy(), self_actions.cpu().numpy(), self_poses.cpu().numpy(), self_memory_c2w.cpu().numpy(), self_frame_idx.cpu().numpy()
         else:
+            self_frames = torch.from_numpy(self_frames)
+            self_actions = torch.from_numpy(self_actions).to(device)
+            self_poses = torch.from_numpy(self_poses).to(device)
+            self_memory_c2w = torch.from_numpy(self_memory_c2w).to(device)
+            self_frame_idx = torch.from_numpy(self_frame_idx).to(device)
+            curr_actions = curr_actions.to(device)
+
             last_frame = self_frames[-1].clone()
-            self_poses = self_poses.to(device)
-            self_memory_c2w = self_memory_c2w.to(device)
-            self_frame_idx = self_frame_idx.to(device)
             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)
+            new_pose_condition_offset = self.pose_prediction_model(last_frame.to(device), curr_actions[None], last_pose_condition)
             
             new_pose_condition_offset[:,3:] = torch.round(new_pose_condition_offset[:,3:])
             new_pose_condition = last_pose_condition + new_pose_condition_offset
             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_actions = torch.cat([self_actions, curr_actions[None, None]])
+            self_poses = torch.cat([self_poses, new_pose_condition[None]])
             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_memory_c2w = torch.cat([self_memory_c2w, new_c2w_mat[None]])
             self_frame_idx = torch.cat([self_frame_idx, torch.tensor([[context_frames_idx]]).to(device)])
 
-        conditions = self.actions.clone()
+        conditions = self_actions.clone()
         pose_conditions = self_poses.clone()
         c2w_mat = self_memory_c2w .clone()
         frame_idx = self_frame_idx.clone()
@@ -903,7 +910,8 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
         xs_pred = self.decode(xs_pred[n_context_frames:].to(device)).cpu()
 
-        return xs_pred[-1,0].cpu(), self_frames.cpu(), self_poses.cpu(), self_memory_c2w.cpu(), self_frame_idx.cpu()
+        return xs_pred[-1,0].cpu().numpy(), self_frames.cpu().numpy(), self_actions.cpu().numpy(), \
+            self_poses.cpu().numpy(), self_memory_c2w.cpu().numpy(), self_frame_idx.cpu().numpy()
 
 
     def reset(self):

app.py

@@ -177,30 +177,33 @@ load_custom_checkpoint(algo=worldmem.pose_prediction_model, checkpoint_path=cfg.
 worldmem.to("cuda").eval()
 
 
-actions = torch.zeros((1, 25))
-poses = torch.zeros((1, 5))
+actions = np.zeros((1, 25), dtype=np.float32)
+poses = np.zeros((1, 5), dtype=np.float32)
 
 memory_frames.append(load_image_as_tensor(DEFAULT_IMAGE))
 
 self_frames = None
+self_actions = 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, self_frames, self_poses, self_memory_c2w, self_frame_idx):
-    new_frame, self_frames, self_poses, self_memory_c2w, self_frame_idx = worldmem.interactive(first_frame,
+def run_interactive(first_frame, action, first_pose, curr_frame, device, self_frames, self_actions, 
+                            self_poses, self_memory_c2w, self_frame_idx):
+    new_frame, self_frames, self_actions, self_poses, self_memory_c2w, self_frame_idx = worldmem.interactive(first_frame,
                                     action,
                                     first_pose, 
                                     curr_frame, 
                                     device=device,
                                     self_frames=self_frames,
+                                    self_actions=self_actions,
                                     self_poses=self_poses,
                                     self_memory_c2w=self_memory_c2w,
                                     self_frame_idx=self_frame_idx)
-    # return new_frame, self_frames, self_poses, self_memory_c2w, self_frame_idx
-    return self_frames[:,:,0,0,0]
+
+    return new_frame, self_frames, self_actions, self_poses, self_memory_c2w, self_frame_idx
 
 def set_denoising_steps(denoising_steps, sampling_timesteps_state):
     worldmem.sampling_timesteps = denoising_steps
@@ -215,6 +218,7 @@ def generate(keys):
     global input_history
     global memory_curr_frame
     global self_frames
+    global self_actions
     global self_poses
     global self_memory_c2w
     global self_frame_idx
@@ -222,12 +226,13 @@ def generate(keys):
     for i in range(len(actions)):
         memory_curr_frame += 1
 
-        new_frame, self_frames, self_poses, self_memory_c2w, self_frame_idx = run_interactive(memory_frames[0],
+        new_frame, self_frames, self_actions, self_poses, self_memory_c2w, self_frame_idx = run_interactive(memory_frames[0],
                                       actions[i],
                                       None, 
                                       memory_curr_frame, 
                                       device=device,
                                       self_frames=self_frames,
+                                      self_actions=self_actions,
                                       self_poses=self_poses,
                                       self_memory_c2w=self_memory_c2w,
                                       self_frame_idx=self_frame_idx)
@@ -254,6 +259,7 @@ def reset():
     global input_history
     global memory_frames
     global self_frames
+    global self_actions
     global self_poses
     global self_memory_c2w
     global self_frame_idx
@@ -263,16 +269,17 @@ def reset():
     self_memory_c2w = None
     self_frame_idx = None
     memory_frames = []
-    memory_frames.append(load_image_as_tensor(DEFAULT_IMAGE))
+    memory_frames.append(load_image_as_tensor(DEFAULT_IMAGE).numpy())
     memory_curr_frame = 0
     input_history = ""
 
-    self_frames = run_interactive(memory_frames[0],
+    new_frame, self_frames, self_actions, self_poses, self_memory_c2w, self_frame_idx = run_interactive(memory_frames[0],
                                 actions[0],
                                 poses[0],
                                 memory_curr_frame,
                                 device=device,
                                 self_frames=self_frames,
+                                self_actions=self_actions,
                                 self_poses=self_poses,
                                 self_memory_c2w=self_memory_c2w,
                                 self_frame_idx=self_frame_idx)