Spaces:

hugging-apps
/

pixtext

Running on Zero

Upload app.py with huggingface_hub

e3f529c verified 4 days ago

9.42 kB

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)

	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)