app.py

from __future__ import annotations

import json
import os
import statistics
from pathlib import Path
from typing import Any

import gradio as gr

from ocr_workbench.client import WorkerError, check_health, normalize_endpoint, run_page
from ocr_workbench.documents import (
    DocumentError,
    cleanup_stale_runs,
    gallery_items,
    normalize_gradio_path,
    prepare_document,
)
from ocr_workbench.export import build_exports
from ocr_workbench.registry import ModelSpec, load_registry


REGISTRY = load_registry()
MAX_INPUT_MB = int(os.getenv("MAX_INPUT_MB", "80"))
DEFAULT_MAX_PAGES = int(os.getenv("DEFAULT_MAX_PAGES", "8"))
ABSOLUTE_MAX_PAGES = int(os.getenv("ABSOLUTE_MAX_PAGES", "40"))
STALE_RUN_HOURS = float(os.getenv("STALE_RUN_HOURS", "12"))


try:
    import spaces
except Exception:
    spaces = None


def _gateway_probe() -> str:
    return "ready"


if spaces is not None:
    gateway_zero_gpu_probe = spaces.GPU(duration=1)(_gateway_probe)
else:
    gateway_zero_gpu_probe = _gateway_probe


def _spec(model_id: str) -> ModelSpec:
    try:
        return REGISTRY[model_id]
    except KeyError as exc:
        raise gr.Error(f"Unknown model: {model_id}") from exc


def _model_info(spec: ModelSpec) -> str:
    endpoint = spec.endpoint() or "未設定"
    return (
        f"### {spec.label}\n\n"
        f"{spec.description}\n\n"
        f"**Worker URL:** `{endpoint}`  \n"
        f"**環境変数:** `{spec.endpoint_env}`  \n"
        f"**画像出力:** {spec.result_note}"
    )


def on_model_change(model_id: str) -> tuple[str, str, str, int, int, int, int, bool, str]:
    spec = _spec(model_id)
    return (
        spec.default_prompt,
        spec.endpoint(),
        _model_info(spec),
        spec.default_max_tokens,
        min(spec.default_max_pages, ABSOLUTE_MAX_PAGES),
        spec.default_dpi,
        spec.default_request_timeout,
        spec.default_layout_as_thought,
        spec.default_image_mode,
    )


def health_check(model_id: str, endpoint_override: str) -> str:
    spec = _spec(model_id)
    endpoint = normalize_endpoint(endpoint_override) or spec.endpoint()
    try:
        payload = check_health(endpoint)
    except WorkerError as exc:
        return f"❌ **Health check failed:** {exc}"
    return "✅ **Worker reachable**\n\n```json\n" + json.dumps(payload, ensure_ascii=False, indent=2) + "\n```"


def run_ocr(
    file_value: object,
    model_id: str,
    prompt: str,
    endpoint_override: str,
    page_selection: str,
    max_pages: int,
    dpi: int,
    max_new_tokens: int,
    layout_as_thought: bool,
    unlimited_image_mode: str,
    request_timeout: int,
    progress: gr.Progress = gr.Progress(track_tqdm=False),
) -> tuple[
    list[tuple[str, str]],
    list[tuple[str, str]],
    str,
    str,
    str,
    str,
    str,
]:
    cleanup_stale_runs(STALE_RUN_HOURS)
    source_path = normalize_gradio_path(file_value)
    spec = _spec(model_id)
    endpoint = normalize_endpoint(endpoint_override) or spec.endpoint()
    if not endpoint:
        raise gr.Error(
            f"{spec.label} のWorker URLが未設定です。Space Variable "
            f"`{spec.endpoint_env}` または画面のWorker URL欄を設定してください。"
        )

    max_pages = min(max(1, int(max_pages)), ABSOLUTE_MAX_PAGES)
    progress(0.03, desc="入力を展開しています")
    try:
        run_dir, pages = prepare_document(
            source_path,
            dpi=int(dpi),
            selection=page_selection,
            max_pages=max_pages,
            max_input_mb=MAX_INPUT_MB,
        )
    except DocumentError as exc:
        raise gr.Error(str(exc)) from exc

    source_gallery = gallery_items(pages)
    options: dict[str, Any] = {
        "max_new_tokens": int(max_new_tokens),
        "layout_as_thought": bool(layout_as_thought),
        "image_mode": unlimited_image_mode,
    }

    responses: list[dict[str, Any]] = []
    elapsed_values: list[float] = []
    for index, page in enumerate(pages, start=1):
        progress(
            0.08 + 0.82 * ((index - 1) / len(pages)),
            desc=f"Page {page.page_number} / {pages[-1].page_number}: {spec.label}",
        )
        try:
            response = run_page(
                endpoint=endpoint,
                model_id=model_id,
                page_path=page.image_path,
                prompt=prompt,
                options=options,
                timeout_seconds=float(request_timeout),
            )
        except WorkerError as exc:
            raise gr.Error(f"Page {page.page_number} failed: {exc}") from exc
        responses.append(response)
        elapsed = response.get("metrics", {}).get("elapsed_seconds")
        if isinstance(elapsed, (int, float)):
            elapsed_values.append(float(elapsed))

    progress(0.94, desc="結果をまとめています")
    annotated, markdown, text, raw_json, archive = build_exports(
        run_dir=run_dir,
        model_id=model_id,
        model_label=spec.label,
        pages=pages,
        responses=responses,
    )
    total_elapsed = sum(elapsed_values)
    median_elapsed = statistics.median(elapsed_values) if elapsed_values else 0.0
    warning_count = sum(len(response.get("warnings", []) or []) for response in responses)
    status = (
        f"完了: {spec.label} / {len(pages)} page(s). "
        f"Worker inference total {total_elapsed:.1f}s, median {median_elapsed:.1f}s/page. "
        f"Warnings: {warning_count}."
    )
    progress(1.0, desc="完了")
    return source_gallery, annotated, markdown, text, raw_json, archive, status


model_choices = [(spec.label, spec.id) for spec in REGISTRY.values()]
default_model_id = next(iter(REGISTRY))
default_spec = REGISTRY[default_model_id]

CSS = """
#status-box textarea {font-family: ui-monospace, SFMono-Regular, Menlo, monospace;}
.result-gallery {min-height: 460px;}
"""

with gr.Blocks(title="OCR Model Workbench") as demo:
    gr.Markdown(
        "# OCR Model Workbench\n"
        "PDF・画像を共通UIから最新OCRモデルへ送り、ページ画像、可視化画像、Markdown、テキスト、JSONを比較します。"
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=5):
            input_file = gr.File(
                label="PDF / 画像",
                file_count="single",
                type="filepath",
                file_types=["image", ".pdf"],
            )
            model = gr.Dropdown(
                choices=model_choices,
                value=default_model_id,
                label="モデル",
            )
            prompt = gr.Textbox(
                value=default_spec.default_prompt,
                label="プロンプト（VLM系のみ）",
                lines=4,
            )
            with gr.Accordion("ページ・推論設定", open=False):
                page_selection = gr.Textbox(
                    label="PDFページ指定",
                    placeholder="空欄=先頭から。例: 1-3,5",
                    value="",
                )
                max_pages = gr.Slider(
                    minimum=1,
                    maximum=ABSOLUTE_MAX_PAGES,
                    step=1,
                    value=min(default_spec.default_max_pages, DEFAULT_MAX_PAGES, ABSOLUTE_MAX_PAGES),
                    label="最大ページ数",
                )
                dpi = gr.Slider(96, 300, value=default_spec.default_dpi, step=12, label="PDF rasterize DPI")
                max_new_tokens = gr.Slider(
                    256,
                    32768,
                    value=default_spec.default_max_tokens,
                    step=256,
                    label="生成上限（max_new_tokens / max_length）",
                )
                layout_as_thought = gr.Checkbox(
                    label="Qianfan Layout-as-Thought",
                    value=default_spec.default_layout_as_thought,
                )
                unlimited_image_mode = gr.Radio(
                    choices=["gundam", "base"],
                    value=default_spec.default_image_mode,
                    label="Unlimited-OCR image mode",
                )
                request_timeout = gr.Slider(
                    60,
                    1800,
                    value=default_spec.default_request_timeout,
                    step=30,
                    label="1ページのタイムアウト（秒）",
                )
            with gr.Accordion("Worker接続", open=False):
                endpoint_override = gr.Textbox(
                    value=default_spec.endpoint(),
                    label="Worker URL（空欄ならSpace Variableを使用）",
                    placeholder="https://username-space-name.hf.space",
                )
                health_button = gr.Button("Health check")
            run_button = gr.Button("OCRを実行", variant="primary")
            status = gr.Textbox(label="Status", interactive=False, elem_id="status-box")

        with gr.Column(scale=7):
            model_info = gr.Markdown(_model_info(default_spec))
            health_output = gr.Markdown()
            with gr.Tabs():
                with gr.Tab("入力ページ"):
                    source_gallery = gr.Gallery(
                        label="Source pages",
                        columns=2,
                        height=520,
                        elem_classes=["result-gallery"],
                    )
                with gr.Tab("結果画像"):
                    result_gallery = gr.Gallery(
                        label="Annotated / normalized result images",
                        columns=2,
                        height=520,
                        elem_classes=["result-gallery"],
                    )
                with gr.Tab("Markdown"):
                    markdown_output = gr.Markdown()
                with gr.Tab("Text"):
                    text_output = gr.Textbox(lines=28, buttons=["copy"])
                with gr.Tab("JSON"):
                    json_output = gr.Code(language="json", lines=28)
                with gr.Tab("Download"):
                    download_output = gr.File(label="全結果ZIP")

    model.change(
        fn=on_model_change,
        inputs=model,
        outputs=[
            prompt,
            endpoint_override,
            model_info,
            max_new_tokens,
            max_pages,
            dpi,
            request_timeout,
            layout_as_thought,
            unlimited_image_mode,
        ],
    )
    health_button.click(
        fn=health_check,
        inputs=[model, endpoint_override],
        outputs=health_output,
    )
    run_button.click(
        fn=run_ocr,
        inputs=[
            input_file,
            model,
            prompt,
            endpoint_override,
            page_selection,
            max_pages,
            dpi,
            max_new_tokens,
            layout_as_thought,
            unlimited_image_mode,
            request_timeout,
        ],
        outputs=[
            source_gallery,
            result_gallery,
            markdown_output,
            text_output,
            json_output,
            download_output,
            status,
        ],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2, max_size=12).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        css=CSS,
    )