Spylif's picture
Update app.py
4ec9249 verified
Raw
History Blame Contribute Delete
918 Bytes
from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import torch
import requests
from io import BytesIO
app = FastAPI()
# use TensorFlow format -- Akash 22 april 25
model = CLIPModel.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
processor = CLIPProcessor.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
class ImageInput(BaseModel):
image_url: str
@app.post("/vector")
def vectorize_image(data: ImageInput):
try:
response = requests.get(data.image_url)
image = Image.open(BytesIO(response.content)).convert("RGB")
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
features = model.get_image_features(**inputs)
return {"vector": features[0].tolist()}
except Exception as e:
return {"error": str(e)}