import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM

device = "cuda" if torch.cuda.is_available() else "cpu"

florence_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Florence-2-base",
    trust_remote_code=True
).to(device).eval()

florence_processor = AutoProcessor.from_pretrained(
    "microsoft/Florence-2-base",
    trust_remote_code=True
)

def generate_caption(image):
    if image.mode != "RGB":
        image = image.convert("RGB")

    inputs = florence_processor(
        text="<MORE_DETAILED_CAPTION>",
        images=image,
        return_tensors="pt"
    ).to(device)

    generated_ids = florence_model.generate(
        input_ids=inputs["input_ids"],
        pixel_values=inputs["pixel_values"],
        max_new_tokens=1024,
        early_stopping=False,
        do_sample=False,
        num_beams=3,
    )

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

    parsed_answer = florence_processor.post_process_generation(
        generated_text,
        task="<MORE_DETAILED_CAPTION>",
        image_size=(image.width, image.height)
    )

    prompt = parsed_answer["<MORE_DETAILED_CAPTION>"]
    print("\n\nGeneration completed!:" + prompt)
    return prompt

gr.Interface(
    generate_caption,
    inputs=gr.Image(label="Input Image", type="pil", image_mode="RGB"),
    outputs=gr.Textbox(label="Output Prompt", lines=2, show_copy_button=True),
    deep_link=False
).launch()