mirror of
https://github.com/immich-app/immich.git
synced 2025-01-06 11:56:46 +01:00
28 lines
707 B
Python
28 lines
707 B
Python
|
from pathlib import Path
|
||
|
|
|
||
|
|
from PIL.Image import Image
|
||
|
|
from sentence_transformers import SentenceTransformer
|
||
|
|
|
||
|
|
from ..schemas import ModelType
|
||
|
|
from .base import InferenceModel
|
||
|
|
|
||
|
|
|
||
|
|
class CLIPSTEncoder(InferenceModel):
|
||
|
|
_model_type = ModelType.CLIP
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
model_name: str,
|
||
|
|
cache_dir: Path | None = None,
|
||
|
|
**model_kwargs,
|
||
|
|
):
|
||
|
|
super().__init__(model_name, cache_dir)
|
||
|
|
self.model = SentenceTransformer(
|
||
|
|
self.model_name,
|
||
|
|
cache_folder=self.cache_dir.as_posix(),
|
||
|
|
**model_kwargs,
|
||
|
|
)
|
||
|
|
|
||
|
|
def predict(self, image_or_text: Image | str) -> list[float]:
|
||
|
|
return self.model.encode(image_or_text).tolist()
|