| from functools import lru_cache |
| from typing import Iterable, List, Tuple |
|
|
| import torch |
| from PIL import Image |
|
|
|
|
| BATCH_SIZE = 50 |
| _CACHE_MODEL = None |
| _CACHE_PROCESSOR = None |
|
|
|
|
| def configure_cache(model, processor) -> None: |
| """Bind the shared model and processor for cached inference.""" |
| global _CACHE_MODEL, _CACHE_PROCESSOR |
| _CACHE_MODEL = model |
| _CACHE_PROCESSOR = processor |
|
|
|
|
| def preprocess_image(img_path: str) -> Image.Image: |
| img = Image.open(img_path) |
| img = img.convert("RGB") |
| img.thumbnail((448, 448)) |
| return img |
|
|
|
|
| def _ensure_tuple(labels: Iterable[str]) -> Tuple[str, ...]: |
| if isinstance(labels, tuple): |
| return labels |
| return tuple(labels) |
|
|
|
|
| @lru_cache(maxsize=5) |
| def _cached_logits(image_path: str, label_tuple: Tuple[str, ...]) -> List[float]: |
| if _CACHE_MODEL is None or _CACHE_PROCESSOR is None: |
| raise RuntimeError("Cache manager not configured with model and processor.") |
|
|
| device = next(_CACHE_MODEL.parameters()).device |
| dtype = next(_CACHE_MODEL.parameters()).dtype |
| image = preprocess_image(image_path) |
| logits: List[float] = [] |
|
|
| with torch.no_grad(): |
| for start in range(0, len(label_tuple), BATCH_SIZE): |
| batch = list(label_tuple[start : start + BATCH_SIZE]) |
| inputs = _CACHE_PROCESSOR(images=image, text=batch, return_tensors="pt", padding=True) |
|
|
| prepared = {} |
| for key, value in inputs.items(): |
| if torch.is_tensor(value): |
| moved = value.to(device) |
| if torch.is_floating_point(moved): |
| moved = moved.to(dtype=dtype) |
| prepared[key] = moved |
| else: |
| prepared[key] = value |
|
|
| outputs = _CACHE_MODEL(**prepared) |
| logits.extend(outputs.logits_per_image[0].detach().cpu().tolist()) |
|
|
| image.close() |
| scores = torch.softmax(torch.tensor(logits), dim=0).tolist() |
| return scores |
|
|
|
|
| def cached_inference(image_path: str, labels: Iterable[str]) -> List[float]: |
| label_tuple = _ensure_tuple(labels) |
| return _cached_logits(image_path, label_tuple) |
|
|
