new

pipeline/kiss3d_wrapper.py

@@ -137,19 +137,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
     # load llm
     llm_configs = config_.get('llm', None)
     if llm_configs is not None and False:
@@ -503,8 +504,8 @@ class kiss3d_wrapper(object):
             'num_inference_steps': num_inference_steps,
             'guidance_scale': 3.5,
             'num_images_per_prompt': 1,
-            'width': 2048//2,
-            'height': 1024//2,
+            'width': 2048,
+            'height': 1024,
             'output_type': 'np',
             'generator': generator,
             'joint_attention_kwargs': {"scale": lora_scale}

pipeline/pipeline_config/default.yaml

@@ -14,20 +14,20 @@ multiview:
   unet: "./checkpoint/zero123++/flexgen_19w.ckpt"
   num_inference_steps: 50
   seed: 42
-  device: 'cuda:1'
+  device: 'cuda:0'
 
 reconstruction:
   model_config: "./models/lrm/config/PRM_inference.yaml"
   base_model: "./checkpoint/lrm/final_ckpt.ckpt"
-  device: 'cuda:1'
+  device: 'cuda:0'
 
 caption:
   base_model: "multimodalart/Florence-2-large-no-flash-attn"
-  device: 'cuda:1'
+  device: 'cuda:0'
 
 llm:
   base_model: "Qwen/Qwen2-7B-Instruct"
-  device: 'cuda:1'
+  device: 'cuda:0'
 
 use_zero_gpu: false # for huggingface demo only
 3d_bundle_templates: './init_3d_Bundle'