Update app.py

app.py

@@ -1,235 +1,71 @@
-# import gradio as gr
-# from huggingface_hub import InferenceClient
-
-
-# def respond(
-#     message,
-#     history: list[dict[str, str]],
-#     system_message,
-#     max_tokens,
-#     temperature,
-#     top_p,
-#     hf_token: gr.OAuthToken,
-# ):
-#     """
-#     For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-#     """
-#     client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
-
-#     messages = [{"role": "system", "content": system_message}]
-
-#     messages.extend(history)
-
-#     messages.append({"role": "user", "content": message})
-
-#     response = ""
-
-#     for message in client.chat_completion(
-#         messages,
-#         max_tokens=max_tokens,
-#         stream=True,
-#         temperature=temperature,
-#         top_p=top_p,
-#     ):
-#         choices = message.choices
-#         token = ""
-#         if len(choices) and choices[0].delta.content:
-#             token = choices[0].delta.content
-
-#         response += token
-#         yield response
-
-
-# """
-# For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-# """
-# chatbot = gr.ChatInterface(
-#     respond,
-#     type="messages",
-#     additional_inputs=[
-#         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-#         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-#         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-#         gr.Slider(
-#             minimum=0.1,
-#             maximum=1.0,
-#             value=0.95,
-#             step=0.05,
-#             label="Top-p (nucleus sampling)",
-#         ),
-#     ],
-# )
-
-# with gr.Blocks() as demo:
-#     with gr.Sidebar():
-#         gr.LoginButton()
-#     chatbot.render()
-
-
-# if __name__ == "__main__":
-#     demo.launch()
-
-# Medical Image Chatbot Application
-# app.py
-from ultralytics import YOLO
-# Imports for the Vision-Language Model (VLM)
-from transformers import AutoProcessor, AutoModelForCausalLM
 import gradio as gr
-from PIL import Image
-import torch
-import os 
-
-# Global variable to store the exact error if model loading fails
-LLAVA_LOAD_ERROR = None
-
-# -----------------------------
-# 1. Load models
-# -----------------------------
-# YOLO Model Loading (requires 'best.pt' in the same directory)
-try:
-    yolo_model = YOLO("best.pt")
-except FileNotFoundError:
-    print("WARNING: 'best.pt' not found. App will use a dummy classification for structure testing.")
-    # Dummy class for structure testing if best.pt is missing
-    class DummyYOLO:
-        def __call__(self, image):
-            class_names = {0: "Normal", 1: "Pneumonia", 2: "Fracture"}
-            class DummyResult:
-                def __init__(self):
-                    self.names = class_names
-                    # Simulating a top-1 prediction for class 1 (Pneumonia)
-                    self.probs = type('Obj', (object,), {'top1': 1}) 
-            return [DummyResult()]
-    yolo_model = DummyYOLO()
-
-# Load LLaVA-Med VLM (Processor and Model)
-LLAVA_MODEL_ID = "microsoft/llava-med-v1.5-mistral-7b"
-print(f"Loading VLM: {LLAVA_MODEL_ID}. This may take some time...")
-
-try:
-    # Determine device: Use GPU if available, otherwise CPU
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {device}")
-
-    llava_processor = AutoProcessor.from_pretrained(LLAVA_MODEL_ID)
-    
-    # Load model with optimizations for large models
-    llava_model = AutoModelForCausalLM.from_pretrained(
-        LLAVA_MODEL_ID,
-        torch_dtype=torch.float16, 
-        device_map="auto" # Tries to intelligently place model parts (necessary for large models)
-    )
-    # Ensure the model is moved to the determined device
-    llava_model.to(device) 
+from huggingface_hub import InferenceClient
+
+
+def respond(
+    message,
+    history: list[dict[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    hf_token: gr.OAuthToken,
+):
+    """
+    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+    """
+    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+
+    messages = [{"role": "system", "content": system_message}]
+
+    messages.extend(history)
+
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+
+    for message in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        choices = message.choices
+        token = ""
+        if len(choices) and choices[0].delta.content:
+            token = choices[0].delta.content
+
+        response += token
+        yield response
+
+
+"""
+For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+"""
+chatbot = gr.ChatInterface(
+    respond,
+    type="messages",
+    additional_inputs=[
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
+    ],
+)
+
+with gr.Blocks() as demo:
+    with gr.Sidebar():
+        gr.LoginButton()
+    chatbot.render()
 
-except Exception as e:
-    # Store the specific error message to display it in the UI
-    LLAVA_LOAD_ERROR = str(e)
-    print(f"CRITICAL ERROR: Failed to load LLaVA-Med model. Error: {LLAVA_LOAD_ERROR}")
-    llava_processor = None
-    llava_model = None
 
-
-# -----------------------------
-# 2. Define main logic
-# -----------------------------
-def medical_chat(image_path, question):
-    if image_path is None:
-        return "⚠️ Please upload a medical image before asking a question."
-
-    # Load image from filepath provided by Gradio (type="filepath" in gr.Image)
-    try:
-        image = Image.open(image_path)
-    except Exception:
-        return "❌ Could not open the image file from the path provided by Gradio."
-
-    # YOLO classification
-    try:
-        results = yolo_model(image)
-        if hasattr(results[0], 'probs') and results[0].probs is not None and results[0].names:
-            predicted_class = results[0].names[int(results[0].probs.top1)]
-        else:
-            predicted_class = "Unknown class (YOLO detection failed to find a simple classification result)."
-    except Exception as e:
-        predicted_class = f"YOLO analysis failed (Error: {e})."
-        print(f"YOLO Error: {e}")
-
-    yolo_output = f"The YOLO model classified the image as: **{predicted_class}**."
-
-    # Direct YOLO question
-    if "output of yolo" in question.lower():
-        return yolo_output
-
-    # Explanation using LLaVA-Med (VLM)
-    if llava_model is None or llava_processor is None:
-        # Return the specific error details and troubleshooting tips
-        troubleshooting_tips = (
-            "**Troubleshooting:**\n"
-            "1. **Memory:** LLaVA-Med V1.5 is a 7B parameter model, requiring significant RAM/VRAM. \n"
-            "2. **Dependencies:** Ensure all libraries in `requirements.txt` (`accelerate`, `bitsandbytes`, `safetensors`) are installed.\n"
-            "3. **Alternative:** If running on a CPU-only or low-memory machine, consider switching to a smaller text-only model like BioGPT-Large."
-        )
-        return (
-            f"{yolo_output}\n\n---\n\n"
-            f"❌ **CRITICAL ERROR: LLaVA-Med model failed to load.**\n"
-            f"**Specific Error:** {LLAVA_LOAD_ERROR or 'No specific error message captured.'}\n\n"
-            f"{troubleshooting_tips}"
-        )
-
-    # LLaVA Prompting (multimodal input structure)
-    llava_prompt = (
-        f"USER: <image>\n"
-        f"The medical image was classified by a separate model as '{predicted_class}'. "
-        f"Based on the visual evidence in the image and this classification, {question} ASSISTANT:"
-    )
-
-    try:
-        # Prepare inputs and move to device
-        inputs = llava_processor(text=llava_prompt, images=image, return_tensors="pt")
-        inputs = {k: v.to(llava_model.device) for k, v in inputs.items()}
-
-        # Generate response
-        output = llava_model.generate(
-            **inputs,
-            max_new_tokens=200,
-            do_sample=True,
-            temperature=0.7,
-        )
-
-        # Decode and clean up output
-        answer = llava_processor.decode(output[0], skip_special_tokens=True)
-        # We only want the ASSISTANT's response
-        clean_answer = answer.split("ASSISTANT:")[-1].strip() 
-
-    except Exception as e:
-        clean_answer = f"LLaVA-Med failed during response generation. Error: {e}"
-
-    return f"{yolo_output}\n\n---\n\n**Explanation (Powered by LLaVA-Med VLM):** {clean_answer}"
-
-# -----------------------------
-# 3. Build custom UI using Blocks (Custom multimodal layout)
-# -----------------------------
-with gr.Blocks(title="🧠 YOLO + Medical Chatbot (LLaVA-Med)") as demo:
-    gr.Markdown("# 🩺 Medical Image Analyzer & Chatbot (YOLO + LLaVA-Med VLM)")
-    gr.Markdown("Upload an image. YOLO classifies it, and LLaVA-Med uses the image and classification result to provide an expert-level explanation.")
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.Image(type="filepath", label="🩻 Upload Medical Image", height=300)
-            text_input = gr.Textbox(
-                label="💬 Ask your question",
-                placeholder="e.g., What is the output of YOLO? or Explain this disease and its symptoms."
-            )
-            submit_btn = gr.Button("Analyze & Ask", variant="primary")
-
-        with gr.Column(scale=2):
-            output_text = gr.Textbox(label="🧠 Combined Response", lines=10, interactive=False)
-
-    # Trigger the analysis function when the button is clicked
-    submit_btn.click(fn=medical_chat, inputs=[image_input, text_input], outputs=output_text)
-
-# -----------------------------
-# 4. Launch the app
-# -----------------------------
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()
+