|
import requests |
|
from PIL import Image |
|
from transformers import AutoProcessor, Blip2ForConditionalGeneration |
|
import torch |
|
|
|
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(inp): |
|
inp = transforms.ToTensor()(inp).unsqueeze(0) |
|
with torch.no_grad(): |
|
prediction = torch.nn.functional.softmax(model(inp)[0], dim=0) |
|
confidences = {labels[i]: float(prediction[i]) for i in range(1000)} |
|
return confidences |
|
|
|
demo = gr.Interface(fn=predict, |
|
inputs=gr.inputs.Image(type="pil"), |
|
outputs=gr.outputs.Label(num_top_classes=3) |
|
) |
|
|
|
def predict(imageurl): |
|
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() |