Update src/ai_processor.py

src/ai_processor.py

@@ -1,21 +1,11 @@
 import os
-# Force CPU-only until we enter the GPU context
-os.environ['CUDA_VISIBLE_DEVICES'] = ''
-
-import torch
-# Prevent any CUDA initialization in the main process
-torch.cuda.is_available = lambda: False
-
 import io
 import base64
 import logging
-import cv2
 import numpy as np
+import cv2
 from PIL import Image
 from datetime import datetime
-from transformers import pipeline
-from ultralytics import YOLO
-from tensorflow.keras.models import load_model
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
@@ -37,6 +27,8 @@ default_system_prompt = (
     "patient context."
 )
 
+# No torch or transformers-related imports at top-level!
+
 @spaces.GPU(enable_queue=True, duration=120)
 def generate_medgemma_report(
     patient_info: str,
@@ -46,14 +38,22 @@ def generate_medgemma_report(
     segmentation_image_path: str,
     max_new_tokens: int = None
 ) -> str:
-    """Runs on GPU. Lazy-loads the MedGemma pipeline and returns the markdown report."""
+    # --- All GPU-related imports and model loading here! ---
+    import torch
+    from transformers import pipeline
+    from PIL import Image
+
+    # System prompt as before
+    global default_system_prompt
+
+    # Lazy-load MedGemma pipeline on GPU
     if not hasattr(generate_medgemma_report, "_pipe"):
         try:
             cfg = Config()
             generate_medgemma_report._pipe = pipeline(
                 'image-text-to-text',
                 model='google/medgemma-4b-it',
-                device='auto',
+                device='cuda',  # Explicitly on GPU
                 torch_dtype='auto',
                 offload_folder='offload',
                 token=cfg.HF_TOKEN
@@ -65,13 +65,13 @@ def generate_medgemma_report(
 
     pipe = generate_medgemma_report._pipe
 
-    # Assemble messages
+    # Compose messages
     msgs = [
         {'role': 'system', 'content': [{'type': 'text', 'text': default_system_prompt}]},
         {'role': 'user',   'content': []},
     ]
 
-    # Attach images
+    # Attach images if available
     for path in (detection_image_path, segmentation_image_path):
         if path and os.path.exists(path):
             msgs[1]['content'].append({'type': 'image', 'image': Image.open(path)})
@@ -96,7 +96,7 @@ class AIProcessor:
         self.px_per_cm = self.config.PIXELS_PER_CM
         self._initialize_models()
         self._load_knowledge_base()
-
+    
     def _initialize_models(self):
         """Load all CPU-only models here."""
         # Set HuggingFace token
@@ -106,6 +106,7 @@ class AIProcessor:
 
         # YOLO detection (CPU-only)
         try:
+            from ultralytics import YOLO
             self.models_cache['det'] = YOLO(self.config.YOLO_MODEL_PATH)
             logging.info("✅ YOLO model loaded (CPU only)")
         except Exception as e:
@@ -114,6 +115,7 @@ class AIProcessor:
 
         # Segmentation model (CPU)
         try:
+            from tensorflow.keras.models import load_model
             self.models_cache['seg'] = load_model(self.config.SEG_MODEL_PATH, compile=False)
             logging.info("✅ Segmentation model loaded (CPU)")
         except Exception as e:
@@ -121,6 +123,7 @@ class AIProcessor:
 
         # Classification pipeline (CPU)
         try:
+            from transformers import pipeline
             self.models_cache['cls'] = pipeline(
                 'image-classification',
                 model='Hemg/Wound-classification',
@@ -241,6 +244,7 @@ class AIProcessor:
     ) -> str:
         det = visual_results.get('detection_image_path', '')
         seg = visual_results.get('segmentation_image_path', '')
+        # This GPU call is safe: it triggers all CUDA/model code *inside* the decorator context.
         report = generate_medgemma_report(
             patient_info, visual_results, guideline_context,
             det, seg, max_new_tokens
@@ -324,4 +328,4 @@ class AIProcessor:
             return {'risk_score': risk_score, 'risk_level': level, 'risk_factors': risk_factors}
         except Exception as e:
             logging.error(f"Risk assessment error: {e}")
-            return {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []}
+            return {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []}