|
import requests |
|
from PIL import Image |
|
from transformers import AutoProcessor, Blip2ForConditionalGeneration |
|
import torch |
|
import gradio as gr |
|
|
|
processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b") |
|
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16) |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
model.to(device) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def predict(imageurl): |
|
image = Image.open(requests.get(imageurl, stream=True).raw).convert('RGB') |
|
inputs = processor(image, return_tensors="pt") |
|
generated_ids = model.generate(**inputs, max_new_tokens=20) |
|
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() |
|
return('caption: '+generated_text) |
|
|
|
demo = gr.Interface(fn=predict, |
|
inputs="text", |
|
outputs=gr.outputs.Label(num_top_classes=3) |
|
) |
|
|
|
demo.launch() |