srivatsavdamaraju's picture
Update app.py
4ab173c verified
raw
history blame
1.23 kB
import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer
import gradio as gr
# Check if GPU is available and set the device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Load the model and tokenizer
model = AutoModel.from_pretrained(
'openbmb/MiniCPM-V',
trust_remote_code=True,
torch_dtype=torch.bfloat16 if device == 'cuda' else torch.float32
)
model = model.eval().to(device)
tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V', trust_remote_code=True)
# Function to process the image and question
def predict(image, question):
image = image.convert('RGB')
msgs = [{'role': 'user', 'content': [image, question]}]
res = model.chat(image=None, msgs=msgs, tokenizer=tokenizer)
generated_text = ""
for new_text in res:
generated_text += new_text
return generated_text
# Set up the Gradio interface
iface = gr.Interface(
fn=predict,
inputs=[
gr.Image(type="pil", label="Upload an Image"),
gr.Textbox(label="Ask a Question")
],
outputs="text",
title="Image Question Answering",
description="Upload an image and ask a question about its content."
)
# Launch the Gradio app
iface.launch()