"""PixText Studio: one image, editable OCR regions, real GPU refinement.""" from __future__ import annotations import hashlib import io import json import logging import math import os import shutil import subprocess import sys import tempfile import threading import time import urllib.request import uuid import zipfile from concurrent.futures import ThreadPoolExecutor from functools import lru_cache from pathlib import Path from fastapi import FastAPI, File, HTTPException, UploadFile from fastapi.responses import FileResponse, HTMLResponse from fastapi.staticfiles import StaticFiles from PIL import Image, ImageOps, UnidentifiedImageError from pydantic import BaseModel, Field ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(ROOT / "vendor")) RUNTIME = Path(os.environ.get("PIXTEXT_RUNTIME", tempfile.gettempdir())) / "pixtext-studio" RUNTIME.mkdir(parents=True, exist_ok=True) SESSION_TTL = 3600 MAX_UPLOAD = 16 * 1024 * 1024 Image.MAX_IMAGE_PIXELS = 16_777_216 POOL = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pixtext") LOCK = threading.RLock() SESSIONS: dict[str, dict] = {} JOBS: dict[str, dict] = {} PIPELINE = None PIPELINE_VARIANT = None EXAMPLES = { "english": {"label": "English editorial", "subtitle": "Four words, one poster"}, "chinese": {"label": "Chinese New Year", "subtitle": "Character-level corrections"}, } EXAMPLE_SETTINGS = { "variant": "meanflow", "steps": 4, "seed": 42, "guidance": 2.6, "text_threshold": 0.5, } if os.environ.get("PIXTEXT_GRADIO") == "1": from gradio import Server app = Server(title="PixText Studio", docs_url=None, redoc_url=None) else: app = FastAPI(title="PixText Studio", docs_url=None, redoc_url=None) app.mount("/static", StaticFiles(directory=ROOT / "static"), name="static") class ExampleFiles(StaticFiles): async def get_response(self, path, scope): response = await super().get_response(path, scope) if path.endswith("-preview.webp") and response.status_code in (200, 304): response.headers["Cache-Control"] = "public, max-age=31536000, immutable" return response app.mount("/examples", ExampleFiles(directory=ROOT / "examples"), name="examples") @lru_cache(maxsize=2) def example_preview(example: str): path = ROOT / "examples" / f"{example}-preview.webp" version = hashlib.sha256(path.read_bytes()).hexdigest()[:12] return f"/examples/{path.name}?v={version}" class ExampleRequest(BaseModel): example: str class Region(BaseModel): id: str = Field(min_length=1, max_length=80) bbox: tuple[float, float, float, float] text: str = Field(max_length=64) enabled: bool = True class RefineRequest(BaseModel): regions: list[Region] = Field(max_length=256) variant: str = "meanflow" steps: int = Field(default=4, ge=1, le=40) seed: int = Field(default=42, ge=0, le=2**32 - 1) guidance: float = Field(default=2.6, ge=0, le=10) class DetectRequest(BaseModel): text_threshold: float = Field(default=0.5, ge=0.1, le=0.95) link_threshold: float = Field(default=0.2, ge=0.1, le=0.95) def session_for(sid: str): with LOCK: value = SESSIONS.get(sid) if value is None: raise HTTPException(404, "This image session has expired. Please open the image again.") value["touched"] = time.monotonic() return value def make_session(image: Image.Image, example: str | None = None): if image.width < 16 or image.height < 16 or max(image.size) > 4096: raise HTTPException(422, "Use an image between 16 and 4096 pixels on each side.") if image.width * image.height > Image.MAX_IMAGE_PIXELS: raise HTTPException(422, "Use an image with no more than 16 megapixels.") with LOCK: for sid, value in list(SESSIONS.items()): active = value.get("job") and JOBS.get(value["job"], {}).get("status") in ( "queued", "running" ) if time.monotonic() - value["touched"] > SESSION_TTL and not active: shutil.rmtree(value["directory"], ignore_errors=True) SESSIONS.pop(sid) for jid in [k for k, j in JOBS.items() if j["session"] == sid]: JOBS.pop(jid) if len(SESSIONS) >= 64: raise HTTPException(503, "The studio is busy. Please try again shortly.") sid = uuid.uuid4().hex directory = RUNTIME / sid directory.mkdir() if example: shutil.copyfile(ROOT / "examples" / f"{example}.png", directory / "source.png") shutil.copyfile(ROOT / "examples" / f"{example}-preview.webp", directory / "preview.webp") else: image.save(directory / "source.png") image.save(directory / "preview.webp", quality=94, method=4) SESSIONS[sid] = { "directory": directory, "size": image.size, "example": example, "touched": time.monotonic(), "job": None, } return {"id": sid, "width": image.width, "height": image.height, "image": example_preview(example) if example else f"/api/sessions/{sid}/preview.webp", "example": example} def update_job(jid: str, **values): with LOCK: JOBS[jid].update(values) def submit(sid: str, kind: str, function, *args): session = session_for(sid) with LOCK: current = JOBS.get(session.get("job")) if current and current["status"] in ("queued", "running"): if current["kind"] == kind: return {"job_id": current["id"]} raise HTTPException(409, "Wait for the current operation to finish.") if sum(j["status"] in ("queued", "running") for j in JOBS.values()) >= 12: raise HTTPException(503, "The queue is full. Please try again shortly.") jid = uuid.uuid4().hex JOBS[jid] = {"id": jid, "session": sid, "kind": kind, "status": "queued", "progress": 0, "message": "Waiting for the model", "created": time.time()} session["job"] = jid def run(): started = time.perf_counter() update_job(jid, status="running", message="Preparing", progress=0.02) try: result = function(session, jid, *args) result["elapsed"] = round(time.perf_counter() - started, 2) update_job(jid, status="done", progress=1, message="Complete", result=result) except Exception as exc: logging.exception("PixText %s job failed", kind) if type(exc).__name__ == "OutOfMemoryError": message = "The GPU ran out of memory. Try refining fewer regions at a time." elif "CUDA GPU" in str(exc): message = "Refinement needs a CUDA GPU. Detection and box editing are ready to use." elif "checksum" in str(exc).lower(): message = "A model file failed its checksum check. Please retry or contact the Space owner." elif "RAQM" in str(exc): message = "Text rendering needs a Pillow build with RAQM. Please contact the Space owner." else: message = ("Detection" if kind == "detect" else "Refinement") + " could not finish. Please retry or contact the Space owner." update_job(jid, status="error", message=message) POOL.submit(run) return {"job_id": jid} def download_checked(url: str, path: Path, expected: str): if path.exists() and hashlib.sha256(path.read_bytes()).hexdigest() == expected: return path.parent.mkdir(parents=True, exist_ok=True) partial = path.with_suffix(path.suffix + ".partial") try: with urllib.request.urlopen(url, timeout=90) as response, partial.open("wb") as stream: shutil.copyfileobj(response, stream) if hashlib.sha256(partial.read_bytes()).hexdigest() != expected: raise RuntimeError("An OCR model checksum did not match. Please try the download again.") partial.replace(path) finally: partial.unlink(missing_ok=True) def ocr_root(jid: str): existing = os.environ.get("PIXTEXT_OCR_ROOT") if existing: return Path(existing) root = Path(os.environ.get("HF_HOME", str(Path.home() / ".cache/huggingface"))) / "pixtext_ocr" craft = root / "craft_mlt_25k.pth" craft_sha = "4a5efbfb48b4081100544e75e1e2b57f8de3d84f213004b14b85fd4b3748db17" if not craft.exists() or hashlib.sha256(craft.read_bytes()).hexdigest() != craft_sha: update_job(jid, message="Downloading CRAFT detection model", progress=0.06) root.mkdir(parents=True, exist_ok=True) archive = root / "craft.zip.partial" try: with urllib.request.urlopen("https://github.com/JaidedAI/EasyOCR/releases/download/pre-v1.1.6/craft_mlt_25k.zip", timeout=90) as response, archive.open("wb") as stream: shutil.copyfileobj(response, stream) with zipfile.ZipFile(archive) as zipped: entry = next(n for n in zipped.namelist() if Path(n).name == "craft_mlt_25k.pth") content = zipped.read(entry) if hashlib.sha256(content).hexdigest() != craft_sha: raise RuntimeError("The CRAFT checksum did not match. Please try again.") partial = craft.with_suffix(".partial") partial.write_bytes(content) partial.replace(craft) finally: archive.unlink(missing_ok=True) base = "https://huggingface.co/PaddlePaddle/PP-OCRv6_medium_rec_onnx/resolve/50c7eacafc52fa7bcf4194e8cd08e46f8558504b/" for name, sha in { "inference.onnx": "9c09abf0957f7968c7586464b7397b84ad2387a0497a351af40e9acc71b673ba", "inference.yml": "991b700facf5b50a7de193468207d5f4255b538dde0d312ae3b7c7a9b6873129", }.items(): update_job(jid, message="Preparing PP-OCRv6 recognition model", progress=0.12) download_checked(base + name, root / "PP-OCRv6_medium_rec_onnx" / name, sha) return root def detect(session, jid, request: DetectRequest): root = ocr_root(jid) update_job(jid, message="Detecting characters and words", progress=0.3) if os.environ.get("PIXTEXT_GRADIO") == "1": output = session["directory"] / f"{jid}-ocr.json" env = os.environ.copy() env.update({"SPACES_ZERO_GPU": "0", "CUDA_VISIBLE_DEVICES": "", "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1"}) completed = subprocess.run( [sys.executable, str(ROOT / "ocr_worker.py"), str(session["directory"] / "source.png"), str(root), str(output), str(request.text_threshold), str(request.link_threshold)], env=env, capture_output=True, text=True, timeout=120, ) if completed.returncode: logging.error("OCR worker failed: %s", completed.stderr[-4000:]) raise RuntimeError("OCR could not finish") result = json.loads(output.read_text()) else: import torch from pixtext_ocr import detect_words threads = int(os.environ.get("PIXTEXT_CPU_THREADS", "4")) if torch.get_num_threads() != threads: torch.set_num_threads(threads) with Image.open(session["directory"] / "source.png") as image: result = detect_words(image, root, request.text_threshold, request.link_threshold) return {"regions": result["regions"], "engine": result["engine"]} def refine(session, jid, request: RefineRequest, progress_report=None): global PIPELINE, PIPELINE_VARIANT import torch from pixtext import PixTextPipeline, TextEdit report = progress_report or (lambda **values: update_job(jid, **values)) if not torch.cuda.is_available(): raise RuntimeError("Refinement needs a CUDA GPU. Detection and box editing are ready to use.") threads = int(os.environ.get("PIXTEXT_CPU_THREADS", "4")) if torch.get_num_threads() != threads: torch.set_num_threads(threads) if os.environ.get("PIXTEXT_GPU_MEMORY_FRACTION"): torch.cuda.set_per_process_memory_fraction(float(os.environ["PIXTEXT_GPU_MEMORY_FRACTION"])) if PIPELINE_VARIANT != request.variant: if PIPELINE is not None: PIPELINE = None torch.cuda.empty_cache() name = "PixText-MeanFlow" if request.variant == "meanflow" else "PixText-Teacher" report(message=f"Loading {name}", progress=0.04) if os.environ.get("PIXTEXT_MODEL_ROOT"): package = Path(os.environ["PIXTEXT_MODEL_ROOT"]) / name else: from huggingface_hub import snapshot_download package = Path(snapshot_download("ATH-MaaS/PixText", allow_patterns=[f"models/{name}/*"], local_dir=RUNTIME / "model-packages")) / "models" / name PIPELINE = PixTextPipeline.from_pretrained(package, device="cuda", dtype=torch.bfloat16) PIPELINE_VARIANT = request.variant rows = [row for row in request.regions if row.enabled] edits = [TextEdit(tuple(row.bbox), row.text.strip()) for row in rows] directory = session["directory"] def progress(done, total): report(progress=0.15 + 0.8 * done / total, message=f"Refining ยท step {done} of {total}", current=done, total=total) result = PIPELINE.edit(directory / "source.png", edits, seed=request.seed, steps=request.steps, guidance_scale=request.guidance, sample_id="studio", progress_callback=progress) report(message="Preparing region comparisons", progress=0.97) result.image.save(directory / f"{jid}.png") result.image.save(directory / f"{jid}.webp", quality=94, method=4) crops = [] with Image.open(directory / "source.png") as before: for i, row in enumerate(rows): x0, y0, x1, y1 = row.bbox padding = max(8, (y1 - y0) * 0.2) box = (max(0, math.floor(x0 - padding)), max(0, math.floor(y0 - padding)), min(before.width, math.ceil(x1 + padding)), min(before.height, math.ceil(y1 + padding))) for label, image in (("before", before), ("after", result.image)): image.crop(box).save(directory / f"{jid}-{i}-{label}.webp", quality=96, method=4) crops.append({"id": row.id, "text": row.text, "bbox": row.bbox, "before": f"{jid}-{i}-before.webp", "after": f"{jid}-{i}-after.webp"}) metadata = {"regions": [row.model_dump() for row in rows], "model": request.variant, "steps": request.steps, "seed": request.seed, "guidance": request.guidance} (directory / f"{jid}.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2)) return {**metadata, "image": f"{jid}.png", "preview": f"{jid}.webp", "regions": crops, "metadata": f"{jid}.json"} @app.get("/") def index(): html = (ROOT / "static/index.html").read_text() script = (ROOT / "static/app.js").read_text() html = html.replace('', "") return HTMLResponse(html.replace("", "")) @app.get("/api/config") def config(): return {"examples": EXAMPLES, "models": ["meanflow", "teacher"], "transport": "gradio" if os.environ.get("PIXTEXT_GRADIO") == "1" else "http"} @app.post("/api/session") def example_session(request: ExampleRequest): if request.example not in EXAMPLES: raise HTTPException(404, "Example not found") with Image.open(ROOT / "examples" / f"{request.example}.png") as image: session = make_session(image, request.example) return {**session, **example_preset(request.example)} @app.post("/api/upload") async def upload(file: UploadFile = File(...)): content = await file.read(MAX_UPLOAD + 1) await file.close() if len(content) > MAX_UPLOAD: raise HTTPException(413, "Choose an image smaller than 16 MB.") try: with Image.open(io.BytesIO(content)) as opened: if opened.format not in ("PNG", "JPEG", "WEBP"): raise HTTPException(422, "Choose a PNG, JPEG or WebP image.") image = ImageOps.exif_transpose(opened).convert("RGB") return make_session(image) except (UnidentifiedImageError, OSError, Image.DecompressionBombError): raise HTTPException(422, "This image could not be opened. Try a PNG, JPEG or WebP file.") @app.get("/api/sessions/{sid}/targets") def example_targets(sid: str): session = session_for(sid) if not session["example"]: raise HTTPException(404, "No example targets for this image") return example_preset(session["example"]) def example_preset(example: str): rows = json.loads((ROOT / "examples" / f"{example}.json").read_text()) return {"regions": [{**row, "id": f"example-{i}", "enabled": True, "origin": "example", "recognized_text": "Example target"} for i, row in enumerate(rows)], "settings": dict(EXAMPLE_SETTINGS)} @app.post("/api/sessions/{sid}/detect") def start_detect(sid: str, request: DetectRequest): return submit(sid, "detect", detect, request) @app.post("/api/sessions/{sid}/refine") def start_refine(sid: str, request: RefineRequest): if os.environ.get("PIXTEXT_GRADIO") == "1": raise HTTPException(409, "Use the studio's queued refinement API.") validate_refine(sid, request) return submit(sid, "refine", refine, request) def validate_refine(sid: str, request: RefineRequest): session = session_for(sid) if request.variant not in ("meanflow", "teacher"): raise HTTPException(422, "Choose MeanFlow or Teacher.") selected = [r for r in request.regions if r.enabled] if not selected or len(selected) > 24: raise HTTPException(422, "Select between 1 and 24 regions to refine.") for region in selected: x0, y0, x1, y1 = region.bbox if not all(math.isfinite(v) for v in region.bbox) or not ( 0 <= x0 < x1 <= session["size"][0] and 0 <= y0 < y1 <= session["size"][1] ): raise HTTPException(422, "A box is outside the image. Adjust it before refining.") if not region.text.strip(): raise HTTPException(422, "Enter the target text for every selected region.") return session @app.get("/api/sessions/{sid}/job") def current_job(sid: str): session = session_for(sid) jid = session.get("job") if not jid: raise HTTPException(404, "No active job") return job(sid, jid) @app.get("/api/sessions/{sid}/jobs/{jid}") def job(sid: str, jid: str): session_for(sid) with LOCK: value = JOBS.get(jid) if value is None or value["session"] != sid: raise HTTPException(404, "Job not found") return {k: v for k, v in value.items() if k != "session"} @app.get("/api/sessions/{sid}/{filename}") def session_file(sid: str, filename: str): session = session_for(sid) # File names are server-generated. A session cannot address another session's data. if Path(filename).name != filename or not filename.endswith((".png", ".webp", ".json")): raise HTTPException(404, "File not found") path = session["directory"] / filename if not path.is_file(): raise HTTPException(404, "File not found") return FileResponse(path, headers={"Cache-Control": "private, no-store"})