update

Dockerfile

@@ -2,4 +2,4 @@ FROM python:3.9
 
 COPY . .
 
-CMD ["python", "main.py"]
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -1 +1,81 @@
-print("hello World")
+import base64
+import logging
+from io import BytesIO
+from typing import List
+import clip
+import numpy as np
+import torch
+from PIL import Image
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+logging.basicConfig(
+    format='%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s',
+    level=logging.DEBUG,
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+model, preprocess = clip.load("models/ViT-B-32.pt")
+model.cpu().eval()
+
+app = FastAPI()
+
+
+class Texts(BaseModel):
+    texts_source: List[str]
+    texts_target: List[str]
+
+
+class Images(BaseModel):
+    images: List[str]
+    texts: List[str]
+
+
+class CLIPResponse(BaseModel):
+    similarity: List[List[float]]
+
+
+@app.post("/clip_image_to_text", response_model=CLIPResponse, tags=["CLIP"])
+def clip_image_to_text(data: Images) -> CLIPResponse:
+    preprocessed_images = []
+    for image in data.images:
+        if 'base64,' not in image:
+            raise HTTPException(422, "Image must be in base64")
+        image = BytesIO(base64.b64decode(image.split('base64,')[-1]))
+        image = Image.open(image)
+        preprocessed_images.append(preprocess(image))
+
+    image_input = torch.tensor(np.stack(preprocessed_images)).cpu()
+    text_tokens = clip.tokenize(["This is the " + desc for desc in data.texts]).cpu()
+
+    with torch.no_grad():
+        image_features = model.encode_image(image_input).float()
+        text_features = model.encode_text(text_tokens).float()
+
+    image_features /= image_features.norm(dim=-1, keepdim=True)
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+    similarity = text_features.cpu().numpy() @ image_features.cpu().numpy().T
+    logging.debug(f"Similarity: {similarity}")
+    return CLIPResponse(similarity=similarity.tolist())
+
+
+@app.post("/clip_text_to_text", response_model=CLIPResponse, tags=["CLIP"])
+def clip_text_to_text(data: Texts) -> CLIPResponse:
+    text_input = clip.tokenize([f"This is {text}" for text in data.texts_source]).cpu()
+    text_output = clip.tokenize(data.texts_target).cpu()
+
+    with torch.no_grad():
+        input_features = model.encode_text(text_input).float()
+        output_features = model.encode_text(text_output).float()
+
+    input_features /= input_features.norm(dim=-1, keepdim=True)
+    output_features /= output_features.norm(dim=-1, keepdim=True)
+
+    similarity = output_features.cpu().numpy() @ input_features.cpu().numpy().T
+    logging.debug(f"Similarity: {similarity}")
+    return CLIPResponse(similarity=similarity.tolist())
+
+
+@app.get("/ping", tags=["TEST"])
+def ping():
+    return "pong"
+

requirements.txt

@@ -0,0 +1,4 @@
+fastapi==0.103.0
+uvicorn==0.23.2
+pydantic==2.3.0
+clip @ git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33