Update app.py

app.py

@@ -1,37 +1,32 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import pipeline
 from PIL import Image
 import torch
 import os
 import spaces
 
-# --- Initialize the Model and Processor Directly ---
-print("Loading MedGemma model and processor...")
-model_id = "google/medgemma-4b-it"
-model = None
-processor = None
-model_loaded = False
-
+# --- Initialize the Model Pipeline (As per your working example) ---
+print("Loading MedGemma model...")
 try:
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
+    pipe = pipeline(
+        "image-text-to-text",
+        model="google/medgemma-4b-it",
         torch_dtype=torch.bfloat16,
         device_map="auto",
         token=os.environ.get("HF_TOKEN")
     )
-    processor = AutoProcessor.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
     model_loaded = True
-    print("Model and processor loaded successfully!")
+    print("Model loaded successfully!")
 except Exception as e:
     model_loaded = False
     print(f"Error loading model: {e}")
 
-# --- Core Analysis Function (Final Corrected Version) ---
+# --- Core Analysis Function (Using the logic from your working example) ---
 @spaces.GPU()
 def analyze_symptoms(symptom_image: Image.Image, symptoms_text: str):
     """
-    Analyzes symptoms using the definitive two-step templating and processing method
-    required by modern multimodal chat models.
+    Analyzes user's symptoms using the definitive calling convention demonstrated
+    in the working X-ray analyzer example.
     """
     if not model_loaded:
         return "Error: The AI model could not be loaded. Please check the Space logs."
@@ -41,59 +36,40 @@ def analyze_symptoms(symptom_image: Image.Image, symptoms_text: str):
         return "Please describe your symptoms or upload an image for analysis."
 
     try:
-        # --- STEP 1: Build the structured messages list ---
-        system_instruction = (
-            "You are an expert, empathetic AI medical assistant. "
-            "Analyze the potential medical condition based on the following information. "
-            "Provide a list of possible conditions, your reasoning, and a clear, actionable next-steps plan."
+        # --- DEFINITIVE MESSAGE CONSTRUCTION (from your example) ---
+        system_prompt = (
+            "You are an expert, empathetic AI medical assistant. Analyze the potential "
+            "medical condition based on the following information. Provide a list of "
+            "possible conditions, your reasoning, and a clear, actionable next-steps plan. "
+            "Start your analysis by describing the user-provided information."
         )
 
-        # The 'content' for a user's turn is a LIST of dictionaries.
-        user_content_list = []
-        if symptom_image:
-            # Add a placeholder dictionary for the image.
-            user_content_list.append({"type": "image"})
+        user_content = []
+        # The user's prompt text is always present.
+        user_content.append({"type": "text", "text": symptoms_text})
         
-        # Add the dictionary for the text.
-        text_content = f"{symptoms_text}\n\n{system_instruction}"
-        user_content_list.append({"type": "text", "text": text_content})
-
+        # The actual PIL image object is added to the content list if it exists.
+        if symptom_image:
+            user_content.append({"type": "image", "image": symptom_image})
+            
         messages = [
-            {"role": "user", "content": user_content_list}
+            {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
+            {"role": "user", "content": user_content}
         ]
-
-        # --- STEP 2: Generate the prompt string using the official template ---
-        # This will correctly create a string with all special tokens, including <image>.
-        prompt = processor.tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-
-        # --- STEP 3: Process the prompt string and image together ---
-        # This is where the prompt's <image> token is linked to the actual image data.
-        inputs = processor(
-            text=prompt,
-            images=symptom_image, # This can be None for text-only cases
-            return_tensors="pt"
-        ).to(model.device)
-
-        # Generation parameters
-        generate_kwargs = {
+        
+        generation_args = {
             "max_new_tokens": 512,
             "do_sample": True,
             "temperature": 0.7,
         }
-
-        print("Generating model output with the definitive two-step process...")
         
-        # Generate the response
-        generate_ids = model.generate(**inputs, **generate_kwargs)
-
-        # Decode only the newly generated tokens
-        input_token_len = inputs["input_ids"].shape[-1]
-        result = processor.batch_decode(generate_ids[:, input_token_len:], skip_special_tokens=True)[0]
+        # --- DEFINITIVE PIPELINE CALL (from your example) ---
+        # The entire messages structure is passed to the `text` argument.
+        output = pipe(text=messages, **generation_args)
         
+        # The result is the 'content' of the last generated message.
+        result = output[0]["generated_text"][-1]["content"]
+
         disclaimer = "\n\n***Disclaimer: I am an AI assistant and not a medical professional. This is not a diagnosis. Please consult a doctor for any health concerns.***"
         
         return result.strip() + disclaimer
@@ -102,7 +78,7 @@ def analyze_symptoms(symptom_image: Image.Image, symptoms_text: str):
         print(f"An exception occurred during analysis: {type(e).__name__}: {e}")
         return f"An error occurred during analysis. Please check the logs for details: {str(e)}"
 
-# --- Gradio Interface (No changes needed) ---
+# --- Gradio Interface (Your original, no changes needed) ---
 with gr.Blocks(theme=gr.themes.Soft(), title="AI Symptom Analyzer") as demo:
     gr.HTML("""
         <div style="text-align: center; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 2rem; border-radius: 10px; margin-bottom: 2rem;">
@@ -134,10 +110,13 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI Symptom Analyzer") as demo:
                 label="AI Analysis", lines=25, show_copy_button=True, placeholder="Analysis results will appear here...")
 
     def clear_all():
-        return None, "", ""
+        # This function should return values for all outputs cleared by the button
+        return None, ""
 
+    # The clear button now correctly clears the image and text input.
     analyze_btn.click(fn=analyze_symptoms, inputs=[image_input, symptoms_input], outputs=output_text)
-    clear_btn.click(fn=clear_all, outputs=[image_input, symptoms_input, output_text])
+    clear_btn.click(fn=lambda: (None, "", ""), outputs=[image_input, symptoms_input, output_text])
+
 
 if __name__ == "__main__":
     print("Starting Gradio interface...")