import math import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 (must precede torch) import torch # noqa: E402 import gradio as gr # noqa: E402 from huggingface_hub import snapshot_download # noqa: E402 from PIL import Image, ImageDraw, ImageOps # noqa: E402 from pixtext import PixTextPipeline, TextEdit # noqa: E402 REPO_ID = "ATH-MaaS/PixText" # local_dir gives real files (the package loader rejects HF-cache symlinks) root = snapshot_download(REPO_ID, allow_patterns=["models/*"], local_dir="pixtext_weights") MODELS = { "MeanFlow (4 steps, fast)": PixTextPipeline.from_pretrained( os.path.join(root, "models", "PixText-MeanFlow"), device="cuda", dtype=torch.bfloat16 ), "Teacher (25 steps)": PixTextPipeline.from_pretrained( os.path.join(root, "models", "PixText-Teacher"), device="cuda", dtype=torch.bfloat16 ), } DEFAULT_MODEL = "MeanFlow (4 steps, fast)" def parse_edits(edits_text: str) -> list[TextEdit]: """Parse lines of `x0, y0, x1, y1, text` into TextEdit objects.""" edits = [] for n, line in enumerate((edits_text or "").splitlines(), 1): line = line.strip() if not line or line.startswith("#"): continue parts = line.split(",", 4) if len(parts) != 5: raise gr.Error(f"Line {n}: expected `x0, y0, x1, y1, text`, got: {line!r}") try: box = tuple(float(p) for p in parts[:4]) except ValueError: raise gr.Error(f"Line {n}: box coordinates must be numbers: {line!r}") text = parts[4].strip() try: edits.append(TextEdit(bbox=box, text=text)) except ValueError as e: raise gr.Error(f"Line {n}: {e}") return edits def draw_boxes(image, edits_text, pending=None): """Draw the edit boxes (and a pending first corner) on top of the image.""" if image is None: return None img = ImageOps.exif_transpose(image).convert("RGB").copy() d = ImageDraw.Draw(img) w = max(2, round(min(img.size) / 300)) try: edits = parse_edits(edits_text) except gr.Error: edits = [] for e in edits: x0, y0, x1, y1 = e.bbox d.rectangle([x0, y0, x1, y1], outline=(0, 220, 255), width=w) if pending is not None: x, y = pending r = 3 * w d.ellipse([x - r, y - r, x + r, y + r], fill=(255, 60, 60)) return img def on_click(image, edits_text, pending, new_text, evt: gr.SelectData): """Two clicks on the preview define a box; the second click appends an edit line.""" if image is None: return edits_text, None, None x, y = evt.index[0], evt.index[1] if pending is None: return edits_text, (x, y), draw_boxes(image, edits_text, (x, y)) x0, y0 = pending bx0, bx1 = sorted((x0, x)) by0, by1 = sorted((y0, y)) if bx1 - bx0 < 2 or by1 - by0 < 2: return edits_text, None, draw_boxes(image, edits_text) text = (new_text or "").strip() or "text" line = f"{bx0}, {by0}, {bx1}, {by1}, {text}" edits_text = (edits_text.rstrip() + "\n" + line) if (edits_text or "").strip() else line return edits_text, None, draw_boxes(image, edits_text) def _gpu_duration(image, edits_text, model_name=DEFAULT_MODEL, steps=0, *args, **kwargs): # Measured on ZeroGPU: MeanFlow ~4 s per 12-edit pass, Teacher (25 steps) ~12 s per pass. try: n = max(1, len(parse_edits(edits_text))) except Exception: n = 1 passes = math.ceil(n / 12) meanflow = "MeanFlow" in str(model_name) eff_steps = int(steps) if steps and int(steps) > 0 else (4 if meanflow else 25) est = 6 + passes * eff_steps * (0.9 if meanflow else 0.5) return int(min(180, max(15, math.ceil(est * 1.15)))) @spaces.GPU(duration=_gpu_duration) def refine( image: Image.Image, edits_text: str, model_name: str = DEFAULT_MODEL, steps: int = 0, guidance_scale: float = 3.3, seed: int = 42, progress=gr.Progress(track_tqdm=True), ): """Re-render text inside boxes of an image with PixText. Args: image: The input image containing text to refine. edits_text: One edit per line, `x0, y0, x1, y1, target text`, in original-image pixels. model_name: Which PixText package to use (MeanFlow is 4-step, Teacher is 25-step). steps: Sampling steps; 0 uses the model's default (4 for MeanFlow, 25 for Teacher). guidance_scale: Glyph guidance scale (default 3.3). seed: Random seed. Returns: A (input, output) image pair for comparison. """ if image is None: raise gr.Error("Please upload an image.") edits = parse_edits(edits_text) if not edits: raise gr.Error("Add at least one edit (click two corners on the preview, or type a line).") image = ImageOps.exif_transpose(image).convert("RGB") for e in edits: if e.bbox[2] > image.width or e.bbox[3] > image.height: raise gr.Error(f"Box {e.bbox} lies outside the {image.width}x{image.height} image.") pipe = MODELS.get(model_name, MODELS[DEFAULT_MODEL]) try: result = pipe.edit( image, edits, seed=int(seed), steps=int(steps) if int(steps) > 0 else None, guidance_scale=float(guidance_scale), ) except ValueError as e: raise gr.Error(str(e)) return (image, result.image) CSS = """ #col-container { max-width: 1200px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ TAGLINE_EDITS = """382, 405, 396, 431, d 449, 410, 574, 444, Excellence 595, 396, 680, 439, Quality""" POSTER_EDITS = """457, 29, 506.2, 94.2, 喜 503, 29.1, 559.4, 92.5, 发 749.5, 32.7, 804.1, 87.3, 甜 846.4, 30.1, 897.4, 88.1, 初 896.6, 30.8, 953, 92.4, 恋 400.7, 360.6, 455.3, 417, 鞭 455.1, 360.6, 509.7, 418.6, 炮 739.8, 362.1, 797.8, 421.9, 给 843.9, 360.5, 898.5, 420.3, 解 898.1, 360.4, 957.9, 422, 馋 506.3, 696.7, 560.9, 751.3, 满 558.1, 698.9, 616.1, 755.3, 满 807.5, 698.2, 853.3, 756.2, 畜""" with gr.Blocks(title="PixText") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """# PixText ✍️ Glyph-conditioned **local text refinement in pixel space**: mark the regions to fix, give the target text, and PixText regenerates just those regions (Chinese & English), leaving every other pixel untouched. [Model](https://huggingface.co/ATH-MaaS/PixText) · [Code](https://github.com/ATH-MaaS/PixText) · [Paper](https://arxiv.org/abs/2610.18284)""" ) pending = gr.State(None) with gr.Row(): with gr.Column(): image = gr.Image(label="Input image", type="pil", height=360) preview = gr.Image( label="Edit regions — click two opposite corners to add a box", type="pil", interactive=False, height=360, ) with gr.Row(): new_text = gr.Textbox( label="Text for the next box you draw", placeholder="e.g. Excellence or 喜", scale=3, ) clear_btn = gr.Button("Clear boxes", scale=1) edits_text = gr.Textbox( label="Edits (one per line: x0, y0, x1, y1, target text)", lines=5, placeholder="120, 80, 160, 125, 鹿", ) run = gr.Button("Refine text", variant="primary") with gr.Column(): output = gr.ImageSlider(label="Before / after", type="pil", format="png", height=560) with gr.Accordion("Advanced settings", open=False): model_name = gr.Radio( list(MODELS.keys()), value=DEFAULT_MODEL, label="Model" ) steps = gr.Slider( 0, 50, value=0, step=1, label="Steps (0 = model default: 4 / 25)" ) guidance_scale = gr.Slider( 1.0, 8.0, value=3.3, step=0.1, label="Glyph guidance scale" ) seed = gr.Number(value=42, precision=0, label="Seed") gr.Examples( examples=[ ["examples/tagline.png", TAGLINE_EDITS], ["examples/poster.png", POSTER_EDITS], ], inputs=[image, edits_text], outputs=[output], fn=refine, cache_examples=True, cache_mode="lazy", ) image.change(draw_boxes, [image, edits_text], preview, api_visibility="private").then( lambda: None, None, pending, api_visibility="private" ) edits_text.change(draw_boxes, [image, edits_text], preview, api_visibility="private") preview.select(on_click, [image, edits_text, pending, new_text], [edits_text, pending, preview], api_visibility="private") clear_btn.click(lambda: ("", None), None, [edits_text, pending], api_visibility="private") run.click( refine, [image, edits_text, model_name, steps, guidance_scale, seed], output, api_name="refine", ) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)