Upd MedSwin loader with GPU

ui.py

@@ -1,11 +1,13 @@
 """Gradio UI setup"""
 import time
 import gradio as gr
+import spaces
 from config import TITLE, DESCRIPTION, CSS, MEDSWIN_MODELS, DEFAULT_MEDICAL_MODEL
 from indexing import create_or_update_index
 from pipeline import stream_chat
 from voice import transcribe_audio, generate_speech
 from models import initialize_medical_model, is_model_loaded, get_model_loading_state, set_model_loading_state
+from logger import logger
 
 
 def create_demo():
@@ -258,6 +260,29 @@ def create_demo():
                     outputs=[agentic_thoughts_box, show_thoughts_state]
                 )
                 
+                # GPU-decorated function to load any model (for user selection)
+                @spaces.GPU(max_duration=120)
+                def load_model_with_gpu(model_name):
+                    """Load medical model (GPU-decorated for ZeroGPU compatibility)"""
+                    try:
+                        if not is_model_loaded(model_name):
+                            logger.info(f"Loading medical model: {model_name}...")
+                            set_model_loading_state(model_name, "loading")
+                            try:
+                                initialize_medical_model(model_name)
+                                logger.info(f"✅ Medical model {model_name} loaded successfully!")
+                                return "✅ The model has been loaded successfully", True
+                            except Exception as e:
+                                logger.error(f"Failed to load medical model {model_name}: {e}")
+                                set_model_loading_state(model_name, "error")
+                                return f"❌ Error loading model: {str(e)[:100]}", False
+                        else:
+                            logger.info(f"Medical model {model_name} is already loaded")
+                            return "✅ The model has been loaded successfully", True
+                    except Exception as e:
+                        logger.error(f"Error loading model {model_name}: {e}")
+                        return f"❌ Error: {str(e)[:100]}", False
+                
                 def load_model_and_update_status(model_name):
                     """Load model and update status, return status text and whether model is ready"""
                     try:
@@ -270,14 +295,9 @@ def create_demo():
                         elif state == "error":
                             return "❌ Error loading model. Please try again.", False
                         
-                        # Start loading
-                        set_model_loading_state(model_name, "loading")
-                        try:
-                            initialize_medical_model(model_name)
-                            return "✅ The model has been loaded successfully", True
-                        except Exception as e:
-                            set_model_loading_state(model_name, "error")
-                            return f"❌ Error loading model: {str(e)[:100]}", False
+                        # Use GPU-decorated function to load the model
+                        status_text, is_ready = load_model_with_gpu(model_name)
+                        return status_text, is_ready
                     except Exception as e:
                         return f"❌ Error: {str(e)[:100]}", False
                 
@@ -293,6 +313,29 @@ def create_demo():
                     else:
                         return "⚠️ Model not loaded. Click to load or it will load on first use.", False
                 
+                # GPU-decorated function to load model on startup
+                @spaces.GPU(max_duration=120)
+                def load_default_model_on_startup():
+                    """Load default medical model on startup (GPU-decorated for ZeroGPU compatibility)"""
+                    try:
+                        if not is_model_loaded(DEFAULT_MEDICAL_MODEL):
+                            logger.info(f"Loading default medical model on startup: {DEFAULT_MEDICAL_MODEL}...")
+                            set_model_loading_state(DEFAULT_MEDICAL_MODEL, "loading")
+                            try:
+                                initialize_medical_model(DEFAULT_MEDICAL_MODEL)
+                                logger.info(f"✅ Default medical model {DEFAULT_MEDICAL_MODEL} loaded successfully on startup!")
+                                return f"✅ {DEFAULT_MEDICAL_MODEL} loaded successfully"
+                            except Exception as e:
+                                logger.error(f"Failed to load default medical model on startup: {e}")
+                                set_model_loading_state(DEFAULT_MEDICAL_MODEL, "error")
+                                return f"❌ Error loading model: {str(e)[:100]}"
+                        else:
+                            logger.info(f"Default medical model {DEFAULT_MEDICAL_MODEL} is already loaded")
+                            return f"✅ {DEFAULT_MEDICAL_MODEL} is ready"
+                    except Exception as e:
+                        logger.error(f"Error in model loading startup: {e}")
+                        return f"⚠️ Startup loading error: {str(e)[:100]}"
+                
                 # Initialize status on load
                 def init_model_status():
                     status_text, is_ready = check_model_status(DEFAULT_MEDICAL_MODEL)
@@ -314,9 +357,9 @@ def create_demo():
                     outputs=[model_status, submit_button, message_input]
                 )
                 
-                # Initialize status
+                # Load default model on startup (GPU-decorated function)
                 demo.load(
-                    fn=init_model_status,
+                    fn=load_default_model_on_startup,
                     outputs=[model_status]
                 )