mirror of
https://github.com/MrSleeps/jellyfin-scene-search.git
synced 2026-11-15 22:41:28 +00:00
Sleeps
112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
|
|
|
|
DEFAULT_CAPTION_PROMPT_PATH = Path(__file__).parent / "prompts" / "caption.txt"
|
|
DEFAULT_SPECIALIST_PROMPT_PATH = Path(__file__).parent / "prompts" / "caption_specialist.txt"
|
|
|
|
|
|
def load_default_caption_prompt() -> str:
|
|
prompt = DEFAULT_CAPTION_PROMPT_PATH.read_text(encoding="utf-8").strip()
|
|
if not prompt:
|
|
raise RuntimeError(f"Default caption prompt is empty: {DEFAULT_CAPTION_PROMPT_PATH}")
|
|
return prompt
|
|
|
|
|
|
def ensure_prompt_file(prompt_path: Path, default_path: Path) -> str:
|
|
prompt_path = prompt_path.expanduser().resolve()
|
|
if not prompt_path.exists():
|
|
prompt_path.parent.mkdir(parents=True, exist_ok=True)
|
|
prompt_path.write_text(
|
|
default_path.read_text(encoding="utf-8").strip() + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
prompt = prompt_path.read_text(encoding="utf-8").strip()
|
|
if not prompt:
|
|
raise RuntimeError(f"Caption prompt is empty: {prompt_path}")
|
|
return prompt
|
|
|
|
|
|
def ensure_caption_prompt_file(prompt_path: Path) -> str:
|
|
return ensure_prompt_file(prompt_path, DEFAULT_CAPTION_PROMPT_PATH)
|
|
|
|
|
|
def ensure_specialist_prompt_file(prompt_path: Path) -> str:
|
|
return ensure_prompt_file(prompt_path, DEFAULT_SPECIALIST_PROMPT_PATH)
|
|
|
|
|
|
class OllamaCaptioner:
|
|
def __init__(
|
|
self,
|
|
url: str,
|
|
model: str,
|
|
timeout_seconds: int = 180,
|
|
prompt: str | None = None,
|
|
max_tokens: int = 120,
|
|
) -> None:
|
|
self.url = url.rstrip("/")
|
|
self.model = model
|
|
self.timeout_seconds = timeout_seconds
|
|
self.prompt = prompt or load_default_caption_prompt()
|
|
self.max_tokens = max_tokens
|
|
self.session = requests.Session()
|
|
|
|
def caption(self, image_path: Path, context_hint: str | None = None) -> str:
|
|
image_b64 = base64.b64encode(image_path.read_bytes()).decode("ascii")
|
|
prompt = self.prompt
|
|
if context_hint:
|
|
prompt = (
|
|
f"{prompt}\n\nMetadata context for this video: {context_hint}\n"
|
|
"Use this only as a vocabulary/context hint. Do not claim a tagged "
|
|
"detail is visible unless it is actually visible in the frame."
|
|
)
|
|
response = self.session.post(
|
|
f"{self.url}/api/chat",
|
|
json={
|
|
"model": self.model,
|
|
"stream": False,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": prompt,
|
|
"images": [image_b64],
|
|
}
|
|
],
|
|
"options": {
|
|
"temperature": 0.0,
|
|
"num_predict": self.max_tokens,
|
|
},
|
|
},
|
|
timeout=self.timeout_seconds,
|
|
)
|
|
if not response.ok:
|
|
raise RuntimeError(
|
|
f"Ollama {response.status_code}: {response.text[:500]}"
|
|
)
|
|
|
|
data = response.json()
|
|
content = (data.get("message") or {}).get("content")
|
|
if not isinstance(content, str) or not content.strip():
|
|
raise RuntimeError(f"Ollama returned no caption: {data}")
|
|
|
|
return " ".join(content.strip().split())
|
|
|
|
def unload(self) -> None:
|
|
response = self.session.post(
|
|
f"{self.url}/api/generate",
|
|
json={
|
|
"model": self.model,
|
|
"keep_alive": 0,
|
|
},
|
|
timeout=self.timeout_seconds,
|
|
)
|
|
if not response.ok:
|
|
raise RuntimeError(
|
|
f"Ollama unload {response.status_code}: {response.text[:500]}"
|
|
)
|