Update src/ai_processor.py

src/ai_processor.py

@@ -1,11 +1,9 @@
-import os
 import io
 import base64
 import logging
 import cv2
 import numpy as np
 from PIL import Image
-import torch
 from datetime import datetime
 from transformers import pipeline
 from ultralytics import YOLO
@@ -40,26 +38,12 @@ class AIProcessor:
         self._initialize_models()
 
     def _initialize_models(self):
-        """Initialize AI models; only MedGemma uses GPU."""
-        # Set HuggingFace token
+        """Initialize CPU-only AI models; MedGemma is loaded on demand within GPU context."""
+        # Set HuggingFace token early
         if self.config.HF_TOKEN:
             HfFolder.save_token(self.config.HF_TOKEN)
             logging.info("HuggingFace token set successfully")
 
-        # MedGemma pipeline on GPU
-        try:
-            self.models_cache['medgemma_pipe'] = pipeline(
-                'image-text-to-text',
-                model='google/medgemma-4b-it',
-                device='cuda',
-                torch_dtype=torch.bfloat16,
-                offload_folder='offload',
-                token=self.config.HF_TOKEN
-            )
-            logging.info("✅ MedGemma pipeline loaded on GPU")
-        except Exception as e:
-            logging.warning(f"MedGemma pipeline not available: {e}")
-
         # YOLO detection on CPU
         try:
             self.models_cache['det'] = YOLO(self.config.YOLO_MODEL_PATH)
@@ -86,7 +70,7 @@ class AIProcessor:
         except Exception as e:
             logging.warning(f"Wound classification model not available: {e}")
 
-        # Embedding for knowledge base
+        # Embedding for knowledge base on CPU
         try:
             self.models_cache['embedding_model'] = HuggingFaceEmbeddings(
                 model_name='sentence-transformers/all-MiniLM-L6-v2',
@@ -190,21 +174,43 @@ class AIProcessor:
     @spaces.GPU(enable_queue=True, duration=120)
     def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
         """Run MedGemma on GPU; return markdown report."""
+        # lazy-load MedGemma pipeline here to avoid CUDA init in main process
         if 'medgemma_pipe' not in self.models_cache:
-            return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+            try:
+                self.models_cache['medgemma_pipe'] = pipeline(
+                    'image-text-to-text',
+                    model='google/medgemma-4b-it',
+                    device='cuda',
+                    torch_dtype=torch.bfloat16,
+                    offload_folder='offload',
+                    token=self.config.HF_TOKEN
+                )
+                logging.info("✅ MedGemma pipeline loaded on GPU")
+            except Exception as e:
+                logging.warning(f"MedGemma pipeline not available: {e}")
+                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+
         # build messages
-        msgs = [{ 'role':'system', 'content':[{'type':'text','text': default_system_prompt}] },
-                { 'role':'user', 'content':[]}]
+        msgs = [
+            {'role':'system','content':[{'type':'text','text':default_system_prompt}]},
+            {'role':'user','content':[]}
+        ]
         # images
-        if image_pil: msgs[1]['content'].append({'type':'image','image':image_pil})
+        if image_pil:
+            msgs[1]['content'].append({'type':'image','image':image_pil})
         for key in ('detection_image_path','segmentation_image_path'):
             p = visual_results.get(key)
             if p and os.path.exists(p):
-                msgs[1]['content'].append({'type':'image', 'image': Image.open(p)})
-        # text prompt stub (expand as needed)
-        prompt = f"## Patient\n{patient_info}\n## Visual Type: {visual_results['wound_type']}"
+                msgs[1]['content'].append({'type':'image','image':Image.open(p)})
+        # text prompt
+        prompt = f"## Patient\n{patient_info}\n## Wound Type: {visual_results['wound_type']}"
         msgs[1]['content'].append({'type':'text','text':prompt})
-        out = self.models_cache['medgemma_pipe'](text=msgs, max_new_tokens=max_new_tokens or self.config.MAX_NEW_TOKENS)
+
+        out = self.models_cache['medgemma_pipe'](
+            text=msgs,
+            max_new_tokens=max_new_tokens or self.config.MAX_NEW_TOKENS,
+            do_sample=False
+        )
         report = out[0]['generated_text'][-1].get('content','')
         return report or self._generate_fallback_report(patient_info, visual_results, guideline_context)
 
@@ -220,10 +226,15 @@ class AIProcessor:
         fn = f"{datetime.now():%Y%m%d_%H%M%S}.png"
         path = os.path.join(self.config.UPLOADS_DIR, fn)
         image_pil.convert('RGB').save(path)
-        if self.config.HF_TOKEN and self.config.DATASET_ID:
+        if self.config.HF_TOKEN and hasattr(self.config,'DATASET_ID') and self.config.DATASET_ID:
             try:
                 api = HfApi()
-                api.upload_file(path_or_fileobj=path, path_in_repo=f"images/{fn}", repo_id=self.config.DATASET_ID, repo_type='dataset')
+                api.upload_file(
+                    path_or_fileobj=path,
+                    path_in_repo=f"images/{fn}",
+                    repo_id=self.config.DATASET_ID,
+                    repo_type='dataset'
+                )
             except Exception as e:
                 logging.warning(f"HF upload failed: {e}")
         return path
@@ -236,17 +247,18 @@ class AIProcessor:
             info = ", ".join(f"{k}:{v}" for k,v in questionnaire_data.items() if v)
             gc = self.query_guidelines(info)
             report = self.generate_final_report(info, vis, gc, image_pil)
-            return {'success':True, 'visual_analysis':vis, 'report':report, 'saved_image_path':saved}
+            return {'success':True,'visual_analysis':vis,'report':report,'saved_image_path':saved}
         except Exception as e:
             logging.error(f"Pipeline error: {e}")
-            return {'success':False, 'error':str(e)}
-
+            return {'success':False,'error':str(e)}
 
     def analyze_wound(self, image, questionnaire_data):
         """Legacy wrapper."""
-        if isinstance(image, str): image = Image.open(image)
+        if isinstance(image,str):
+            image = Image.open(image)
         return self.full_analysis_pipeline(image, questionnaire_data)
 
+        
     def _assess_risk_legacy(self, questionnaire_data):
         """Legacy risk assessment for backward compatibility"""
         risk_factors = []