inference.py

"""Local inference for Anime_Image2Prompt.

The model and processor are loaded only from the directory containing this
file. No base model, adapter, or network download is required.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

import torch
from PIL import Image, ImageOps
from transformers import AutoModelForMultimodalLM, AutoProcessor


MODEL_DIR = Path(__file__).resolve().parent
INFERENCE_CONFIG = MODEL_DIR / "inference_config.json"
TAG_FIELDS = ("general_tags", "character_tags", "copyright_tags")
REQUIRED_MODEL_FILES = (
    "config.json",
    "model.safetensors",
    "processor_config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "chat_template.jinja",
    "inference_config.json",
)


def _validate_model_directory(model_dir: Path) -> None:
    missing = [name for name in REQUIRED_MODEL_FILES if not (model_dir / name).is_file()]
    if missing:
        raise FileNotFoundError(
            f"Incomplete Anime_Image2Prompt package; missing: {', '.join(missing)}"
        )


def _select_runtime(device: str) -> tuple[str, torch.dtype]:
    cuda_available = torch.cuda.is_available()
    if device == "cuda" and not cuda_available:
        raise RuntimeError(
            "CUDA was requested but is unavailable. Install a CUDA-enabled PyTorch "
            "build and verify the NVIDIA driver, or use --device cpu."
        )

    use_cuda = device == "cuda" or (device == "auto" and cuda_available)
    if not use_cuda:
        return "cpu", torch.float32

    supports_bf16 = getattr(torch.cuda, "is_bf16_supported", lambda: False)()
    return "cuda", torch.bfloat16 if supports_bf16 else torch.float16


def _parse_model_json(text: str) -> dict[str, list[str]]:
    start = text.find("{")
    if start < 0:
        raise ValueError(f"The model did not return JSON. Raw output: {text!r}")

    try:
        payload, _ = json.JSONDecoder().raw_decode(text[start:])
    except json.JSONDecodeError as error:
        raise ValueError(f"The model returned invalid JSON. Raw output: {text!r}") from error

    if not isinstance(payload, dict) or set(payload) != set(TAG_FIELDS):
        actual = tuple(payload.keys()) if isinstance(payload, dict) else type(payload).__name__
        raise ValueError(f"Unexpected model output fields: {actual!r}")

    normalized: dict[str, list[str]] = {}
    for field in TAG_FIELDS:
        value = payload[field]
        if not isinstance(value, list) or len(value) > 1:
            raise ValueError(f"{field} must be [] or a one-string list: {value!r}")
        if value:
            if not isinstance(value[0], str):
                raise ValueError(f"{field} must contain a string: {value!r}")
            tags = value[0].split(",")
            if any(not tag for tag in tags):
                raise ValueError(f"{field} contains an empty tag: {value!r}")
            if len(tags) != len(set(tags)):
                raise ValueError(f"{field} contains duplicate tags: {value!r}")
        normalized[field] = value
    return normalized


def _move_inputs(inputs: dict[str, Any], device: torch.device) -> dict[str, Any]:
    return {
        key: value.to(device) if hasattr(value, "to") else value
        for key, value in inputs.items()
    }


class AnimeImage2Prompt:
    """Reusable local inference session for Anime_Image2Prompt."""

    def __init__(self, device: str = "auto", low_vram: bool = False) -> None:
        _validate_model_directory(MODEL_DIR)
        runtime, dtype = _select_runtime(device)
        self.runtime = runtime
        self.dtype = dtype
        self.config = json.loads(INFERENCE_CONFIG.read_text(encoding="utf-8"))

        if runtime == "cuda":
            device_map: str | dict[str, str] = "auto"
        else:
            device_map = {"": "cpu"}

        self.processor = AutoProcessor.from_pretrained(
            MODEL_DIR,
            local_files_only=True,
        )
        if low_vram:
            pixels = self.config["image_pixels"]
            self.processor.image_processor.size = {
                "shortest_edge": pixels["min"],
                "longest_edge": pixels["low_vram_max"],
            }

        self.model = AutoModelForMultimodalLM.from_pretrained(
            MODEL_DIR,
            dtype=dtype,
            device_map=device_map,
            attn_implementation="sdpa",
            local_files_only=True,
        ).eval()

    @property
    def device(self) -> torch.device:
        return self.model.device

    def predict(
        self,
        image: str | Path | Image.Image,
        max_new_tokens: int = 1024,
    ) -> dict[str, list[str]]:
        """Generate structured prompt tags for one image."""
        if isinstance(image, Image.Image):
            prepared_image = ImageOps.exif_transpose(image).convert("RGB")
        else:
            image_path = Path(image).expanduser().resolve()
            if not image_path.is_file():
                raise FileNotFoundError(f"Input image does not exist: {image_path}")
            with Image.open(image_path) as source:
                prepared_image = ImageOps.exif_transpose(source).convert("RGB")
                prepared_image.load()

        messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": self.config["system_prompt"]}],
            },
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": prepared_image},
                    {"type": "text", "text": self.config["user_prompt"]},
                ],
            },
        ]
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            enable_thinking=False,
            return_dict=True,
            return_tensors="pt",
        )
        inputs = _move_inputs(inputs, self.device)

        with torch.inference_mode():
            output_ids = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                eos_token_id=self.config["assistant_end_token_id"],
                pad_token_id=self.processor.tokenizer.pad_token_id,
            )

        prompt_length = inputs["input_ids"].shape[1]
        generated_ids = output_ids[:, prompt_length:]
        raw_text = self.processor.batch_decode(
            generated_ids,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )[0]
        return _parse_model_json(raw_text)

    @staticmethod
    def to_prompt(payload: dict[str, list[str]]) -> str:
        """Flatten structured output to a comma-and-space-separated prompt."""
        tags: list[str] = []
        for field in TAG_FIELDS:
            values = payload.get(field, [])
            if values:
                tags.extend(values[0].split(","))
        return ", ".join(tags)


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate Danbooru-style prompt tags with Anime_Image2Prompt."
    )
    parser.add_argument("image", type=Path, help="path to one input image")
    parser.add_argument(
        "--output",
        type=Path,
        help="optional output file; the result is always printed to stdout",
    )
    parser.add_argument(
        "--format",
        choices=("json", "prompt"),
        default="json",
        help="output format (default: json)",
    )
    parser.add_argument(
        "--device",
        choices=("auto", "cuda", "cpu"),
        default="auto",
        help="inference device (default: auto)",
    )
    parser.add_argument(
        "--low-vram",
        action="store_true",
        help="reduce image tokens to lower peak VRAM usage",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=1024,
        help="maximum generated tokens (default: 1024)",
    )
    return parser.parse_args()


def main() -> int:
    args = _parse_args()
    if args.max_new_tokens < 1:
        raise ValueError("--max-new-tokens must be greater than zero")

    runtime, dtype = _select_runtime(args.device)
    if runtime == "cuda":
        hardware = torch.cuda.get_device_name(0)
        print(f"Loading Anime_Image2Prompt on {hardware} ({dtype})...", file=sys.stderr)
    else:
        print("Loading Anime_Image2Prompt on CPU (this may be slow)...", file=sys.stderr)

    tagger = AnimeImage2Prompt(device=args.device, low_vram=args.low_vram)
    print("Generating prompt...", file=sys.stderr)
    result = tagger.predict(args.image, max_new_tokens=args.max_new_tokens)

    if args.format == "json":
        rendered = json.dumps(result, ensure_ascii=False, indent=2) + "\n"
    else:
        rendered = tagger.to_prompt(result) + "\n"

    if args.output:
        output_path = args.output.expanduser().resolve()
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(rendered, encoding="utf-8", newline="\n")
        print(f"Saved: {output_path}", file=sys.stderr)

    print(rendered, end="")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (FileNotFoundError, RuntimeError, ValueError) as error:
        print(f"Error: {error}", file=sys.stderr)
        raise SystemExit(1) from error