Spaces:
Build error
Build error
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() | |