Deploy OCR Model Workbench

app.py

@@ -61,13 +61,16 @@ def _model_info(spec: ModelSpec) -> str:
     )
 
 
-def on_model_change(model_id: str) -> tuple[str, str, str, int, bool, str]:
+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,
     )
@@ -221,10 +224,10 @@ with gr.Blocks(title="OCR Model Workbench") as demo:
                     minimum=1,
                     maximum=ABSOLUTE_MAX_PAGES,
                     step=1,
-                    value=min(DEFAULT_MAX_PAGES, ABSOLUTE_MAX_PAGES),
+                    value=min(default_spec.default_max_pages, DEFAULT_MAX_PAGES, ABSOLUTE_MAX_PAGES),
                     label="最大ページ数",
                 )
-                dpi = gr.Slider(96, 300, value=180, step=12, label="PDF rasterize DPI")
+                dpi = gr.Slider(96, 300, value=default_spec.default_dpi, step=12, label="PDF rasterize DPI")
                 max_new_tokens = gr.Slider(
                     256,
                     32768,
@@ -244,7 +247,7 @@ with gr.Blocks(title="OCR Model Workbench") as demo:
                 request_timeout = gr.Slider(
                     60,
                     1800,
-                    value=600,
+                    value=default_spec.default_request_timeout,
                     step=30,
                     label="1ページのタイムアウト（秒）",
                 )
@@ -293,6 +296,9 @@ with gr.Blocks(title="OCR Model Workbench") as demo:
             endpoint_override,
             model_info,
             max_new_tokens,
+            max_pages,
+            dpi,
+            request_timeout,
             layout_as_thought,
             unlimited_image_mode,
         ],

ocr_workbench/client.py

@@ -4,6 +4,7 @@ import json
 import os
 import inspect
 import time
+from concurrent.futures import TimeoutError
 from pathlib import Path
 from typing import Any
 
@@ -43,10 +44,13 @@ def check_health(endpoint: str, timeout_seconds: float = 20.0) -> dict[str, Any]
     if not endpoint:
         raise WorkerError("Worker URL is not configured.")
     try:
-        payload = _client(endpoint).predict(api_name="/health")
+        job = _client(endpoint).submit(api_name="/health")
+        payload = job.result(timeout=timeout_seconds)
         if not isinstance(payload, dict):
             raise WorkerError("Worker health response is not a JSON object.")
         return payload
+    except TimeoutError as exc:
+        raise WorkerError(f"Worker health check timed out after {timeout_seconds:.0f}s.") from exc
     except Exception as exc:
         raise WorkerError(f"Worker health check failed: {exc}") from exc
 
@@ -71,7 +75,7 @@ def run_page(
         if delay:
             time.sleep(delay)
         try:
-            payload = _client(endpoint).predict(
+            job = _client(endpoint).submit(
                 handle_file(str(page_path)),
                 model_id,
                 prompt or "",
@@ -79,6 +83,7 @@ def run_page(
                 os.getenv("WORKER_API_TOKEN", "").strip(),
                 api_name="/ocr",
             )
+            payload = job.result(timeout=timeout_seconds)
             if not isinstance(payload, dict):
                 raise WorkerError("Worker response is not a JSON object.")
             required = {"model", "text", "markdown", "metrics"}
@@ -86,6 +91,14 @@ def run_page(
             if missing:
                 raise WorkerError(f"Worker response is missing fields: {missing}")
             return payload
+        except TimeoutError as exc:
+            last_error = WorkerError(f"Worker request timed out after {timeout_seconds:.0f}s.")
+            try:
+                job.cancel()
+            except Exception:
+                pass
+            if attempt >= len(retry_delays):
+                break
         except (ValueError, OSError, WorkerError, Exception) as exc:
             last_error = exc
             if attempt >= len(retry_delays):

ocr_workbench/registry.py

@@ -15,6 +15,9 @@ class ModelSpec:
     description: str
     result_note: str
     default_max_tokens: int = 4096
+    default_max_pages: int = 8
+    default_dpi: int = 180
+    default_request_timeout: int = 600
     default_layout_as_thought: bool = False
     default_image_mode: str = "gundam"
 
@@ -32,13 +35,16 @@ _BUILTIN_MODELS: list[ModelSpec] = [
         id="paddleocr-vl-1.6",
         label="PaddleOCR-VL 1.6",
         endpoint_env="PADDLEOCR_VL_WORKER_URL",
-        default_prompt="Parse this document to Markdown.",
+        default_prompt="OCR:",
         description=(
-            "Compact document parser for text, layout, tables, formulas, charts and seals. "
-            "The Storage Bucket can be mounted on this worker as a persistent model/cache volume."
+            "Compact OCR/document VLM. This worker uses the Transformers/PyTorch backend so "
+            "inference runs inside ZeroGPU."
         ),
-        result_note="Returns PaddleOCR export images plus Markdown/JSON where available.",
-        default_max_tokens=4096,
+        result_note="Returns model text plus the normalized input image in ZeroGPU mode.",
+        default_max_tokens=512,
+        default_max_pages=1,
+        default_dpi=120,
+        default_request_timeout=900,
     ),
     ModelSpec(
         id="qianfan-ocr",
@@ -105,6 +111,9 @@ def _custom_models() -> list[ModelSpec]:
                 description=str(item.get("description", "Custom OCR worker.")),
                 result_note=str(item.get("result_note", "Uses the common OCR worker response schema.")),
                 default_max_tokens=int(item.get("default_max_tokens", 4096)),
+                default_max_pages=int(item.get("default_max_pages", 8)),
+                default_dpi=int(item.get("default_dpi", 180)),
+                default_request_timeout=int(item.get("default_request_timeout", 600)),
                 default_layout_as_thought=bool(item.get("default_layout_as_thought", False)),
                 default_image_mode=str(item.get("default_image_mode", "gundam")),
             )