update

common.py

@@ -88,12 +88,14 @@ os.environ["GRADIO_ANALYTICS_ENABLED"] = "false"
 os.environ.setdefault("OPENAI_API_KEY", "sk-placeholder")
 MAX_SEED = 100000
 
+# Global variables for lazy initialization
+_RBG_REMOVER = None
+_RBG14_REMOVER = None
+
 # DELIGHT = DelightingModel()
 # IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
 # IMAGESR_MODEL = ImageStableSR()
 if os.getenv("GRADIO_APP").startswith("imageto3d"):
-    RBG_REMOVER = RembgRemover()
-    RBG14_REMOVER = BMGG14Remover()
     SAM_PREDICTOR = SAMPredictor(model_type="vit_h", device="cpu")
     if "sam3d" in os.getenv("GRADIO_APP"):
         PIPELINE = Sam3dInference(device="cuda")
@@ -111,8 +113,6 @@ if os.getenv("GRADIO_APP").startswith("imageto3d"):
     )
     os.makedirs(TMP_DIR, exist_ok=True)
 elif os.getenv("GRADIO_APP").startswith("textto3d"):
-    RBG_REMOVER = RembgRemover()
-    RBG14_REMOVER = BMGG14Remover()
     if "sam3d" in os.getenv("GRADIO_APP"):
         PIPELINE = Sam3dInference(device="cuda")
     else:
@@ -161,11 +161,15 @@ def end_session(req: gr.Request) -> None:
         shutil.rmtree(user_dir)
 
 
+@spaces.GPU
 def preprocess_image_fn(
     image: str | np.ndarray | Image.Image,
     rmbg_tag: str = "rembg",
     preprocess: bool = True,
 ) -> tuple[Image.Image, Image.Image]:
+    """Preprocess image with lazy model initialization to avoid CUDA init at import time."""
+    global _RBG_REMOVER, _RBG14_REMOVER
+
     if isinstance(image, str):
         image = Image.open(image)
     elif isinstance(image, np.ndarray):
@@ -173,7 +177,16 @@ def preprocess_image_fn(
 
     image_cache = image.copy()  # resize_pil(image.copy(), 1024)
 
-    bg_remover = RBG_REMOVER if rmbg_tag == "rembg" else RBG14_REMOVER
+    # Lazy initialization - models are created on first call within @spaces.GPU context
+    if rmbg_tag == "rembg":
+        if _RBG_REMOVER is None:
+            _RBG_REMOVER = RembgRemover()
+        bg_remover = _RBG_REMOVER
+    else:
+        if _RBG14_REMOVER is None:
+            _RBG14_REMOVER = BMGG14Remover()
+        bg_remover = _RBG14_REMOVER
+
     image = bg_remover(image)
     image = keep_largest_connected_component(image)