cool-api / app.py
preston-cell's picture
Update app.py
40c03fb verified
raw
history blame
569 Bytes
import gradio as gr
from transformers import AutoProcessor, AutoModelForImageTextToText
import torch
model_path = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
processor = AutoProcessor.from_pretrained(model_path)
model = AutoModelForImageTextToText.from_pretrained(
model_path,
torch_dtype=torch.bfloat16,
_attn_implementation="flash_attention_2"
).to("cuda")
def launch(input):
out = model.generate(**input)
return(out)
iface = gr.Interface(launch,
inputs=gr.Image(type='pil'),
outputs="text")
iface.launch()