new

pipeline/kiss3d_wrapper.py

@@ -141,20 +141,20 @@ def init_wrapper_from_config(config_path):
     caption_model = None
 
     # load reconstruction model
-    # logger.info('==> Loading reconstruction model ...')
-    # recon_device = config_['reconstruction'].get('device', 'cpu')
-    # recon_model_config = OmegaConf.load(config_['reconstruction']['model_config'])
-    # recon_model = instantiate_from_config(recon_model_config.model_config)
-    # # load recon model checkpoint
-    # model_ckpt_path = hf_hub_download(repo_id="LTT/PRM", filename="final_ckpt.ckpt", repo_type="model")
-    # state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']
-    # state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}
-    # recon_model.load_state_dict(state_dict, strict=True)
-    # recon_model.to(recon_device)
-    # recon_model.eval()
+    logger.info('==> Loading reconstruction model ...')
+    recon_device = config_['reconstruction'].get('device', 'cpu')
+    recon_model_config = OmegaConf.load(config_['reconstruction']['model_config'])
+    recon_model = instantiate_from_config(recon_model_config.model_config)
+    # load recon model checkpoint
+    model_ckpt_path = hf_hub_download(repo_id="LTT/PRM", filename="final_ckpt.ckpt", repo_type="model")
+    state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']
+    state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}
+    recon_model.load_state_dict(state_dict, strict=True)
+    recon_model.to(recon_device)
+    recon_model.eval()
     # logger.warning(f"GPU memory allocated after load reconstruction model on {recon_device}: {torch.cuda.memory_allocated(device=recon_device) / 1024**3} GB")
-    recon_model = None
-    recon_model_config = None
+    # recon_model = None
+    # recon_model_config = None
     # load llm
     llm_configs = config_.get('llm', None)
     if llm_configs is not None and False: