n

pipeline/kiss3d_wrapper.py

@@ -75,8 +75,8 @@ def init_wrapper_from_config(config_path):
     else:
         flux_pipe = FluxImg2ImgPipeline.from_pretrained(flux_base_model_pth, torch_dtype=dtype_[flux_dtype], token=access_token)
     # flux_pipe.enable_vae_slicing()
-    # flux_pipe.enable_xformers_memory_efficient_attention()
-    flux_pipe.enable_sequential_cpu_offload()
+    flux_pipe.enable_xformers_memory_efficient_attention()
+    # flux_pipe.enable_sequential_cpu_offload()
     # load flux model and controlnet
     if flux_controlnet_pth is not None and False:
         flux_controlnet = FluxControlNetModel.from_pretrained(flux_controlnet_pth, torch_dtype=torch.bfloat16)
@@ -139,20 +139,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: