grounding_demo/native_runtime.py

"""Checkpoint-native inference. No case catalog, references, or HTTP fallback here.

This module deliberately avoids importing torch at import time: spaces must install
its CUDA hooks first. The GPU boundary lives in zerogpu_backend, not in this module.
"""
from __future__ import annotations

import hashlib
import json
import os
import re
from pathlib import Path


class ConfigurationError(RuntimeError):
    pass


class IncompleteGeneration(RuntimeError):
    pass


def resolve_snapshot() -> Path:
    local = os.getenv("GROUNDING_MODEL_PATH", "").strip()
    repo = os.getenv("GROUNDING_MODEL_REPO", "").strip()
    if local and repo:
        raise ConfigurationError("Configure either GROUNDING_MODEL_PATH or GROUNDING_MODEL_REPO, not both.")
    if local:
        return Path(local).expanduser().resolve(strict=True)
    if not repo:
        raise ConfigurationError("Configure GROUNDING_MODEL_REPO and GROUNDING_MODEL_REVISION first.")
    revision = os.getenv("GROUNDING_MODEL_REVISION", "").strip()
    if not re.fullmatch(r"[0-9a-fA-F]{40}", revision):
        raise ConfigurationError("GROUNDING_MODEL_REVISION must be the full 40-character commit SHA of reviewed model code.")
    from huggingface_hub import snapshot_download

    # The custom processor uses local open(). Download a pinned snapshot first;
    # passing a Hub repo ID directly to this processor does not work.
    return Path(snapshot_download(
        repo_id=repo, revision=revision, token=os.getenv("HF_TOKEN") or None,
        allow_patterns=["*.json", "*.py", "*.jinja", "*.txt", "*.safetensors"],
    ))


def inspect_contract(root: Path, *, require_weights: bool = True) -> dict:
    """Validate the public GroundingPI release without executing model code."""
    required = [
        "config.json", "preprocessor_config.json", "generation_config.json",
        "tokenizer.json", "tokenizer_config.json", "chat_template.jinja",
        "configuration_groundingpi.py", "configuration_groundingpi_vision.py",
        "modeling_groundingpi.py", "modeling_groundingpi_vision.py",
        "processing_groundingpi.py", "image_processing_groundingpi.py", "media_utils.py",
        "streammind_gate.py",
    ]
    missing = [name for name in required if not (root / name).is_file()]
    if missing:
        raise ConfigurationError(f"Model package is missing files: {', '.join(missing)}")
    config = json.loads((root / "config.json").read_text())
    if config.get("model_type") != "groundingpi":
        raise ConfigurationError("This demo requires the GroundingPI/GroundingPI model package.")
    if config.get("vision_config", {}).get("model_type") != "groundingpi_vision":
        raise ConfigurationError("The model does not use the expected Kimi-K3 ViT.")
    if not config.get("auto_map", {}).get("AutoModelForImageTextToText"):
        raise ConfigurationError("The model is missing the AutoModelForImageTextToText mapping.")
    if not config.get("auto_map", {}).get("AutoProcessor"):
        raise ConfigurationError("The model is missing its custom AutoProcessor mapping.")
    tokenizer = json.loads((root / "tokenizer.json").read_text())
    vocab = dict(tokenizer.get("model", {}).get("vocab", {}))
    vocab.update({item["content"]: item["id"] for item in tokenizer.get("added_tokens", [])})
    gam_tokens = [f"<{n}>" for n in range(1000)] + [
        "<|box_start|>", "<|box_end|>", "<|object_ref_start|>", "<|object_ref_end|>",
    ]
    if any(token not in vocab for token in gam_tokens):
        raise ConfigurationError("The tokenizer is missing native GAM coordinate or grounding tokens.")
    if vocab.get("<|image_pad|>") != config.get("image_token_id"):
        raise ConfigurationError("The image token ID does not match the model configuration.")
    generation = json.loads((root / "generation_config.json").read_text())
    eos = generation.get("eos_token_id")
    eos = eos if isinstance(eos, list) else [eos]
    if not eos or any(not isinstance(token, int) for token in eos):
        raise ConfigurationError("generation_config has no valid end-of-sequence token.")
    shards = []
    if require_weights:
        index = root / "model.safetensors.index.json"
        if index.is_file():
            shards = sorted(set(json.loads(index.read_text())["weight_map"].values()))
        elif (root / "model.safetensors").is_file():
            shards = ["model.safetensors"]
        if not shards or any(
            Path(name).name != name or not name.endswith(".safetensors")
            or not (root / name).is_file() or (root / name).stat().st_size == 0
            for name in shards
        ):
            raise ConfigurationError("Model weight shards are incomplete. Upload the weights as well as the tokenizer and configuration.")
    prep = json.loads((root / "preprocessor_config.json").read_text())
    return {
        "architecture": config["architectures"], "vision": "groundingpi_vision",
        "config_sha256": hashlib.sha256((root / "config.json").read_bytes()).hexdigest(),
        "processor_sha256": hashlib.sha256((root / "preprocessor_config.json").read_bytes()).hexdigest(),
        "media_proc_cfg": prep.get("media_proc_cfg"), "gam_coordinate_tokens": 1000,
        "eos_token_ids": eos, "weight_shards": shards,
        "scope": "package_contract_only_not_gpu_inference",
    }


def user_messages(prompt: str) -> list[dict]:
    text = prompt.replace("<image>", "").strip()
    # Exactly one image. Do not let text inject another image placeholder or role.
    forbidden = ["<|image_pad|>", "<|video_pad|>", "<|image_placeholder|>",
                 "<|vision_start|>", "<|vision_end|>", "<|im_start|>", "<|im_end|>"]
    if not text or any(token in text for token in forbidden):
        raise ValueError("The prompt is empty or contains reserved image/conversation control tokens.")
    return [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": text}]}]


def decode_answer(processor, generated_ids: list[int], eos_ids: list[int], limit: int) -> str:
    # Keep GAM coordinate tokens. skip_special_tokens=True would erase evidence.
    if not generated_ids:
        raise IncompleteGeneration("The model returned no answer.")
    if generated_ids[-1] not in eos_ids:
        raise IncompleteGeneration(
            f"Generation did not reach EOS ({len(generated_ids)}/{limit} tokens). Truncated output is not accepted as a valid result."
        )
    ids = list(generated_ids)
    while ids and ids[-1] in eos_ids:
        ids.pop()
    text = processor.decode(ids, skip_special_tokens=False, spaces_between_special_tokens=False)
    if not text.strip():
        raise IncompleteGeneration("The model returned EOS without a grounding result.")
    return text


class NativeRuntime:
    def __init__(self, model, processor, eos_ids: list[int]):
        self.model, self.processor, self.eos_ids = model, processor, eos_ids

    @classmethod
    def load(cls, root: Path) -> "NativeRuntime":
        contract = inspect_contract(root)
        import torch
        from transformers import AutoModelForImageTextToText, AutoProcessor

        processor = AutoProcessor.from_pretrained(str(root), trust_remote_code=True)
        model = AutoModelForImageTextToText.from_pretrained(
            str(root), trust_remote_code=True, dtype=torch.bfloat16,
            attn_implementation="sdpa",
        ).eval().to("cuda")
        return cls(model, processor, contract["eos_token_ids"])

    def prepare(self, image_path: str | Path, prompt: str, limit: int):
        from PIL import Image

        path = Path(image_path)
        if path.stat().st_size > 30 * 1024 * 1024:
            raise ValueError("Image exceeds 30 MiB. Reduce the file size and retry.")
        if not 1 <= limit <= 2048:
            raise ConfigurationError("GROUNDING_MAX_TOKENS must be between 1 and 2048.")
        with Image.open(path) as source:
            if source.width * source.height > 50_000_000:
                raise ValueError("Image exceeds 50 megapixels. Reduce its dimensions and retry.")
            # Preserve RGBA: the checkpoint applies its own alpha background policy.
            image = source.copy()
        rendered = self.processor.apply_chat_template(
            user_messages(prompt), tokenize=False, add_generation_prompt=True,
        )
        inputs = self.processor(text=[rendered], images=[image], return_tensors="pt")
        prompt_tokens = int(inputs["input_ids"].shape[-1])
        if prompt_tokens + limit > 8192:
            raise ValueError("Image and prompt exceed the 8192-token context budget. Shorten the prompt; inputs will not be silently truncated.")
        if "pixel_values" not in inputs or "image_grid_thw" not in inputs:
            raise ConfigurationError("The processor returned no image pixels or grid. Inference has been stopped.")
        return inputs, prompt_tokens

    def generate_ids(self, inputs, limit: int) -> list[int]:
        """Must run inside @spaces.GPU on ZeroGPU. No loading or HTTP here."""
        import torch

        device_inputs = {key: value.to("cuda") for key, value in inputs.items()}
        with torch.inference_mode():
            output = self.model.generate(
                **device_inputs, max_new_tokens=limit, do_sample=False,
                use_cache=True, eos_token_id=self.eos_ids,
            )
        start = inputs["input_ids"].shape[-1]
        # No GPU tensor escapes the temporary GPU lease.
        return output[0, start:].detach().cpu().tolist()