File size: 1,312 Bytes
c4d212f d6ef205 0845f18 d6ef205 0845f18 d6ef205 |
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 |
import torch
from PIL import Image
from transformers import Blip2Processor, Blip2ForConditionalGeneration, BitsAndBytesConfig
# Load model and processor
device = "cuda" if torch.cuda.is_available() else "cpu"
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
model = Blip2ForConditionalGeneration.from_pretrained(
"Salesforce/blip2-flan-t5-xl", device_map="auto"
)
def get_image_answer(image: Image.Image, question: str) -> str:
if image.mode != "RGB":
image = image.convert("RGB")
inputs = processor(images=image, text=question, return_tensors="pt")
for key in inputs:
if inputs[key].dtype in [torch.float32, torch.float64]:
# Cast only float tensors (like pixel values) to float16 if on CUDA
inputs[key] = inputs[key].to(device, torch.float16 if device == "cuda" else torch.float32)
else:
# Leave token inputs (e.g., input_ids) as integers
inputs[key] = inputs[key].to(device)
print("Prompt Passed to VLM:", f"Question: {question} Answer:")
output_ids = model.generate(**inputs)
answer = processor.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
print("Model Response:", answer)
return answer
|