mirror of
https://github.com/MrSleeps/jellyfin-scene-search.git
synced 2026-11-15 22:41:28 +00:00
Sleeps
98 lines
3.8 KiB
Python
98 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from transformers import AutoModel, AutoProcessor
|
|
|
|
|
|
class SiglipEmbedder:
|
|
def __init__(self, model_id: str, device: str = "cuda") -> None:
|
|
if device == "cuda" and not torch.cuda.is_available():
|
|
device = "cpu"
|
|
self.device = torch.device(device)
|
|
self.model_id = model_id
|
|
self.processor = AutoProcessor.from_pretrained(model_id)
|
|
self.model = AutoModel.from_pretrained(model_id).to(self.device).eval()
|
|
self._dim: int | None = None
|
|
|
|
@staticmethod
|
|
def _extract_features(output: object) -> torch.Tensor:
|
|
"""Return the embedding tensor across Transformers 4.x/5.x APIs."""
|
|
if isinstance(output, torch.Tensor):
|
|
return output
|
|
|
|
# Newer Transformers versions can return BaseModelOutputWithPooling.
|
|
for attr in ("image_embeds", "text_embeds", "pooler_output"):
|
|
value = getattr(output, attr, None)
|
|
if isinstance(value, torch.Tensor):
|
|
return value
|
|
|
|
if isinstance(output, (tuple, list)):
|
|
for value in output:
|
|
if isinstance(value, torch.Tensor) and value.ndim == 2:
|
|
return value
|
|
|
|
raise TypeError(
|
|
f"Unsupported feature output from Transformers: {type(output).__name__}. "
|
|
"Expected a Tensor or an object containing pooler_output/image_embeds/text_embeds."
|
|
)
|
|
|
|
@staticmethod
|
|
def _normalise(features: torch.Tensor) -> torch.Tensor:
|
|
return F.normalize(features.float(), p=2, dim=-1)
|
|
|
|
def embed_images(self, images: Sequence[Image.Image]) -> np.ndarray:
|
|
inputs = self.processor(images=list(images), return_tensors="pt")
|
|
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
|
with torch.inference_mode():
|
|
output = self.model.get_image_features(**inputs)
|
|
features = self._normalise(self._extract_features(output))
|
|
arr = features.cpu().numpy().astype(np.float32)
|
|
self._dim = int(arr.shape[-1])
|
|
return arr
|
|
|
|
def embed_text(self, text: str) -> np.ndarray:
|
|
inputs = self.processor(text=[text], padding="max_length", return_tensors="pt")
|
|
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
|
with torch.inference_mode():
|
|
output = self.model.get_text_features(**inputs)
|
|
features = self._normalise(self._extract_features(output))
|
|
arr = features[0].cpu().numpy().astype(np.float32)
|
|
self._dim = int(arr.shape[-1])
|
|
return arr
|
|
|
|
def match_probabilities(
|
|
self,
|
|
text_vector: np.ndarray,
|
|
image_vectors: np.ndarray,
|
|
) -> np.ndarray:
|
|
"""Apply SigLIP2's learned logit scale/bias to stored cosine embeddings."""
|
|
text = torch.as_tensor(text_vector, dtype=torch.float32, device=self.device)
|
|
images = torch.as_tensor(image_vectors, dtype=torch.float32, device=self.device)
|
|
|
|
text = F.normalize(text, p=2, dim=-1)
|
|
images = F.normalize(images, p=2, dim=-1)
|
|
cosine = images @ text
|
|
|
|
logit_scale = getattr(self.model, "logit_scale", None)
|
|
logit_bias = getattr(self.model, "logit_bias", None)
|
|
if logit_scale is None or logit_bias is None:
|
|
raise RuntimeError(
|
|
"Loaded model does not expose SigLIP logit_scale/logit_bias; "
|
|
"cannot calculate model-native match scores."
|
|
)
|
|
|
|
with torch.inference_mode():
|
|
logits = cosine * logit_scale.float().exp() + logit_bias.float()
|
|
probabilities = torch.sigmoid(logits)
|
|
|
|
return probabilities.detach().cpu().numpy().astype(np.float32)
|
|
|
|
@property
|
|
def dimension(self) -> int | None:
|
|
return self._dim
|