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)}