app.py

import torch
import gradio as gr
from transformers import pipeline

import os
os.environ["CUDA_HOME"] = "/usr/local/cuda"


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

pipe = pipeline(
    task="image-to-text",
    model="mistralai/Mistral-Small-3.2-24B-Instruct-2506",
    framework="pt",
    torch_dtype=torch.float16 if device=="cuda" else torch.float32,
    device=0 if device=="cuda" else -1,
)

def run_vision(img):
    if img is None:
        return "Upload an image!"
    result = pipe(img)
    return result[0]["generated_text"]

with gr.Blocks() as demo:
    gr.Markdown("# Image-to-Text with Mistral‑Small‑3.2")
    image_input = gr.Image(type="pil")
    text_output = gr.Textbox()
    run_btn = gr.Button("Generate")
    run_btn.click(run_vision, inputs=image_input, outputs=text_output)

if __name__ == "__main__":
    demo.launch()