Spaces:
Runtime error
Runtime error
File size: 1,944 Bytes
c7f5de3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
import torch
from PIL import Image
class InstructBlip:
device = "cuda" if torch.cuda.is_available() else "cpu"
def __init__(self, model_pretrain:str = "Salesforce/instructblip-vicuna-7b"):
self.model = InstructBlipForConditionalGeneration.from_pretrained(model_pretrain
, device_map={"": 0}, torch_dtype=torch.float16)
self.processor = InstructBlipProcessor.from_pretrained(model_pretrain)
def image_captioning(self, image: Image.Image) -> str:
prompt = "What are the features of this picture?"
inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device)
outputs = self.model.generate(
**inputs,
do_sample=False,
num_beams=5,
max_length=256,
min_length=1,
top_p=0.9,
repetition_penalty=1.5,
length_penalty=1.0,
temperature=1,
)
generated_text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()
return generated_text
def visual_question_answering(self, image: Image.Image, prompt: str) -> str:
inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(device)
outputs = self.model.generate(
**inputs,
do_sample=False,
num_beams=5,
max_length=256,
min_length=1,
top_p=0.9,
repetition_penalty=1.5,
length_penalty=1.0,
temperature=1,
)
generated_text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()
return generated_text
|