Update app.py

app.py

@@ -17,19 +17,13 @@ import logging
 import json
 import gc
 
-from utils.detector import ArabicSignDetector
-from utils.translator import MedicalTranslator
-from utils.medical_agent import MedicalAgent
-from utils.speech import SpeechProcessor
-from utils.sign_generator import SignGenerator
-
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = Flask(__name__)
 CORS(app)
 
-# Global instances
+# Global instances - will be initialized lazily
 detector = None
 translator = None
 medical_agent = None
@@ -55,14 +49,16 @@ def setup_environment():
     
     return device
 
-def initialize_models():
-    global detector, translator, medical_agent, speech_processor, sign_generator
-
-    logger.info("🔄 Starting sequential model loading...")
+def initialize_essential_models():
+    """Initialize only essential models to avoid OOM"""
+    global detector, speech_processor, sign_generator
+    
+    logger.info("🔄 Initializing essential models only...")
     
     try:
-        # Step 1: Load detector first
+        # Step 1: Load detector first (most critical)
         logger.info("📥 Step 1: Loading YOLO detector...")
+        from utils.detector import ArabicSignDetector
         detector = ArabicSignDetector()
         logger.info("✅ YOLO Detector loaded")
         
@@ -72,42 +68,75 @@ def initialize_models():
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         
-        # Step 2: Load translator
-        logger.info("📥 Step 2: Loading translator...")
-        translator = MedicalTranslator()
-        logger.info("✅ Medical Translator loaded")
-        
-        # Clear memory
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
-        # Step 3: Load medical agent (HuatuoGPT)
-        logger.info("📥 Step 3: Loading HuatuoGPT medical agent...")
-        medical_agent = MedicalAgent()
-        logger.info("✅ HuatuoGPT Medical Agent loaded")
-        
-        # Clear memory
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
-        # Step 4: Load speech processor
-        logger.info("📥 Step 4: Loading speech processor...")
+        # Step 2: Load lightweight models
+        logger.info("📥 Step 2: Loading speech processor...")
+        from utils.speech import SpeechProcessor
         speech_processor = SpeechProcessor()
         logger.info("✅ Speech Processor loaded")
         
-        # Step 5: Load sign generator
-        logger.info("📥 Step 5: Loading sign generator...")
+        # Step 3: Load sign generator
+        logger.info("📥 Step 3: Loading sign generator...")
+        from utils.sign_generator import SignGenerator
         sign_generator = SignGenerator()
         logger.info("✅ Sign Generator loaded")
         
-        logger.info("🎉 All models loaded sequentially!")
+        logger.info("🎉 Essential models loaded! Heavy models will load on demand.")
         
     except Exception as e:
-        logger.error(f"❌ Sequential loading failed: {e}")
-        # Continue with partial loading
-        pass
+        logger.error(f"❌ Essential models loading failed: {e}")
+        raise
+
+def get_translator():
+    """Lazy loader for translator"""
+    global translator
+    if translator is None:
+        try:
+            logger.info("🔄 Lazy loading translator...")
+            from utils.translator import MedicalTranslator
+            translator = MedicalTranslator()
+            logger.info("✅ Translator loaded")
+        except Exception as e:
+            logger.error(f"❌ Translator loading failed: {e}")
+            # Fallback translator
+            class FallbackTranslator:
+                def ar_to_en(self, text): return f"[EN] {text}"
+                def en_to_ar(self, text): return f"[AR] {text}"
+            translator = FallbackTranslator()
+    return translator
+
+def get_medical_agent():
+    """Lazy loader for medical agent with lighter model"""
+    global medical_agent
+    if medical_agent is None:
+        try:
+            logger.info("🔄 Lazy loading medical agent...")
+            # Try to import the lite version first
+            try:
+                from utils.medical_agent_lite import LiteMedicalAgent
+                medical_agent = LiteMedicalAgent()
+                logger.info("✅ Lite Medical Agent loaded")
+            except ImportError:
+                # Fallback to original with error handling
+                from utils.medical_agent import MedicalAgent
+                medical_agent = MedicalAgent()
+                logger.info("✅ Original Medical Agent loaded")
+        except Exception as e:
+            logger.error(f"❌ Medical agent loading failed: {e}")
+            # Ultimate fallback
+            class UltimateFallbackAgent:
+                def __init__(self):
+                    self.sessions = {}
+                def process_input(self, text, session_id):
+                    return {
+                        'response': 'Please describe your medical concern?',
+                        'question_count': 1,
+                        'state': 'questioning',
+                        'workflow_used': False
+                    }
+                def process_doctor_input(self, text):
+                    return "Please describe your symptoms?"
+            medical_agent = UltimateFallbackAgent()
+    return medical_agent
 
 @app.route('/')
 def index():
@@ -115,7 +144,7 @@ def index():
     try:
         return send_file('index.html')
     except Exception as e:
-        return "Medical Agent API with HuatuoGPT is running! Add index.html for web interface."
+        return "Medical Sign Language API is running! Add index.html for web interface."
 
 @app.route('/ui')
 def serve_ui():
@@ -126,30 +155,26 @@ def serve_ui():
 def health_check():
     return jsonify({
         "status": "healthy", 
-        "models_loaded": bool(detector and medical_agent),
-        "medical_ai": "HuatuoGPT",
-        "message": "Medical AI system operational"
+        "models_loaded": bool(detector),
+        "essential_models": "YOLO, Speech, Sign",
+        "heavy_models": "Load on demand",
+        "message": "System operational with lazy loading"
     })
 
 @app.route('/debug-model')
 def debug_model():
     """Debug endpoint to check model status"""
     detector_status = {
-        'model_loaded': detector.model is not None if detector else False,
-        'model_path': 'unknown',
-        'classes_loaded': False,
-        'total_classes': 0
+        'model_loaded': detector is not None and detector.model is not None,
+        'translator_loaded': translator is not None,
+        'medical_agent_loaded': medical_agent is not None,
+        'speech_loaded': speech_processor is not None,
+        'sign_loaded': sign_generator is not None,
     }
     
-    if detector and detector.model:
-        detector_status['model_path'] = 'best.pt'
-        detector_status['classes_loaded'] = True
-        detector_status['total_classes'] = len(detector.model.names)
-        detector_status['class_names'] = dict(detector.model.names)
-    
     return jsonify({
-        'detector_status': detector_status,
-        'message': 'Model debug information'
+        'models_status': detector_status,
+        'message': 'Lazy loading enabled for heavy models'
     })
 
 @app.route('/debug-files')
@@ -222,19 +247,21 @@ def process_sign_language():
         arabic_text = detection_result['arabic_text']
         logger.info(f"📝 Detected Arabic: {arabic_text}")
 
-        # Translate to English for medical agent
-        english_text = translator.ar_to_en(arabic_text)
+        # Lazy load translator
+        translator_instance = get_translator()
+        english_text = translator_instance.ar_to_en(arabic_text)
         logger.info(f"🌐 Translated to English: {english_text}")
 
-        # Process with HuatuoGPT medical agent
-        agent_response = medical_agent.process_input(
+        # Lazy load medical agent
+        medical_agent_instance = get_medical_agent()
+        agent_response = medical_agent_instance.process_input(
             english_text,
             session_id=session_id
         )
-        logger.info(f"🤖 HuatuoGPT response: {agent_response}")
+        logger.info(f"🤖 Medical response: {agent_response}")
 
         # Translate response back to Arabic
-        arabic_response = translator.en_to_ar(agent_response['response'])
+        arabic_response = translator_instance.en_to_ar(agent_response['response'])
         logger.info(f"🌐 Translated to Arabic: {arabic_response}")
 
         # Generate sign animation for the response
@@ -270,7 +297,7 @@ def process_sign_language():
             'conversation_state': agent_response.get('state', 'questioning'),
             'session_id': session_id,
             'workflow_used': agent_response.get('workflow_used', False),
-            'medical_ai': 'HuatuoGPT'
+            'medical_ai': 'Medical AI'
         }
 
         # Add TTS audio if available
@@ -311,12 +338,14 @@ def process_audio():
         doctor_text = speech_processor.speech_to_text(audio_path)
         logger.info(f"🎤 Doctor said: {doctor_text}")
 
-        # Process with medical agent
-        patient_question = medical_agent.process_doctor_input(doctor_text)
-        logger.info(f"🤖 HuatuoGPT rephrased: {patient_question}")
+        # Lazy load medical agent
+        medical_agent_instance = get_medical_agent()
+        patient_question = medical_agent_instance.process_doctor_input(doctor_text)
+        logger.info(f"🤖 Medical rephrased: {patient_question}")
 
-        # Translate to Arabic
-        arabic_question = translator.en_to_ar(patient_question)
+        # Lazy load translator
+        translator_instance = get_translator()
+        arabic_question = translator_instance.en_to_ar(patient_question)
         logger.info(f"🌐 Translated to Arabic: {arabic_question}")
 
         # Generate sign data for the question
@@ -347,7 +376,7 @@ def process_audio():
             'patient_question_arabic': arabic_question,
             'sign_data': sign_data,
             'session_id': session_id,
-            'medical_ai': 'HuatuoGPT'
+            'medical_ai': 'Medical AI'
         }
 
         if tts_audio:
@@ -404,7 +433,7 @@ def conversation_status():
         'success': True,
         'session_id': session_id,
         'max_questions': 3,
-        'medical_ai': 'HuatuoGPT',
+        'medical_ai': 'Medical AI',
         'system_ready': all([
             detector is not None,
             translator is not None, 
@@ -422,8 +451,9 @@ def reset_conversation():
         session_id = data.get('session_id', 'default_session')
         
         # Reset session in medical agent
-        if hasattr(medical_agent, 'sessions') and session_id in medical_agent.sessions:
-            del medical_agent.sessions[session_id]
+        medical_agent_instance = get_medical_agent()
+        if hasattr(medical_agent_instance, 'sessions') and session_id in medical_agent_instance.sessions:
+            del medical_agent_instance.sessions[session_id]
             logger.info(f"🔄 Medical conversation reset for session: {session_id}")
         else:
             logger.info(f"🔄 New session started: {session_id}")
@@ -485,12 +515,12 @@ def serve_static(filename):
     except:
         return "File not found", 404
 
-@spaces.GPU
+@spaces.GPU(enable_zero_gpu=True)
 def create_app():
-    """Application factory pattern with GPU declaration"""
-    print("🚀 Initializing Medical Sign Language App with GPU support...")
+    """Application factory pattern with GPU declaration and ZeroGPU"""
+    print("🚀 Initializing Medical Sign Language App with ZeroGPU support...")
     setup_environment()
-    initialize_models()
+    initialize_essential_models()  # Only load essential models
     return app
 
 if __name__ == '__main__':