rag_server/applications/clip_embedding/embedding.py

import io
import urllib.request
from typing import List

import torch
from PIL import Image

from .clip_model import get_model

# init model
model, processor, DEVICE, DTYPE, MAX_BATCH = get_model()


def _normalize(x: torch.Tensor) -> torch.Tensor:
    return x / (x.norm(dim=-1, keepdim=True) + 1e-12)

def _to_list(x: torch.Tensor):
    return x.detach().cpu().tolist()

async def embed_image_url(img_url_list: List[str]):
    images = []
    for u in img_url_list:
        with urllib.request.urlopen(u, timeout=15) as r:
            img = Image.open(io.BytesIO(r.read())).convert("RGB")
            images.append(img)

    outputs = []
    for chunk_start in range(0, len(images), MAX_BATCH):
        chunk = images[chunk_start:chunk_start + MAX_BATCH]
        inputs = processor(images=chunk, return_tensors="pt")
        inputs = {k: v.to(DEVICE, dtype=DTYPE) if hasattr(v, "to") else v for k, v in inputs.items()}
        feats = model.get_image_features(**inputs)
        feats = _normalize(feats)
        outputs.extend(_to_list(feats))

    return outputs