Update src/ai_processor.py

src/ai_processor.py

@@ -3,14 +3,12 @@
 # Turn on deep logging: export LOGLEVEL=DEBUG SMARTHEAL_DEBUG=1
 
 import os
-import time
 import logging
 from datetime import datetime
 from typing import Optional, Dict, List, Tuple
 
-# ---- Environment defaults ----
+# ---- Environment defaults (do NOT globally hint CUDA here) ----
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
 LOGLEVEL = os.getenv("LOGLEVEL", "INFO").upper()
 SMARTHEAL_DEBUG = os.getenv("SMARTHEAL_DEBUG", "0") == "1"
 
@@ -35,6 +33,7 @@ from spaces import GPU as _SPACES_GPU
 def smartheal_gpu_stub(ping: int = 0) -> str:
     return "ready"
 
+# ---- Paths / constants ----
 UPLOADS_DIR = "uploads"
 os.makedirs(UPLOADS_DIR, exist_ok=True)
 
@@ -54,15 +53,37 @@ SEG_THRESH = float(os.getenv("SEG_THRESH", "0.5"))
 models_cache: Dict[str, object] = {}
 knowledge_base_cache: Dict[str, object] = {}
 
-# ---------- Lazy imports ----------
+# ---------- Utilities to prevent CUDA in main process ----------
+from contextlib import contextmanager
+
+@contextmanager
+def _no_cuda_env():
+    """
+    Mask GPUs so any library imported/constructed in the main process
+    cannot see CUDA (required for Spaces Stateless GPU).
+    """
+    prev = os.environ.get("CUDA_VISIBLE_DEVICES")
+    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+    try:
+        yield
+    finally:
+        if prev is None:
+            os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        else:
+            os.environ["CUDA_VISIBLE_DEVICES"] = prev
+
+# ---------- Lazy imports (wrapped where needed) ----------
 def _import_ultralytics():
-    from ultralytics import YOLO
+    # Prevent Ultralytics from probing CUDA on import
+    with _no_cuda_env():
+        from ultralytics import YOLO
     return YOLO
 
 def _import_tf_loader():
     import tensorflow as tf
     try:
-        tf.config.set_visible_devices([], "GPU")  # keep TF on CPU
+        # Keep TF on CPU only
+        tf.config.set_visible_devices([], "GPU")
     except Exception:
         pass
     from tensorflow.keras.models import load_model
@@ -122,6 +143,27 @@ Keep to 220–300 words. Do NOT provide diagnosis. Avoid contraindicated advice.
 """
 
 # ---------- VLM (MedGemma replaced with Qwen2-VL) ----------
+@_SPACES_GPU(enable_queue=True)
+def _vlm_infer_gpu(messages, model_id: str, max_new_tokens: int, token: Optional[str]):
+    """
+    Runs entirely inside a Spaces GPU worker. It's the ONLY place we allow CUDA init.
+    """
+    from transformers import pipeline
+    pipe = pipeline(
+        task="image-text-to-text",
+        model=model_id,
+        device_map="auto",            # CUDA init happens here, safely in GPU worker
+        token=token,
+        trust_remote_code=True,
+        model_kwargs={"low_cpu_mem_usage": True},
+    )
+    out = pipe(text=messages, max_new_tokens=max_new_tokens, do_sample=False, temperature=0.2)
+    try:
+        txt = out[0]["generated_text"][-1].get("content", "")
+    except Exception:
+        txt = out[0].get("generated_text", "")
+    return (txt or "").strip() or "⚠️ Empty response"
+
 def generate_medgemma_report(  # kept name so callers don't change
     patient_info: str,
     visual_results: Dict,
@@ -131,6 +173,7 @@ def generate_medgemma_report(  # kept name so callers don't change
 ) -> str:
     """
     MedGemma replacement using Qwen/Qwen2-VL-2B-Instruct via image-text-to-text.
+    Loads & runs ONLY inside a GPU worker to satisfy Stateless GPU constraints.
     """
     if os.getenv("SMARTHEAL_ENABLE_VLM", "1") != "1":
         return "⚠️ VLM disabled"
@@ -138,20 +181,6 @@ def generate_medgemma_report(  # kept name so callers don't change
     model_id = os.getenv("SMARTHEAL_VLM_MODEL", "Qwen/Qwen2-VL-2B-Instruct")
     max_new_tokens = max_new_tokens or int(os.getenv("SMARTHEAL_VLM_MAX_TOKENS", "600"))
 
-    try:
-        from transformers import pipeline
-        pipe = pipeline(
-            task="image-text-to-text",
-            model=model_id,
-            device_map=None,              # keep CPU by default for Spaces stability
-            token=HF_TOKEN,
-            trust_remote_code=True,
-            model_kwargs={"low_cpu_mem_usage": True},
-        )
-    except Exception as e:
-        logging.error(f"❌ Could not load VLM ({model_id}): {e}")
-        return "⚠️ VLM error"
-
     uprompt = SMARTHEAL_USER_PREFIX.format(
         patient_info=patient_info,
         wound_type=visual_results.get("wound_type", "Unknown"),
@@ -163,34 +192,28 @@ def generate_medgemma_report(  # kept name so callers don't change
         guideline_context=(guideline_context or "")[:900],
     )
 
+    messages = [
+        {"role": "system", "content": [{"type": "text", "text": SMARTHEAL_SYSTEM_PROMPT}]},
+        {"role": "user", "content": [
+            {"type": "image", "image": image_pil},
+            {"type": "text",  "text": uprompt},
+        ]},
+    ]
+
     try:
-        messages = [
-            {"role": "system", "content": [{"type": "text", "text": SMARTHEAL_SYSTEM_PROMPT}]},
-            {"role": "user", "content": [
-                {"type": "image", "image": image_pil},
-                {"type": "text", "text": uprompt},
-            ]},
-        ]
-        out = pipe(text=messages, max_new_tokens=max_new_tokens, do_sample=False, temperature=0.2)
-        if out and len(out) > 0:
-            try:
-                text = out[0]["generated_text"][-1].get("content", "")
-            except Exception:
-                text = out[0].get("generated_text", "")
-            text = (text or "").strip()
-            return text if text else "⚠️ Empty response"
-        return "⚠️ No output generated"
+        # IMPORTANT: do not import transformers or touch CUDA here. Only call the GPU worker.
+        return _vlm_infer_gpu(messages, model_id, max_new_tokens, HF_TOKEN)
     except Exception as e:
-        logging.error(f"❌ VLM generation error: {e}")
+        logging.error(f"VLM call failed: {e}")
         return "⚠️ VLM error"
 
-UPLOADS_DIR = "uploads"
-os.makedirs(UPLOADS_DIR, exist_ok=True)
-
 # ---------- Initialize CPU models ----------
 def load_yolo_model():
     YOLO = _import_ultralytics()
-    return YOLO(YOLO_MODEL_PATH)
+    # Construct model with CUDA masked to avoid auto-selecting cuda:0
+    with _no_cuda_env():
+        model = YOLO(YOLO_MODEL_PATH)
+    return model
 
 def load_segmentation_model():
     load_model = _import_tf_loader()
@@ -216,7 +239,7 @@ def initialize_cpu_models() -> None:
     if "det" not in models_cache:
         try:
             models_cache["det"] = load_yolo_model()
-            logging.info("✅ YOLO loaded (CPU)")
+            logging.info("✅ YOLO loaded (CPU; CUDA masked in main)")
         except Exception as e:
             logging.error(f"YOLO load failed: {e}")
 
@@ -653,6 +676,7 @@ class AIProcessor:
             det_model = self.models_cache.get("det")
             if det_model is None:
                 raise RuntimeError("YOLO model not loaded")
+            # Force CPU inference and avoid CUDA touch
             results = det_model.predict(image_cv, verbose=False, device="cpu")
             if (not results) or (not getattr(results[0], "boxes", None)) or (len(results[0].boxes) == 0):
                 try:
@@ -797,12 +821,9 @@ class AIProcessor:
             vs = self.knowledge_base_cache.get("vector_store")
             if not vs:
                 return "Knowledge base is not available."
-            try:
-                retriever = vs.as_retriever(search_kwargs={"k": 5})
-                docs = retriever.get_relevant_documents(query)
-            except Exception:
-                retriever = vs.as_retriever(search_kwargs={"k": 5})
-                docs = retriever.invoke(query)
+            retriever = vs.as_retriever(search_kwargs={"k": 5})
+            # Modern API (avoid get_relevant_documents deprecation)
+            docs = retriever.invoke(query)
             lines: List[str] = []
             for d in docs:
                 src = (d.metadata or {}).get("source", "N/A")