new

pipeline/kiss3d_wrapper.py

@@ -102,34 +102,37 @@ def init_wrapper_from_config(config_path):
     # TODO: load pulid model
 
     # init multiview model
-    logger.info('==> Loading multiview diffusion model ...')
-    multiview_device = config_['multiview'].get('device', 'cpu')
-    multiview_pipeline = DiffusionPipeline.from_pretrained(
-        config_['multiview']['base_model'], 
-        custom_pipeline=config_['multiview']['custom_pipeline'],
-        torch_dtype=torch.float16,
-    )
-    multiview_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
-        multiview_pipeline.scheduler.config, timestep_spacing='trailing'
-    )
+    # logger.info('==> Loading multiview diffusion model ...')
+    # multiview_device = config_['multiview'].get('device', 'cpu')
+    # multiview_pipeline = DiffusionPipeline.from_pretrained(
+    #     config_['multiview']['base_model'], 
+    #     custom_pipeline=config_['multiview']['custom_pipeline'],
+    #     torch_dtype=torch.float16,
+    # )
+    # multiview_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+    #     multiview_pipeline.scheduler.config, timestep_spacing='trailing'
+    # )
     
-    # unet_ckpt_path = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="flexgen_19w.ckpt", repo_type="model", token=access_token)
-    unet_ckpt_path = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="flexgen.ckpt", repo_type="model", token=access_token)
-    if unet_ckpt_path is not None:
-        state_dict = torch.load(unet_ckpt_path, map_location='cpu')
-        # state_dict = {k[10:]: v for k, v in state_dict.items() if k.startswith('unet.unet.')}
-        multiview_pipeline.unet.load_state_dict(state_dict, strict=True)
+    # # unet_ckpt_path = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="flexgen_19w.ckpt", repo_type="model", token=access_token)
+    # unet_ckpt_path = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="flexgen.ckpt", repo_type="model", token=access_token)
+    # if unet_ckpt_path is not None:
+    #     state_dict = torch.load(unet_ckpt_path, map_location='cpu')
+    #     # state_dict = {k[10:]: v for k, v in state_dict.items() if k.startswith('unet.unet.')}
+    #     multiview_pipeline.unet.load_state_dict(state_dict, strict=True)
+
+    # multiview_pipeline.to(multiview_device)
+    # logger.warning(f"GPU memory allocated after load multiview model on {multiview_device}: {torch.cuda.memory_allocated(device=multiview_device) / 1024**3} GB")
+    multiview_pipeline = None
 
-    multiview_pipeline.to(multiview_device)
-    logger.warning(f"GPU memory allocated after load multiview model on {multiview_device}: {torch.cuda.memory_allocated(device=multiview_device) / 1024**3} GB")
 
     # load caption model
-    logger.info('==> Loading caption model ...')
-    caption_device = config_['caption'].get('device', 'cpu')
-    caption_model = AutoModelForCausalLM.from_pretrained(config_['caption']['base_model'], \
-                    torch_dtype=torch.bfloat16, trust_remote_code=True).to(caption_device)
-    caption_processor = AutoProcessor.from_pretrained(config_['caption']['base_model'], trust_remote_code=True)
-    logger.warning(f"GPU memory allocated after load caption model on {caption_device}: {torch.cuda.memory_allocated(device=caption_device) / 1024**3} GB")
+    # logger.info('==> Loading caption model ...')
+    # caption_device = config_['caption'].get('device', 'cpu')
+    # caption_model = AutoModelForCausalLM.from_pretrained(config_['caption']['base_model'], \
+    #                 torch_dtype=torch.bfloat16, trust_remote_code=True).to(caption_device)
+    # caption_processor = AutoProcessor.from_pretrained(config_['caption']['base_model'], trust_remote_code=True)
+    # logger.warning(f"GPU memory allocated after load caption model on {caption_device}: {torch.cuda.memory_allocated(device=caption_device) / 1024**3} GB")
+    caption_processor = None
 
     # load reconstruction model
     logger.info('==> Loading reconstruction model ...')
@@ -147,7 +150,7 @@ def init_wrapper_from_config(config_path):
 
     # load llm
     llm_configs = config_.get('llm', None)
-    if llm_configs is not None:
+    if llm_configs is not None and False:
         logger.info('==> Loading LLM ...')
         llm_device = llm_configs.get('device', 'cpu')
         llm, llm_tokenizer = load_llm_model(llm_configs['base_model'])