should work

app.py

@@ -16,44 +16,67 @@ logger = logging.getLogger(__name__)
 model = None
 processor = None
 device = None
+model_loaded = False
 
 def load_model():
-    """Load the AI model once at startup"""
-    global model, processor, device
+    """Load the AI model exactly like in Colab"""
+    global model, processor, device, model_loaded
     
-    logger.info("Loading AI model...")
-    
-    # Get Hugging Face token from environment
-    hf_token = os.environ.get('HF_TOKEN')
+    logger.info("Loading AI model (Colab style)...")
     
+    # === Load AI Model === (exactly like Colab)
     model_id = "mychen76/paligemma-3b-mix-448-med_30k-ct-brain"
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     dtype = torch.float16 if torch.cuda.is_available() else torch.float32
     
     logger.info(f"Using device: {device}")
+    logger.info(f"Using dtype: {dtype}")
+    logger.info(f"CUDA available: {torch.cuda.is_available()}")
     
     try:
-        # Load processor and model with authentication token
-        processor = AutoProcessor.from_pretrained(
-            model_id, 
-            token=hf_token
-        )
+        # Load exactly like Colab (no token, no trust_remote_code)
+        logger.info("Loading processor...")
+        processor = AutoProcessor.from_pretrained(model_id)
+        
+        logger.info("Loading model...")
         model = PaliGemmaForConditionalGeneration.from_pretrained(
             model_id, 
-            torch_dtype=dtype,
-            token=hf_token
+            torch_dtype=dtype
         ).to(device).eval()
         
         logger.info("Model loaded successfully!")
+        model_loaded = True
         return True
         
     except Exception as e:
         logger.error(f"Error loading model: {e}")
+        logger.error(f"Error type: {type(e)}")
+        model_loaded = False
         return False
 
+def run_model(img):
+    """Run model inference exactly like Colab"""
+    prompt = "<image> Findings:"
+    inputs = processor(images=img, text=prompt, return_tensors="pt").to(device, dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+    generated_ids = model.generate(**inputs, max_new_tokens=100)
+    result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return result
+
 def analyze_brain_scan(image, patient_name="", patient_age="", symptoms=""):
     """Analyze brain scan image and return medical findings"""
     try:
+        if not model_loaded or model is None:
+            return """
+## ⚠️ Model Loading Error
+
+The AI model is not available. This could be due to:
+- Model loading issues
+- Memory limitations
+- Network connectivity
+
+Please check the logs or try refreshing.
+"""
+        
         if image is None:
             return "Please upload a brain scan image."
         
@@ -61,22 +84,8 @@ def analyze_brain_scan(image, patient_name="", patient_age="", symptoms=""):
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image).convert("RGB")
         
-        # Run AI inference
-        prompt = "<image> Findings:"
-        inputs = processor(
-            images=image, 
-            text=prompt, 
-            return_tensors="pt"
-        ).to(device, dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
-        
-        with torch.no_grad():
-            generated_ids = model.generate(**inputs, max_new_tokens=100)
-            
-        result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        
-        # Clean up the result
-        if result.startswith(prompt):
-            result = result[len(prompt):].strip()
+        # Run AI inference using Colab method
+        result = run_model(image)
         
         # Format the response
         timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
@@ -106,6 +115,9 @@ Always consult with qualified medical professionals for actual diagnosis.
 def create_api_response(image, patient_name="", patient_age="", symptoms=""):
     """Create API-compatible response for integration"""
     try:
+        if not model_loaded or model is None:
+            return {"error": "Model not loaded"}
+        
         if image is None:
             return {"error": "No image provided"}
         
@@ -113,24 +125,10 @@ def create_api_response(image, patient_name="", patient_age="", symptoms=""):
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image).convert("RGB")
         
-        # Run AI inference
-        prompt = "<image> Findings:"
-        inputs = processor(
-            images=image, 
-            text=prompt, 
-            return_tensors="pt"
-        ).to(device, dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+        # Run AI inference using Colab method
+        result = run_model(image)
         
-        with torch.no_grad():
-            generated_ids = model.generate(**inputs, max_new_tokens=100)
-            
-        result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        
-        # Clean up the result
-        if result.startswith(prompt):
-            result = result[len(prompt):].strip()
-        
-        # Create API response
+        # Create API response (matching your original format)
         response = {
             "prediction": result,
             "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
@@ -141,7 +139,8 @@ def create_api_response(image, patient_name="", patient_age="", symptoms=""):
             },
             "model_info": {
                 "model_id": "mychen76/paligemma-3b-mix-448-med_30k-ct-brain",
-                "device": str(device)
+                "device": str(device),
+                "model_loaded": model_loaded
             }
         }
         
@@ -151,9 +150,23 @@ def create_api_response(image, patient_name="", patient_age="", symptoms=""):
         logger.error(f"API error: {e}")
         return {"error": f"Analysis failed: {str(e)}"}
 
+def get_model_status():
+    """Get current model status"""
+    return f"""
+## 🤖 Model Status
+
+- **Model Loaded**: {model_loaded}
+- **Device**: {device}
+- **CUDA Available**: {torch.cuda.is_available()}
+- **Model Object**: {type(model).__name__ if model else 'None'}
+- **Processor Object**: {type(processor).__name__ if processor else 'None'}
+- **PyTorch Version**: {torch.__version__}
+"""
+
 # Load model at startup
-logger.info("Initializing Brain CT Analyzer...")
-if load_model():
+logger.info("Initializing Brain CT Analyzer (Colab Style)...")
+load_success = load_model()
+if load_success:
     logger.info("Model loaded successfully!")
 else:
     logger.error("Failed to load model!")
@@ -169,6 +182,12 @@ with gr.Blocks(title="Brain CT Analyzer", theme=gr.themes.Soft()) as demo:
     **⚠️ Important:** This is for educational/research purposes only. Always consult qualified medical professionals.
     """)
     
+    # Model status section
+    with gr.Accordion("🔧 Model Status", open=not model_loaded):
+        status_output = gr.Markdown(value=get_model_status())
+        refresh_btn = gr.Button("🔄 Refresh Status")
+        refresh_btn.click(fn=get_model_status, outputs=status_output)
+    
     with gr.Row():
         with gr.Column(scale=1):
             image_input = gr.Image(
@@ -195,13 +214,14 @@ with gr.Blocks(title="Brain CT Analyzer", theme=gr.themes.Soft()) as demo:
             analyze_btn = gr.Button(
                 "🔍 Analyze Brain Scan",
                 variant="primary",
-                size="lg"
+                size="lg",
+                interactive=model_loaded
             )
         
         with gr.Column(scale=1):
             result_output = gr.Markdown(
                 label="Analysis Results",
-                value="Upload an image and click 'Analyze Brain Scan' to see results."
+                value="Upload an image and click 'Analyze Brain Scan' to see results." if model_loaded else "⚠️ Model not loaded. Check status above."
             )
     
     # API endpoint simulation
@@ -209,19 +229,20 @@ with gr.Blocks(title="Brain CT Analyzer", theme=gr.themes.Soft()) as demo:
         api_output = gr.JSON(label="API Response Format")
     
     # Event handlers
-    analyze_btn.click(
-        fn=analyze_brain_scan,
-        inputs=[image_input, patient_name, patient_age, symptoms],
-        outputs=result_output
-    )
-    
-    analyze_btn.click(
-        fn=create_api_response,
-        inputs=[image_input, patient_name, patient_age, symptoms],
-        outputs=api_output
-    )
+    if model_loaded:
+        analyze_btn.click(
+            fn=analyze_brain_scan,
+            inputs=[image_input, patient_name, patient_age, symptoms],
+            outputs=result_output
+        )
+        
+        analyze_btn.click(
+            fn=create_api_response,
+            inputs=[image_input, patient_name, patient_age, symptoms],
+            outputs=api_output
+        )
     
-    # Example images (if you have any)
+    # Instructions
     gr.Markdown("""
     ## 📋 Usage Instructions:
     1. Upload a brain CT scan image (JPEG or PNG)
@@ -231,6 +252,9 @@ with gr.Blocks(title="Brain CT Analyzer", theme=gr.themes.Soft()) as demo:
     
     ## 🔗 Integration:
     This interface can be integrated with your medical app using the Gradio API.
+    
+    ## ✅ Based on Working Colab Code:
+    This version uses the exact same model loading and inference code as your working Google Colab notebook.
     """)
 
 if __name__ == "__main__":
@@ -238,4 +262,4 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=7860,
         share=True
-    ) 
+    )