Upd models loader

ui.py

@@ -328,10 +328,18 @@ def create_demo():
                                 status_lines.append(f"❌ MedSwin ({model_name}): error loading")
                             else:
                                 # Use GPU-decorated function to load the model
-                                status_text, is_ready = load_model_with_gpu(model_name)
-                                if is_ready:
-                                    status_lines.append(f"✅ MedSwin ({model_name}): loaded and ready")
-                                else:
+                                try:
+                                    result = load_model_with_gpu(model_name)
+                                    if result and isinstance(result, tuple) and len(result) == 2:
+                                        status_text, is_ready = result
+                                        if is_ready:
+                                            status_lines.append(f"✅ MedSwin ({model_name}): loaded and ready")
+                                        else:
+                                            status_lines.append(f"⏳ MedSwin ({model_name}): loading...")
+                                    else:
+                                        status_lines.append(f"⏳ MedSwin ({model_name}): loading...")
+                                except Exception as e:
+                                    logger.error(f"Error calling load_model_with_gpu: {e}")
                                     status_lines.append(f"⏳ MedSwin ({model_name}): loading...")
                         
                         # TTS model status
@@ -346,15 +354,15 @@ def create_demo():
                         # ASR (Whisper) model status
                         if WHISPER_AVAILABLE:
                             if config.global_whisper_model is not None:
-                                status_lines.append("✅ ASR (Whisper large-v3-turbo): loaded and ready")
+                                status_lines.append("✅ ASR (Whisper): loaded and ready")
                             else:
-                                status_lines.append("⚠️ ASR (Whisper large-v3-turbo): not loaded")
+                                status_lines.append("⚠️ ASR (Whisper): not loaded")
                         else:
                             status_lines.append("❌ ASR: library not available")
-                            
-                            status_text = "\n".join(status_lines)
-                            is_ready = is_model_loaded(model_name)
-                            return status_text, is_ready
+                        
+                        status_text = "\n".join(status_lines)
+                        is_ready = is_model_loaded(model_name)
+                        return status_text, is_ready
                     except Exception as e:
                         return f"❌ Error: {str(e)[:100]}", False
                 
@@ -452,23 +460,57 @@ def create_demo():
                 
                 # Initialize status on load
                 def init_model_status():
-                    status_text, is_ready = check_model_status(DEFAULT_MEDICAL_MODEL)
-                    return status_text
+                    try:
+                        result = check_model_status(DEFAULT_MEDICAL_MODEL)
+                        if result and isinstance(result, tuple) and len(result) == 2:
+                            status_text, is_ready = result
+                            return status_text
+                        else:
+                            return "⚠️ Unable to check model status"
+                    except Exception as e:
+                        logger.error(f"Error in init_model_status: {e}")
+                        return f"⚠️ Error: {str(e)[:100]}"
                 
                 # Update status when model selection changes
                 def update_model_status_on_change(model_name):
-                    status_text, is_ready = check_model_status(model_name)
-                    return status_text
+                    try:
+                        result = check_model_status(model_name)
+                        if result and isinstance(result, tuple) and len(result) == 2:
+                            status_text, is_ready = result
+                            return status_text
+                        else:
+                            return "⚠️ Unable to check model status"
+                    except Exception as e:
+                        logger.error(f"Error in update_model_status_on_change: {e}")
+                        return f"⚠️ Error: {str(e)[:100]}"
                 
                 # Handle model selection change
                 def on_model_change(model_name):
-                    status_text, is_ready = load_model_and_update_status(model_name)
-                    submit_enabled = is_ready
-                    return (
-                        status_text,
-                        gr.update(interactive=submit_enabled),
-                        gr.update(interactive=submit_enabled)
-                    )
+                    try:
+                        result = load_model_and_update_status(model_name)
+                        if result and isinstance(result, tuple) and len(result) == 2:
+                            status_text, is_ready = result
+                            submit_enabled = is_ready
+                            return (
+                                status_text,
+                                gr.update(interactive=submit_enabled),
+                                gr.update(interactive=submit_enabled)
+                            )
+                        else:
+                            error_msg = "⚠️ Unable to load model status"
+                            return (
+                                error_msg,
+                                gr.update(interactive=False),
+                                gr.update(interactive=False)
+                            )
+                    except Exception as e:
+                        logger.error(f"Error in on_model_change: {e}")
+                        error_msg = f"⚠️ Error: {str(e)[:100]}"
+                        return (
+                            error_msg,
+                            gr.update(interactive=False),
+                            gr.update(interactive=False)
+                        )
                 
                 # Update status display periodically or on model status changes
                 def refresh_model_status(model_name):
@@ -494,8 +536,19 @@ def create_demo():
                     outputs=None
                 )
                 # Finally update status to show all models
+                def update_status_after_load():
+                    try:
+                        result = check_model_status(DEFAULT_MEDICAL_MODEL)
+                        if result and isinstance(result, tuple) and len(result) == 2:
+                            return result[0]
+                        else:
+                            return "⚠️ Unable to check model status"
+                    except Exception as e:
+                        logger.error(f"Error updating status after load: {e}")
+                        return f"⚠️ Error: {str(e)[:100]}"
+                
                 demo.load(
-                    fn=lambda: check_model_status(DEFAULT_MEDICAL_MODEL)[0],
+                    fn=update_status_after_load,
                     inputs=None,
                     outputs=[model_status]
                 )
@@ -511,9 +564,23 @@ def create_demo():
                         # Check if model is loaded
                         if not is_model_loaded(medical_model_name):
                             # Try to load it
-                            status_text, is_ready = load_model_and_update_status(medical_model_name)
-                            if not is_ready:
-                                error_msg = "⚠️ Model is not ready. Please wait for the model to finish loading before sending messages."
+                            try:
+                                result = load_model_and_update_status(medical_model_name)
+                                if result and isinstance(result, tuple) and len(result) == 2:
+                                    status_text, is_ready = result
+                                    if not is_ready:
+                                        error_msg = "⚠️ Model is not ready. Please wait for the model to finish loading before sending messages."
+                                        updated_history = history + [{"role": "assistant", "content": error_msg}]
+                                        yield updated_history, ""
+                                        return
+                                else:
+                                    error_msg = "⚠️ Unable to check model status. Please try again."
+                                    updated_history = history + [{"role": "assistant", "content": error_msg}]
+                                    yield updated_history, ""
+                                    return
+                            except Exception as e:
+                                logger.error(f"Error loading model: {e}")
+                                error_msg = f"⚠️ Error loading model: {str(e)[:200]}. Please try again."
                                 updated_history = history + [{"role": "assistant", "content": error_msg}]
                                 yield updated_history, ""
                                 return