# 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 = """ """ 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()