update

viewcrafter.py

@@ -72,10 +72,14 @@ class ViewCrafter:
         view_masks = renderer(point_cloud_mask)
         return images, view_masks
     
-    def run_render(self, pcd, imgs,masks, H, W, camera_traj,num_views):
+    def run_render(self, pcd, imgs,masks, H, W, camera_traj,num_views,use_cpu=False):
+        if use_cpu:
+            device = torch.device("cpu")
+        else:
+            device = self.device
         render_setup = setup_renderer(camera_traj, image_size=(H,W))
         renderer = render_setup['renderer']
-        render_results, viewmask = self.render_pcd(pcd, imgs, masks, num_views,renderer,self.device)
+        render_results, viewmask = self.render_pcd(pcd, imgs, masks, num_views,renderer,device)
         return render_results, viewmask
 
     
@@ -119,7 +123,7 @@ class ViewCrafter:
             ## FIXME hard coded candidate view数量, 以left为例,第一次迭代从[左,左上]中选取, 从第二次开始可以从[左,左上,左下]中选取
             num_candidates = 2
             candidate_poses,thetas,phis = generate_candidate_poses(c2ws, H, W, focals, principal_points, self.opts.d_theta[0], self.opts.d_phi[0],num_candidates, self.device)
-            _, viewmask = self.run_render([pcd[-1]], [imgs[-1]],masks, H, W, candidate_poses,num_candidates)
+            _, viewmask = self.run_render([pcd[-1]], [imgs[-1]],masks, H, W, candidate_poses,num_candidates,use_cpu=True)
             nbv_id = torch.argmin(viewmask.sum(dim=[1,2,3])).item()
             save_image( viewmask.permute(0,3,1,2), os.path.join(self.opts.save_dir,f"candidate_mask0_nbv{nbv_id}.png"), normalize=True, value_range=(0, 1))
             theta_nbv = thetas[nbv_id]
@@ -139,11 +143,14 @@ class ViewCrafter:
                     r = [float(i) for i in lines[2].split()]
             else: 
                 phi, theta, r = self.gradio_traj
-            camera_traj,num_views = generate_traj_txt(c2ws, H, W, focals, principal_points, phi, theta, r,self.opts.video_length, self.device,viz_traj=True, save_dir = self.opts.save_dir)
+            device = torch.device("cpu")
+            camera_traj,num_views = generate_traj_txt(c2ws, H, W, focals, principal_points, phi, theta, r,self.opts.video_length, device,viz_traj=True, save_dir = self.opts.save_dir)
+            # camera_traj,num_views = generate_traj_txt(c2ws, H, W, focals, principal_points, phi, theta, r,self.opts.video_length, self.device,viz_traj=True, save_dir = self.opts.save_dir)
         else:
             raise KeyError(f"Invalid Mode: {self.opts.mode}")
 
-        render_results, viewmask = self.run_render([pcd[-1]], [imgs[-1]],masks, H, W, camera_traj,num_views)
+        render_results, viewmask = self.run_render([pcd[-1]], [imgs[-1]],masks, H, W, camera_traj,num_views,use_cpu=True)
+        render_results = render_results.to(self.device)
         render_results = F.interpolate(render_results.permute(0,3,1,2), size=(576, 1024), mode='bilinear', align_corners=False).permute(0,2,3,1)
         render_results[0] = self.img_ori
         if self.opts.mode == 'single_view_txt':