Florence2-LaTex / app.py
Chan-Y's picture
Update app.py
1128ae6 verified
raw
history blame
995 Bytes
import gradio as gr
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_name = "Chan-Y/Florence-2-LaTex-v2"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).to(device)
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
def process_image(image):
if image:
inputs = processor(images=image, return_tensors="pt").to(device)
outputs = model.generate(**inputs)
text = processor.decode(outputs[0], skip_special_tokens=True)
return text
return "No image provided."
with gr.Blocks() as demo:
with gr.Row():
image_input = gr.Image(label="Upload an Image", type="pil")
process_button = gr.Button("Process Image")
output = gr.Textbox(label="Output Message")
process_button.click(process_image, inputs=image_input, outputs=output)
demo.launch()