Create app.py

app.py

@@ -0,0 +1,94 @@
+import gradio as gr
+import torch
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForCausalLM
+
+# Use CPU as requested
+device = "cpu"
+
+def load_vlm(model_name):
+    """Helper to load model and processor."""
+    try:
+        print(f"Loading {model_name}...")
+        model = AutoModelForCausalLM.from_pretrained(
+            f'microsoft/{model_name}', 
+            trust_remote_code=True
+        ).to(device).eval()
+        processor = AutoProcessor.from_pretrained(
+            f'microsoft/{model_name}', 
+            trust_remote_code=True
+        )
+        return model, processor
+    except Exception as e:
+        print(f"Error loading {model_name}: {e}")
+        return None, None
+
+# Load both models
+model_base, proc_base = load_vlm('Florence-2-base')
+model_large, proc_large = load_vlm('Florence-2-large')
+
+def describe_image(uploaded_image, model_choice):
+    if uploaded_image is None:
+        return "Please upload an image."
+
+    # Select model based on UI choice
+    if model_choice == "Florence-2-base":
+        model, processor = model_base, proc_base
+    else:
+        model, processor = model_large, proc_large
+
+    if model is None:
+        return f"{model_choice} failed to load."
+
+    if not isinstance(uploaded_image, Image.Image):
+        uploaded_image = Image.fromarray(uploaded_image)
+
+    # Core generation logic
+    inputs = processor(text="<MORE_DETAILED_CAPTION>", images=uploaded_image, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=1024,
+            num_beams=3,
+            do_sample=False,
+        )
+
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    result = processor.post_process_generation(
+        generated_text,
+        task="<MORE_DETAILED_CAPTION>",
+        image_size=(uploaded_image.width, uploaded_image.height)
+    )
+    
+    return result["<MORE_DETAILED_CAPTION>"]
+
+# Simplified Gradio Layout
+css = ".submit-btn { background-color: #4682B4 !important; color: white !important; }"
+
+with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
+    gr.Markdown("# **Florence-2 Models Image Captions**")
+    gr.Markdown("> Select the model to use. **Base** is faster; **Large** is more accurate.")
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Upload Image", type="pil")
+            model_choice = gr.Radio(
+                choices=["Florence-2-base", "Florence-2-large"], 
+                label="Model Choice", 
+                value="Florence-2-base"
+            )
+            generate_btn = gr.Button("Generate Caption", elem_classes="submit-btn")
+            
+        with gr.Column():
+            output = gr.Textbox(label="Generated Caption", lines=6, interactive=True)
+
+    generate_btn.click(
+        fn=describe_image, 
+        inputs=[image_input, model_choice], 
+        outputs=output
+    )
+
+if __name__ == "__main__":
+    demo.launch()