Update src/ai_processor.py

src/ai_processor.py

@@ -42,7 +42,7 @@ class AIProcessor:
                     model="google/medgemma-4b-it",
                     torch_dtype=torch.bfloat16,
                     offload_folder="offload",
-                    device_map="auto",
+                    device_map="cuda",
                     token=self.config.HF_TOKEN
                 )
                 logging.info("✅ MedGemma pipeline loaded successfully")
@@ -246,7 +246,7 @@ class AIProcessor:
             logging.error(f"Guidelines query error: {e}")
             return f"Error querying guidelines: {str(e)}"
 
-    
+    @spaces.GPU(enable_queue=True, duration=120) 
     def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
         """Generate comprehensive medical report using MedGemma"""
         try:
@@ -438,7 +438,7 @@ Based on the visual analysis, this appears to be a {visual_results.get('wound_ty
             logging.error(f"Image saving error: {e}")
             return None
 
-    @spaces.GPU(enable_queue=True, duration=120) 
+    
     def full_analysis_pipeline(self, image, questionnaire_data):
         """Complete analysis pipeline with real-time models"""
         try: