Remove create model at import time

demo_gradio.py

@@ -14,32 +14,61 @@ import numpy as np
 import colorsys
 
 
-# Load model (once, to avoid reloading)
-@spaces.GPU
-def load_model():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    args = get_argparser().parse_args()
-    args.zero_shot = True
-    model = DataParallel(build_model(args).to(device))
-    WEIGHTS_PATH = hf_hub_download(
-        repo_id="jerpelhan/geco2-assets",
-        filename="weights/CNTQG_multitrain_ca44.pth",
-        repo_type="dataset",
-    )
+_MODEL = None
+_ARGS = None
+_WEIGHTS_PATH = None
+
+def _get_args():
+    global _ARGS
+    if _ARGS is None:
+        args = get_argparser().parse_args()
+        args.zero_shot = True
+        _ARGS = args
+    return _ARGS
+
+def _get_weights_path():
+    global _WEIGHTS_PATH
+    if _WEIGHTS_PATH is None:
+        _WEIGHTS_PATH = hf_hub_download(
+            repo_id="jerpelhan/geco2-assets",
+            filename="weights/CNTQG_multitrain_ca44.pth",
+            repo_type="dataset",
+        )
+    return _WEIGHTS_PATH
+
+def get_model_on_device(device: torch.device):
+    """
+    Lazily build and load model, then move to the requested device.
+    IMPORTANT: model is constructed/loaded without initializing CUDA in the main process.
+    This function will be called from inside the @spaces.GPU worker.
+    """
+    global _MODEL
+    if _MODEL is None:
+        args = _get_args()
+
+        # Build on CPU first to avoid CUDA init in the wrong process
+        model = build_model(args)
+        model = DataParallel(model)  # wrap before loading; matches your original
 
-    model.load_state_dict(torch.load(WEIGHTS_PATH, map_location="cpu", weights_only=True)["model"], strict=False)
-    model.eval()
-    return model, device
+        weights_path = _get_weights_path()
+        ckpt = torch.load(weights_path, map_location="cpu", weights_only=True)
+        state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
+        model.load_state_dict(state, strict=False)
 
+        model.eval()
+        _MODEL = model
 
-model, device = load_model()
+    # Ensure correct device for this invocation
+    _MODEL = _MODEL.to(device)
+    return _MODEL
 
 
 # **Function to Process Image Once**
 @spaces.GPU
 def process_image_once(inputs, enable_mask):
     model.module.return_masks = enable_mask
-
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = get_model_on_device(device)
     image = inputs["image"]
     drawn_boxes = inputs["points"]
     image_tensor = torch.tensor(image).to(device)