handler.py

"""HuggingFace Inference Endpoint handler for OmniShotCut.

Deploy this as a *custom* Inference Endpoint (dedicated GPU). The model weights
are pulled from the authors' hub repo on first load, so this endpoint repo only
needs this file plus requirements.txt — no weights to host.

Request body (JSON):
    {
      "inputs": "<video url | base64-encoded mp4>",
      "parameters": {"mode": "default"}      # or "clean_shot"
    }
Raw binary bodies (Content-Type: video/mp4) are also accepted.

Response (JSON):
    default    -> {"shots": [[s,e],...], "intra_labels": [...], "inter_labels": [...]}
    clean_shot -> {"shots": [[s,e],...]}
"""
import base64
import os
import tempfile
import urllib.request
from typing import Any, Dict

import omnishotcut

_CKPT = "uva-cv-lab/OmniShotCut"
_FILENAME = "OmniShotCut_ckpt.pth"


class EndpointHandler:
    def __init__(self, path: str = "") -> None:
        # Weights auto-download from the HuggingFace hub and cache in the image.
        # Loading here (once, at cold start) keeps every request warm.
        self.model = omnishotcut.load(_CKPT, filename=_FILENAME)

    # -- HF calls this per request -------------------------------------------
    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inp = data.get("inputs", data)
        params = data.get("parameters") or {}
        mode = params.get("mode", "default")
        if mode not in ("default", "clean_shot"):
            return {"error": f"unknown mode {mode!r}; use 'default' or 'clean_shot'"}

        path, is_tmp = self._materialize(inp)
        try:
            if mode == "clean_shot":
                ranges = self.model.inference(path, mode="clean_shot")
                return {"mode": mode, "shots": _ints(ranges)}
            ranges, intra, inter = self.model.inference(path, mode="default")
            return {
                "mode": mode,
                "shots": _ints(ranges),
                "intra_labels": [str(x) for x in intra],
                "inter_labels": [str(x) for x in inter],
            }
        except Exception as exc:  # surface a clean error instead of a 500 with no body
            return {"error": f"{type(exc).__name__}: {exc}"}
        finally:
            if is_tmp and path and os.path.exists(path):
                os.remove(path)

    # -- resolve the request payload to a local mp4 path ---------------------
    def _materialize(self, inp: Any):
        if isinstance(inp, (bytes, bytearray)):
            return self._write_tmp(bytes(inp)), True
        if isinstance(inp, str):
            if inp.startswith(("http://", "https://")):
                fn = _tmp_name()
                urllib.request.urlretrieve(inp, fn)  # endpoint fetches the video
                return fn, True
            return self._write_tmp(base64.b64decode(inp)), True  # base64 mp4
        raise ValueError("inputs must be a video URL, base64 mp4 string, or raw bytes")

    @staticmethod
    def _write_tmp(raw: bytes) -> str:
        fn = _tmp_name()
        with open(fn, "wb") as f:
            f.write(raw)
        return fn


def _tmp_name() -> str:
    fd, fn = tempfile.mkstemp(suffix=".mp4")
    os.close(fd)
    return fn


def _ints(ranges) -> list:
    return [[int(a), int(b)] for a, b in ranges]