import gradio as gr
import torch
from PIL import Image
import os
from transformers import AutoProcessor, LlavaNextForConditionalGeneration

# Model config
MODEL_ID = "OralGPT/OralGPT-Omni-7B-Instruct" # Based on README news

def load_model():
    print(f"Loading model {MODEL_ID}...")
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    model = LlavaNextForConditionalGeneration.from_pretrained(
        MODEL_ID, 
        torch_dtype=torch.float16, 
        low_cpu_mem_usage=True,
        device_map="auto"
    )
    return processor, model

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

def predict(image, text):
    global processor, model
    if processor is None or model is None:
        processor, model = load_model()

    if image is None:
        return "Please upload a dental image."

    # Prepare prompt following LLaVA-v1.6 / LLaVA-Next style (which OralGPT-Omni uses)
    # Note: OralGPT-Omni is based on Qwen/Vicuna usually.
    # The prompt template might need adjustment based on the specific base model.
    
    prompt = f"A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. USER: <image>\n{text} ASSISTANT:"
    
    inputs = processor(text=prompt, images=image, return_tensors="pt").to(model.device, torch.float16)
    
    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=512)
    
    response = processor.decode(output[0], skip_special_tokens=True)
    
    # Post-process to extract only the assistant's part
    if "ASSISTANT:" in response:
        response = response.split("ASSISTANT:")[-1].strip()
    
    return response

# Gradio Interface
description = """
# 🦷 OralGPT-Omni: Digital Dentistry Assistant
Upload a dental image (X-ray, photograph, etc.) and ask a question.
OralGPT is specialized in dental multimodal analysis including Panoramic X-rays, Periapical radiographs, and more.
"""

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Dental Image"),
        gr.Textbox(label="Question", placeholder="Describe this dental X-ray...")
    ],
    outputs=gr.Textbox(label="OralGPT Response"),
    title="OralGPT-Omni Demo",
    description=description,
    examples=[
        ["OralGPT/assets/mmoral-logo.png", "What is shown in this logo?"]
    ]
)

if __name__ == "__main__":
    demo.launch()