feat(backend): add lazy model registry and downloader

rohan-pandeyy · rahulharpal1603 · commit 894daca00dc2 · 2026-05-11T11:04:01.000+05:30
diff --git a/.gitignore b/.gitignore
@@ -9,6 +9,7 @@ pnpm-debug.log*
 lerna-debug.log*
 
 backend/app/models/image-generation/*
+backend/app/models/ONNX_Exports/
 
 node_modules
 
diff --git a/backend/.gitignore b/backend/.gitignore
@@ -31,6 +31,7 @@ MANIFEST
 #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 *.manifest
 *.spec
+!PictoPy.spec
 
 # Installer logs
 pip-log.txt
diff --git a/backend/app/models/FaceNet.py b/backend/app/models/FaceNet.py
@@ -1,29 +1,69 @@
-# app/facenet/FaceNet.py
-
+import os
+import threading
 import onnxruntime
+from app.models.model_registry import MODEL_REGISTRY
 from app.utils.FaceNet import FaceNet_util_normalize_embedding
 from app.utils.ONNX import ONNX_util_get_execution_providers
 from app.logging.setup_logging import get_logger
 
-# Initialize logger
 logger = get_logger(__name__)
 
 
 class FaceNet:
     def __init__(self, model_path):
-        self.session = onnxruntime.InferenceSession(
-            model_path, providers=ONNX_util_get_execution_providers()
-        )
-        self.input_tensor_name = self.session.get_inputs()[0].name
-        self.output_tensor_name = self.session.get_outputs()[0].name
+        self.model_path = model_path
+        self._session: onnxruntime.InferenceSession | None = None
+        self.input_tensor_name: str | None = None
+        self.output_tensor_name: str | None = None
+        self._lock = threading.Lock()
+
+    def get_session(self) -> onnxruntime.InferenceSession:
+        # Fast path: capture reference before returning so close() on another
+        # thread between the check and the return cannot cause a None return.
+        session = self._session
+        if session is not None:
+            return session
+
+        with self._lock:
+            if self._session is None:
+                if not os.path.exists(self.model_path):
+                    model_key = None
+                    for key, spec in MODEL_REGISTRY.items():
+                        if spec["filename"] in self.model_path:
+                            model_key = key
+                            break
+                    model_name = (
+                        model_key if model_key else os.path.basename(self.model_path)
+                    )
+                    raise RuntimeError(
+                        f"Model '{model_name}' is not installed. "
+                        "Please install it from Settings → AI Models before using this feature."
+                    )
+
+                self._session = onnxruntime.InferenceSession(
+                    self.model_path, providers=ONNX_util_get_execution_providers()
+                )
+                self.input_tensor_name = self._session.get_inputs()[0].name
+                self.output_tensor_name = self._session.get_outputs()[0].name
+
+            # Capture inside the lock before releasing — prevents close() on
+            # another thread from nulling _session between lock release and return.
+            session = self._session
+
+        return session
 
     def get_embedding(self, preprocessed_image):
-        result = self.session.run(
+        session = self.get_session()
+        result = session.run(
             [self.output_tensor_name], {self.input_tensor_name: preprocessed_image}
         )[0]
         embedding = result[0]
         return FaceNet_util_normalize_embedding(embedding)
 
     def close(self):
-        del self.session
-        logger.info("FaceNet model session closed.")
+        with self._lock:
+            if self._session is not None:
+                self._session = None
+                self.input_tensor_name = None
+                self.output_tensor_name = None
+                logger.info("FaceNet model session closed.")
diff --git a/backend/app/models/YOLO.py b/backend/app/models/YOLO.py
@@ -19,45 +19,77 @@ def __init__(self, path, conf_threshold=0.7, iou_threshold=0.5):
         self.model_path = path
         self.conf_threshold = conf_threshold
         self.iou_threshold = iou_threshold
-        # Create ONNX session once and reuse it
-        self.session = onnxruntime.InferenceSession(
-            self.model_path, providers=ONNX_util_get_execution_providers()
-        )
-
-        # Initialize model info
-        self.get_input_details()
-        self.get_output_details()
+        self._session = None
+        import threading
+        self._lock = threading.Lock()
+
+    def get_session(self):
+        session = self._session
+        if session is not None:
+            return session
+
+        with self._lock:
+            if self._session is None:
+                import os
+                from app.models.model_registry import MODEL_REGISTRY
+
+                if not os.path.exists(self.model_path):
+                    model_key = None
+                    for key, spec in MODEL_REGISTRY.items():
+                        if spec["filename"] in self.model_path:
+                            model_key = key
+                            break
+                    model_name = model_key if model_key else os.path.basename(self.model_path)
+                    raise RuntimeError(
+                        f"Model '{model_name}' is not installed. "
+                        "Please install it from Settings → AI Models before using this feature."
+                    )
+
+                self._session = onnxruntime.InferenceSession(
+                    self.model_path, providers=ONNX_util_get_execution_providers()
+                )
+                # Initialize model info once session is created
+                self.get_input_details()
+                self.get_output_details()
+
+            return self._session
 
     def __call__(self, image):
         return self.detect_objects(image)
 
     def close(self):
-        del self.session  # Clean up the ONNX session
+        with self._lock:
+            if self._session is not None:
+                self._session = None
         logger.info("YOLO model session closed.")
 
     @log_memory_usage
     def detect_objects(self, image):
+        session = self.get_session()
         input_tensor = self.prepare_input(image)
-        outputs = self.inference(input_tensor)
+        outputs = self.inference(input_tensor, session=session)
         self.boxes, self.scores, self.class_ids = self.process_output(outputs)
         return self.boxes, self.scores, self.class_ids
 
-    def inference(self, input_tensor):
-        time.perf_counter()
-        outputs = self.session.run(
+    def inference(self, input_tensor, session=None):
+        start = time.perf_counter()
+        if session is None:
+            session = self.get_session()
+        outputs = session.run(
             self.output_names, {self.input_names[0]: input_tensor}
         )
+        logger.debug("Inference completed in %.4fs", time.perf_counter() - start)
         return outputs
 
     def get_input_details(self):
-        model_inputs = self.session.get_inputs()
+        model_inputs = self._session.get_inputs()
         self.input_names = [inp.name for inp in model_inputs]
         self.input_shape = model_inputs[0].shape
         self.input_height = self.input_shape[2]
         self.input_width = self.input_shape[3]
 
     def get_output_details(self):
-        model_outputs = self.session.get_outputs()
+        model_outputs = self._session.get_outputs()
         self.output_names = [out.name for out in model_outputs]
 
     def prepare_input(self, image):
diff --git a/backend/app/models/model_registry.py b/backend/app/models/model_registry.py
@@ -0,0 +1,108 @@
+from typing import TypedDict, Literal
+import json
+import os
+import sys
+from platformdirs import user_data_dir
+
+FeatureType = Literal["object_detection", "face_detection", "face_embedding"]
+TierType = Literal["nano", "small", "medium", "required"]
+
+
+class ModelSpec(TypedDict):
+    filename: str
+    url: str
+    sha256: str
+    size_mb: float
+    feature: FeatureType
+    tier: TierType
+
+
+MODEL_REGISTRY: dict[str, ModelSpec] = {
+    "yolo_nano": ModelSpec(
+        filename="YOLOv11_Nano.onnx",
+        url="https://github.com/AOSSIE-Org/PictoPy/releases/download/models-v1.0/YOLOv11_Nano.onnx",
+        sha256="64e0a360a854cd1cebf697e38967adb73f24a2c41a86379f8a0bfae1b0c6af0b",
+        size_mb=10.2,
+        feature="object_detection",
+        tier="nano",
+    ),
+    "yolo_nano_face": ModelSpec(
+        filename="YOLOv11_Nano_Face.onnx",
+        url="https://github.com/AOSSIE-Org/PictoPy/releases/download/models-v1.0/YOLOv11_Nano_Face.onnx",
+        sha256="1d09cb0f31d46700a3e80838623aeafa125f2cd0d1c9a12f0b0853bf64b0a83d",
+        size_mb=10.1,
+        feature="face_detection",
+        tier="nano",
+    ),
+    "yolo_small": ModelSpec(
+        filename="YOLOv11_Small.onnx",
+        url="https://github.com/AOSSIE-Org/PictoPy/releases/download/models-v1.0/YOLOv11_Small.onnx",
+        sha256="8bfa953dbe93bef33b09be00da01cb4727f25416cb5d2fdb2a9f1083283b8aaa",
+        size_mb=36.2,
+        feature="object_detection",
+        tier="small",
+    ),
+    "yolo_small_face": ModelSpec(
+        filename="YOLOv11_Small_Face.onnx",
+        url="https://github.com/AOSSIE-Org/PictoPy/releases/download/models-v1.0/YOLOv11_Small_Face.onnx",
+        sha256="213333bbecd049c8e1d8bc63b4799df08d2ec52c8f8a6737b37b75ba839d7c03",
+        size_mb=36.1,
+        feature="face_detection",
+        tier="small",
+    ),
+    "yolo_medium": ModelSpec(
+        filename="YOLOv11_Medium.onnx",
+        url="https://github.com/AOSSIE-Org/PictoPy/releases/download/models-v1.0/YOLOv11_Medium.onnx",
+        sha256="f53a0bf6a141788f329ab92b438045f0ebac73ff10285b9ef0d06551cdbd01ea",
+        size_mb=76.9,
+        feature="object_detection",
+        tier="medium",
+    ),
+    "yolo_medium_face": ModelSpec(
+        filename="YOLOv11_Medium_Face.onnx",
+        url="https://github.com/AOSSIE-Org/PictoPy/releases/download/models-v1.0/YOLOv11_Medium_Face.onnx",
+        sha256="fd3ab085d776ee1ce59fd47def82c2abf49bbfde3a10a1b5b660b76cb41a6912",
+        size_mb=76.6,
+        feature="face_detection",
+        tier="medium",
+    ),
+    "facenet": ModelSpec(
+        filename="FaceNet_128D.onnx",
+        url="https://github.com/AOSSIE-Org/PictoPy/releases/download/models-v1.0/FaceNet_128D.onnx",
+        sha256="c37946ea8cce94141777dcbcccb8a61786c9a7c4f0c9f40471bd27029fa664ed",
+        size_mb=87.0,
+        feature="face_embedding",
+        tier="required",
+    ),
+}
+
+TIER_MODELS: dict[str, list[str]] = {
+    "nano":   ["yolo_nano", "yolo_nano_face"],
+    "small":  ["yolo_small", "yolo_small_face"],
+    "medium": ["yolo_medium", "yolo_medium_face"],
+    "required": ["facenet"], # Required model; not user-selectable
+}
+
+USER_DATA_MODELS = os.path.join(user_data_dir("PictoPy"), "models")
+LOCAL_ONNX_EXPORTS = os.path.join(os.path.dirname(__file__), "ONNX_Exports")
+
+
+def ensure_model_exports_directory() -> None:
+    """Create the active model exports directory if it does not exist."""
+    if getattr(sys, 'frozen', False):
+        os.makedirs(USER_DATA_MODELS, exist_ok=True)
+    else:
+        os.makedirs(LOCAL_ONNX_EXPORTS, exist_ok=True)
+
+
+def get_model_path(key: str) -> str:
+    filename = MODEL_REGISTRY[key]["filename"]
+    ensure_model_exports_directory()
+    
+    # In production (compiled by PyInstaller), use the platform-appropriate user data directory.
+    if getattr(sys, 'frozen', False):
+        return os.path.normpath(os.path.join(USER_DATA_MODELS, filename))
+        
+    # In development, strictly use the local repo folder
+    return os.path.normpath(os.path.join(LOCAL_ONNX_EXPORTS, filename))
+
diff --git a/backend/app/utils/model_bootstrap.py b/backend/app/utils/model_bootstrap.py
@@ -0,0 +1,34 @@
+import asyncio
+import logging
+
+from app.database.metadata import db_get_metadata
+from app.models.model_registry import TIER_MODELS
+from app.utils.model_downloader import ensure_model
+
+logger = logging.getLogger(__name__)
+
+
+def _get_preferred_yolo_tier() -> str:
+    metadata = db_get_metadata() or {}
+    user_preferences = metadata.get("user_preferences") or {}
+    tier = user_preferences.get("YOLO_model_size", "small")
+
+    if tier not in TIER_MODELS:
+        return "small"
+
+    return tier
+
+
+async def _ensure_ai_tagging_models_async() -> None:
+    tier = _get_preferred_yolo_tier()
+    model_keys = list(TIER_MODELS.get(tier, TIER_MODELS["small"]))
+    if "facenet" not in model_keys:
+        model_keys.append("facenet")
+
+    for model_key in model_keys:
+        logger.info("Ensuring AI tagging model is available: %s", model_key)
+        await ensure_model(model_key)
+
+
+def ensure_ai_tagging_models() -> None:
+    asyncio.run(_ensure_ai_tagging_models_async())
diff --git a/backend/app/utils/model_downloader.py b/backend/app/utils/model_downloader.py
diff --git a/backend/requirements.txt b/backend/requirements.txt
