Update app.py

app.py

@@ -1,7 +1,7 @@
-# app.py - MedGemma with Fixed Authentication
+# app.py - Working MedGemma with Correct Implementation
 import gradio as gr
 import torch
-from transformers import AutoProcessor, AutoModelForImageTextToText
+from transformers import AutoProcessor, AutoModelForImageTextToText, pipeline
 from PIL import Image
 import os
 import logging
@@ -30,13 +30,14 @@ def authenticate_hf():
 # Model configuration
 MODEL_ID = "google/medgemma-4b-it"
 
-# Global variables for model and processor
+# Global variables
 model = None
 processor = None
+pipeline_model = None
 
 def load_model():
-    """Load model and processor with authentication"""
-    global model, processor
+    """Load MedGemma model using the recommended approach"""
+    global model, processor, pipeline_model
     
     try:
         # First authenticate
@@ -45,33 +46,43 @@ def load_model():
             logger.error("❌ Authentication required for MedGemma")
             return False
             
-        logger.info(f"Loading model: {MODEL_ID}")
-        
-        # Check if CUDA is available
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device: {device}")
+        logger.info(f"Loading MedGemma: {MODEL_ID}")
+        
+        # Method 1: Try using pipeline (recommended by HuggingFace)
+        try:
+            logger.info("Attempting to load using pipeline...")
+            pipeline_model = pipeline(
+                "image-text-to-text", 
+                model=MODEL_ID,
+                torch_dtype=torch.float32,
+                device_map="auto" if torch.cuda.is_available() else None,
+                trust_remote_code=True
+            )
+            logger.info("✅ Pipeline model loaded successfully!")
+            return True
+        except Exception as e:
+            logger.warning(f"Pipeline loading failed: {e}")
+            
+        # Method 2: Try direct model loading
+        logger.info("Attempting direct model loading...")
         
-        # Load processor first
-        logger.info("Loading processor...")
+        # Load processor
         processor = AutoProcessor.from_pretrained(
-            MODEL_ID, 
+            MODEL_ID,
             trust_remote_code=True,
             token=True
         )
-        logger.info("✅ Processor loaded successfully")
+        logger.info("✅ Processor loaded")
         
-        # Load model with authentication
-        logger.info("Loading model...")
+        # Load model
         model = AutoModelForImageTextToText.from_pretrained(
             MODEL_ID,
-            torch_dtype=torch.float32,  # Use float32 for CPU compatibility
-            device_map=None,  # Let PyTorch handle device placement
+            torch_dtype=torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
             trust_remote_code=True,
-            low_cpu_mem_usage=True,
             token=True
         )
         logger.info("✅ Model loaded successfully!")
-        
         return True
         
     except Exception as e:
@@ -85,19 +96,19 @@ model_loaded = load_model()
 
 def analyze_medical_image(image, clinical_question, patient_history=""):
     """Analyze medical image with clinical context"""
-    global model, processor
+    global model, processor, pipeline_model
     
     # Check if model is loaded
-    if not model_loaded or model is None or processor is None:
+    if not model_loaded:
         return """❌ **Model Loading Issue**
 
-The model failed to load properly. This could be due to:
+MedGemma failed to load. This is likely due to:
 
-1. **Memory constraints**: The model requires significant RAM
-2. **Hardware limitations**: Consider upgrading to GPU hardware
-3. **Temporary issue**: Try refreshing the page
+1. **Transformers version**: Make sure you're using transformers >= 4.52.0
+2. **Authentication**: Ensure HF_TOKEN is properly set
+3. **Model compatibility**: MedGemma requires the latest transformers library
 
-**Current Status**: Model loading failed - please try again or contact support."""
+**Status**: Model loading failed. Please try refreshing the page or contact support."""
     
     if image is None:
         return "⚠️ Please upload a medical image first."
@@ -106,65 +117,72 @@ The model failed to load properly. This could be due to:
         return "⚠️ Please provide a clinical question."
     
     try:
-        # Prepare the conversation
-        messages = [
-            {
-                "role": "system",
-                "content": [{"type": "text", "text": "You are MedGemma, an expert medical AI assistant specialized in medical image analysis. Provide detailed, structured analysis while emphasizing that this is for educational purposes only and should not replace professional medical diagnosis."}]
-            }
-        ]
-        
-        # Build user message content
-        user_content = []
-        
-        # Add patient history if provided
-        if patient_history.strip():
-            user_content.append({"type": "text", "text": f"Patient History: {patient_history}\n\n"})
-        
-        # Add the clinical question
-        user_content.append({"type": "text", "text": f"Clinical Question: {clinical_question}"})
-        
-        # Add the image
-        user_content.append({"type": "image", "image": image})
-        
-        messages.append({
-            "role": "user",
-            "content": user_content
-        })
-        
-        # Process inputs
-        logger.info("Processing input...")
-        inputs = processor.apply_chat_template(
-            messages, 
-            add_generation_prompt=True, 
-            tokenize=True,
-            return_dict=True, 
-            return_tensors="pt"
-        )
-        
-        # Move to appropriate device if model is on GPU
-        if torch.cuda.is_available() and next(model.parameters()).is_cuda:
-            device = next(model.parameters()).device
-            inputs = {k: v.to(device) for k, v in inputs.items()}
-        
-        input_len = inputs["input_ids"].shape[-1]
-        
-        # Generate response
-        logger.info("Generating response...")
-        with torch.inference_mode():
-            generation = model.generate(
-                **inputs, 
-                max_new_tokens=1000,  # Reduced for stability
-                do_sample=True,
-                temperature=0.3,
-                top_p=0.95,
-                repetition_penalty=1.1,
-                pad_token_id=processor.tokenizer.eos_token_id if hasattr(processor, 'tokenizer') else None
+        # Method 1: Use pipeline if available
+        if pipeline_model is not None:
+            logger.info("Using pipeline for analysis...")
+            
+            # Prepare message in the format expected by pipeline
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": image},
+                        {"type": "text", "text": f"Patient History: {patient_history}\n\nClinical Question: {clinical_question}\n\nAs MedGemma, provide a detailed medical analysis of this image for educational purposes only."}
+                    ]
+                }
+            ]
+            
+            # Generate response using pipeline
+            result = pipeline_model(messages, max_new_tokens=1000)
+            
+            # Extract response text
+            response = result[0]['generated_text'] if isinstance(result, list) else result['generated_text']
+            
+        # Method 2: Use direct model if pipeline failed
+        elif model is not None and processor is not None:
+            logger.info("Using direct model for analysis...")
+            
+            # Prepare messages for direct model
+            messages = [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": "You are MedGemma, an expert medical AI assistant. Provide detailed medical analysis for educational purposes only."}]
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": f"Patient History: {patient_history}\n\nClinical Question: {clinical_question}"},
+                        {"type": "image", "image": image}
+                    ]
+                }
+            ]
+            
+            # Process inputs
+            inputs = processor.apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+                tokenize=True,
+                return_dict=True,
+                return_tensors="pt"
             )
-            generation = generation[0][input_len:]
+            
+            # Generate response
+            with torch.inference_mode():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=1000,
+                    do_sample=True,
+                    temperature=0.3,
+                    top_p=0.9
+                )
+            
+            # Decode response
+            response = processor.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
         
-        # Decode response
-        response = processor.decode(generation, skip_special_tokens=True)
+        else:
+            return "❌ No model available for analysis. Please try refreshing the page."
+        
+        # Clean up response
         response = response.strip()
         
         # Add medical disclaimer
@@ -187,31 +205,18 @@ The model failed to load properly. This could be due to:
         logger.error(f"❌ Error in analysis: {str(e)}")
         import traceback
         logger.error(f"Full traceback: {traceback.format_exc()}")
-        return f"❌ Analysis failed: {str(e)}\n\nPlease try with a different image or question."
+        return f"❌ Analysis failed: {str(e)}\n\nPlease try again with a different image or question."
 
 # Create Gradio interface
 def create_interface():
     with gr.Blocks(
-        title="MedGemma Medical Image Analysis", 
+        title="MedGemma Medical Analysis", 
         theme=gr.themes.Soft(),
         css="""
-        .gradio-container {
-            max-width: 1200px !important;
-        }
-        .disclaimer {
-            background-color: #fef2f2;
-            border: 1px solid #fecaca;
-            border-radius: 8px;
-            padding: 16px;
-            margin: 16px 0;
-        }
-        .success {
-            background-color: #f0f9ff;
-            border: 1px solid #bae6fd;
-            border-radius: 8px;
-            padding: 16px;
-            margin: 16px 0;
-        }
+        .gradio-container { max-width: 1200px !important; }
+        .disclaimer { background-color: #fef2f2; border: 1px solid #fecaca; border-radius: 8px; padding: 16px; margin: 16px 0; }
+        .success { background-color: #f0f9ff; border: 1px solid #bae6fd; border-radius: 8px; padding: 16px; margin: 16px 0; }
+        .warning { background-color: #fffbeb; border: 1px solid #fed7aa; border-radius: 8px; padding: 16px; margin: 16px 0; }
         """
     ) as demo:
         
@@ -227,21 +232,22 @@ def create_interface():
         
         # Status display
         if model_loaded:
-            gr.Markdown("""
+            method = "Pipeline" if pipeline_model else "Direct Model"
+            gr.Markdown(f"""
             <div class="success">
-            ✅ <strong>SYSTEM READY</strong><br>
-            MedGemma model is loaded and authenticated. You can now analyze medical images.
+            ✅ <strong>MEDGEMMA READY</strong><br>
+            Model loaded successfully using {method} method. Ready for medical image analysis.
             </div>
             """)
         else:
             gr.Markdown("""
-            <div class="disclaimer">
-            ⚠️ <strong>SYSTEM LOADING</strong><br>
-            MedGemma model is still loading. Please wait a few moments and refresh the page.
+            <div class="warning">
+            ⚠️ <strong>MODEL LOADING FAILED</strong><br>
+            MedGemma failed to load. Please ensure you have the latest transformers library and proper authentication.
             </div>
             """)
         
-        # Warning banner
+        # Medical disclaimer
         gr.Markdown("""
         <div class="disclaimer">
         ⚠️ <strong>IMPORTANT MEDICAL DISCLAIMER</strong><br>
@@ -251,66 +257,64 @@ def create_interface():
         """)
         
         with gr.Row():
-            # Left column - Inputs
+            # Left column
             with gr.Column(scale=1):
-                gr.Markdown("## 📤 Upload Medical Image")
+                gr.Markdown("## 📤 Medical Image Upload")
                 
                 image_input = gr.Image(
                     label="Medical Image", 
                     type="pil",
-                    height=300,
-                    sources=["upload", "clipboard"]
+                    height=300
                 )
                 
                 clinical_question = gr.Textbox(
                     label="Clinical Question *",
-                    placeholder="Examples:\n• Describe the findings in this chest X-ray\n• What pathological changes are visible?\n• Provide differential diagnosis\n• Identify any abnormalities",
+                    placeholder="Examples:\n• Describe findings in this chest X-ray\n• What pathological changes are visible?\n• Provide differential diagnosis\n• Identify abnormalities",
                     lines=4
                 )
                 
                 patient_history = gr.Textbox(
                     label="Patient History (Optional)",
-                    placeholder="e.g., 65-year-old male with chronic cough and dyspnea",
+                    placeholder="e.g., 65-year-old male with chronic cough",
                     lines=2
                 )
                 
                 with gr.Row():
                     clear_btn = gr.Button("🗑️ Clear", variant="secondary")
-                    analyze_btn = gr.Button("🔍 Analyze Image", variant="primary", size="lg")
-                
-                # System status
-                auth_status = "✅ Authenticated" if model_loaded else "🔄 Loading"
-                model_status = "✅ Ready" if model_loaded else "🔄 Loading"
+                    analyze_btn = gr.Button("🔍 Analyze", variant="primary", size="lg")
                 
+                # System info
                 gr.Markdown(f"""
-                **Authentication:** {auth_status}  
-                **Model Status:** {model_status}  
-                **Device:** {'CUDA' if torch.cuda.is_available() else 'CPU'}
+                **Status:** {'✅ Ready' if model_loaded else '❌ Failed'}  
+                **Method:** {'Pipeline' if pipeline_model else 'Direct' if model else 'None'}  
+                **Device:** {'CUDA' if torch.cuda.is_available() else 'CPU'}  
+                **Transformers:** {getattr(__import__('transformers'), '__version__', 'Unknown')}
                 """)
                 
-            # Right column - Output
+            # Right column
             with gr.Column(scale=1):
-                gr.Markdown("## 📋 Medical Analysis")
+                gr.Markdown("## 📋 Medical Analysis Results")
                 
                 output = gr.Textbox(
-                    label="AI Analysis Results",
+                    label="AI Medical Analysis",
                     lines=20,
                     show_copy_button=True,
-                    placeholder="Upload a medical image and ask a clinical question to get started..."
+                    placeholder="Upload a medical image and ask a clinical question..." if model_loaded else "Model unavailable - please check system status"
                 )
         
-        # Example cases
-        with gr.Accordion("📚 Example Cases", open=False):
-            examples = gr.Examples(
-                examples=[
-                    [
-                        "https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png",
-                        "Analyze this chest X-ray for any abnormal findings. Comment on heart size, lung fields, and overall anatomy.",
-                        "Adult patient with respiratory symptoms"
-                    ]
-                ],
-                inputs=[image_input, clinical_question, patient_history]
-            )
+        # Examples
+        if model_loaded:
+            with gr.Accordion("📚 Example Cases", open=False):
+                examples = gr.Examples(
+                    examples=[
+                        [
+                            "https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png",
+                            "Analyze this chest X-ray systematically. Comment on heart size, lung fields, and any abnormalities.",
+                            "Adult patient with respiratory symptoms"
+                        ]
+                    ],
+                    inputs=[image_input, clinical_question, patient_history]
+                )
         
         # Event handlers
         analyze_btn.click(
@@ -320,11 +324,8 @@ def create_interface():
             show_progress=True
         )
         
-        def clear_all():
-            return None, "", "", ""
-        
         clear_btn.click(
-            fn=clear_all,
+            fn=lambda: (None, "", "", ""),
             outputs=[image_input, clinical_question, patient_history, output]
         )
         
@@ -333,13 +334,12 @@ def create_interface():
         ---
         ### 🔬 About MedGemma
         
-        MedGemma-4B is Google's specialized medical AI model for educational medical image analysis.
-        It demonstrates strong performance across radiology, pathology, dermatology, and ophthalmology.
+        MedGemma-4B is Google's specialized medical AI model requiring transformers >= 4.52.0.
         
         ### 🔒 Privacy & Ethics
-        - Real-time processing with no data retention
-        - Designed for educational and research use only
-        - No PHI or patient data should be uploaded
+        - Real-time processing, no data storage
+        - Educational and research purposes only
+        - No patient data should be uploaded
         
         **Model:** Google MedGemma-4B | **License:** Apache 2.0
         """)