Spaces:
Sleeping
Sleeping
File size: 3,222 Bytes
84ee555 | 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 | from functools import lru_cache
from typing import Any
import open_clip
from PIL import Image
import numpy as np
class VisionLanguageModel:
"""
A class that wraps the OpenCLIP model for vision-language tasks.
Args:
model_name (str, optional): The name of the OpenCLIP model. Defaults to "ViT-L-14".
pretrained (str, optional): The name of the pretrained model. Defaults to "laion2B-s32B-b82K".
device (str, optional): The device to use for computations. Defaults to "cuda".
Attributes:
device (str): The device used for computations.
model: The OpenCLIP model instance.
processor: The image processor instance.
tokenizer: The tokenizer instance.
"""
def __init__(
self,
model_name: str = "ViT-B-16",
pretrained: str = "laion2B-s34B-b88K",
device: str = "cuda",
):
"""
Initializes the VisionLanguageModel instance.
Args:
model_name (str, optional): The name of the OpenCLIP model. Defaults to "ViT-B-16".
pretrained (str, optional): The name of the pretrained model. Defaults to "laion2B-s34B-b88K".
device (str, optional): The device to use for computations. Defaults to "cuda".
"""
self.device: str = device
self._load_model(model_name, pretrained)
@lru_cache(maxsize=1)
def _load_model(self, model_name: str, pretrained: str):
"""
Loads the OpenCLIP model and tokenizer.
Args:
model_name (str): The name of the OpenCLIP model.
pretrained (str): The name of the pretrained model.
"""
self.model, _, self.preprocess = open_clip.create_model_and_transforms(
model_name=model_name, pretrained=pretrained, device=self.device
)
self.model.eval()
self.tokenizer = open_clip.get_tokenizer(model_name=model_name)
def get_embedding(self, input: str | Image.Image | Any) -> np.ndarray:
"""
Computes the embedding for the given input.
Args:
input (str | Image.Image | Any): The input to compute the embedding for.
Can be a string, a PIL Image, or any other type (which will raise a ValueError).
Returns:
np.ndarray: The computed embedding as a NumPy array.
Raises:
ValueError: If the input type is not supported.
Examples:
>>> model = VisionLanguageModel()
>>> text_embedding = model.get_embedding("This is a sample text")
>>> image = Image.open("image.jpg").convert("RGB")
>>> image_embedding = model.get_embedding(image)
"""
if isinstance(input, str):
tokens = self.tokenizer(input).to(self.device)
vector = self.model.encode_text(tokens)
elif isinstance(input, Image.Image):
image_input = self.preprocess(input).unsqueeze(0).to(self.device)
vector = self.model.encode_image(image_input)
else:
raise ValueError(f"Invalid input type: {type(input)}")
vector /= vector.norm(dim=-1, keepdim=True)
return vector.cpu().detach().numpy().astype("float32")
|