blip2-debunker / test.py
zgjiangtoby's picture
init
d710f4c
from PIL import Image
import requests
from transformers import Blip2Processor, Blip2ForConditionalGeneration
import torch
from peft import LoraConfig, get_peft_model, PeftModel
config = LoraConfig(
r=16,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
)
device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "./blip2_fakenews_all"
# device_map = {"": 0}
device_map = "auto"
processor = Blip2Processor.from_pretrained("blip2")
model = Blip2ForConditionalGeneration.from_pretrained("blip2",
load_in_8bit=True,
device_map=device_map)
model = PeftModel.from_pretrained(model, model_name)
model = get_peft_model(model, config)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
prompt = "Question: Is this real or fake? Answer: real. Question: Why? "
inputs = processor(images=image, text=prompt, return_tensors="pt").to(device, torch.float16)
generated_ids = model.generate(**inputs)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
print(generated_text)