omniret/embedding.py

from __future__ import annotations

import json
import tempfile
import wave
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from urllib.parse import urlparse
from urllib.request import urlopen

import imageio_ffmpeg
import numpy as np
import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from PIL import Image
from safetensors.torch import load_model as load_safetensors_model

from .config import OmniRetModelConfig
from .modeling import OmniRetModel

_MEDIA_FIELDS = ("audio", "image", "video")
_INPUT_FIELDS = {"instruction", "text", "audio", "image", "video", "max_frames"}


@dataclass(frozen=True)
class _Input:
    instruction: str
    text: str
    modality: str | None
    media: str | Path | None
    max_frames: int


def _validate_input(value: dict[str, Any], index: int) -> _Input:
    if not isinstance(value, dict):
        raise TypeError(f"input {index} must be a dictionary")
    unknown = sorted(set(value) - _INPUT_FIELDS)
    if unknown:
        raise ValueError(f"input {index} has unknown keys: {unknown}")
    present = [field for field in _MEDIA_FIELDS if value.get(field) is not None]
    if len(present) > 1:
        raise ValueError(f"input {index} has multiple media fields")
    text = str(value.get("text") or "").strip()
    if not text and not present:
        raise ValueError(f"input {index} has no text or media")
    media = value.get(present[0]) if present else None
    if media is not None and not isinstance(media, (str, Path)):
        raise TypeError(f"input {index} {present[0]} must be a path or HTTP(S) URL")
    max_frames = value.get("max_frames", 8)
    if not isinstance(max_frames, int) or not 1 <= max_frames <= 8:
        raise ValueError(f"input {index} max_frames must be between 1 and 8")
    return _Input(
        instruction=str(value.get("instruction") or "").strip(),
        text=text,
        modality=present[0] if present else None,
        media=media,
        max_frames=max_frames,
    )


def _validate_inputs(inputs: list[dict[str, Any]]) -> list[_Input]:
    if not isinstance(inputs, list):
        raise TypeError("inputs must be a list")
    if not inputs:
        raise ValueError("inputs must not be empty")
    return [_validate_input(value, index) for index, value in enumerate(inputs)]


def _format_input(value: _Input) -> str:
    pieces = []
    if value.modality:
        pieces.append(f"{value.modality.title()}: <{value.modality}>")
    if value.text:
        pieces.append(value.text)
    body = "\n".join(pieces)
    return f"Instruct: {value.instruction}\nQuery:\n{body}" if value.instruction else body


@contextmanager
def _local_path(value: str | Path, index: int, field: str):
    text = str(value)
    if text.startswith(("http://", "https://")):
        suffix = Path(urlparse(text).path).suffix
        try:
            with urlopen(text, timeout=30) as response, tempfile.NamedTemporaryFile(suffix=suffix) as handle:
                payload = response.read(512 * 1024 * 1024 + 1)
                if len(payload) > 512 * 1024 * 1024:
                    raise ValueError(f"input {index} {field} download exceeds 512 MiB")
                handle.write(payload)
                handle.flush()
                yield Path(handle.name)
        except OSError as error:
            raise ValueError(f"input {index} could not download {field}: {error}") from error
        return
    path = Path(value)
    if not path.is_file():
        raise FileNotFoundError(f"input {index} {field} file does not exist: {path}")
    yield path


def _load_wav(path: Path, index: int) -> np.ndarray:
    with wave.open(str(path)) as handle:
        if (handle.getnchannels(), handle.getsampwidth(), handle.getframerate()) != (1, 2, 16000):
            raise ValueError(f"input {index} expected mono 16-bit 16 kHz PCM WAV")
        return np.frombuffer(handle.readframes(handle.getnframes()), dtype=np.int16).astype(np.float32) / 32768


def _load_video(path: Path, max_frames: int, index: int) -> list[Image.Image]:
    reader = imageio_ffmpeg.read_frames(path)
    try:
        metadata = next(reader)
    finally:
        reader.close()
    total = max(1, int(float(metadata["fps"]) * float(metadata["duration"])))
    indices = np.linspace(0, total - 1, min(max_frames, total)).round().astype(int).tolist()
    select = "+".join(f"eq(n\\,{frame})" for frame in indices)
    reader = imageio_ffmpeg.read_frames(path, output_params=["-vf", f"select={select}", "-vsync", "0"])
    try:
        metadata = next(reader)
        frames = [Image.frombytes("RGB", metadata["size"], frame) for frame in reader]
    finally:
        reader.close()
    if not frames:
        raise ValueError(f"input {index} video contains no decodable frames")
    return frames


def _load_media(value: str | Path, modality: str, max_frames: int, index: int):
    with _local_path(value, index, modality) as path:
        try:
            if modality == "audio":
                return _load_wav(path, index)
            if modality == "video":
                return _load_video(path, max_frames, index)
            with Image.open(path) as image:
                return image.convert("RGB")
        except (OSError, RuntimeError, ValueError, wave.Error) as error:
            if isinstance(error, ValueError) and str(error).startswith(f"input {index}"):
                raise
            raise ValueError(f"input {index} could not decode {modality}: {error}") from error


def _model_dir(model_name_or_path: str | Path) -> Path:
    path = Path(model_name_or_path)
    return path if path.is_dir() else Path(snapshot_download(repo_id=str(model_name_or_path)))


class OmniRetEmbedder:
    def __init__(
        self,
        model_name_or_path: str | Path,
        torch_dtype: torch.dtype | None = None,
        attn_implementation: str | None = None,
        device: str | torch.device | None = None,
    ) -> None:
        model_dir = _model_dir(model_name_or_path)
        metadata = json.loads((model_dir / "config.json").read_text())
        bases = {
            name: snapshot_download(repo_id=value["repo_id"], revision=value["revision"])
            for name, value in metadata["base_models"].items()
        }
        config = OmniRetModelConfig(
            text_model_name=bases["text"],
            vision_model_name=bases["vision"],
            audio_model_name=bases["audio"],
            attn_implementation=attn_implementation,
        )
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        dtype = torch_dtype or (torch.bfloat16 if self.device.type == "cuda" else torch.float32)
        self.model = OmniRetModel.from_pretrained(config).to(dtype=dtype)
        load_safetensors_model(self.model, str(model_dir / "model.safetensors"), strict=True)
        self.model = self.model.to(device=self.device).eval()

    @torch.inference_mode()
    def process(self, inputs: list[dict[str, Any]]) -> torch.Tensor:
        values = _validate_inputs(inputs)
        texts = [_format_input(value) for value in values]
        modalities = [value.modality or "image" for value in values]
        media = [
            None if value.media is None else _load_media(value.media, modalities[index], value.max_frames, index)
            for index, value in enumerate(values)
        ]
        raw = self.model.encode_raw_media_batch(media, modalities) if any(item is not None for item in media) else None
        tokens, media_mask = raw if isinstance(raw, tuple) else (raw, None)
        embeddings, _ = self.model.encode_batch(
            texts,
            tokens,
            modalities,
            media_mask,
            exclude_instruction_prefix=any(value.instruction for value in values),
        )
        return F.normalize(embeddings.float(), dim=-1)