new

pipeline/kiss3d_wrapper.py

@@ -67,7 +67,7 @@ def init_wrapper_from_config(config_path):
     flux_dtype = config_['flux'].get('dtype', 'bf16')
     flux_controlnet_pth = config_['flux'].get('controlnet', None)
     # flux_lora_pth = config_['flux'].get('lora', None)
-    flux_lora_pth = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="rgb_normal_large.safetensors", repo_type="model", token=access_token)
+    flux_lora_pth = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="rgb_normal.safetensors", repo_type="model", token=access_token)
     flux_redux_pth = config_['flux'].get('redux', None)
     # taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype_[flux_dtype]).to(flux_device)
     if flux_base_model_pth.endswith('safetensors'):
@@ -102,23 +102,23 @@ def init_wrapper_from_config(config_path):
     # logger.warning(f"GPU memory allocated after load flux model on {flux_device}: {torch.cuda.memory_allocated(device=flux_device) / 1024**3} GB")
 
     # 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.ckpt", repo_type="model", token=access_token)
-    # if unet_ckpt_path is not None:
-    #     state_dict = torch.load(unet_ckpt_path, map_location='cpu')
-    #     multiview_pipeline.unet.load_state_dict(state_dict, strict=True)
+    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')
+        multiview_pipeline.unet.load_state_dict(state_dict, strict=True)
 
-    # multiview_pipeline.to(multiview_device)
+    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