Spaces:
Running
Running
| """PixelGen demo Space — a text-to-pixel-art model with 1,257 hand-written | |
| parameters and zero training steps. | |
| The `pixelgen/` package and the `weights/pixelgen.pt` checkpoint are vendored | |
| from https://huggingface.co/DJLougen/pixelgen (MIT). The demo mirrors the | |
| reference implementation (`PixelGen.sprite` / `generate.py`) 1:1: same | |
| tokenizer, same forward pass, same nearest-neighbour upscale — it only adds a | |
| UI. The model is tiny and CPU-bound, so inference runs directly in the web | |
| process; `@spaces.GPU` stays as a portable no-op in case the hardware is ever | |
| moved to ZeroGPU. | |
| """ | |
| import spaces # imported first; a no-op off ZeroGPU | |
| import gradio as gr | |
| from PIL import Image | |
| from pixelgen import weights as weights_io | |
| from pixelgen.model import PixelGen | |
| from pixelgen.palette import COLOR_WORDS, RAMP_NAMES | |
| from pixelgen.templates import CLASSES | |
| from pixelgen.tokenizer import CLASS_SYNONYMS | |
| MODEL_ID = "DJLougen/pixelgen" | |
| DEFAULT_SEED = 0 | |
| DEFAULT_SCALE = 8 # 16 x 8 -> 128x128 PNG, the model's default output | |
| def _load_model() -> PixelGen: | |
| """Load the shipped checkpoint; fall back to rebuilding weights in code. | |
| `PixelGen()` sets every weight deterministically in `__init__`, so both | |
| paths produce bit-identical sprites (the repo's round-trip test covers | |
| this) — the fallback only guards against a missing/corrupt file. | |
| """ | |
| try: | |
| model = weights_io.load() # resolves to weights/pixelgen.pt | |
| print(f"loaded {MODEL_ID} from weights/pixelgen.pt") | |
| except Exception as exc: # pragma: no cover - robustness fallback | |
| print(f"checkpoint load failed ({exc!r}); rebuilding weights from source") | |
| model = PixelGen() | |
| return model | |
| MODEL = _load_model() | |
| N_PARAMS = sum(p.numel() for p in MODEL.parameters()) | |
| def generate(prompt: str, seed: int = DEFAULT_SEED, scale: int = DEFAULT_SCALE): | |
| """Generate a PICO-8 pixel-art sprite from a short text prompt. | |
| Args: | |
| prompt: A colour word (red, blue, green, yellow, purple, orange, pink, | |
| brown, white — plus synonyms such as gold, crimson, emerald) and a | |
| class word (potion, sword, shield, ship, invader, slime, skull, | |
| tree, key, heart, star — plus synonyms such as alien, blade). | |
| For example "red potion" or "golden key". Unknown words are | |
| ignored; a prompt with no known class renders the default sprite | |
| (a green potion). | |
| seed: Selects the hand-marked "?" pixels, giving small variations of | |
| the same sprite across seeds. | |
| scale: Nearest-neighbour upscale factor; the output PNG is | |
| 16*scale x 16*scale pixels (default 8 -> 128x128). | |
| Returns: | |
| (sprite, seed_strip): the sprite as a PIL image, plus the same prompt | |
| rendered at seeds 0-7 to show the seed variation. | |
| """ | |
| seed, scale = int(seed), int(scale) | |
| sprite = Image.fromarray(MODEL.generate([prompt], [seed], scale=scale)[0], "RGB") | |
| strip_arr = MODEL.generate([prompt] * 8, list(range(8)), scale=scale) | |
| strip = [Image.fromarray(arr, "RGB") for arr in strip_arr] | |
| return sprite, strip | |
| def build_prompt(body: str, colour: str) -> str: | |
| """Compose a prompt like 'red potion' from the two dropdowns.""" | |
| return f"{colour} {body}" | |
| def vocab_markdown() -> str: | |
| """Render the model's fixed 59-word vocabulary as reference tables.""" | |
| cls_rows = "\n".join(f"| {c} | {', '.join(ws)} |" for c, ws in CLASS_SYNONYMS.items()) | |
| by_ramp: dict = {} | |
| for word, ramp in COLOR_WORDS.items(): | |
| by_ramp.setdefault(ramp, []).append(word) | |
| col_rows = "\n".join(f"| {r} | {', '.join(ws)} |" for r, ws in by_ramp.items()) | |
| return ( | |
| "Prompts are a bag of words: every known word counts, everything else is " | |
| "ignored. A prompt with **no known class word** falls back to the default " | |
| "class (**potion**) and one with **no known colour word** to the default " | |
| "ramp (**green**) — so `\"dragon\"` gives you a green potion.\n\n" | |
| "### Sprite classes\n" | |
| "| class | words |\n|---|---|\n" | |
| f"{cls_rows}\n\n" | |
| "### Colour ramps\n" | |
| "| colour | words |\n|---|---|\n" | |
| f"{col_rows}" | |
| ) | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| .gradio-container img { image-rendering: pixelated; } | |
| """ | |
| EXAMPLES = [ | |
| ["red potion", 0], | |
| ["golden key", 3], | |
| ["blue invader", 0], | |
| ["purple slime", 4], | |
| ] | |
| with gr.Blocks(title="PixelGen") as demo: | |
| gr.Markdown( | |
| f"# 🕹️ PixelGen\n" | |
| f"A text-to-pixel-art model with **{N_PARAMS:,} parameters — zero of them " | |
| f"trained**. Every weight was typed in by hand: prompts like " | |
| f"`\"red potion\"` or `\"blue invader\"` become 16×16 PICO-8 sprites, " | |
| f"upscaled to crisp pixel art. " | |
| f"[Model card](https://huggingface.co/{MODEL_ID}) · " | |
| f"[weight-space programming](https://rikiyat.github.io/blog/weight-space-programming/)" | |
| ) | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| prompt = gr.Textbox( | |
| label="Prompt", value="red potion", placeholder="e.g. blue invader" | |
| ) | |
| with gr.Accordion("…or build one", open=False): | |
| body_dd = gr.Dropdown(CLASSES, value="potion", label="sprite class") | |
| colour_dd = gr.Dropdown(RAMP_NAMES, value="red", label="colour") | |
| for dd in (body_dd, colour_dd): | |
| dd.change(build_prompt, [body_dd, colour_dd], prompt) | |
| seed = gr.Slider(0, 99, value=DEFAULT_SEED, step=1, label="seed") | |
| scale = gr.Slider( | |
| 4, 16, value=DEFAULT_SCALE, step=1, | |
| label="scale (output is 16×scale px)", | |
| ) | |
| go = gr.Button("Generate", variant="primary") | |
| with gr.Column(scale=1): | |
| out = gr.Image(label="sprite", type="pil", height=300) | |
| gr.Markdown( | |
| "**Same prompt, seeds 0–7** — a seed only toggles the hand-marked " | |
| "`?` pixels, so seeds give small variations of the same sprite." | |
| ) | |
| strip = gr.Gallery(columns=8, height=150, label="seeds") | |
| with gr.Accordion("Prompt vocabulary", open=False): | |
| gr.Markdown(vocab_markdown()) | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=[prompt, seed], | |
| outputs=[out, strip], | |
| fn=generate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| go.click(generate, [prompt, seed, scale], [out, strip], api_name="generate") | |
| prompt.submit(generate, [prompt, seed, scale], [out, strip]) | |
| demo.load(generate, [prompt, seed, scale], [out, strip]) | |
| if __name__ == "__main__": | |
| # Gradio 6: theme/css are launch() parameters, not Blocks constructor ones. | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |