src/unlimited_ocr_rdna4/cli.py

from __future__ import annotations

import argparse
import json
import sys
import traceback
from pathlib import Path

from .constants import MODEL_REVISION, VERSION
from .errors import AppError, ModelIntegrityError
from .infer import RunOptions, default_output_path, run_inference
from .model_store import prepare_model, verify_model
from .paths import default_cache_dir, default_model_dir
from .runtime import inspect_runtime, runtime_issues


def bounded_int(minimum: int, maximum: int):
    def parse(value: str) -> int:
        number = int(value)
        if not minimum <= number <= maximum:
            raise argparse.ArgumentTypeError(f"must be between {minimum} and {maximum}")
        return number

    return parse


def _add_output_flags(parser: argparse.ArgumentParser, *, suppress_defaults: bool = False) -> None:
    default = argparse.SUPPRESS if suppress_defaults else False
    parser.add_argument("--json", action="store_true", default=default, help="emit structured JSON")
    parser.add_argument("-q", "--quiet", action="store_true", default=default, help="suppress progress diagnostics")
    parser.add_argument("--debug", action="store_true", default=default, help="show tracebacks for unexpected failures")
    parser.add_argument(
        "--no-color", action="store_true", default=default, help="disable color (currently the default)"
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="unlimited-ocr-rdna4",
        description="Run Baidu Unlimited-OCR locally on a single AMD RDNA 4 GPU.",
        epilog="Examples: unlimited-ocr-rdna4 prepare; unlimited-ocr-rdna4 run --input page.png",
    )
    _add_output_flags(parser)
    parser.add_argument("--version", action="version", version=VERSION)
    subparsers = parser.add_subparsers(dest="command")

    prepare = subparsers.add_parser("prepare", help="download and verify the pinned Baidu model")
    _add_output_flags(prepare, suppress_defaults=True)
    prepare.add_argument("--model-dir", type=Path, default=default_model_dir(), help="prepared model destination")
    prepare.add_argument("--cache-dir", type=Path, default=default_cache_dir(), help="persistent download cache")
    prepare.add_argument("-n", "--dry-run", action="store_true", help="show the plan without downloading")

    doctor = subparsers.add_parser("doctor", help="inspect ROCm, PyTorch, RDNA 4, and model readiness")
    _add_output_flags(doctor, suppress_defaults=True)
    doctor.add_argument("--device", help="ROCm ordinal or stable ROCr UUID; omit to inspect all visible GPUs")
    doctor.add_argument("--model-dir", type=Path, default=default_model_dir(), help="prepared model directory")
    doctor.add_argument("--require-model", action="store_true", help="fail when the prepared model is absent")

    run = subparsers.add_parser("run", help="parse one image or PDF into Markdown")
    _add_output_flags(run, suppress_defaults=True)
    run.add_argument("--input", required=True, type=Path, help="input image or PDF")
    run.add_argument("-o", "--output", type=Path, help="Markdown output path")
    run.add_argument(
        "--device",
        default=None,
        help="ROCm ordinal or stable ROCr UUID (default: UNLIMITED_OCR_DEVICE or 0)",
    )
    run.add_argument("--model-dir", type=Path, default=default_model_dir(), help="prepared model directory")
    run.add_argument("--mode", choices=("gundam", "base"), default="gundam", help="single-page image profile")
    run.add_argument("--max-length", type=bounded_int(512, 32768), default=4096, help="total sequence limit")
    run.add_argument("--dpi", type=bounded_int(72, 400), default=200, help="PDF rendering resolution")
    run.add_argument("--start-page", type=bounded_int(1, 100000), default=1, help="first PDF page, 1-based")
    run.add_argument("--max-pages", type=bounded_int(1, 100), default=20, help="maximum PDF pages per run")
    run.add_argument(
        "--max-page-pixels",
        type=bounded_int(1_000_000, 200_000_000),
        default=60_000_000,
        help="maximum rendered pixels for one PDF page",
    )
    run.add_argument(
        "--max-total-pixels",
        type=bounded_int(1_000_000, 2_000_000_000),
        default=400_000_000,
        help="maximum aggregate rendered pixels for a PDF run",
    )
    run.add_argument(
        "--max-page-rendered-mib",
        type=bounded_int(1, 2048),
        default=512,
        help="maximum temporary size of one rendered PDF page",
    )
    run.add_argument(
        "--max-rendered-mib",
        type=bounded_int(1, 8192),
        default=2048,
        help="maximum aggregate size of rendered PDF pages",
    )
    run.add_argument("--prompt", default="<image>document parsing.", help="model prompt")
    run.add_argument("-f", "--force", action="store_true", help="replace an existing output file")
    return parser


def _emit(payload: dict, *, as_json: bool, quiet: bool = False) -> None:
    if as_json:
        print(json.dumps(payload, ensure_ascii=False, sort_keys=True))
        return
    if quiet:
        primary = payload.get("output") or payload.get("model_dir") or payload.get("status")
        if primary is not None:
            print(primary)
        return
    for key, value in payload.items():
        if isinstance(value, (dict, list, tuple)):
            print(f"{key}={json.dumps(value, ensure_ascii=False, sort_keys=True)}")
        else:
            print(f"{key}={value}")


def _command_prepare(args: argparse.Namespace) -> int:
    if not args.quiet:
        action = "Would prepare" if args.dry_run else "Preparing"
        print(f"{action} {MODEL_REVISION} at {args.model_dir}", file=sys.stderr, flush=True)
    status = prepare_model(args.model_dir, args.cache_dir, dry_run=args.dry_run)
    payload = {"schema_version": 1, **status.__dict__, "status": "dry-run" if args.dry_run else "ok"}
    _emit(payload, as_json=args.json, quiet=args.quiet)
    return 0


def _command_doctor(args: argparse.Namespace) -> int:
    runtime = inspect_runtime(device=args.device, require_single=False, validate=False)
    issues = runtime_issues(runtime, require_single=False)
    model_payload: dict[str, object]
    try:
        model = verify_model(args.model_dir.expanduser().resolve(), full_weight_hash=False)
        model_payload = {"prepared": True, "model_dir": model.model_dir, "revision": model.revision}
    except ModelIntegrityError as exc:
        if args.require_model:
            raise
        model_payload = {"prepared": False, "model_dir": str(args.model_dir), "detail": str(exc)}
    payload = {
        "schema_version": 1,
        "status": "ok" if not issues else "error",
        "runtime": runtime.to_dict(),
        "runtime_issues": issues,
        "model": model_payload,
    }
    _emit(payload, as_json=args.json, quiet=args.quiet)
    return 0 if not issues else 3


def _command_run(args: argparse.Namespace) -> int:
    import os

    device = args.device or os.environ.get("UNLIMITED_OCR_DEVICE")
    visibility_is_preconfigured = any(
        name in os.environ for name in ("ROCR_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES")
    )
    if device is None and not visibility_is_preconfigured:
        device = "0"
    runtime = inspect_runtime(device=device, require_single=True)
    input_path = args.input.expanduser()
    output = args.output.expanduser() if args.output else default_output_path(input_path)
    if not args.quiet:
        print(f"Loading pinned model on RDNA 4 device {device}; inference is offline", file=sys.stderr, flush=True)
    result = run_inference(
        RunOptions(
            input_path=input_path,
            output_path=output,
            model_dir=args.model_dir,
            mode=args.mode,
            max_length=args.max_length,
            dpi=args.dpi,
            start_page=args.start_page,
            max_pages=args.max_pages,
            max_page_pixels=args.max_page_pixels,
            max_total_pixels=args.max_total_pixels,
            max_page_rendered_bytes=args.max_page_rendered_mib * 1024**2,
            max_rendered_bytes=args.max_rendered_mib * 1024**2,
            prompt=args.prompt,
            force=args.force,
            quiet=args.quiet,
        ),
        runtime,
    )
    for warning in result.warnings:
        print(f"warning: {warning}", file=sys.stderr)
    payload = {"status": "ok", **result.to_dict()}
    _emit(payload, as_json=args.json, quiet=args.quiet)
    return 0


def run(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if not args.command:
        parser.print_help(sys.stderr)
        return 2
    try:
        if args.command == "prepare":
            return _command_prepare(args)
        if args.command == "doctor":
            return _command_doctor(args)
        if args.command == "run":
            return _command_run(args)
        parser.error(f"unknown command: {args.command}")
    except KeyboardInterrupt:
        print("interrupted", file=sys.stderr)
        return 130
    except AppError as exc:
        if args.json:
            print(
                json.dumps({"schema_version": 1, "status": "error", "error": str(exc), "exit_code": exc.exit_code}),
                file=sys.stderr,
            )
        else:
            print(f"error: {exc}", file=sys.stderr)
        return exc.exit_code
    except Exception as exc:
        if args.debug:
            traceback.print_exc()
        else:
            print(f"error: unexpected failure: {exc}; rerun with --debug", file=sys.stderr)
        return 1
    return 1


def main() -> None:
    raise SystemExit(run())