import gradio as gr
from transformers import AutoProcessor, IdeficsForVisionText2Text
from peft import PeftModel
from PIL import Image
import torch
import re

# Set up the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load the processor
processor = AutoProcessor.from_pretrained("idefics9b-lora-roco")

# Load base model (no quantization)
base_model = IdeficsForVisionText2Text.from_pretrained(
    "HuggingFaceM4/idefics-9b-instruct",
    torch_dtype=torch.float32  # Use float32 for CPU
).to(device)

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "idefics9b-lora-roco").to(device)
model.eval()

# Ensure image is RGB
def convert_to_rgb(image):
    if image.mode == "RGB":
        return image
    return image.convert("RGB")

# Inference function
def generate_caption(image: Image.Image) -> str:
    prompt = "User: What do you see in this picture?\nAssistant:"
    image = convert_to_rgb(image)

    inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
    bad_words_ids = processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
    eos_token_id = processor.tokenizer("<end_of_utterance>", add_special_tokens=False).input_ids[0]

    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            eos_token_id=eos_token_id,
            bad_words_ids=bad_words_ids,
            max_new_tokens=100,
            do_sample=False,
        )

    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]

    # Clean the generated text
    clean_text = re.sub(r"<.*?>", "", generated_text).strip()

    # Extract only after 'Assistant:'
    match = re.search(r"Assistant:\s*(.*)", clean_text, re.IGNORECASE | re.DOTALL)
    if match:
        clean_text = match.group(1).strip()

    return clean_text

# Launch Gradio interface
demo = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil", label="Upload Medical Image", image_mode='RGB', width=512, height=512),
    outputs=gr.Textbox(label="Generated Caption"),
    title="IDEFICS-9B Medical Image Captioning",
    description="Upload a medical image and get the caption using the fine-tuned IDEFICS-9B model (with LoRA).",
    theme="default"
)

demo.launch()