GlitchJesus/RecursionLabImageAI

1import gradio as gr2from PIL import Image3import torch4from transformers import AutoProcessor, AutoModelForCausalLM5 6# Load Florence-2 (runs on CPU, free tier compatible)7model_id = "microsoft/Florence-2-large"8device = "cuda" if torch.cuda.is_available() else "cpu"9dtype = torch.float16 if torch.cuda.is_available() else torch.float3210 11print(f"Loading model on {device}...")12model = AutoModelForCausalLM.from_pretrained(13    model_id,14    torch_dtype=dtype,15    trust_remote_code=True16).to(device)17processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)18print("Model loaded.")19 20def analyze_image(image, prompt):21    if image is None:22        return "No image uploaded."23    24    if not prompt:25        prompt = "<MORE_DETAILED_CAPTION>"26    27    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, dtype)28    29    with torch.no_grad():30        generated_ids = model.generate(31            input_ids=inputs["input_ids"],32            pixel_values=inputs["pixel_values"],33            max_new_tokens=512,34            do_sample=False35        )36    37    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]38    parsed = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))39    40    # Return the first value from parsed dict41    if isinstance(parsed, dict):42        return list(parsed.values())[0]43    return str(parsed)44 45# Available tasks for Florence-246TASKS = [47    "<CAPTION>",48    "<DETAILED_CAPTION>",49    "<MORE_DETAILED_CAPTION>",50    "<OCR>",51    "<OCR_WITH_REGION>",52    "<OBJECT_DETECTION>",53    "<REGION_TO_CATEGORY>",54    "<REGION_TO_DESCRIPTION>",55]56 57with gr.Blocks(title="Vision Analyzer") as demo:58    gr.Markdown("# Image Understanding")59    gr.Markdown("Upload an image and select what you want to extract from it.")60    61    with gr.Row():62        with gr.Column():63            image_input = gr.Image(type="pil", label="Upload Image")64            task_dropdown = gr.Dropdown(choices=TASKS, value="<MORE_DETAILED_CAPTION>", label="Analysis Type")65            text_prompt = gr.Textbox(label="Or enter custom prompt (overrides dropdown)", placeholder="Describe what you see...", lines=2)66            analyze_btn = gr.Button("Analyze")67        with gr.Column():68            output = gr.Textbox(label="Result", lines=15, show_copy_button=True)69    70    analyze_btn.click(fn=analyze_image, inputs=[image_input, text_prompt], outputs=output)71    72    gr.Markdown("---")73    gr.Markdown("Powered by Microsoft Florence-2-large on HuggingFace free tier.")74 75demo.launch(server_name="0.0.0.0", server_port=7860)76