separate generation process

app.py

@@ -35,7 +35,6 @@ GRADIO_OBJ_SHADING_PATH = 'gradio_output_shading.obj'
 #opt = tyro.cli(AllConfigs)
 
 ckpt_path = hf_hub_download(repo_id="rgxie/LDM", filename="LDM6v01.ckpt")
-#ckpt_path = '/ssd3/xrg/tensor23d/pretrained/last_6view_0610_14.ckpt'
 
 opt = Options(
     input_size=512, 
@@ -83,8 +82,6 @@ if opt.resume is not None:
     print(f'[INFO] load resume success!')
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-print('first device')
-print(device)
 model = model.half().to(device)
 model.eval()
 
@@ -141,18 +138,14 @@ pipe_image_plus = pipe_image_plus.to(device)
 # load rembg
 bg_remover = rembg.new_session()
 
-# process function
-@spaces.GPU
-def process(condition_input_image, prompt, prompt_neg='', input_elevation=0, input_num_steps=30, input_seed=42, mv_moedl_option=None):
 
+@spaces.GPU
+def generate_mv(condition_input_image, prompt, prompt_neg='', input_elevation=0, input_num_steps=30, input_seed=42, mv_moedl_option=None):
     # seed
     kiui.seed_everything(input_seed)
 
     os.makedirs(os.path.join(opt.workspace, "gradio"), exist_ok=True)
     output_video_path = os.path.join(opt.workspace,"gradio", GRADIO_VIDEO_PATH)
-    output_obj_rgb_path = os.path.join(opt.workspace,"gradio", GRADIO_OBJ_PATH)
-    output_obj_albedo_path = os.path.join(opt.workspace,"gradio", GRADIO_OBJ_ALBEDO_PATH)
-    output_obj_shading_path = os.path.join(opt.workspace,"gradio", GRADIO_OBJ_SHADING_PATH)
     
     # text-conditioned
     if condition_input_image is None:
@@ -199,7 +192,15 @@ def process(condition_input_image, prompt, prompt_neg='', input_elevation=0, inp
             mv_image_grid = rearrange(mv_image, 'c (n h) (m w) -> (m h) (n w) c', n=3, m=2).numpy()
             mv_image = rearrange(mv_image, 'c (n h) (m w) -> (n m) h w c', n=3, m=2).numpy()
             input_image = mv_image
-               
+    return mv_image_grid, processed_image, input_image 
+
+@spaces.GPU
+def generate_3d(input_image, condition_input_image, mv_moedl_option=None, input_seed=42):
+    kiui.seed_everything(input_seed)
+    
+    output_obj_rgb_path = os.path.join(opt.workspace,"gradio", GRADIO_OBJ_PATH)
+    output_obj_albedo_path = os.path.join(opt.workspace,"gradio", GRADIO_OBJ_ALBEDO_PATH)
+    output_obj_shading_path = os.path.join(opt.workspace,"gradio", GRADIO_OBJ_SHADING_PATH)
     # generate gaussians
      # [4, 256, 256, 3], float32
     input_image = torch.from_numpy(input_image).permute(0, 3, 1, 2).float().to(device) # [4, 3, 256, 256]
@@ -274,7 +275,8 @@ def process(condition_input_image, prompt, prompt_neg='', input_elevation=0, inp
             save_obj(vertices, faces, vertex_colors[2], output_obj_shading_path)
         
 
-    return mv_image_grid, processed_image, output_obj_rgb_path, output_obj_albedo_path, output_obj_shading_path
+    return output_obj_rgb_path, output_obj_albedo_path, output_obj_shading_path
+
 
 # gradio UI
 
@@ -385,7 +387,13 @@ with block:
                 )
 
             
-        button_gen.click(process, inputs=[condition_input_image, input_text, input_neg_text, input_elevation, input_num_steps, input_seed,mv_moedl_option], outputs=[mv_image_grid,processed_image, output_obj_rgb_path, output_obj_albedo_path, output_obj_shading_path])
-    
-    
+        input_image = gr.State()
+        button_gen.click(fn=generate_mv, inputs=[condition_input_image, input_text, input_neg_text, input_elevation, input_num_steps, input_seed, mv_moedl_option], 
+                         outputs=[mv_image_grid, processed_image, input_image],).success(
+                            fn=generate_3d,
+                            inputs=[input_image, condition_input_image, mv_moedl_option, input_seed], 
+                            outputs=[output_obj_rgb_path, output_obj_albedo_path, output_obj_shading_path] ,
+                         )
+        
+        
 block.launch(server_name="0.0.0.0", share=False)

core/models.py

@@ -50,13 +50,9 @@ class LTRFM_NeRF(nn.Module):
         )
         
         aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]]).cuda()
-        print(aabb.device)
         grid_size = torch.tensor([opt.splat_size, opt.splat_size, opt.splat_size]).cuda()
         near_far =torch.tensor([opt.znear, opt.zfar]).cuda()
         
-        print(device)
-        print('aabb')
-        print(aabb.device)
         # tensorf Renderer
         self.tensorRF = TensorVMSplit_NeRF(aabb, grid_size, density_n_comp=opt.density_n_comp,appearance_n_comp=opt.app_n_comp,app_dim=opt.app_dim,\
             density_dim=opt.density_dim,near_far=near_far, shadingMode=opt.shadingMode, pos_pe=opt.pos_pe, view_pe=opt.view_pe, fea_pe=opt.fea_pe)
@@ -127,9 +123,6 @@ class LTRFM_NeRF(nn.Module):
         xyz_samples=xyz_samples.permute(0,2,3,4,1)
         xyz_samples=xyz_samples.view(1,-1,1,3)
         
-        print("app and xyz")
-        print(planes['app_planes'].device)
-        print(xyz_samples.device)
 
         grid_out = self.tensorRF.predict_sdf(planes,xyz_samples)
         grid_out['sigma']=grid_out['sigma'].view(grid_size,grid_size,grid_size).float()
@@ -191,68 +184,6 @@ class LTRFM_NeRF(nn.Module):
         return vertices, faces, uvs, mesh_tex_idx, [texture_map,texture_map_albedo]
 
     
-    def forward(self, data, step_ratio=1):
-        # data: output of the dataloader
-        # return: loss
-        #self.set_beta(data['t'])
-        results = {}
-        loss = 0
-
-        images = data['input'] # [B, 4, 9, h, W], input features
-        
-        # use the first view to predict gaussians
-        svd_volume = self.forward_svd_volume(images,data) # [B, N, 14]
-
-        results['svd_volume'] = svd_volume
-
-        # always use white bg
-        bg_color = torch.ones(3, dtype=torch.float32).to(device)
-        
-        # use the other views for rendering and supervision
-        results = self.tensorRF(svd_volume, data['all_rays_o'], data['all_rays_d'],is_train=True, bg_color=bg_color, N_samples=self.opt.n_sample)
-        pred_shading = results['image'] # [B, V, C, output_size, output_size]
-        pred_alphas = results['alpha'] # [B, V, 1, output_size, output_size]
-        pred_albedos = results['albedo'] # [B, V, C, output_size, output_size]
-        
-        pred_images = pred_shading*pred_albedos
-
-        results['images_pred'] = pred_images
-        results['alphas_pred'] = pred_alphas
-        results['pred_albedos'] = pred_albedos
-        results['pred_shading'] = pred_shading
-        
-
-        gt_images = data['images_output'] # [B, V, 3, output_size, output_size], ground-truth novel views
-        gt_albedos = data['albedos_output'] # [B, V, 3, output_size, output_size], ground-truth novel views
-        gt_masks = data['masks_output'] # [B, V, 1, output_size, output_size], ground-truth masks
-
-        gt_images = gt_images * gt_masks + bg_color.view(1, 1, 3, 1, 1) * (1 - gt_masks)
-        gt_albedos = gt_albedos * gt_masks + bg_color.view(1, 1, 3, 1, 1) * (1 - gt_masks)
-
-        loss_mse = F.mse_loss(pred_images, gt_images) + F.mse_loss(pred_alphas, gt_masks) + F.mse_loss(pred_albedos, gt_albedos)
-        loss = loss + loss_mse
-
-        # eikonal_loss = ((results['eik_grads'].norm(2, dim=1) - 1) ** 2).mean()
-        # loss = loss+ 0.1*eikonal_loss
-
-        if self.opt.lambda_lpips > 0:
-            loss_lpips = self.lpips_loss(
-                F.interpolate(gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False), 
-                F.interpolate(pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False),
-            ).mean()
-            results['loss_lpips'] = loss_lpips
-            loss = loss + self.opt.lambda_lpips * loss_lpips
-            
-        results['loss'] = loss
-
-        # metric
-        with torch.no_grad():
-            psnr = -10 * torch.log10(torch.mean((pred_images.detach() - gt_images) ** 2))
-            results['psnr'] = psnr
-
-        return results
-    
-    
     def render_frame(self, data):
         # data: output of the dataloader
         # return: loss
@@ -631,70 +562,6 @@ class LTRFM_Mesh(nn.Module):
         }
         return out
     
-    def forward(self, data, step_ratio=1):
-        # data: output of the dataloader
-        # return: loss
-
-        results = {}
-        loss = 0
-
-        images = data['input'] # [B, 4, 9, h, W], input features
-        
-        # use the first view to predict gaussians
-        svd_volume = self.forward_svd_volume(images,data) # [B, N, 14]
-        
-        results['svd_volume'] = svd_volume
-        
-        # return the rendered images
-        results = self.forward_geometry(svd_volume, data['w2c'], self.opt.output_size)
-
-
-        # always use white bg
-        bg_color = torch.ones(3, dtype=torch.float32).to(device)
-        
-        # use the other views for rendering and supervision
-        #results = self.tensorRF(svd_volume, data['all_rays_o'], data['all_rays_d'],is_train=True, bg_color=bg_color, N_samples=self.opt.n_sample)
-        
-        
-        pred_shading = results['image'] # [B, V, C, output_size, output_size]
-        pred_alphas = results['mask'] # [B, V, 1, output_size, output_size]
-        pred_albedos = results['albedo'] # [B, V, C, output_size, output_size]
-        
-        pred_images=pred_shading*pred_albedos
-
-        results['images_pred'] = pred_images
-        results['alphas_pred'] = pred_alphas
-        results['pred_albedos'] = pred_albedos
-        results['pred_shading'] = pred_shading
-        
-
-        gt_images = data['images_output'] # [B, V, 3, output_size, output_size], ground-truth novel views
-        gt_albedos = data['albedos_output'] # [B, V, 3, output_size, output_size], ground-truth novel views
-        gt_masks = data['masks_output'] # [B, V, 1, output_size, output_size], ground-truth masks
-
-        gt_images = gt_images * gt_masks + bg_color.view(1, 1, 3, 1, 1) * (1 - gt_masks)
-        gt_albedos = gt_albedos * gt_masks + bg_color.view(1, 1, 3, 1, 1) * (1 - gt_masks)
-
-        loss_mse = F.mse_loss(pred_images, gt_images) + F.mse_loss(pred_alphas, gt_masks) + F.mse_loss(pred_albedos, gt_albedos)
-        loss = loss + loss_mse
-
-        if self.opt.lambda_lpips > 0:
-            loss_lpips = self.lpips_loss(
-                F.interpolate(gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False), 
-                F.interpolate(pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False),
-            ).mean()
-            results['loss_lpips'] = loss_lpips
-            loss = loss + self.opt.lambda_lpips * loss_lpips
-            
-        results['loss'] = loss
-
-        # metric
-        with torch.no_grad():
-            psnr = -10 * torch.log10(torch.mean((pred_images.detach() - gt_images) ** 2))
-            results['psnr'] = psnr
-
-        return results
-    
     
     def render_frame(self, data):
         # data: output of the dataloader

core/tensorBase.py

@@ -205,23 +205,7 @@ class TensorBase(torch.nn.Module):
         self.near_far = near_far
         self.step_ratio = 0.9 #step_ratio原作0.5
 
-        print("aabb", self.aabb.view(-1))
-        print("grid size", gridSize)
-        self.aabbSize = self.aabb[1] - self.aabb[0]
-        print(self.aabbSize.device)
-        self.invaabbSize = 2.0/self.aabbSize
-        print(self.invaabbSize.device)
-        # self.invaabbSize = self.invaabbSize.to(self.aabb.device)
-        # print(self.invaabbSize.device)
-        self.gridSize= gridSize.float()
-        self.units=self.aabbSize / (self.gridSize-1)
-        self.stepSize=torch.mean(self.units)*self.step_ratio  # TBD step_ratio? why so small 0.5
-        self.aabbDiag = torch.sqrt(torch.sum(torch.square(self.aabbSize)))
-        self.nSamples=int((self.aabbDiag / self.stepSize).item()) + 1
-        print("sampling step size: ", self.stepSize)
-        print("sampling number: ", self.nSamples)
-
-        # self.update_stepSize(gridSize)
+        self.update_stepSize(gridSize)
 
         self.matMode = [[0,1], [0,2], [1,2]]
         self.vecMode =  [2, 1, 0]
@@ -252,14 +236,8 @@ class TensorBase(torch.nn.Module):
         print(self.renderModule)
 
     def update_stepSize(self, gridSize):
-        print("aabb", self.aabb.view(-1))
-        print("grid size", gridSize)
         self.aabbSize = self.aabb[1] - self.aabb[0]
-        print(self.aabbSize.device)
         self.invaabbSize = 2.0/self.aabbSize
-        print(self.invaabbSize.device)
-        # self.invaabbSize = self.invaabbSize.to(self.aabb.device)
-        # print(self.invaabbSize.device)
         self.gridSize= gridSize.float()
         self.units=self.aabbSize / (self.gridSize-1)
         self.stepSize=torch.mean(self.units)*self.step_ratio  # TBD step_ratio? why so small 0.5
@@ -281,14 +259,8 @@ class TensorBase(torch.nn.Module):
         pass
     
     def normalize_coord(self, xyz_sampled):
-        print("debug")
-        
-        print(xyz_sampled.device)
-        print(self.aabb[0].device)
-        print(self.invaabbSize.device)
         if xyz_sampled.device!=self.invaabbSize.device:
             self.invaabbSize=self.invaabbSize.to(xyz_sampled.device)
-        
         return (xyz_sampled-self.aabb[0]) * self.invaabbSize - 1
 
     def get_optparam_groups(self, lr_init_spatial = 0.02, lr_init_network = 0.001):