models/vision/minicpm_wrapper.py

"""MiniCPM-V 4.6 wrapper. Loads the model and runs inference on film scans."""

from __future__ import annotations

import json
import logging
import os
import re
from typing import Any

from config import CHECKPOINT_DIR, get_vision_config, require_gpu_for_inference
from models.vision.prompts import DETECTION_PROMPT_INT

logger = logging.getLogger(__name__)

DETECTION_PROMPT = DETECTION_PROMPT_INT


def _resolve_model_path() -> str:
    """Pick configured fine-tuned model or public base model."""
    cfg = get_vision_config()
    explicit = os.getenv("HALIDE_VISION_MODEL_ID")
    if explicit:
        logger.info("Using explicit vision model %s", explicit)
        return explicit

    if cfg.use_finetuned:
        local_candidates = [
            cfg.local_model_path,
            CHECKPOINT_DIR / "minicpm-v-4.6-merged",
        ]
        seen: set[str] = set()
        for path in local_candidates:
            key = str(path.resolve())
            if key in seen:
                continue
            seen.add(key)
            if path.exists() and (path / "config.json").exists():
                logger.info("Using local fine-tuned vision model at %s", path)
                return str(path)
        logger.info("Using fine-tuned vision model repo %s", cfg.finetuned_model_id)
        return cfg.finetuned_model_id

    logger.info("Using base vision model %s", cfg.base_model_id)
    return cfg.base_model_id


class MiniCPMVDetector:
    """Lazy-loading wrapper around MiniCPM-V 4.6 for film defect detection."""

    def __init__(self, model_path: str | None = None) -> None:
        self._model_path = model_path or _resolve_model_path()
        self._model: Any = None
        self._processor: Any = None
        self._dtype: Any = None
        self._device: str = "cpu"

    @property
    def model_path(self) -> str:
        return self._model_path

    def load(self) -> None:
        if self._model is not None:
            return
        require_gpu_for_inference("vision")
        import torch
        from transformers import AutoModelForImageTextToText, AutoProcessor

        logger.info("Loading MiniCPM-V 4.6 from %s", self._model_path)
        self._processor = AutoProcessor.from_pretrained(
            self._model_path, trust_remote_code=True
        )
        self._dtype = _select_cuda_dtype(torch)
        self._model = AutoModelForImageTextToText.from_pretrained(
            self._model_path,
            torch_dtype=self._dtype,
            device_map="auto",
            trust_remote_code=True,
        )
        self._device = str(next(self._model.parameters()).device)
        logger.info("Model loaded on %s with dtype %s", self._device, self._dtype)

    def detect(self, image: Any) -> dict:
        """Run defect detection on a PIL image. Returns parsed JSON dict."""
        import torch

        if self._model is None:
            self.load()

        cfg = get_vision_config()
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": DETECTION_PROMPT},
                ],
            }
        ]

        inputs = _apply_chat_template(
            self._processor,
            messages,
            downsample_mode=cfg.downsample_mode,
            max_slice_nums=cfg.max_slice_nums,
        ).to(self._device)

        with torch.inference_mode():
            generated = self._model.generate(
                **inputs,
                downsample_mode=cfg.downsample_mode,
                max_new_tokens=cfg.max_new_tokens,
                do_sample=False,
            )

        trimmed = [out[len(inp):] for inp, out in zip(inputs.input_ids, generated)]
        text = self._processor.batch_decode(
            trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )[0]

        return _parse_defect_json(text)

    def close(self) -> None:
        if self._model is not None:
            del self._model
            self._model = None
        if self._processor is not None:
            del self._processor
            self._processor = None


def _parse_defect_json(text: str) -> dict:
    """Extract and parse the first JSON object from model output."""
    text = text.strip()
    if text.startswith("```"):
        text = re.sub(r"^```(?:json)?\s*", "", text, flags=re.IGNORECASE)
        text = re.sub(r"\s*```$", "", text)
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            return {"defects": parsed}
        if isinstance(parsed, dict):
            return parsed
        return {"defects": [], "_raw": text, "_parse_error": "json_not_object"}
    except json.JSONDecodeError:
        pass

    match = re.search(r"\{[\s\S]*\}", text)
    if not match:
        logger.warning("No JSON found in model output: %r", text[:200])
        return {"defects": [], "_raw": text, "_parse_error": "no_json_object"}
    try:
        parsed = json.loads(match.group(0))
        if isinstance(parsed, dict):
            return parsed
        return {"defects": [], "_raw": text, "_parse_error": "json_not_object"}
    except json.JSONDecodeError as exc:
        fragments = _parse_defect_fragments(text)
        if fragments:
            logger.warning(
                "Salvaged %s defect fragments from malformed JSON: %s",
                len(fragments),
                exc,
            )
            return {
                "defects": fragments,
                "_parse_error": str(exc),
                "_parse_warning": "salvaged_defect_fragments",
            }
        logger.warning("JSON parse error: %s; raw: %r", exc, text[:200])
        return {"defects": [], "_raw": text, "_parse_error": str(exc)}


def _parse_defect_fragments(text: str) -> list[dict[str, Any]]:
    """Recover complete defect objects from truncated JSON arrays."""
    fragments: list[dict[str, Any]] = []
    for match in re.finditer(r"\{[^{}]*\"label\"[^{}]*\"bbox\"\s*:\s*\[[^\]]+\][^{}]*\}", text):
        try:
            candidate = json.loads(match.group(0))
        except json.JSONDecodeError:
            continue
        if isinstance(candidate, dict):
            fragments.append(candidate)
    return fragments


def _apply_chat_template(
    processor: Any,
    messages: list[dict],
    *,
    downsample_mode: str,
    max_slice_nums: int,
) -> Any:
    """Call MiniCPM chat template across Transformers API variants."""
    kwargs = {
        "tokenize": True,
        "add_generation_prompt": True,
        "return_dict": True,
        "return_tensors": "pt",
    }
    try:
        return processor.apply_chat_template(
            messages,
            **kwargs,
            downsample_mode=downsample_mode,
            max_slice_nums=max_slice_nums,
        )
    except TypeError:
        return processor.apply_chat_template(
            messages,
            **kwargs,
            processor_kwargs={
                "downsample_mode": downsample_mode,
                "max_slice_nums": max_slice_nums,
            },
        )


def _select_cuda_dtype(torch_module: Any) -> Any:
    major, _minor = torch_module.cuda.get_device_capability()
    if major >= 8:
        return torch_module.bfloat16
    return torch_module.float16


_default_detector: MiniCPMVDetector | None = None


def get_detector() -> MiniCPMVDetector:
    global _default_detector
    if _default_detector is None:
        _default_detector = MiniCPMVDetector()
    return _default_detector