app.py

import os
import sys
from typing import Iterable
import gradio as gr
import torch
import requests
from PIL import Image
from transformers import AutoProcessor, Florence2ForConditionalGeneration
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

# --- Theme and CSS Definition ---

colors.steel_blue = colors.Color(
    name="steel_blue",
    c50="#EBF3F8", c100="#D3E5F0", c200="#A8CCE1", c300="#7DB3D2",
    c400="#529AC3", c500="#4682B4", c600="#3E72A0", c700="#36638C",
    c800="#2E5378", c900="#264364", c950="#1E3450",
)

class SteelBlueTheme(Soft):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.steel_blue,
        neutral_hue: colors.Color | str = colors.slate,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue,
            text_size=text_size, font=font, font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
            slider_color="*secondary_500",
            slider_color_dark="*secondary_600",
            block_title_text_weight="600",
            block_border_width="3px",
            block_shadow="*shadow_drop_lg",
            button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px",
            color_accent_soft="*primary_100",
            block_label_background_fill="*primary_200",
        )

steel_blue_theme = SteelBlueTheme()

css = """
#main-title h1 {
    font-size: 2.3em !important;
}
#output-title h2 {
    font-size: 2.1em !important;
}
"""

# --- Model Loading ---

# Note: This will load all four models into memory, which can be very resource-intensive.
# Ensure you have enough VRAM/RAM available. `device_map="auto"` will help distribute them.

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

models = {}
processors = {}

print("Loading Florence-2 models... This may take a while.")
for name, repo_id in MODEL_IDS.items():
    print(f"Loading {name}...")
    # Load the model with bfloat16 precision and automatic device mapping
    model = Florence2ForConditionalGeneration.from_pretrained(
        repo_id,
        dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True
    )
    processor = AutoProcessor.from_pretrained(repo_id, trust_remote_code=True)
    models[name] = model
    processors[name] = processor
    print(f"✅ Finished loading {name}.")

print("\n🎉 All models loaded successfully!")


# --- Inference Function ---

def run_florence2_inference(model_name: str, image: Image.Image, task_prompt: str,
                            max_new_tokens: int = 1024, num_beams: int = 3):
    """
    Runs inference using the selected Florence-2 model.
    """
    if image is None:
        return "Please upload an image to get started."

    # Select the model and processor based on user choice
    model = models[model_name]
    processor = processors[model_name]

    # Prepare inputs for the model
    inputs = processor(text=task_prompt, images=image, return_tensors="pt").to(model.device, torch.bfloat16)

    # Generate output IDs
    generated_ids = model.generate(
        input_ids=inputs["input_ids"],
        pixel_values=inputs["pixel_values"],
        max_new_tokens=max_new_tokens,
        num_beams=num_beams,
        do_sample=False # For more deterministic output
    )

    # Decode the generated IDs to text
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]

    # Post-process the generation to get a structured answer
    # Important: The task prompt must be passed to the post-processing function
    image_size = image.size
    parsed_answer = processor.post_process_generation(
        generated_text, task=task_prompt, image_size=image_size
    )

    return parsed_answer


# --- Gradio Interface Definition ---

# Define a list of tasks that Florence-2 can perform
florence_tasks = [
    "<OD>", "<CAPTION>", "<DETAILED_CAPTION>", "<MORE_DETAILED_CAPTION>",
    "<DENSE_REGION_CAPTION>", "<REGION_PROPOSAL>", "<OCR>", "<OCR_WITH_REGION>"
]

# Download an example image for the interface
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
example_image = Image.open(requests.get(url, stream=True).raw).convert("RGB")


with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
    gr.Markdown("# **Florence-2 Vision Models**", elem_id="main-title")
    gr.Markdown("Select a model, upload an image, choose a task, and click Submit to see the results.")

    with gr.Row():
        with gr.Column(scale=2):
            image_upload = gr.Image(type="pil", label="Upload Image", value=example_image)
            task_prompt = gr.Dropdown(
                label="Select Task",
                choices=florence_tasks,
                value="<OD>"
            )
            model_choice = gr.Radio(
                choices=list(MODEL_IDS.keys()),
                label="Select Model",
                value="Florence-2-large-ft"
            )
            image_submit = gr.Button("Submit", variant="primary")

            with gr.Accordion("Advanced options", open=False):
                max_new_tokens = gr.Slider(
                    label="Max New Tokens", minimum=128, maximum=2048, step=128, value=1024
                )
                num_beams = gr.Slider(
                    label="Number of Beams", minimum=1, maximum=10, step=1, value=3
                )

        with gr.Column(scale=3):
            gr.Markdown("## Output", elem_id="output-title")
            parsed_output = gr.JSON(label="Parsed Answer", interactive=False)

    image_submit.click(
        fn=run_florence2_inference,
        inputs=[model_choice, image_upload, task_prompt, max_new_tokens, num_beams],
        outputs=[parsed_output]
    )

if __name__ == "__main__":
    demo.queue().launch(show_error=True, debug=True)