update

asset3d_gen/models/delight_model.py

@@ -28,6 +28,7 @@ class DelightingModel(object):
         device: str = "cuda",
         seed: int = 0,
     ) -> None:
+        self.model_path = model_path
         self.image_guide_scale = image_guide_scale
         self.text_guide_scale = text_guide_scale
         self.num_infer_step = num_infer_step
@@ -37,29 +38,30 @@ class DelightingModel(object):
         )
         self.seed = seed
         self.device = device
-
-        if model_path is None:
-            suffix = "hunyuan3d-delight-v2-0"
-            model_path = snapshot_download(
-                repo_id="tencent/Hunyuan3D-2", allow_patterns=f"{suffix}/*"
+        self.pipeline = None # lazy load model adapt to @spaces.GPU
+        
+    def _lazy_init_pipeline(self):
+        if self.pipeline is None:
+            model_path = self.model_path
+            if model_path is None:
+                suffix = "hunyuan3d-delight-v2-0"
+                model_path = snapshot_download(
+                    repo_id="tencent/Hunyuan3D-2", allow_patterns=f"{suffix}/*"
+                )
+                model_path = os.path.join(model_path, suffix)
+
+            pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+                model_path,
+                torch_dtype=torch.float16,
+                safety_checker=None,
             )
-            model_path = os.path.join(model_path, suffix)
-
-        pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
-            model_path,
-            torch_dtype=torch.float16,
-            safety_checker=None,
-        )
-        pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
-            pipeline.scheduler.config
-        )
-        pipeline.set_progress_bar_config(disable=True)
+            pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+                pipeline.scheduler.config
+            )
+            pipeline.set_progress_bar_config(disable=True)
 
-        pipeline.to(self.device, torch.float16)
-        # pipeline.enable_model_cpu_offload()
-        # pipeline.enable_xformers_memory_efficient_attention()
-        
-        self.pipeline = pipeline
+            pipeline.to(self.device, torch.float16)
+            self.pipeline = pipeline
 
     def recenter_image(
         self, image: Image.Image, border_ratio: float = 0.2
@@ -110,6 +112,8 @@ class DelightingModel(object):
         preprocess: bool = False,
         target_wh: tuple[int, int] = None,
     ) -> Image.Image:
+        self._lazy_init_pipeline()
+        
         if isinstance(image, str):
             image = Image.open(image)
         elif isinstance(image, np.ndarray):
@@ -153,7 +157,7 @@ if __name__ == "__main__":
     delighting_model = DelightingModel(
         # model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
     )
-    image_path = "scripts/apps/assets/example_image/room_bottle_002.jpeg"
+    image_path = "/home/users/xinjie.wang/xinjie/asset3d-gen/scripts/apps/assets/example_image/room_bottle_002.jpeg"
     image = delighting_model(
         image_path, preprocess=True, target_wh=(512, 512)
     )  # noqa