fix synchronize session tensor access and model deletion lifecycle

- snapshot tensor names to avoid close() invalidating session.run() inputs.
- lock model deletion to prevent delete-while-in-use races

Author: rohan-pandeyy

backend/app/models/FaceNet.py

@@ -25,10 +25,17 @@ def __init__(self, model_path):
         self.output_tensor_name: str | None = None
         self._lock = threading.Lock()
 
-    def get_session(self) -> onnxruntime.InferenceSession:
+    def get_session(self) -> tuple[onnxruntime.InferenceSession, str, str]:
         session = self._session
         if session is not None:
-            return session
+            input_name = self.input_tensor_name
+            output_name = self.output_tensor_name
+            if input_name is None or output_name is None:
+                raise RuntimeError(
+                    f"Model session for '{self._model_key}' was closed while "
+                    "get_session() was executing."
+                )
+            return session, input_name, output_name
 
         with self._lock:
             if self._session is None:
@@ -50,19 +57,31 @@ def get_session(self) -> onnxruntime.InferenceSession:
                     self.model_path, providers=ONNX_util_get_execution_providers()
                 )
                 if self._model_key is not None and not self._session_registered:
-                    mark_model_session_active(self._model_key)
+                    try:
+                        mark_model_session_active(self._model_key)
+                    except RuntimeError:
+                        self._session = None
+                        raise
                     self._session_registered = True
                 self.input_tensor_name = self._session.get_inputs()[0].name
                 self.output_tensor_name = self._session.get_outputs()[0].name
 
             session = self._session
+            input_name = self.input_tensor_name
+            output_name = self.output_tensor_name
+
+            if session is None or input_name is None or output_name is None:
+                raise RuntimeError(
+                    f"Model session for '{self._model_key}' was closed while "
+                    "get_session() was executing."
+                )
 
-        return session
+        return session, input_name, output_name
 
     def get_embedding(self, preprocessed_image):
-        session = self.get_session()
+        session, input_tensor_name, output_tensor_name = self.get_session()
         result = session.run(
-            [self.output_tensor_name], {self.input_tensor_name: preprocessed_image}
+            [output_tensor_name], {input_tensor_name: preprocessed_image}
         )[0]
         embedding = result[0]
         return FaceNet_util_normalize_embedding(embedding)

backend/app/models/YOLO.py

@@ -30,6 +30,7 @@ def __init__(self, path, conf_threshold=0.7, iou_threshold=0.5):
         self.iou_threshold = iou_threshold
         self._session = None
         import threading
+
         self._lock = threading.Lock()
 
     def get_session(self):
@@ -47,7 +48,9 @@ def get_session(self):
                         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)
+                    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."
@@ -57,7 +60,11 @@ def get_session(self):
                     self.model_path, providers=ONNX_util_get_execution_providers()
                 )
                 if self._model_key is not None and not self._session_registered:
-                    mark_model_session_active(self._model_key)
+                    try:
+                        mark_model_session_active(self._model_key)
+                    except RuntimeError:
+                        self._session = None
+                        raise
                     self._session_registered = True
                 # Initialize model info once session is created
                 self.get_input_details()
@@ -95,9 +102,7 @@ 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}
-        )
+        outputs = session.run(self.output_names, {self.input_names[0]: input_tensor})
         logger.debug("Inference completed in %.4fs", time.perf_counter() - start)
         return outputs
 

backend/app/models/model_registry.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from typing import TypedDict, Literal
-import json
 import os
 import sys
 from platformdirs import user_data_dir
@@ -79,10 +78,10 @@ class ModelSpec(TypedDict):
 }
 
 TIER_MODELS: dict[str, list[str]] = {
-    "nano":   ["yolo_nano", "yolo_nano_face"],
-    "small":  ["yolo_small", "yolo_small_face"],
+    "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
+    "required": ["facenet"],  # Required model; not user-selectable
 }
 
 USER_DATA_MODELS = os.path.join(user_data_dir("PictoPy"), "models")
@@ -91,7 +90,7 @@ class ModelSpec(TypedDict):
 
 def ensure_model_exports_directory() -> None:
     """Create the active model exports directory if it does not exist."""
-    if getattr(sys, 'frozen', False):
+    if getattr(sys, "frozen", False):
         os.makedirs(USER_DATA_MODELS, exist_ok=True)
     else:
         os.makedirs(LOCAL_ONNX_EXPORTS, exist_ok=True)
@@ -100,11 +99,11 @@ def ensure_model_exports_directory() -> None:
 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):
+    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))
 
@@ -115,4 +114,3 @@ def get_model_key_from_path(model_path: str) -> str | None:
         if spec["filename"].lower() == target_filename:
             return key
     return None
-

backend/app/models/session_registry.py

@@ -3,11 +3,16 @@
 import threading
 
 _active_sessions: dict[str, int] = {}
+_models_pending_deletion: set[str] = set()
 _registry_lock = threading.Lock()
 
 
 def mark_model_session_active(model_key: str) -> None:
     with _registry_lock:
+        if model_key in _models_pending_deletion:
+            raise RuntimeError(
+                f"Model '{model_key}' is being deleted; cannot start a new session."
+            )
         _active_sessions[model_key] = _active_sessions.get(model_key, 0) + 1
 
 
@@ -23,3 +28,17 @@ def mark_model_session_inactive(model_key: str) -> None:
 def get_active_session_count(model_key: str) -> int:
     with _registry_lock:
         return _active_sessions.get(model_key, 0)
+
+
+def try_mark_model_for_deletion(model_key: str) -> int | None:
+    with _registry_lock:
+        count = _active_sessions.get(model_key, 0)
+        if count > 0:
+            return count
+        _models_pending_deletion.add(model_key)
+        return None
+
+
+def release_model_deletion_mark(model_key: str) -> None:
+    with _registry_lock:
+        _models_pending_deletion.discard(model_key)

backend/app/routes/models.py

@@ -2,14 +2,17 @@
 import json
 import uuid
 import asyncio
-from typing import Dict, Optional
+from typing import Dict
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from fastapi import APIRouter, HTTPException, status
 from pydantic import BaseModel
 from fastapi.responses import StreamingResponse
 from app.models.model_registry import MODEL_REGISTRY, TIER_MODELS, get_model_path
-from app.models.session_registry import get_active_session_count
+from app.models.session_registry import (
+    try_mark_model_for_deletion,
+    release_model_deletion_mark,
+)
 from app.utils.hardware_detect import get_hardware_info
 from app.utils.model_downloader import ensure_model
 import logging
@@ -113,8 +116,10 @@ async def delete_model(model_key: str):
         )
 
     path = get_model_path(model_key)
-    active_session_count = get_active_session_count(model_key)
-    if active_session_count > 0:
+
+    # Check no sessions are active and reserve the model for deletion.
+    active_session_count = try_mark_model_for_deletion(model_key)
+    if active_session_count is not None:
         raise HTTPException(
             status_code=status.HTTP_409_CONFLICT,
             detail=(
@@ -123,21 +128,27 @@ async def delete_model(model_key: str):
             ),
         )
 
-    if os.path.exists(path):
-        try:
-            await asyncio.to_thread(os.remove, path)
+    try:
+        if os.path.exists(path):
+            try:
+                await asyncio.to_thread(os.remove, path)
+                return {
+                    "success": True,
+                    "message": f"Model {model_key} deleted successfully.",
+                }
+            except Exception as e:
+                logger.error(f"Failed to delete model {model_key}: {e}")
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail=f"Failed to delete model: {str(e)}",
+                )
+        else:
             return {
                 "success": True,
-                "message": f"Model {model_key} deleted successfully.",
+                "message": f"Model {model_key} already not present.",
             }
-        except Exception as e:
-            logger.error(f"Failed to delete model {model_key}: {e}")
-            raise HTTPException(
-                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                detail=f"Failed to delete model: {str(e)}",
-            )
-    else:
-        return {"success": True, "message": f"Model {model_key} already not present."}
+    finally:
+        release_model_deletion_mark(model_key)
 
 
 @router.post("/setup")

backend/app/utils/model_bootstrap.py

@@ -31,4 +31,4 @@ async def _ensure_ai_tagging_models_async() -> None:
 
 
 def ensure_ai_tagging_models() -> None:
-    asyncio.run(_ensure_ai_tagging_models_async())
+    asyncio.run(_ensure_ai_tagging_models_async())