Spaces:
Runtime error
Runtime error
File size: 1,370 Bytes
2a146da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | from __future__ import annotations
from typing import Iterable
import torch
from PIL import Image
from pathoverse.models.registry import registry
class ModelEngine:
"""
High-level inference interface for PathoVerse foundation models.
"""
def __init__(
self,
model_name: str,
device: str = "cuda",
batch_size: int = 8,
) -> None:
self.model_name = model_name
self.device = device
self.batch_size = batch_size
self.adapter = registry.create(
model_name,
device=device,
batch_size=batch_size,
)
def load(self) -> None:
self.adapter.load()
def encode_images(
self,
images: Iterable[Image.Image],
) -> torch.Tensor:
images = list(images)
if not images:
raise ValueError("No images supplied.")
outputs = []
for start in range(
0,
len(images),
self.batch_size,
):
batch = images[start : start + self.batch_size]
inputs = self.adapter.preprocess(batch)
embeddings = self.adapter.encode(inputs)
outputs.append(
embeddings.detach().cpu()
)
return torch.cat(outputs, dim=0)
def info(self):
return self.adapter.info() |