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

torch.manual_seed(42)

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

MODEL_IDS = {
    "Florence-2-base": "florence-community/Florence-2-base",
    "Florence-2-large": "florence-community/Florence-2-large",
}

TASK = "<MORE_DETAILED_CAPTION>"

models = {}
processors = {}

for name, repo_id in MODEL_IDS.items():
    try:
        models[name] = (
            Florence2ForConditionalGeneration.from_pretrained(repo_id, torch_dtype=dtype)
            .to(device)
            .eval()
        )
        processors[name] = AutoProcessor.from_pretrained(repo_id)
    except Exception as e:
        print(f"Error loading {name}: {e}")
        models[name] = None
        processors[name] = None


def describe_image(uploaded_image, model_choice):
    """
    Generates a detailed description of the input image using the selected model.

    Args:
        uploaded_image (PIL.Image.Image | np.ndarray): The image to describe.
        model_choice (str): "Florence-2-base" or "Florence-2-large".

    Returns:
        str: A detailed textual description of the image or an error message.
    """
    if uploaded_image is None:
        return "Please upload an image."

    if model_choice not in MODEL_IDS:
        return "Invalid model choice."

    model = models.get(model_choice)
    processor = processors.get(model_choice)
    if model is None or processor is None:
        return f"{model_choice} failed to load."

    if not isinstance(uploaded_image, Image.Image):
        uploaded_image = Image.fromarray(uploaded_image)
    uploaded_image = uploaded_image.convert("RGB")

    inputs = processor(text=TASK, images=uploaded_image, return_tensors="pt").to(device, dtype)

    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=1024,
            early_stopping=False,
            do_sample=False,
            num_beams=3,
        )

    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
    processed_description = processor.post_process_generation(
        generated_text,
        task=TASK,
        image_size=(uploaded_image.width, uploaded_image.height),
    )
    image_description = processed_description[TASK]
    print("\nImage description generated!:", image_description)
    return image_description


# Description for the interface
description = "> Select the model to use for generating the image description. 'Base' is smaller and faster, while 'Large' is more accurate but slower."
if device == "cpu":
    description += " Note: Running on CPU, which may be slow for large models."

# Define examples
examples = [
    ["images/2.jpeg", "Florence-2-large"],
    ["images/1.jpeg", "Florence-2-base"],
    ["images/3.jpeg", "Florence-2-large"],
    ["images/4.jpeg", "Florence-2-large"],
]

css = """
.submit-btn {
    background-color: #4682B4 !important;
    color: white !important;
}
.submit-btn:hover {
    background-color: #87CEEB !important;
}
"""

with gr.Blocks() as demo:
    gr.Markdown("# **Florence-2 Models Image Captions**")
    gr.Markdown(description)
    with gr.Row():
        # Left column: Input image and Generate button
        with gr.Column():
            image_input = gr.Image(label="Upload Image", type="pil")
            generate_btn = gr.Button("Generate Caption", elem_classes="submit-btn")
        # Right column: Model choice and output
        with gr.Column():
            model_choice = gr.Radio(
                list(MODEL_IDS.keys()),
                label="Model Choice",
                value="Florence-2-base",
            )
            output = gr.Textbox(label="Generated Caption", lines=4, interactive=True)

    # Examples (both inputs wired so the model choice in each example is applied)
    gr.Examples(examples=examples, inputs=[image_input, model_choice])

    # Connect the button to the function
    generate_btn.click(fn=describe_image, inputs=[image_input, model_choice], outputs=output)

# Launch the interface
demo.launch(
    theme="bethecloud/storj_theme",
    css=css,
    debug=True,
    mcp_server=True,
    ssr_mode=False,
    show_error=True,
)