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 = [ "", "", "", "", "", "", "", "" ] # 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="" ) 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)