app.py

# app.py - MedGemma with Authentication
import gradio as gr
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText
from PIL import Image
import os
import logging
from huggingface_hub import login

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Authenticate with Hugging Face
def authenticate_hf():
    """Authenticate with Hugging Face using token"""
    try:
        hf_token = os.getenv('HF_TOKEN')
        if hf_token:
            login(token=hf_token)
            logger.info("✅ Authenticated with Hugging Face")
            return True
        else:
            logger.warning("⚠️ No HF_TOKEN found in environment")
            return False
    except Exception as e:
        logger.error(f"❌ Authentication failed: {e}")
        return False

# Model configuration
MODEL_ID = "google/medgemma-4b-it"

# Global variables for model and processor
model = None
processor = None

def load_model():
    """Load model and processor with authentication"""
    global model, processor
    
    try:
        # First authenticate
        auth_success = authenticate_hf()
        if not auth_success:
            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}")
        
        # Load model with authentication
        model = AutoModelForImageTextToText.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
            device_map="auto" if device == "cuda" else None,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
            token=True  # Use authenticated token
        )
        
        processor = AutoProcessor.from_pretrained(
            MODEL_ID, 
            trust_remote_code=True,
            token=True  # Use authenticated token
        )
        
        logger.info("✅ Model loaded successfully!")
        return True
        
    except Exception as e:
        logger.error(f"❌ Error loading model: {str(e)}")
        return False

# Initialize model at startup
model_loaded = load_model()

def analyze_medical_image(image, clinical_question, patient_history=""):
    """Analyze medical image with clinical context"""
    global model, processor
    
    # Check if model is loaded
    if not model_loaded or model is None or processor is None:
        return """❌ **Model Authentication Issue**

MedGemma requires authentication. Please ensure:

1. **HF_TOKEN is set**: The Space owner needs to add their Hugging Face token to Space Settings → Repository secrets
2. **Model access approved**: Make sure you have access to MedGemma at https://huggingface.co/google/medgemma-4b-it
3. **Space restart**: After adding the token, restart the Space

**Current Status**: Authentication failed - model cannot load without proper token."""
    
    if image is None:
        return "⚠️ Please upload a medical image first."
    
    if not clinical_question.strip():
        return "⚠️ Please provide a clinical question."
    
    try:
        # Prepare the conversation with proper medical context
        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. Be thorough but clear in your explanations."}]
            }
        ]
        
        # 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
        inputs = processor.apply_chat_template(
            messages, 
            add_generation_prompt=True, 
            tokenize=True,
            return_dict=True, 
            return_tensors="pt"
        )
        
        # Move to appropriate device
        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 with appropriate settings
        with torch.inference_mode():
            generation = model.generate(
                **inputs, 
                max_new_tokens=1500,
                do_sample=True,
                temperature=0.3,  # Lower temperature for more focused medical analysis
                top_p=0.95,
                repetition_penalty=1.1,
                pad_token_id=processor.tokenizer.eos_token_id
            )
            generation = generation[0][input_len:]
        
        # Decode response
        response = processor.decode(generation, skip_special_tokens=True)
        
        # Clean up response
        response = response.strip()
        
        # Add structured disclaimer
        disclaimer = """

---
### ⚠️ MEDICAL DISCLAIMER
**This analysis is for educational and research purposes only.**
- This AI assistant is not a substitute for professional medical advice
- Always consult qualified healthcare professionals for diagnosis and treatment
- Do not make medical decisions based solely on this analysis
- In case of medical emergency, contact emergency services immediately
---
        """
        
        return response + disclaimer
        
    except Exception as e:
        logger.error(f"Error in analyze_medical_image: {str(e)}")
        return f"❌ Error processing request: {str(e)}\n\nPlease try again or contact support if the issue persists."

# Create Gradio interface
def create_interface():
    with gr.Blocks(
        title="MedGemma Medical Image 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;
        }
        .auth-warning {
            background-color: #fffbeb;
            border: 1px solid #fed7aa;
            border-radius: 8px;
            padding: 16px;
            margin: 16px 0;
        }
        """
    ) as demo:
        
        # Header
        gr.Markdown("""
        # 🏥 MedGemma Medical Image Analysis
        
        **Advanced Medical AI Assistant powered by Google's MedGemma-4B**
        
        This tool can analyze various medical imaging modalities including:
        - 🫁 **Chest X-rays** - Pneumonia, COVID-19, lung pathology
        - 🧠 **CT Scans** - Brain, chest, abdomen imaging
        - 🔬 **Histopathology** - Microscopic tissue analysis
        - 👁️ **Ophthalmology** - Retinal imaging, eye conditions
        - 🩺 **Dermatology** - Skin lesions and conditions
        """)
        
        # Authentication status
        if not model_loaded:
            gr.Markdown("""
            <div class="auth-warning">
            🔐 <strong>AUTHENTICATION REQUIRED</strong><br>
            MedGemma model requires authentication. Please:
            <ol>
            <li>Ensure you have access to the model at <a href="https://huggingface.co/google/medgemma-4b-it">MedGemma page</a></li>
            <li>Add your HF_TOKEN to Space Settings → Repository secrets</li>
            <li>Restart the Space</li>
            </ol>
            </div>
            """)
        
        # Warning banner
        gr.Markdown("""
        <div class="disclaimer">
        ⚠️ <strong>IMPORTANT MEDICAL DISCLAIMER</strong><br>
        This tool is for <strong>educational and research purposes only</strong>. 
        Do not upload real patient data or use for actual medical diagnosis. 
        Always consult qualified healthcare professionals.
        </div>
        """)
        
        with gr.Row():
            # Left column - Inputs
            with gr.Column(scale=1):
                gr.Markdown("## 📤 Upload & Configure")
                
                image_input = gr.Image(
                    label="Medical Image", 
                    type="pil",
                    height=350,
                    sources=["upload", "clipboard"]
                )
                
                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 based on imaging\n• Identify any abnormalities present",
                    lines=4,
                    max_lines=6
                )
                
                patient_history = gr.Textbox(
                    label="Patient History (Optional)",
                    placeholder="Example: 65-year-old male presenting with chronic cough, shortness of breath, and chest pain. History of smoking for 30 years.",
                    lines=3,
                    max_lines=5
                )
                
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
                    analyze_btn = gr.Button("🔍 Analyze Image", variant="primary", size="lg")
                
                # Model status
                auth_status = "✅ Authenticated" if model_loaded else "🔐 Authentication Required"
                model_status = "✅ Loaded" if model_loaded else "❌ Not Loaded"
                
                gr.Markdown(f"""
                **Authentication:** {auth_status}  
                **Model Status:** {model_status}  
                **Model:** {MODEL_ID}  
                **Device:** {'CUDA' if torch.cuda.is_available() else 'CPU'}
                """)
                
            # Right column - Output
            with gr.Column(scale=1):
                gr.Markdown("## 📋 Medical Analysis Results")
                
                output = gr.Textbox(
                    label="AI Medical Analysis",
                    lines=25,
                    max_lines=35,
                    show_copy_button=True,
                    placeholder="Upload an image and ask a clinical question to get started..."
                )
        
        # Example cases section
        gr.Markdown("## 💡 Example Use Cases")
        
        with gr.Accordion("Click to see 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 and describe any abnormal findings. Comment on the heart size, lung fields, and overall chest anatomy.",
                        "Adult patient presenting with respiratory symptoms including cough and shortness of breath."
                    ],
                    [
                        None,
                        "What pathological changes are visible in this medical image? Provide a structured analysis including anatomical observations and potential diagnoses.",
                        "Patient with acute onset symptoms"
                    ],
                    [
                        None,
                        "Perform a systematic review of this imaging study. Include: 1) Technical quality assessment, 2) Normal anatomical structures, 3) Abnormal findings, 4) Clinical significance.",
                        ""
                    ],
                    [
                        None,
                        "Compare the findings in this image to normal anatomy. What are the key differences and what might they suggest clinically?",
                        "Follow-up imaging for known condition"
                    ]
                ],
                inputs=[image_input, clinical_question, patient_history],
                label="Click any example to load it"
            )
        
        # Event handlers
        analyze_btn.click(
            fn=analyze_medical_image,
            inputs=[image_input, clinical_question, patient_history],
            outputs=output,
            show_progress=True
        )
        
        def clear_all():
            return None, "", "", ""
        
        clear_btn.click(
            fn=clear_all,
            outputs=[image_input, clinical_question, patient_history, output]
        )
        
        # Footer information
        gr.Markdown("""
        ---
        ### 🔬 About MedGemma
        
        MedGemma is Google's specialized medical AI model trained on medical imaging and clinical text. 
        **Note**: This model requires authentication and access approval.
        
        ### 🔒 Privacy & Data Policy
        - **No data storage**: Images and text are processed in real-time and not saved
        - **No patient data**: Use only synthetic, anonymized, or educational images  
        - **Educational use**: This tool is designed for learning and research purposes
        
        ### 📞 Support
        For technical issues or questions, please create an issue in the [Hugging Face Space repository](https://huggingface.co/spaces).
        
        **Model**: Google MedGemma-4B | **Framework**: Transformers + Gradio | **License**: Apache 2.0
        """)
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )