File size: 5,937 Bytes
f338ee9 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """
Adapted from https://github.com/yuvalkirstain/PickScore. Originally MIT License, Copyright (c) 2021.
"""
from transformers import AutoProcessor, AutoModel
from PIL import Image
import torch
import numpy as np
from transformers import AutoProcessor, AutoModel
from datasets import load_from_disk, load_dataset
import torch
from PIL import Image
from io import BytesIO
from tqdm.auto import tqdm
import sys
def open_image(image):
if isinstance(image, bytes):
image = Image.open(BytesIO(image))
elif isinstance(image, str):
image = Image.open(image)
image = image.convert("RGB")
return image
class PickScorer(torch.nn.Module):
def __init__(self, processor_name_or_path, model_pretrained_name_or_path, device='cuda'):
super().__init__()
self.processor = AutoProcessor.from_pretrained(processor_name_or_path)
self.model = AutoModel.from_pretrained(model_pretrained_name_or_path).to(device)
self.device = device
self.eval()
@torch.no_grad()
def __call__(self, prompt, images):
# preprocess
image_inputs = self.processor(
images=images,
padding=True,
truncation=True,
max_length=77,
return_tensors="pt",
).to(self.device)
text_inputs = self.processor(
text=prompt,
padding=True,
truncation=True,
max_length=77,
return_tensors="pt",
).to(self.device)
with torch.no_grad():
# embed
image_embs = self.model.get_image_features(**image_inputs)
image_embs = image_embs / torch.norm(image_embs, dim=-1, keepdim=True)
text_embs = self.model.get_text_features(**text_inputs)
text_embs = text_embs / torch.norm(text_embs, dim=-1, keepdim=True)
# score
scores = self.model.logit_scale.exp() * (text_embs @ image_embs.T)[0]
# get probabilities if you have multiple images to choose from
if len(scores) == 1:
probs = scores
else:
probs = torch.softmax(scores, dim=-1)
return probs.cpu().tolist()
def score(self, img_path, prompt):
if isinstance(img_path, list):
result = []
for one_img_path in img_path:
# Load your image and prompt
with torch.no_grad():
# Process the image
if isinstance(one_img_path, str):
image = self.preprocess_val(Image.open(one_img_path)).unsqueeze(0).to(device=self.device, non_blocking=True)
elif isinstance(one_img_path, Image.Image):
image = self.preprocess_val(one_img_path).unsqueeze(0).to(device=self.device, non_blocking=True)
else:
raise TypeError('The type of parameter img_path is illegal.')
# Process the prompt
text = self.tokenizer([prompt]).to(device=self.device, non_blocking=True)
with torch.cuda.amp.autocast():
outputs = self.model(image, text)
image_features, text_features = outputs["image_features"], outputs["text_features"]
logits_per_image = image_features @ text_features.T
pick_score = self.model.logit_scale.exp() * torch.diagonal(logits_per_image).cpu().numpy()
result.append(pick_score[0])
return result
elif isinstance(img_path, str):
# Load your image and prompt
with torch.no_grad():
# Process the image
image = self.preprocess_val(Image.open(img_path)).unsqueeze(0).to(device=self.device, non_blocking=True)
# Process the prompt
text = self.tokenizer([prompt]).to(device=self.device, non_blocking=True)
with torch.cuda.amp.autocast():
outputs = self.model(image, text)
image_features, text_features = outputs["image_features"], outputs["text_features"]
logits_per_image = image_features @ text_features.T
pick_score = self.model.logit_scale.exp() * torch.diagonal(logits_per_image).cpu().numpy()
return [pick_score[0]]
elif isinstance(img_path, Image.Image):
# Load your image and prompt
with torch.no_grad():
# Process the image
image = self.preprocess_val(img_path).unsqueeze(0).to(device=self.device, non_blocking=True)
# Process the prompt
text = self.tokenizer([prompt]).to(device=self.device, non_blocking=True)
with torch.cuda.amp.autocast():
outputs = self.model(image, text)
image_features, text_features = outputs["image_features"], outputs["text_features"]
logits_per_image = image_features @ text_features.T
pick_score = self.model.logit_scale.exp() * torch.diagonal(logits_per_image).cpu().numpy()
return [pick_score[0]]
else:
raise TypeError('The type of parameter img_path is illegal.')
if __name__ == "__main__":
pickscorer = PickScorer(processor_name_or_path="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", model_pretrained_name_or_path="yuvalkirstain/PickScore_v1")
image0 = open_image('./image0.png')
image1 = open_image('./image1.png')
prompt = "photorealistic image of a lone painter standing in a gallery, watching an exhibition of paintings made entirely with AI. In the foreground of the image a robot looks proudly at his art"
probs = pickscorer(prompt, [image0])
probs1 = pickscorer(prompt, [image1])
print(probs)
print(probs1)
|