Spaces:
Running on Zero
Running on Zero
| # 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() | |
| 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() | |