File size: 3,081 Bytes
334dcac e4bcc80 334dcac 55f430c 334dcac 55f430c 334dcac |
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 |
# File name: model.py
import json
import os
import numpy as np
import torch
from starlette.requests import Request
from PIL import Image
import ray
from ray import serve
from clip_retrieval.load_clip import load_clip, get_tokenizer
# from clip_retrieval.clip_client import ClipClient, Modality
@serve.deployment(num_replicas=6, ray_actor_options={"num_cpus": .2, "num_gpus": 0.1})
class CLIPTransform:
def __init__(self):
# os.environ["OMP_NUM_THREADS"] = "20"
# torch.set_num_threads(20)
# Load model
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
self._clip_model="ViT-L/14"
self._clip_model_id ="laion5B-L-14"
self.model, self.preprocess = load_clip(self._clip_model, use_jit=True, device=self.device)
self.tokenizer = get_tokenizer(self._clip_model)
print ("using device", self.device)
def text_to_embeddings(self, prompt):
text = self.tokenizer([prompt]).to(self.device)
with torch.no_grad():
prompt_embededdings = self.model.encode_text(text)
prompt_embededdings /= prompt_embededdings.norm(dim=-1, keepdim=True)
return(prompt_embededdings)
def image_to_embeddings(self, input_im):
input_im = Image.fromarray(input_im)
prepro = self.preprocess(input_im).unsqueeze(0).to(self.device)
with torch.no_grad():
image_embeddings = self.model.encode_image(prepro)
image_embeddings /= image_embeddings.norm(dim=-1, keepdim=True)
return(image_embeddings)
def preprocessed_image_to_emdeddings(self, prepro):
with torch.no_grad():
image_embeddings = self.model.encode_image(prepro)
image_embeddings /= image_embeddings.norm(dim=-1, keepdim=True)
return(image_embeddings)
async def __call__(self, http_request: Request) -> str:
request = await http_request.json()
# print(type(request))
# print(str(request))
# switch based if we are using text or image
embeddings = None
if "text" in request:
prompt = request["text"]
embeddings = self.text_to_embeddings(prompt)
elif "image_url" in request:
image_url = request["image_url"]
# download image from url
import requests
from io import BytesIO
image_bytes = requests.get(image_url).content
input_image = Image.open(BytesIO(image_bytes))
input_image = input_image.convert('RGB')
input_image = np.array(input_image)
embeddings = self.image_to_embeddings(input_image)
elif "preprocessed_image" in request:
prepro = request["preprocessed_image"]
# create torch tensor on the device
prepro = torch.tensor(prepro).to(self.device)
embeddings = self.preprocessed_image_to_emdeddings(prepro)
else:
raise Exception("Invalid request")
return embeddings.cpu().numpy().tolist()
deployment_graph = CLIPTransform.bind()
|