| from typing import List | |
| import torch | |
| from PIL import Image | |
| from transformers import CLIPModel, CLIPProcessor | |
| MODEL_DIM = 512 | |
| class ClipWrapper: | |
| def __init__(self): | |
| self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") | |
| self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") | |
| def images2vec(self, images: List[Image.Image]) -> torch.Tensor: | |
| inputs = self.processor(images=images, return_tensors="pt") | |
| with torch.no_grad(): | |
| model_inputs = {k: v.to(self.model.device) for k, v in inputs.items()} | |
| image_embeds = self.model.vision_model(**model_inputs) | |
| clip_vectors = self.model.visual_projection(image_embeds[1]) | |
| return clip_vectors / clip_vectors.norm(dim=-1, keepdim=True) | |
| def texts2vec(self, texts: List[str]) -> torch.Tensor: | |
| inputs = self.processor(text=texts, return_tensors="pt", padding=True) | |
| with torch.no_grad(): | |
| model_inputs = {k: v.to(self.model.device) for k, v in inputs.items()} | |
| text_embeds = self.model.text_model(**model_inputs) | |
| text_vectors = self.model.text_projection(text_embeds[1]) | |
| return text_vectors / text_vectors.norm(dim=-1, keepdim=True) | |