Launch PixText Studio with OCR, editable regions and ZeroGPU refinement

.dockerignore

@@ -0,0 +1,9 @@
+.git
+.venv
+__pycache__
+*.pyc
+.pytest_cache
+.ruff_cache
+tests
+runtime
+.env*

.gitignore

@@ -0,0 +1,11 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+.venv/
+.env
+.env.*
+runtime/
+deps/
+*.key
+*.pem

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.12-slim
+RUN apt-get update && apt-get install -y --no-install-recommends build-essential libgl1 libglib2.0-0 libraqm0 && rm -rf /var/lib/apt/lists/*
+RUN useradd -m -u 1000 user
+WORKDIR /home/user/app
+COPY requirements.txt ./
+RUN pip install --no-cache-dir -r requirements.txt
+COPY --chown=user:user . .
+USER user
+ENV HOME=/home/user PYTHONUNBUFFERED=1 PYTHONPATH=/home/user/app/vendor HF_HOME=/home/user/.cache/huggingface
+EXPOSE 7860
+CMD ["python", "-m", "uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 PixText Authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,15 +1,36 @@
 ---
 title: PixText Studio
-emoji: 👀
-colorFrom: gray
-colorTo: purple
+emoji: ✍️
+colorFrom: yellow
+colorTo: red
 sdk: gradio
 sdk_version: 6.30.0
-python_version: '3.12'
-app_file: app.py
-pinned: false
+python_version: 3.12
+app_file: space.py
 license: apache-2.0
-short_description: Detect and refine text in images.
+models:
+  - ATH-MaaS/PixText
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PixText Studio
+
+Detect words, adjust their boxes, enter the text you want, and compare the repaired regions. The interface uses the same CRAFT + PP-OCRv6 detection and PixText inference code as the ComfyUI nodes.
+
+Upload an image or choose the English and Chinese examples. Select only the regions you want to repair. You can move or resize any box, or draw a new one. MeanFlow defaults to four steps; Teacher defaults to 25. The result includes an image slider, matching region crops, a PNG download and editable region JSON.
+
+## Run
+
+Deploy this directory as a Gradio Space with ZeroGPU, or run locally on an Ampere-or-newer CUDA GPU:
+
+```sh
+python -m pip install -r requirements.txt
+python space.py
+```
+
+Model packages download from [ATH-MaaS/PixText](https://huggingface.co/ATH-MaaS/PixText) when first used. The optional OCR files download on the first detection run and are checked against release SHA-256 hashes. For an existing installation, set `PIXTEXT_MODEL_ROOT` to the directory containing the `PixText-MeanFlow` and `PixText-Teacher` packages and `PIXTEXT_OCR_ROOT` to the OCR model directory. Both variables are optional.
+
+Jobs are queued one at a time. Images are stored in isolated temporary sessions and removed after one hour of inactivity, when the next session is created. Nothing is uploaded to a dataset or sent to an external inference service.
+
+The vendored inference and OCR code is from [PixText commit 7779419](https://github.com/yaoliliu/PixText/tree/7779419de634d9785ec23558942480db713f0f99). Model packages retain their manifests, font notices and SHA-256 checks. CRAFT weights come from [EasyOCR](https://github.com/JaidedAI/EasyOCR); recognition weights and configuration come from [PaddlePaddle](https://huggingface.co/PaddlePaddle/PP-OCRv6_medium_rec_onnx).
+
+The Space uses Gradio Server to serve the custom editor and queue GPU requests. On ZeroGPU, PyTorch SDPA uses the original PixText token-routing mask without runtime compilation. A Dockerfile is also included for dedicated GPU hosting.

app.py

@@ -0,0 +1,371 @@
+"""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 sys
+import tempfile
+import threading
+import time
+import urllib.request
+import uuid
+import zipfile
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+from fastapi import FastAPI, File, HTTPException, UploadFile
+from fastapi.responses import FileResponse
+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"},
+}
+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")
+app.mount("/examples", StaticFiles(directory=ROOT / "examples"), name="examples")
+
+
+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=3.3, 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()
+        image.save(directory / "source.png")
+        SESSIONS[sid] = {
+            "directory": directory, "size": image.size, "example": example,
+            "touched": time.monotonic(), "job": None,
+        }
+    return {"id": sid, "width": image.width, "height": image.height,
+            "image": f"/api/sessions/{sid}/source.png", "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):
+    import torch
+    from pixtext_ocr import detect_words
+
+    torch.set_num_threads(4)
+    root = ocr_root(jid)
+    update_job(jid, message="Detecting words and reading text", progress=0.3)
+    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.")
+    torch.set_num_threads(4)
+    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}/*"])) / "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")
+    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}.png")
+            crops.append({"id": row.id, "text": row.text, "bbox": row.bbox,
+                          "before": f"{jid}-{i}-before.png", "after": f"{jid}-{i}-after.png"})
+    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", "regions": crops, "metadata": f"{jid}.json"}
+
+
+@app.get("/")
+def index():
+    return FileResponse(ROOT / "static/index.html")
+
+
+@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:
+        return make_session(image.convert("RGB"), 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")
+    rows = json.loads((ROOT / "examples" / f"{session['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)]}
+
+
+@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):
+    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}/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", ".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"})

packages.txt

@@ -0,0 +1,3 @@
+libgl1
+libglib2.0-0
+libraqm0

pixtext_ocr.py

@@ -0,0 +1,183 @@
+"""Local, optional CRAFT word detection and PP-OCRv6 CTC recognition.
+
+Models are installed explicitly; inference never downloads files or calls an OCR service.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from functools import lru_cache
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+
+def model_files(root: Path) -> tuple[Path, Path, Path]:
+    paths = (
+        root / "craft_mlt_25k.pth",
+        root / "PP-OCRv6_medium_rec_onnx/inference.onnx",
+        root / "PP-OCRv6_medium_rec_onnx/inference.yml",
+    )
+    missing = [str(p.name) for p in paths if not p.is_file()]
+    if missing:
+        raise ValueError(
+            "Install the optional OCR models in ComfyUI/models/pixtext_ocr. "
+            "Download links and paths are in the English workflow's OCR setup note. Missing: "
+            + ", ".join(missing)
+        )
+    return paths
+
+
+def model_identity(root: Path):
+    return tuple((str(p), p.stat().st_size, p.stat().st_mtime_ns) for p in model_files(root))
+
+
+def decode_ctc(probabilities: np.ndarray, vocabulary: list[str], valid_steps: int):
+    """Collapse repeats before removing blanks, preserving repeated letters across blanks."""
+    sequence = probabilities[:valid_steps]
+    ids, scores = sequence.argmax(-1), sequence.max(-1)
+    text, confidence, previous = [], [], -1
+    for index, score in zip(ids, scores, strict=True):
+        index = int(index)
+        if index and index != previous:
+            if index >= len(vocabulary):
+                raise ValueError("OCR dictionary does not match the recognition model")
+            text.append(vocabulary[index])
+            confidence.append(float(score))
+        previous = index
+    return "".join(text), float(np.mean(confidence)) if confidence else 0.0
+
+
+def crop_quad(rgb: np.ndarray, quad: np.ndarray):
+    import cv2
+
+    width = max(2, round(max(np.linalg.norm(quad[1] - quad[0]), np.linalg.norm(quad[2] - quad[3]))))
+    height = max(
+        2, round(max(np.linalg.norm(quad[3] - quad[0]), np.linalg.norm(quad[2] - quad[1])))
+    )
+    destination = np.float32([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]])
+    matrix = cv2.getPerspectiveTransform(quad.astype(np.float32), destination)
+    return cv2.warpPerspective(rgb, matrix, (width, height))
+
+
+class WordOCR:
+    def __init__(self, identity):
+        try:
+            import onnxruntime as ort
+            import torch
+            import yaml
+            from easyocr.craft import CRAFT
+            from easyocr.detection import copyStateDict
+        except ImportError as exc:
+            raise RuntimeError(
+                "Install requirements-ocr.txt into ComfyUI's Python environment to use automatic OCR."
+            ) from exc
+        craft_path, rec_path, dictionary = [Path(row[0]) for row in identity]
+        self.detector = CRAFT()
+        self.detector.load_state_dict(
+            copyStateDict(torch.load(craft_path, map_location="cpu", weights_only=True))
+        )
+        self.detector.eval()
+        options = ort.SessionOptions()
+        options.intra_op_num_threads = 4
+        options.inter_op_num_threads = 1
+        self.recognizer = ort.InferenceSession(
+            str(rec_path), sess_options=options, providers=["CPUExecutionProvider"]
+        )
+        config = yaml.safe_load(dictionary.read_text(encoding="utf-8"))
+        chars = config["PostProcess"]["character_dict"]
+        self.vocabulary = (
+            ["blank"] + chars + ([" "] if config["PostProcess"].get("use_space_char", True) else [])
+        )
+        self.input_name = self.recognizer.get_inputs()[0].name
+
+    def recognize(self, crops):
+        import cv2
+
+        results = []
+        # Small batches bound RAM for the large multilingual output vocabulary.
+        for start in range(0, len(crops), 8):
+            group = crops[start : start + 8]
+            widths = [min(3200, max(1, math.ceil(48 * c.shape[1] / c.shape[0]))) for c in group]
+            width = max(320, max(widths))
+            batch = np.zeros((len(group), 3, 48, width), dtype=np.float32)
+            for i, (crop, valid_width) in enumerate(zip(group, widths, strict=True)):
+                # Official PP-OCR preprocessing uses BGR, not RGB.
+                resized = cv2.resize(crop[:, :, ::-1], (valid_width, 48))
+                batch[i, :, :, :valid_width] = (
+                    (resized.astype(np.float32) / 255 - 0.5) / 0.5
+                ).transpose(2, 0, 1)
+            predictions = self.recognizer.run(None, {self.input_name: batch})[0]
+            if predictions.shape[-1] != len(self.vocabulary):
+                raise ValueError("OCR model output and dictionary have different vocabulary sizes")
+            for pred, valid_width in zip(predictions, widths, strict=True):
+                results.append(
+                    decode_ctc(
+                        pred,
+                        self.vocabulary,
+                        min(len(pred), math.ceil(len(pred) * valid_width / width)),
+                    )
+                )
+        return results
+
+    def detect(self, image: Image.Image, text_threshold: float, link_threshold: float):
+        from easyocr.detection import test_net
+
+        rgb = np.array(image.convert("RGB"))
+        boxes, _ = test_net(
+            1280, 1.5, self.detector, rgb, text_threshold, link_threshold, 0.4, False, "cpu"
+        )
+        quads = list(boxes[0])
+        if len(quads) > 256:
+            raise ValueError(
+                "More than 256 words detected. Crop the image or raise the text threshold."
+            )
+        recognition = self.recognize([crop_quad(rgb, q) for q in quads])
+        rows = []
+        for quad, (text, score) in zip(quads, recognition, strict=True):
+            if not text.strip():
+                continue
+            low = np.maximum(0, quad.min(axis=0))
+            high = np.minimum([image.width, image.height], quad.max(axis=0))
+            if np.any(high <= low):
+                continue
+            bbox = [round(float(x), 2) for x in [*low, *high]]
+            key = json.dumps([bbox, text], ensure_ascii=False)
+            rows.append(
+                {
+                    "id": "ocr-" + hashlib.sha256(key.encode()).hexdigest()[:16],
+                    "bbox": bbox,
+                    "text": text,
+                    "recognized_text": text,
+                    "confidence": round(score, 4),
+                    "enabled": False,
+                    "origin": "ocr",
+                }
+            )
+        rows.sort(key=lambda row: (round(row["bbox"][1] / 12), row["bbox"][0]))
+        return rows
+
+
+@lru_cache(maxsize=1)
+def load_ocr(identity):
+    return WordOCR(identity)
+
+
+def detect_words(image: Image.Image, root: Path, text_threshold=0.5, link_threshold=0.2):
+    identity = model_identity(root)
+    rows = load_ocr(identity).detect(image, text_threshold, link_threshold)
+    source_id = "pixels:" + hashlib.sha256(image.convert("RGB").tobytes()).hexdigest()
+    document = {
+        "version": 1,
+        "source_id": source_id,
+        "size": list(image.size),
+        "regions": rows,
+        "engine": "CRAFT + PP-OCRv6 medium (CPU)",
+    }
+    document["detection_id"] = hashlib.sha256(
+        json.dumps(document, sort_keys=True, ensure_ascii=False).encode()
+    ).hexdigest()
+    return document

requirements.txt

@@ -0,0 +1,16 @@
+torch==2.8.0
+torchvision==0.23.0
+numpy==2.2.6
+Pillow>=11.3.0
+opencv-python-headless==5.0.0.93
+safetensors>=0.4.5
+easyocr==1.7.2
+onnxruntime>=1.23.2,<2
+PyYAML>=6.0.2
+fastapi==0.162.0
+uvicorn==0.54.0
+python-multipart==0.0.32
+pydantic==2.14.1
+huggingface-hub>=1.16.0,<2
+gradio==6.30.0
+spaces==0.51.3

space.py

@@ -0,0 +1,93 @@
+"""Hugging Face entry point: custom studio UI, Gradio queue, free ZeroGPU."""
+
+import os
+
+os.environ["PIXTEXT_GRADIO"] = "1"
+os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
+
+# Import before torch so ZeroGPU can enable its CUDA staging mode.
+import spaces
+import gradio as gr
+import hashlib
+import json
+import time
+import uuid
+from pathlib import Path
+
+import app as studio
+
+from zerogpu_attention import enable_sdpa
+
+enable_sdpa()
+
+
+def load_models():
+    import torch
+    from pixtext import PixTextPipeline
+    from huggingface_hub import snapshot_download
+
+    torch.set_num_threads(4)
+    if not torch.cuda.is_available() and not spaces.config.Config.zero_gpu:
+        return {}
+    pipelines = {}
+    for variant, name in (("meanflow", "PixText-MeanFlow"), ("teacher", "PixText-Teacher")):
+        if os.environ.get("PIXTEXT_MODEL_ROOT"):
+            package = Path(os.environ["PIXTEXT_MODEL_ROOT"]) / name
+        else:
+            package = Path(snapshot_download("ATH-MaaS/PixText", allow_patterns=[f"models/{name}/*"])) / "models" / name
+        pipeline = PixTextPipeline.from_pretrained(package, device="cpu", dtype=torch.bfloat16)
+        pipeline.model = pipeline.model.to("cuda")
+        pipeline.device = torch.device("cuda")
+        pipelines[variant] = pipeline
+    return pipelines
+
+
+# Stage both models before serving requests, as required by ZeroGPU.
+PIPELINES = load_models()
+
+
+@spaces.GPU(duration=120)
+def gpu_refine(directory: str, size: list[int], jid: str, payload: str):
+    started = time.perf_counter()
+    request = studio.RefineRequest.model_validate_json(payload)
+    studio.PIPELINE = PIPELINES.get(request.variant)
+    studio.PIPELINE_VARIANT = request.variant if studio.PIPELINE is not None else None
+    progress = gr.Progress()
+
+    def report(**values):
+        progress(values.get("progress", 0), desc=values.get("message", "Refining"))
+
+    result = studio.refine({"directory": Path(directory), "size": size}, jid, request, report)
+    result["elapsed"] = round(time.perf_counter() - started, 2)
+    return result
+
+
+@studio.app.api(name="refine_image", concurrency_limit=1, concurrency_id="pixtext-gpu")
+def refine_image(sid: str, payload: str) -> dict:
+    request = studio.RefineRequest.model_validate_json(payload)
+    session = studio.validate_refine(sid, request)
+    canonical = request.model_dump_json()
+    key = hashlib.sha256(canonical.encode()).hexdigest()
+    with studio.LOCK:
+        previous = session.get("last_gradio_result")
+        if previous and previous["key"] == key:
+            job = studio.JOBS.get(previous["jid"])
+            if job and job["status"] == "done":
+                return {"job_id": job["id"], "result": job["result"]}
+        jid = uuid.uuid4().hex
+        studio.JOBS[jid] = {"id": jid, "session": sid, "kind": "refine", "status": "running",
+                            "progress": 0, "message": "Waiting for GPU", "created": time.time()}
+        session["job"] = jid
+    try:
+        result = gpu_refine(str(session["directory"]), list(session["size"]), jid, canonical)
+        studio.update_job(jid, status="done", progress=1, message="Complete", result=result)
+        session["last_gradio_result"] = {"key": key, "jid": jid}
+        return {"job_id": jid, "result": result}
+    except Exception:
+        studio.update_job(jid, status="error", message="Refinement could not finish. Please retry.")
+        raise
+
+
+if __name__ == "__main__":
+    studio.app.launch(server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
+                      server_port=int(os.environ.get("PIXTEXT_PORT", "7860")), show_error=False)

zerogpu_attention.py

@@ -0,0 +1,39 @@
+"""Use PyTorch SDPA with PixText's original token routing, without runtime compilation.
+
+ZeroGPU does not support torch.compile. Only the attention execution changes;
+the dense routing mask comes from the same function as PixText's reference mask.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def enable_sdpa():
+    from pixtext.modeling import attention, transformer
+
+    def forward(self, tokens, position_ids, block_mask, rotary_factors=None):
+        count = tokens.shape[0]
+        if block_mask.dtype != torch.bool or block_mask.shape != (count, count):
+            raise ValueError("Attention routing mask does not match the image tokens")
+        qkv = self.qkv(tokens).reshape(count, 3, self.num_heads, self.head_dim)
+        query, key, value = qkv.permute(1, 2, 0, 3).unbind(dim=0)
+        if rotary_factors is None:
+            rotary_factors = self.rotary.prepare(position_ids)
+        cosine, sine = rotary_factors
+        query = attention._normalize_and_apply_rotary(
+            query, self.q_norm.weight, self.q_norm.eps, cosine, sine
+        ).unsqueeze(0)
+        key = attention._normalize_and_apply_rotary(
+            key, self.k_norm.weight, self.k_norm.eps, cosine, sine
+        ).unsqueeze(0)
+        attended = F.scaled_dot_product_attention(
+            query, key, value.unsqueeze(0), attn_mask=block_mask[None, None], dropout_p=0.0
+        )
+        merged = attended.transpose(1, 2).reshape(count, self.hidden_size)
+        return self.proj_dropout(self.proj(merged))
+
+    attention.GroupSelfAttention.forward = forward
+    # Transformer imports this symbol directly, so replace its reference too.
+    transformer.build_flex_attention_block_mask = attention.build_dense_attention_mask