Spaces:
Running on Zero
Running on Zero
| """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") | |
| 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"} | |
| def index(): | |
| html = (ROOT / "static/index.html").read_text() | |
| script = (ROOT / "static/app.js").read_text() | |
| html = html.replace('<script defer src="/static/app.js"></script>', "") | |
| return HTMLResponse(html.replace("</body>", "<script>\n(() => {\n" + script + "\n})();\n</script></body>")) | |
| def config(): | |
| return {"examples": EXAMPLES, "models": ["meanflow", "teacher"], | |
| "transport": "gradio" if os.environ.get("PIXTEXT_GRADIO") == "1" else "http"} | |
| 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)} | |
| 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.") | |
| 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)} | |
| def start_detect(sid: str, request: DetectRequest): | |
| return submit(sid, "detect", detect, request) | |
| 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 | |
| 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) | |
| 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"} | |
| 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"}) | |