Views
No views yet
1# import
2from transformers import AutoProcessor, AutoModel
3
4# load model
5device = "cuda"
6processor_name_or_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
7model_pretrained_name_or_path = "yuvalkirstain/PickScore_v1"
8
9processor = AutoProcessor.from_pretrained(processor_name_or_path)
10model = AutoModel.from_pretrained(model_pretrained_name_or_path).eval().to(device)
11
12def calc_probs(prompt, images):
13
14 # preprocess
15 image_inputs = processor(
16 images=images,
17 padding=True,
18 truncation=True,
19 max_length=77,
20 return_tensors="pt",
21 ).to(device)
22
23 text_inputs = processor(
24 text=prompt,
25 padding=True,
26 truncation=True,
27 max_length=77,
28 return_tensors="pt",
29 ).to(device)
30
31
32 with torch.no_grad():
33 # embed
34 image_embs = model.get_image_features(**image_inputs)
35 image_embs = image_embs / torch.norm(image_embs, dim=-1, keepdim=True)
36
37 text_embs = model.get_text_features(**text_inputs)
38 text_embs = text_embs / torch.norm(text_embs, dim=-1, keepdim=True)
39
40 # score
41 scores = model.logit_scale.exp() * (text_embs @ image_embs.T)[0]
42
43 # get probabilities if you have multiple images to choose from
44 probs = torch.softmax(scores, dim=-1)
45
46 return probs.cpu().tolist()
47
48pil_images = [Image.open("my_amazing_images/1.jpg"), Image.open("my_amazing_images/2.jpg")]
49prompt = "fantastic, increadible prompt"
50print(calc_probs(prompt, pil_images))1@inproceedings{Kirstain2023PickaPicAO,
2 title={Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation},
3 author={Yuval Kirstain and Adam Polyak and Uriel Singer and Shahbuland Matiana and Joe Penna and Omer Levy},
4 year={2023}
5}