runtime/kev/jevvision.py

"""Convenience loader for the public JEVision bundle.

JEVision is a routed pair of KEV-compatible checkpoints: text-only requests use
the selected text adapter; requests containing images use the visual sidecar.
This is not an AutoModel wrapper and it does not merge the two adapters.
"""
from __future__ import annotations

import base64
import mimetypes
import os
from pathlib import Path
from dataclasses import replace

import torch

from .api import SystemOneRequest
from .checkpoint import Checkpoint, LoadOptions
from .device import default_device
from .serve import SERVE_CONTEXT_TOKENS, Server


TEXT_ADAPTERS = {
    "jevvision-text": "text/jevvision-text",
    "jevbench-m3": "text/jevbench-m3",
    "kev-0.8b": "text/kev-0.8b",
}
MAX_IMAGE_BYTES = 10 * 1024 * 1024
IMAGE_MIME_TYPES = {"image/png", "image/jpeg", "image/webp"}


def image_file_as_data_url(path: str | os.PathLike[str]) -> dict[str, str]:
    """Encode one supported local image in the System One request format."""
    source = Path(path)
    if source.stat().st_size > MAX_IMAGE_BYTES:
        raise ValueError("compressed image exceeds the 10 MiB request limit")
    mime, _ = mimetypes.guess_type(source.name)
    if mime not in IMAGE_MIME_TYPES:
        raise ValueError("image must be PNG, JPEG, or WebP")
    payload = base64.b64encode(source.read_bytes()).decode("ascii")
    return {"data_url": f"data:{mime};base64,{payload}"}


class JEVision:
    """Load and call JEVision locally with KEV/Jev-style typed inputs/outputs."""

    def __init__(self, server: Server):
        self._server = server

    @classmethod
    def from_pretrained(
        cls,
        repo_id: str = "divyanshx11/JEVision",
        *,
        revision: str = "main",
        text_adapter: str = "jevvision-text",
        device: str | None = None,
        cache_dir: str | os.PathLike[str] | None = None,
    ) -> "JEVision":
        """Download the adapter bundle (if needed) and load both routed branches."""
        local = Path(repo_id)
        if local.is_dir():
            root = local
        else:
            from huggingface_hub import snapshot_download

            root = Path(snapshot_download(
                repo_id=repo_id,
                revision=revision,
                cache_dir=str(cache_dir) if cache_dir is not None else None,
                allow_patterns=[
                    "README.md", "LICENSE", "jevvision_manifest.json",
                    "adapter/**", "pointer_head.pt", "report.json", "text/**",
                ],
            ))
        return cls.from_directory(root, text_adapter=text_adapter, device=device)

    @classmethod
    def from_directory(
        cls,
        bundle_dir: str | os.PathLike[str],
        *,
        text_adapter: str = "jevvision-text",
        device: str | None = None,
    ) -> "JEVision":
        """Load a previously downloaded JEVision bundle from disk."""
        import json

        root = Path(bundle_dir)
        if text_adapter not in TEXT_ADAPTERS:
            raise ValueError(f"text_adapter must be one of {tuple(TEXT_ADAPTERS)}")
        manifest_path = root / "jevvision_manifest.json"
        if not manifest_path.is_file():
            raise FileNotFoundError(f"JEVision manifest not found: {manifest_path}")
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        max_context = int(manifest["context"]["max_request_tokens"])
        if SERVE_CONTEXT_TOKENS > max_context:
            raise ValueError(
                f"KEV_SERVE_CONTEXT_TOKENS={SERVE_CONTEXT_TOKENS} exceeds the bundle's tested cap {max_context}"
            )

        text_run = root / TEXT_ADAPTERS[text_adapter]
        if not (text_run / "head.pt").is_file():
            raise FileNotFoundError(f"text adapter checkpoint is incomplete: {text_run}")
        selected_device = device or default_device()
        options = LoadOptions.from_env()
        if selected_device != "cpu" and options.dtype is None:
            options = replace(options, dtype=torch.bfloat16)
        if options.backend is None:
            options = replace(options, backend="auto")

        checkpoint = Checkpoint(str(text_run))
        tokenizer, model = checkpoint.load(selected_device, options)
        from .visual_model import VisualDecisionModel

        visual = VisualDecisionModel.from_run(
            root, selected_device, max_input_tokens=SERVE_CONTEXT_TOKENS
        )
        return cls(Server(checkpoint, tokenizer, model, str(selected_device), visual=visual))

    def system_one(
        self,
        *,
        state,
        questions: dict,
        images: list[dict[str, str]] | None = None,
        model: str = "jev-latest",
    ) -> dict:
        """Return one KEV/System One response dict for the supplied request."""
        request = SystemOneRequest(
            model=model, state=state, questions=questions, images=images
        )
        return self._server.answer(request)

    __call__ = system_one