Blame for app.py

# coding=utf-8
"""OpenHTR-Gen: interactive handwriting text synthesis demo (Space / ZeroGPU)."""
from __future__ import annotations

import os

# ZeroGPU patches CUDA during import. This MUST precede Gradio, engine and
# anything else that may import torch or initialize CUDA.
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
try:
    import spaces  # provided by Hugging Face ZeroGPU
    _HAS_SPACES = True
except ImportError:  # local/non-ZeroGPU execution
    spaces = None
    _HAS_SPACES = False

import random
import threading
import time

import gradio as gr

import engine
import selection


TITLE = "OpenHTR-Gen"
BASE_MODEL = os.environ.get("BASE_MODEL", "black-forest-labs/FLUX.1-Fill-dev")
GPU_DURATION = int(os.environ.get("GPU_DURATION", "120"))
QUANTIZE_8BIT = os.environ.get("HW_QUANTIZE_8BIT", "0").lower() in ("1", "true", "yes", "8bit")
# ZeroGPU packs model weights placed on CUDA at module scope for reuse by GPU
# workers. Keep lazy loading for ordinary local execution unless requested.
EAGER_LOAD = os.environ.get("HW_EAGER_LOAD", "1" if _HAS_SPACES else "0").lower() in (
    "1", "true", "yes"
)
MAX_CHARS = int(os.environ.get("MAX_CHARS", "35"))
SIZE = engine.DEFAULT_SIZE  # 512, matches the training resolution


def _gpu_decorator(duration):
    if spaces is None:
        return lambda fn: fn
    return spaces.GPU(duration=duration)


FONT = engine.load_font()

_pipe = None
_pipe_lock = threading.Lock()


def ensure_pipeline():
    """Initialize once; ZeroGPU prefers calling this at module scope."""
    global _pipe
    if _pipe is not None:
        return _pipe
    with _pipe_lock:
        if _pipe is None:
            t0 = time.time()
            _pipe = engine.load_pipeline(
                base_model=BASE_MODEL,
                lora_path=engine.resolve_lora_path(),
                quantize_8bit=QUANTIZE_8BIT,
                device="cuda",
                zero_gpu=_HAS_SPACES,
            )
            print(f"[load] ready in {time.time() - t0:.1f}s", flush=True)
    return _pipe


def _warm_hf_cache():
    """Pre-download weights so the first GPU call spends less time on I/O."""
    try:
        from huggingface_hub import snapshot_download
        snapshot_download(
            BASE_MODEL,
            allow_patterns=["model_index.json", "transformer/*", "text_encoder_2/*",
                            "tokenizer_2/*", "vae/*", "scheduler/*"],
            token=os.environ.get("HF_TOKEN") or None,
            ignore_patterns=["*.bin", "*.onnx", "*.msgpack"],
        )
        print("[warm] base weights cached")
    except Exception as e:  # pragma: no cover
        print(f"[warm] pre-download failed (will retry on first call): {e}")

    # Fetch the adapter here as well: network I/O should not eat ZeroGPU seconds.
    print(f"[warm] adapter: {engine.resolve_lora_path()}")


if EAGER_LOAD:
    # ZeroGPU CUDA emulation registers the model weights here, before GPU tasks.
    # Do not spawn a competing download thread while loading the same weights.
    ensure_pipeline()
elif _HAS_SPACES:
    # Optional legacy lazy mode: prefetch to keep GPU tasks mostly compute-bound.
    threading.Thread(target=_warm_hf_cache, daemon=True).start()


@_gpu_decorator(GPU_DURATION)
def _infer(image, box, text, steps, guidance, seed, crop_pad):
    pipe = ensure_pipeline()
    return engine.edit_text_region(
        pipe, FONT, image, box, text,
        size=SIZE, steps=steps, guidance_scale=guidance, seed=seed,
        crop_pad_ratio=crop_pad,
    )


def run_demo(editor, text, steps, guidance, seed, crop_pad, pad_px,
             mx0, my0, mx1, my1):
    t_start = time.time()

    bg, mask = selection.editor_mask(editor)
    if bg is None:
        raise gr.Error("Please upload a full-page image first.")

    text = (text or "").strip()
    if not text:
        raise gr.Error("Please enter the text you want to generate.")
    if len(text) > MAX_CHARS:
        text = text[:MAX_CHARS]
        gr.Warning(f"Text is longer than {MAX_CHARS} characters; truncated to: {text}")

    box = selection.manual_box((mx0, my0, mx1, my1), *bg.size)
    box_src = "manual box"
    if box is None:
        box = selection.clamp_box(selection.mask_to_box(mask, pad_px), *bg.size)
        box_src = "brush selection"
    if box is None:
        raise gr.Error("No region selected: paint over the text you want to rewrite, "
                       "or fill in x0/y0/x1/y1 under Advanced settings.")

    if seed is None or int(seed) < 0:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)

    try:
        result = _infer(bg, box, text, int(steps), float(guidance), seed, float(crop_pad))
    except Exception as e:  # noqa: BLE001
        raise gr.Error(f"Generation failed: {type(e).__name__}: {e}") from e

    crop, page = result["crop"], result["page"]
    elapsed = time.time() - t_start
    status = (
        f"**Done** · {elapsed:.1f}s · box ({box_src}) = `{box}` · "
        f"page `{bg.size[0]}×{bg.size[1]}`\n\n"
        f"canvas `{result['canvas'][0]}×{result['canvas'][1]}` · "
        f"crop `{crop.size[0]}×{crop.size[1]}` · seed `{seed}` · steps `{int(steps)}` · "
        f"guidance `{float(guidance):g}`"
    )
    return crop, page, status


CSS = """
footer {visibility: hidden;}
"""

# Gradio's built-in UI strings follow the browser locale. Pin the locale to English
# before the app boots so the whole interface is English for every visitor.
HEAD = """
<script>
  try {
    Object.defineProperty(navigator, "language", { get: () => "en" });
    Object.defineProperty(navigator, "languages", { get: () => ["en"] });
  } catch (e) {}
</script>
"""

with gr.Blocks(title=TITLE, css=CSS, head=HEAD) as demo:
    gr.Markdown(
        f"# {TITLE}\n"
        "Upload a full-page Chinese handwriting image, paint over the text you want to "
        "rewrite, and type the new text. The model re-renders that region in the "
        "handwriting style of the page.\n"
        "> FLUX.1 [dev] is released under a non-commercial license. "
        "Do not use this demo or its outputs commercially."
    )

    with gr.Row():
        with gr.Column(scale=1):
            editor = gr.ImageEditor(
                label="① Upload a page and paint over the text to rewrite",
                type="pil",
                layers=False,
                transforms=[],
                sources=["upload", "clipboard"],
                brush=gr.Brush(colors=["rgb(255,45,45)", "rgb(0,0,0)"], default_size=18),
                height=420,
            )
            text = gr.Textbox(
                label="② Text to generate",
                value="手写文本",
                max_lines=3,
                placeholder="Enter the replacement text",
            )

            with gr.Accordion("Advanced settings", open=False):
                with gr.Row():
                    steps = gr.Slider(4, 30, value=20, step=1, label="Inference steps")
                    guidance = gr.Slider(1, 50, value=30, step=1, label="guidance_scale")
                with gr.Row():
                    seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
                    crop_pad = gr.Slider(0.0, 0.2, value=0.0, step=0.01,
                                         label="Crop padding")
                    pad_px = gr.Slider(0, 20, value=4, step=1,
                                       label="Selection padding (px)")
                gr.Markdown("Manual box (optional; overrides the brush. Pixel coordinates, "
                            "x1 > x0 and y1 > y0).")
                with gr.Row():
                    mx0 = gr.Number(value=0, precision=0, label="x0")
                    my0 = gr.Number(value=0, precision=0, label="y0")
                    mx1 = gr.Number(value=0, precision=0, label="x1")
                    my1 = gr.Number(value=0, precision=0, label="y1")

            btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            out_crop = gr.Image(label="③ Generated crop", type="pil", height=200)
            out_page = gr.Image(label="④ Full-page preview (for reference)",
                                type="pil", height=400)
            status = gr.Markdown()

    btn.click(
        fn=run_demo,
        inputs=[editor, text, steps, guidance, seed, crop_pad, pad_px, mx0, my0, mx1, my1],
        outputs=[out_crop, out_page, status],
    )


if __name__ == "__main__":
    demo.queue(max_size=8).launch()