from __future__ import annotations import ast import asyncio import importlib.metadata import json import queue import re import shutil import subprocess import sys import tempfile import threading import time import traceback import uuid import zipfile from pathlib import Path from typing import Any, Iterable, Iterator, Literal # A one-Python-file Space cannot use requirements.txt. Install only missing or # mismatched model dependencies before importing Torch/Transformers/Gradio. _RUNTIME_REQUIREMENTS = { "gradio": "gradio==6.17.3", "torch": "torch==2.10.0", "torchvision": "torchvision==0.25.0", "transformers": "transformers==4.57.1", "accelerate": "accelerate>=1.10,<2", "safetensors": "safetensors>=0.5", "Pillow": "Pillow==12.1.1", "PyMuPDF": "PyMuPDF==1.27.2.2", "einops": "einops==0.8.2", "addict": "addict==2.4.0", "easydict": "easydict==1.13", "matplotlib": "matplotlib==3.10.8", "psutil": "psutil==7.2.2", } _missing: list[str] = [] for _distribution, _requirement in _RUNTIME_REQUIREMENTS.items(): try: _version = importlib.metadata.version(_distribution) except importlib.metadata.PackageNotFoundError: _missing.append(_requirement) continue if "==" in _requirement and _version != _requirement.split("==", 1)[1]: _missing.append(_requirement) if _missing: print("Installing Space runtime dependencies:", " ".join(_missing), flush=True) subprocess.run( [ sys.executable, "-m", "pip", "install", "--quiet", "--upgrade", "--no-cache-dir", *_missing, ], check=True, ) import fitz import spaces import torch from fastapi import File, Form, HTTPException, UploadFile from fastapi.responses import FileResponse, HTMLResponse from gradio import Server from PIL import Image, ImageDraw, ImageFont from pydantic import BaseModel, Field from transformers import AutoModel, AutoTokenizer, TextStreamer MODEL_ID = "baidu/Unlimited-OCR" DEFAULT_PROMPT = "document parsing." DEFAULT_MULTI_PROMPT = "Multi page parsing." DEFAULT_MAX_TOKENS = 32000 DEFAULT_PDF_DPI = 200 DEFAULT_TIMEOUT = 1200 DEFAULT_BATCH_SIZE = 2 MAX_MULTI_IMAGES = 40 MAX_CONTEXT_LENGTH = 32768 NGRAM_SIZE = 35 NGRAM_WINDOW = 128 NGRAM_WINDOW_MULTI = 1024 IMAGE_MODE_GUNDAM = "gundam" IMAGE_MODE_BASE = "base" SUPPORTED_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} SUPPORTED_EXTENSIONS = SUPPORTED_IMAGE_EXTENSIONS | {".pdf"} THUMBNAIL_MAX_WIDTH = 1024 SCRIPT_DIR = Path(__file__).resolve().parent INDEX_PATH = SCRIPT_DIR / "index.html" SESSION_ROOT = Path(tempfile.gettempdir()) / "unlimited_ocr_space_sessions" SESSION_ROOT.mkdir(parents=True, exist_ok=True) _DET_RE = re.compile(r"<\|det\|>.*?<\|/det\|>", re.DOTALL) _REF_RE = re.compile(r"<\|/?ref\|>") _GROUNDING_RE = re.compile( r"<\|ref\|>(?P.*?)<\|/ref\|>" r".*?<\|det\|>(?P.*?)<\|/det\|>", re.DOTALL, ) _DET_ONLY_RE = re.compile(r"<\|det\|>(?P.*?)<\|/det\|>", re.DOTALL) _REF_DET_RE = re.compile( r"<\|ref\|>.*?<\|/ref\|>\s*<\|det\|>.*?<\|/det\|>", re.DOTALL ) _LAYOUT_ELEMENT_RE = re.compile( r"<\|ref\|>(?P.*?)<\|/ref\|>\s*" r"<\|det\|>.*?<\|/det\|>" r"(?P.*?)(?=<\|ref\|>|$)", re.DOTALL, ) _COORDS_RE = re.compile( r"\[\s*\[[^\]]+\](?:\s*,\s*\[[^\]]+\])*\s*\]" r"|\[\s*-?[\d.]+\s*,\s*-?[\d.]+\s*,\s*-?[\d.]+\s*,\s*-?[\d.]+\s*\]" ) LABEL_COLORS = [ (239, 68, 68), (37, 99, 235), (5, 150, 105), (147, 51, 234), (234, 88, 12), (219, 39, 119), (13, 148, 136), ] # Stable semantic colors across every page and request. The palette order no # longer depends on which label happens to appear first on a page. LABEL_COLOR_BY_TYPE: dict[str, tuple[int, int, int]] = { "text": (5, 150, 105), # green "title": (239, 68, 68), # red "header": (239, 68, 68), # red (det-only output variant) "sub_title": (239, 68, 68), # red "subtitle": (239, 68, 68), # red "list": (37, 99, 235), # blue "image_caption": (13, 148, 136), # teal; must precede image "caption": (13, 148, 136), # teal "image": (147, 51, 234), # purple "figure": (147, 51, 234), # purple "picture": (147, 51, 234), # purple "chart": (147, 51, 234), # purple "seal": (147, 51, 234), # purple "table": (234, 88, 12), # orange "formula": (219, 39, 119), # pink "equation": (219, 39, 119), # pink "footer": (100, 116, 139), # slate "page_number": (100, 116, 139), # slate "page_footnote": (148, 163, 184), # light slate "footnote": (148, 163, 184), # light slate "aside_text": (13, 148, 136), # teal } def _parse_coordinate_array(value: str) -> list[list[float]]: try: parsed = ast.literal_eval(value) except (SyntaxError, ValueError, TypeError): return [] if not isinstance(parsed, list): return [] # Accept both the usual `[[x1, y1, x2, y2]]` form and a single # `[x1, y1, x2, y2]` form emitted by a few model/runtime combinations. if len(parsed) == 4 and all(isinstance(item, (int, float)) for item in parsed): parsed = [parsed] boxes: list[list[float]] = [] for box in parsed: if not isinstance(box, (list, tuple)) or len(box) != 4: continue try: values = [float(item) for item in box] except (TypeError, ValueError): continue if not all(value == value and abs(value) != float("inf") for value in values): continue boxes.append(values) return boxes def _clip_normalized_box(box: Iterable[float]) -> list[int] | None: values = list(box) if len(values) != 4: return None x1, y1, x2, y2 = [max(0, min(1000, round(value))) for value in values] if x2 <= x1 or y2 <= y1: return None return [x1, y1, x2, y2] def parse_grounding(raw_output: str, width: int, height: int) -> list[dict[str, Any]]: """Parse Unlimited-OCR ``/`` spans into normalized and pixel boxes.""" grounding: list[dict[str, Any]] = [] consumed_det_spans: set[tuple[int, int]] = set() def add_item(label: str, coordinate_text: str, span: tuple[int, int]) -> None: coordinate_match = _COORDS_RE.search(coordinate_text or "") if not coordinate_match: return boxes = [_clip_normalized_box(box) for box in _parse_coordinate_array(coordinate_match.group(0))] boxes = [box for box in boxes if box is not None] if not boxes: return pixel_boxes = [ [ round(box[0] * width / 1000), round(box[1] * height / 1000), round(box[2] * width / 1000), round(box[3] * height / 1000), ] for box in boxes ] grounding.append( { "label": label.strip(), "boxes": boxes, "pixel_boxes": pixel_boxes, } ) consumed_det_spans.add(span) for match in _GROUNDING_RE.finditer(raw_output or ""): det_start = raw_output.find("<|det|>", match.start(), match.end()) det_end = raw_output.find("<|/det|>", det_start, match.end()) + len("<|/det|>") add_item(match.group("label"), match.group("coords"), (det_start, det_end)) for match in _DET_ONLY_RE.finditer(raw_output or ""): if any(match.start() >= start and match.end() <= end for start, end in consumed_det_spans): continue body = match.group("body") coords_match = _COORDS_RE.search(body) if not coords_match: continue label = body[: coords_match.start()].strip() add_item(label, coords_match.group(0), (match.start(), match.end())) return grounding def prepare_visualization_grounding( grounding: list[dict[str, Any]], ) -> list[dict[str, Any]]: """Remove list containers and relabel their nested text boxes as `list`.""" list_boxes = [ box for item in grounding if _grounding_label_key(item.get("label")).startswith("list") for box in item.get("boxes") or [] if len(box) == 4 ] def inside_list(box: list[float]) -> bool: if len(box) != 4: return False center_x = (box[0] + box[2]) / 2 center_y = (box[1] + box[3]) / 2 return any( outer[0] <= center_x <= outer[2] and outer[1] <= center_y <= outer[3] for outer in list_boxes ) prepared: list[dict[str, Any]] = [] by_label: dict[str, dict[str, Any]] = {} for item in grounding: source_label = str(item.get("label") or "box").strip() or "box" source_key = _grounding_label_key(source_label) if source_key.startswith("list"): # A list region is only a structural container. Its nested text # detections below are the boxes users actually need to inspect. continue boxes = item.get("boxes") or [] pixel_boxes = item.get("pixel_boxes") or [] for index, box in enumerate(boxes): if len(box) != 4 or index >= len(pixel_boxes) or len(pixel_boxes[index]) != 4: continue output_label = "list" if source_key.startswith("text") and inside_list(box) else source_label output_key = _grounding_label_key(output_label) target = by_label.get(output_key) if target is None: target = {"label": output_label, "boxes": [], "pixel_boxes": []} by_label[output_key] = target prepared.append(target) target["boxes"].append(box) target["pixel_boxes"].append(pixel_boxes[index]) return prepared def strip_grounding(raw_output: str) -> str: """Remove Unlimited-OCR layout labels and coordinate markup from markdown.""" # In Unlimited-OCR's document output, `<|ref|>title/text/image` is # a layout label for the following detection box. Keep the following OCR # content for normal labels, but intentionally drop image content. def replace_layout_element(match: re.Match[str]) -> str: label = match.group("label").strip().casefold() if label in {"image", "figure"}: return "\n\n" return f"\n\n{match.group('content')}" text = _LAYOUT_ELEMENT_RE.sub(replace_layout_element, raw_output or "") # Keep compatibility with det-only output variants. text = _REF_DET_RE.sub("\n\n", text) text = _DET_RE.sub("\n\n", text) text = _REF_RE.sub("", text) text = text.replace("<|grounding|>", "") # Some serving clients serialize newlines as the two literal characters # `\\n`; turn them back into real Markdown line breaks. text = text.replace("\\r\\n", "\n").replace("\\n", "\n") # Models sometimes wrap the complete answer in a Markdown code fence. # Remove only an outer fence so headings/tables render in the Preview tab. text = re.sub(r"^\s*```(?:markdown|md|text)?\s*\n", "", text, flags=re.IGNORECASE) text = re.sub(r"\n\s*```\s*$", "", text) return re.sub(r"\n{3,}", "\n\n", text).strip() def _grounding_label_key(label: Any) -> str: return re.sub(r"\s+", " ", str(label or "box").strip()).casefold() or "box" def _grounding_color_map(grounding: list[dict[str, Any]]) -> dict[str, tuple[int, int, int]]: """Assign a stable semantic color independent of per-page label order.""" colors: dict[str, tuple[int, int, int]] = {} for item in grounding: key = _grounding_label_key(item.get("label")) if key not in colors: # Prefer an exact match. This matters for `image_caption`, which # must not be captured by the more general `image` category. semantic_key = key if key in LABEL_COLOR_BY_TYPE else next( (known for known in LABEL_COLOR_BY_TYPE if key.startswith(f"{known} ")), None, ) if semantic_key: colors[key] = LABEL_COLOR_BY_TYPE[semantic_key] else: # Stable fallback for unexpected labels; unlike hash(), this # value is reproducible across processes and pages. checksum = sum((index + 1) * ord(char) for index, char in enumerate(key)) colors[key] = LABEL_COLORS[checksum % len(LABEL_COLORS)] return colors def draw_grounding(image: Image.Image, grounding: list[dict[str, Any]]) -> Image.Image: """Draw parsed pixel boxes over a page image with a visible translucent fill.""" base = image.convert("RGBA") overlay = Image.new("RGBA", base.size, (0, 0, 0, 0)) overlay_draw = ImageDraw.Draw(overlay) line_width = max(3, min(8, round(min(base.size) / 450))) colors = _grounding_color_map(grounding) try: font = ImageFont.load_default(size=13) except TypeError: font = ImageFont.load_default() boxes_to_draw: list[tuple[tuple[int, int, int], str, tuple[int, int, int, int]]] = [] for item in grounding: color = colors[_grounding_label_key(item.get("label"))] label = str(item.get("label") or "box")[:32] for box in item.get("pixel_boxes") or []: if len(box) != 4: continue x1, y1, x2, y2 = [int(value) for value in box] boxes_to_draw.append((color, label, (x1, y1, x2, y2))) # Draw every fill first so a later overlapping box cannot cover an earlier # border. Then draw a white separator and the semantic-color border. for color, _, box in boxes_to_draw: overlay_draw.rectangle(box, fill=(*color, 48)) for color, _, (x1, y1, x2, y2) in boxes_to_draw: overlay_draw.rectangle( (x1, y1, x2, y2), outline=(255, 255, 255, 245), width=line_width + 2, ) if x2 - x1 > 2 and y2 - y1 > 2: overlay_draw.rectangle( (x1 + 1, y1 + 1, x2 - 1, y2 - 1), outline=(*color, 255), width=line_width, ) output = Image.alpha_composite(base, overlay).convert("RGB") draw = ImageDraw.Draw(output) for color, label, (x1, y1, _, _) in boxes_to_draw: if label: text_bbox = draw.textbbox((x1 + 3, max(0, y1 - 18)), label, font=font) draw.rectangle(text_bbox, fill=color) draw.text((x1 + 3, max(0, y1 - 18)), label, fill="white", font=font) return output def split_multi_page_output(raw_output: str, expected_pages: int) -> list[str]: """Split the model's multi-page response when it emits `` separators.""" if expected_pages <= 1: return [raw_output] parts = [part.strip() for part in re.split(r"\s*\s*", raw_output or "")] parts = [part for part in parts if part] return parts if len(parts) == expected_pages else [] def normalize_page_range(page_from: int | float | None, page_to: int | float | None, total: int) -> tuple[int, int]: if total <= 0: raise ValueError("PDF has no pages") first = max(1, min(int(page_from or 1), total)) last = int(page_to or 0) last = total if last <= 0 else min(last, total) if first > last: raise ValueError(f"Invalid page range: {first} > {last}") return first, last def _save_page_result( session_dir: Path, page: dict[str, Any], raw_output: str, latency: float, warning: str | None = None, ) -> dict[str, Any]: page_index = int(page["index"]) image_path = session_dir / page["filename"] with Image.open(image_path) as opened: image = opened.convert("RGB") grounding = prepare_visualization_grounding( parse_grounding(raw_output, image.width, image.height) ) markdown = strip_grounding(raw_output).replace("", "\n\n---\n\n") prefix = f"page_{page_index + 1:04d}" raw_path = session_dir / f"{prefix}_raw.txt" markdown_path = session_dir / f"{prefix}.md" grounding_path = session_dir / f"{prefix}_grounding.json" raw_path.write_text(raw_output, encoding="utf-8") markdown_path.write_text(markdown, encoding="utf-8") _write_json(grounding_path, grounding) result = { "page_index": page_index, "page_number": page["page_number"], "source_page_number": page.get("source_page_number", page["page_number"]), "source_name": page.get("source_name", ""), "width": image.width, "height": image.height, "raw_output": raw_output, "markdown": markdown, "grounding": grounding, "latency_seconds": round(latency, 3), "warning": warning, } image.close() return result def _ensure_grounding_overlays(session_dir: Path) -> None: """Create downloadable overlays lazily instead of pausing between batches.""" for grounding_path in sorted(session_dir.glob("page_*_grounding.json")): prefix = grounding_path.name.removesuffix("_grounding.json") source_path = session_dir / f"{prefix}.png" overlay_path = session_dir / f"{prefix}_grounding.png" if not source_path.is_file(): continue if overlay_path.is_file() and overlay_path.stat().st_mtime_ns >= max( source_path.stat().st_mtime_ns, grounding_path.stat().st_mtime_ns ): continue grounding = json.loads(grounding_path.read_text(encoding="utf-8")) with Image.open(source_path) as opened: image = opened.convert("RGB") overlay = draw_grounding(image, grounding) try: overlay.save(overlay_path, format="PNG") finally: overlay.close() image.close() def _create_results_zip(session_dir: Path) -> Path: _ensure_grounding_overlays(session_dir) zip_path = session_dir / "unlimited_ocr_streaming_results.zip" with zipfile.ZipFile(zip_path, "w") as archive: for path in sorted(session_dir.iterdir()): if ( path == zip_path or not path.is_file() or path.name.startswith("source_") or path.name.endswith("_thumb.jpg") ): continue # PNG/JPEG/WebP are already compressed. Deflating them again adds a # noticeable CPU pause at the end of long documents for almost no # size benefit. compression = ( zipfile.ZIP_STORED if path.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"} else zipfile.ZIP_DEFLATED ) archive.write(path, path.name, compress_type=compression) return zip_path def _materialize_source_pages( session_dir: Path, source_path: Path, source_name: str, extension: str, start_index: int, pdf_dpi: int, page_from: int, page_to: int, ) -> list[dict[str, Any]]: """Render/save one source incrementally, keeping at most one decoded page.""" pages: list[dict[str, Any]] = [] def save_page(source_page_number: int, image: Image.Image) -> None: page_index = start_index + len(pages) filename = f"page_{page_index + 1:04d}.png" thumbnail_filename = f"page_{page_index + 1:04d}_thumb.jpg" try: image.save(session_dir / filename, format="PNG") thumbnail = image.copy() try: thumbnail.thumbnail( (THUMBNAIL_MAX_WIDTH, max(1, image.height)), Image.Resampling.LANCZOS, ) thumbnail.save( session_dir / thumbnail_filename, format="JPEG", quality=92, subsampling=0, optimize=True, ) finally: thumbnail.close() pages.append( { "index": page_index, "page_number": page_index + 1, "source_page_number": source_page_number, "source_name": source_name, "filename": filename, "width": image.width, "height": image.height, "image_url": f"/api/session/{session_dir.name}/page/{page_index}/image", "thumbnail_url": f"/api/session/{session_dir.name}/page/{page_index}/thumbnail", "thumbnail_filename": thumbnail_filename, } ) finally: image.close() if extension == ".pdf": import fitz with fitz.open(source_path) as document: first, last = normalize_page_range(page_from, page_to, len(document)) matrix = fitz.Matrix(float(pdf_dpi) / 72.0, float(pdf_dpi) / 72.0) for source_page_number in range(first, last + 1): pixmap = document[source_page_number - 1].get_pixmap( matrix=matrix, alpha=False ) image = Image.frombytes( "RGB", (pixmap.width, pixmap.height), pixmap.samples ) save_page(source_page_number, image) else: with Image.open(source_path) as opened: save_page(1, opened.convert("RGB")) return pages class StreamSettings(BaseModel): session_id: str request_id: str = Field(min_length=1, max_length=100) page_indices: list[int] | None = None mode: Literal["gundam", "base"] = IMAGE_MODE_BASE prompt: str = DEFAULT_MULTI_PROMPT api_base: str = "" model_name: str = MODEL_ID max_tokens: int = Field(default=DEFAULT_MAX_TOKENS, ge=1, le=MAX_CONTEXT_LENGTH) request_timeout: int = Field(default=DEFAULT_TIMEOUT, ge=10, le=7200) batch_size: int = Field(default=DEFAULT_BATCH_SIZE, ge=1, le=MAX_MULTI_IMAGES) ngram_enabled: bool = True def normalize_prompt(prompt: str | None, multi_image: bool) -> str: wanted_default = DEFAULT_MULTI_PROMPT if multi_image else DEFAULT_PROMPT value = re.sub(r"^(?:\s*)+", "", (prompt or "").strip()) if not value or value in {DEFAULT_PROMPT, DEFAULT_MULTI_PROMPT}: value = wanted_default # The official infer_multi implementation expects one image marker for the # complete image list. return f"{value}" def _session_dir(session_id: str) -> Path: try: normalized = str(uuid.UUID(session_id)) except ValueError as exc: raise HTTPException(status_code=400, detail="Invalid session ID") from exc path = SESSION_ROOT / normalized if not path.is_dir(): raise HTTPException(status_code=404, detail="Session not found") return path def _read_metadata(session_dir: Path) -> dict[str, Any]: path = session_dir / "document.json" if not path.is_file(): raise HTTPException(status_code=404, detail="Session metadata not found") return json.loads(path.read_text(encoding="utf-8")) def _write_json(path: Path, value: Any) -> None: path.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8") class GenerationCancelled(RuntimeError): pass print(f"Loading {MODEL_ID} tokenizer from Hugging Face Hub...", flush=True) TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) print(f"Loading {MODEL_ID} with ZeroGPU CUDA emulation...", flush=True) MODEL = AutoModel.from_pretrained( MODEL_ID, trust_remote_code=True, use_safetensors=True, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, ).eval().cuda() torch.set_grad_enabled(False) print("Unlimited-OCR model ready.", flush=True) app = Server(debug=True, title="Unlimited-OCR Hugging Face Space") # Hugging Face's Gradio hot-reload launcher looks for the conventional `demo` # symbol. Keep both names pointing at the same Server instance. demo = app _MODEL_LOCK = threading.Lock() _THREAD_CONTEXT = threading.local() _ACTIVE_REQUESTS: dict[str, threading.Event] = {} _SAFE_REQUEST_ID = re.compile(r"^[A-Za-z0-9_-]{1,100}$") class QueueTextStreamer(TextStreamer): def __init__(self, tokenizer: Any, **kwargs: Any) -> None: self._sink = getattr(_THREAD_CONTEXT, "sink", None) self._cancel = getattr(_THREAD_CONTEXT, "cancel", None) self._eos_text = tokenizer.decode( [tokenizer.eos_token_id], skip_special_tokens=False ) kwargs.pop("interval", None) super().__init__(tokenizer, **kwargs) def put(self, value: Any) -> None: if self._cancel is not None and self._cancel.is_set(): raise GenerationCancelled("Generation stopped") super().put(value) def on_finalized_text(self, text: str, stream_end: bool = False) -> None: if self._eos_text: text = text.replace(self._eos_text, "\n") if text and self._sink is not None: self._sink(text) # Unlimited-OCR remote code creates TPSTextStreamer internally. Replacing that # module symbol gives true token streaming without redirecting process stdout. _MODEL_MODULE = sys.modules[MODEL.__class__.__module__] _MODEL_MODULE.TPSTextStreamer = QueueTextStreamer def _run_model_worker( image_paths: list[Path], settings: StreamSettings, cancel: threading.Event, emit: Any, ) -> None: with _MODEL_LOCK: if cancel.is_set(): emit("cancelled", None) emit("done", None) return output_dir = tempfile.mkdtemp(prefix="unlimited_ocr_output_") multi_image = len(image_paths) > 1 _THREAD_CONTEXT.sink = lambda text: emit("chunk", text) _THREAD_CONTEXT.cancel = cancel try: common = { "prompt": normalize_prompt(settings.prompt, multi_image), "output_path": output_dir, "save_results": False, "max_length": min(int(settings.max_tokens), MAX_CONTEXT_LENGTH), "tps_interval": 0, "no_repeat_ngram_size": NGRAM_SIZE if settings.ngram_enabled else 0, "ngram_window": ( NGRAM_WINDOW_MULTI if multi_image else NGRAM_WINDOW ) if settings.ngram_enabled else 0, "temperature": 0.0, } if multi_image: output, _ = MODEL.infer_multi( TOKENIZER, image_files=[str(path) for path in image_paths], image_size=1024, **common, ) emit("final", output) else: is_gundam = settings.mode == IMAGE_MODE_GUNDAM MODEL.infer( TOKENIZER, image_file=str(image_paths[0]), base_size=1024, image_size=640 if is_gundam else 1024, crop_mode=is_gundam, **common, ) except GenerationCancelled: emit("cancelled", None) except BaseException as exc: error_traceback = traceback.format_exc() print( f"[OCR ERROR] request={settings.request_id} " f"pages={settings.page_indices}: {exc}\n{error_traceback}", flush=True, ) emit( "error", { "message": str(exc) or exc.__class__.__name__, "exception_type": exc.__class__.__name__, }, ) finally: ring_window = getattr(MODEL.config, "_ring_window", None) if ring_window is not None: MODEL.config.sliding_window = ring_window for name in ("sink", "cancel"): if hasattr(_THREAD_CONTEXT, name): delattr(_THREAD_CONTEXT, name) shutil.rmtree(output_dir, ignore_errors=True) emit("done", None) @app.api( name="run_ocr_batch", stream_every=0.1, concurrency_limit=1, concurrency_id="unlimited-ocr-gpu", ) @spaces.GPU(duration=60) def run_ocr_batch( session_id: str, page_indices: list[int], mode: str = IMAGE_MODE_BASE, prompt: str = DEFAULT_MULTI_PROMPT, max_tokens: int = DEFAULT_MAX_TOKENS, ngram_enabled: bool = True, request_id: str = "space-request", batch_number: int = 1, total_batches: int = 1, request_timeout: int = DEFAULT_TIMEOUT, ) -> Iterator[dict[str, Any]]: """Run one ZeroGPU batch and stream UI events through Gradio's SSE queue.""" print( f"[OCR START] request={request_id} batch={batch_number}/{total_batches} " f"session={session_id} pages={page_indices} mode={mode} " f"max_tokens={max_tokens}", flush=True, ) session_dir = _session_dir(session_id) pages = _read_metadata(session_dir).get("pages") or [] if not _SAFE_REQUEST_ID.fullmatch(request_id): yield {"event": "error", "data": {"message": "Invalid request ID"}} return if ( not page_indices or len(page_indices) > MAX_MULTI_IMAGES or any(index < 0 or index >= len(pages) for index in page_indices) ): yield {"event": "error", "data": {"message": "Page index out of range"}} return batch = [pages[index] for index in page_indices] paths = [session_dir / page["filename"] for page in batch] settings = StreamSettings( session_id=session_id, request_id=request_id, page_indices=page_indices, mode=mode, prompt=prompt, max_tokens=max_tokens, ngram_enabled=ngram_enabled, request_timeout=request_timeout, ) cancel = threading.Event() _ACTIVE_REQUESTS[request_id] = cancel yield { "event": "batch_started", "data": { "batch": batch_number, "batches": total_batches, "page_indices": page_indices, "prompt": normalize_prompt(prompt, len(paths) > 1), }, } events: queue.Queue[tuple[str, Any]] = queue.Queue() worker = threading.Thread( target=_run_model_worker, args=(paths, settings, cancel, lambda kind, value: events.put((kind, value))), daemon=True, name=f"unlimited-ocr-{request_id}", ) chunks: list[str] = [] final_output = "" characters = 0 started = time.perf_counter() worker.start() try: while worker.is_alive() or not events.empty(): if time.perf_counter() - started >= request_timeout: cancel.set() yield { "event": "error", "data": {"message": f"Inference timed out after {request_timeout}s"}, } return try: kind, value = events.get(timeout=0.05) except queue.Empty: continue if kind == "chunk": delta = str(value) chunks.append(delta) characters += len(delta) yield { "event": "delta", "data": { "batch": batch_number, "page_indices": page_indices, "delta": delta, "characters": characters, }, } elif kind == "final": final_output = str(value or "") elif kind == "error": if isinstance(value, dict): error_message = value.get("message") or "Unknown inference error" exception_type = value.get("exception_type") or "Error" else: error_message = str(value) or "Unknown inference error" exception_type = type(value).__name__ yield { "event": "error", "data": { "message": ( f"Hugging Face inference failed " f"({exception_type}): {error_message}" ), "request_id": request_id, "batch": batch_number, }, } return elif kind == "cancelled": yield {"event": "cancelled", "data": {"message": "Generation stopped"}} return raw_output = "".join(chunks) or final_output latency = time.perf_counter() - started sections = split_multi_page_output(raw_output, len(batch)) warning = None if len(batch) > 1 and not sections: warning = ( "Model did not return one section per input page. " "The combined response is assigned only to the first page." ) pairs = [(batch[0], raw_output)] else: pairs = list(zip(batch, sections or [raw_output], strict=True)) results_path = session_dir / "results.json" if results_path.is_file(): try: results_index = json.loads(results_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): results_index = {} else: results_index = {} for page, page_output in pairs: result = _save_page_result( session_dir, page, page_output, latency, warning ) persisted = dict(result) persisted.pop("raw_output", None) results_index[str(result["page_index"])] = persisted client_result = dict(result) client_result.pop("raw_output", None) yield {"event": "page_result", "data": client_result} _write_json(results_path, results_index) yield { "event": "batch_done", "data": { "batch": batch_number, "batches": total_batches, "latency_seconds": round(latency, 3), "mapped_pages": len(pairs), }, } print( f"[OCR DONE] request={request_id} batch={batch_number}/{total_batches} " f"characters={len(raw_output)} mapped_pages={len(pairs)} " f"latency={latency:.3f}s", flush=True, ) finally: cancel.set() if worker.is_alive(): worker.join(timeout=5) _ACTIVE_REQUESTS.pop(request_id, None) # Direct FastAPI inference is intentionally absent: every CUDA call must run # inside the Gradio API function decorated with spaces.GPU above. @app.get("/", response_class=HTMLResponse) async def homepage() -> HTMLResponse: if not INDEX_PATH.is_file(): raise HTTPException(status_code=500, detail="index.html is missing") defaults = json.dumps({ "apiBase": "", "modelName": MODEL_ID, "maxTokens": DEFAULT_MAX_TOKENS, "pdfDpi": DEFAULT_PDF_DPI, "timeout": DEFAULT_TIMEOUT, "batchSize": DEFAULT_BATCH_SIZE, "backend": "huggingface", }) return HTMLResponse( INDEX_PATH.read_text(encoding="utf-8").replace("__APP_DEFAULTS__", defaults) ) @app.get("/api/hf/status") async def hf_status() -> dict[str, Any]: return { "backend": "huggingface", "ready": True, "model": MODEL_ID, "device": "ZeroGPU (allocated per queued API call)", } @app.post("/api/document") async def upload_document( files: list[UploadFile] = File(...), pdf_dpi: int = Form(DEFAULT_PDF_DPI), page_from: int = Form(1), page_to: int = Form(0), ) -> dict[str, Any]: if not files: raise HTTPException(status_code=400, detail="No files uploaded") if not 72 <= pdf_dpi <= 600: raise HTTPException(status_code=400, detail="PDF DPI must be between 72 and 600") session_id = str(uuid.uuid4()) session_dir = SESSION_ROOT / session_id session_dir.mkdir(parents=True, exist_ok=False) pages: list[dict[str, Any]] = [] try: for source_index, upload in enumerate(files): source_name = Path(upload.filename or f"upload_{source_index + 1}").name extension = Path(source_name).suffix.lower() if extension not in SUPPORTED_EXTENSIONS: raise ValueError(f"Unsupported file type: {extension or 'unknown'}") source_path = session_dir / f"source_{source_index + 1:03d}{extension}" with source_path.open("wb") as destination: while chunk := await upload.read(1024 * 1024): destination.write(chunk) pages.extend(await asyncio.to_thread( _materialize_source_pages, session_dir, source_path, source_name, extension, len(pages), pdf_dpi, page_from, page_to, )) metadata = {"session_id": session_id, "pages": pages} _write_json(session_dir / "document.json", metadata) return metadata except Exception as exc: shutil.rmtree(session_dir, ignore_errors=True) raise HTTPException(status_code=400, detail=str(exc)) from exc @app.get("/api/session/{session_id}/page/{page_index}/image") async def page_image(session_id: str, page_index: int) -> FileResponse: session_dir = _session_dir(session_id) pages = _read_metadata(session_dir).get("pages") or [] if not 0 <= page_index < len(pages): raise HTTPException(status_code=404, detail="Page not found") return FileResponse(session_dir / pages[page_index]["filename"], media_type="image/png") @app.get("/api/session/{session_id}/page/{page_index}/thumbnail") async def page_thumbnail(session_id: str, page_index: int) -> FileResponse: session_dir = _session_dir(session_id) pages = _read_metadata(session_dir).get("pages") or [] if not 0 <= page_index < len(pages): raise HTTPException(status_code=404, detail="Page not found") page = pages[page_index] path = session_dir / page["thumbnail_filename"] return FileResponse(path, media_type="image/jpeg") @app.post("/api/cancel/{request_id}") async def cancel(request_id: str) -> dict[str, bool]: event = _ACTIVE_REQUESTS.get(request_id) if event: event.set() return {"cancelled": bool(event)} @app.get("/api/session/{session_id}/download") async def download(session_id: str) -> FileResponse: session_dir = _session_dir(session_id) archive = await asyncio.to_thread(_create_results_zip, session_dir) return FileResponse( archive, media_type="application/zip", filename="unlimited_ocr_streaming_results.zip", ) @app.delete("/api/session/{session_id}") async def delete_session(session_id: str) -> dict[str, bool]: shutil.rmtree(_session_dir(session_id)) return {"deleted": True} if __name__ == "__main__": app.launch( server_name="0.0.0.0", show_error=True, max_file_size="1gb", allowed_paths=[str(SESSION_ROOT)], )