Update registry.py

registry.py

@@ -17,17 +17,18 @@ def get_model(task: str, model_key: str, device="cpu"):
     """
     logger.info(f"Preparing model wrapper '{model_key}' for task '{task}' on device '{device}'")
 
-    if task == "detection":
-        from models.detection.detector import ObjectDetector
-        return ObjectDetector(model_key=model_key, device=device)
-
-    elif task == "segmentation":
-        from models.segmentation.segmenter import Segmenter
-        return Segmenter(model_key=model_key, device=device)
-
-    elif task == "depth":
-        from models.depth.depth_estimator import DepthEstimator
-        return DepthEstimator(model_key=model_key, device=device)
-
-    else:
-        raise ValueError(f"Unsupported task '{task}'")
+    try:
+        if task == "detection":
+            from models.detection.detector import ObjectDetector
+            return ObjectDetector(model_key=model_key, device=device)
+        elif task == "segmentation":
+            from models.segmentation.segmenter import Segmenter
+            return Segmenter(model_key=model_key, device=device)
+        elif task == "depth":
+            from models.depth.depth_estimator import DepthEstimator
+            return DepthEstimator(model_key=model_key, device=device)
+        else:
+            raise ValueError(f"Unsupported task '{task}'")
+    except Exception as e:
+        logger.error(f"Error loading model '{model_key}' for task '{task}': {e}")
+        raise