Upd first time loader #3

ui.py

@@ -277,7 +277,7 @@ def create_demo():
                             <p style="margin-top: 10px; margin-bottom: 0; font-size: 11px; color: #666;">Click any model name to view details on Hugging Face</p>
                         </div>
                         """
-                    )
+                        )
                 
                 show_thoughts_state = gr.State(value=False)
                 
@@ -377,10 +377,10 @@ def create_demo():
                     if is_model_loaded(model_name):
                         status_lines.append(f"✅ MedSwin ({model_name}): loaded and ready")
                     else:
-                        state = get_model_loading_state(model_name)
-                        if state == "loading":
+                    state = get_model_loading_state(model_name)
+                    if state == "loading":
                             status_lines.append(f"⏳ MedSwin ({model_name}): loading...")
-                        elif state == "error":
+                    elif state == "error":
                             status_lines.append(f"❌ MedSwin ({model_name}): error loading")
                         else:
                             status_lines.append(f"⚠️ MedSwin ({model_name}): not loaded")
@@ -460,6 +460,13 @@ def create_demo():
                         return status_text
                         
                     except Exception as e:
+                        error_msg = str(e)
+                        # Check if it's a ZeroGPU quota/rate limit error - re-raise for retry
+                        if ("429" in error_msg or "Too Many Requests" in error_msg or 
+                            "quota" in error_msg.lower() or "ZeroGPU" in error_msg or
+                            "runnning out" in error_msg.lower() or "running out" in error_msg.lower()):
+                            logger.warning(f"[STARTUP] ZeroGPU quota/rate limit error detected: {error_msg[:100]}")
+                            raise  # Re-raise to trigger retry logic in wrapper
                         logger.error(f"[STARTUP] ❌ Error in model loading startup: {e}")
                         import traceback
                         logger.debug(f"[STARTUP] Full traceback: {traceback.format_exc()}")
@@ -487,7 +494,7 @@ def create_demo():
                         result = check_model_status(model_name)
                         if result and isinstance(result, tuple) and len(result) == 2:
                             status_text, is_ready = result
-                            return status_text
+                    return status_text
                         else:
                             return "⚠️ Unable to check model status"
                     except Exception as e:
@@ -500,12 +507,12 @@ def create_demo():
                         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)
-                            )
+                    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 (
@@ -555,20 +562,59 @@ def create_demo():
                         return f"❌ ASR: error - {str(e)[:100]}"
                 
                 # Load medical model on startup and update status
-                # Use a wrapper to handle GPU context properly
+                # Use a wrapper to handle GPU context properly with retry logic
                 def load_startup_and_update_ui():
-                    """Load model on startup and return status with UI updates"""
-                    try:
-                        status_text = load_medical_model_on_startup()
-                        # Check if model is ready and update submit button state
-                        is_ready = is_model_loaded(DEFAULT_MEDICAL_MODEL)
-                        return status_text, gr.update(interactive=is_ready), gr.update(interactive=is_ready)
-                    except Exception as e:
-                        logger.error(f"Error in load_startup_and_update_ui: {e}")
-                        import traceback
-                        logger.debug(f"Full traceback: {traceback.format_exc()}")
-                        error_msg = f"⚠️ Startup error: {str(e)[:100]}"
-                        return error_msg, gr.update(interactive=False), gr.update(interactive=False)
+                    """Load model on startup with retry logic (max 3 attempts) and return status with UI updates"""
+                    import time
+                    max_retries = 3
+                    base_delay = 5.0  # Start with 5 seconds delay
+                    
+                    for attempt in range(1, max_retries + 1):
+                        try:
+                            logger.info(f"[STARTUP] Attempt {attempt}/{max_retries} to load medical model...")
+                            status_text = load_medical_model_on_startup()
+                            # Check if model is ready and update submit button state
+                            is_ready = is_model_loaded(DEFAULT_MEDICAL_MODEL)
+                            if is_ready:
+                                logger.info(f"[STARTUP] ✅ Model loaded successfully on attempt {attempt}")
+                                return status_text, gr.update(interactive=is_ready), gr.update(interactive=is_ready)
+                            else:
+                                # Model didn't load, but no exception - might be a state issue
+                                logger.warning(f"[STARTUP] Model not ready after attempt {attempt}, but no error")
+                                if attempt < max_retries:
+                                    delay = base_delay * attempt  # Exponential backoff: 5s, 10s, 15s
+                                    logger.info(f"[STARTUP] Retrying in {delay} seconds...")
+                                    time.sleep(delay)
+                                    continue
+                                else:
+                                    return status_text, gr.update(interactive=False), gr.update(interactive=False)
+                        except Exception as e:
+                            error_msg = str(e)
+                            is_quota_error = ("429" in error_msg or "Too Many Requests" in error_msg or 
+                                            "quota" in error_msg.lower() or "ZeroGPU" in error_msg)
+                            
+                            if is_quota_error and attempt < max_retries:
+                                delay = base_delay * attempt  # Exponential backoff: 5s, 10s, 15s
+                                logger.warning(f"[STARTUP] ZeroGPU rate limit/quota error on attempt {attempt}/{max_retries}")
+                                logger.info(f"[STARTUP] Retrying in {delay} seconds...")
+                                time.sleep(delay)
+                                continue
+                            else:
+                                logger.error(f"[STARTUP] Error in load_startup_and_update_ui (attempt {attempt}/{max_retries}): {e}")
+                                import traceback
+                                logger.debug(f"[STARTUP] Full traceback: {traceback.format_exc()}")
+                                
+                                if is_quota_error:
+                                    error_display = "⚠️ ZeroGPU quota/rate limit reached. Please wait or try again later."
+                                else:
+                                    error_display = f"⚠️ Startup error: {str(e)[:100]}"
+                                
+                                if attempt >= max_retries:
+                                    logger.error(f"[STARTUP] Failed after {max_retries} attempts")
+                                    return error_display, gr.update(interactive=False), gr.update(interactive=False)
+                    
+                    # Should not reach here, but just in case
+                    return "⚠️ Startup failed after retries", gr.update(interactive=False), gr.update(interactive=False)
                 
                 demo.load(
                     fn=load_startup_and_update_ui,
@@ -596,25 +642,25 @@ def create_demo():
                                 updated_history = history + [{"role": "assistant", "content": error_msg}]
                                 yield updated_history, ""
                                 return
-                            
+                        
                             # If request is None, create a mock request for compatibility
                             if request is None:
-                                class MockRequest:
-                                    session_hash = "anonymous"
-                                request = MockRequest()
-                            
+                            class MockRequest:
+                                session_hash = "anonymous"
+                            request = MockRequest()
+                        
                             # Let stream_chat handle model loading (it's GPU-decorated and can load on-demand)
-                            for result in stream_chat(
-                                message, history, system_prompt, temperature, max_new_tokens,
-                                top_p, top_k, penalty, retriever_k, merge_threshold,
-                                use_rag, medical_model_name, use_web_search,
-                                enable_clinical_intake, disable_agentic_reasoning, show_thoughts, request
-                            ):
-                                yield result
+                        for result in stream_chat(
+                            message, history, system_prompt, temperature, max_new_tokens,
+                            top_p, top_k, penalty, retriever_k, merge_threshold,
+                            use_rag, medical_model_name, use_web_search,
+                            enable_clinical_intake, disable_agentic_reasoning, show_thoughts, request
+                        ):
+                            yield result
                             # If we get here, stream_chat completed successfully
                             return
                             
-                        except Exception as e:
+                    except Exception as e:
                             error_msg_lower = str(e).lower()
                             is_gpu_error = 'gpu task aborted' in error_msg_lower or 'gpu' in error_msg_lower or 'zerogpu' in error_msg_lower
                             
@@ -636,10 +682,10 @@ def create_demo():
                                 if is_gpu_error:
                                     error_msg = f"⚠️ GPU task was aborted. This can happen if:\n- The request took too long\n- Multiple GPU requests conflicted\n- GPU quota was exceeded\n\nPlease try again or select a different model."
                                 else:
-                                    error_msg = f"⚠️ An error occurred: {str(e)[:200]}"
+                        error_msg = f"⚠️ An error occurred: {str(e)[:200]}"
                                 
-                                updated_history = history + [{"role": "assistant", "content": error_msg}]
-                                yield updated_history, ""
+                        updated_history = history + [{"role": "assistant", "content": error_msg}]
+                        yield updated_history, ""
                                 return
                 
                 submit_button.click(