app.py

"""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('<script defer src="/static/app.js"></script>', "")
    return HTMLResponse(html.replace("</body>", "<script>\n(() => {\n" + script + "\n})();\n</script></body>"))


@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"})