Deploy BitPixelLM Gradio Space

app.py

@@ -1,190 +1,190 @@
-    """
-    BitPixelLM — Hugging Face Gradio Space
-
-    Generates 32x32 pixel art from text prompts using BitPixelLM,
-    a 1.58-bit ternary transformer built from scratch with BitNet b1.58.
-    """
-
-    import sys
-    import json
-    import torch
-    import numpy as np
-    import gradio as gr
-    from pathlib import Path
-    from PIL import Image
-
-    HERE = Path(__file__).parent
-    sys.path.insert(0, str(HERE))
-
-    from model.tokenizer import PaletteTokenizer
-    from model.text_encoder import TextTokenizer, TextEncoder
-    from model.bit_pixel_decoder import BitPixelLMDecoder, BitPixelLM
-
-    # ─── Paths (all relative to this file — bundled in the Space) ──
-    PALETTE_PATH    = HERE / "palette_256.npy"
-    VOCAB_PATH      = HERE / "vocab.json"
-    CHECKPOINT_PATH = HERE / "best.pt"
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = None
-    palette_tok = None
-    text_tok = None
-
-
-    def load_all():
-        global model, palette_tok, text_tok
-
-        palette_tok = PaletteTokenizer(palette_path=str(PALETTE_PATH))
-        with open(VOCAB_PATH) as f:
-            vocab = json.load(f)
-        text_tok = TextTokenizer(vocab)
-
-        ckpt = torch.load(str(CHECKPOINT_PATH), map_location=device, weights_only=False)
-        args = ckpt.get("args", {})
-
-        text_encoder = TextEncoder(
-            vocab_size=text_tok.vocab_size,
-            d_model=args.get("d_model", 256),
-            nhead=args.get("nhead", 8),
-            num_layers=args.get("text_layers", 3),
-            dim_feedforward=args.get("dim_ff", 512),
-            max_seq_len=args.get("max_text_len", 32),
-            dropout=args.get("dropout", 0.1),
-        )
-        pixel_decoder = BitPixelLMDecoder(
-            vocab_size=palette_tok.vocab_size,
-            d_model=args.get("d_model", 256),
-            nhead=args.get("nhead", 8),
-            num_layers=args.get("pixel_layers", 6),
-            dim_feedforward=args.get("dim_ff", 512),
-            img_size=32,
-            dropout=args.get("dropout", 0.1),
-        )
-        m = BitPixelLM(text_encoder, pixel_decoder).to(device)
-        m.load_state_dict(ckpt["model_state_dict"])
-        m.eval()
-        model = m
-        print(f"BitPixelLM loaded | device={device} | vocab={text_tok.vocab_size} words")
-
-
-    def generate(prompt, temperature, top_k, top_p, num_samples, scale):
-        if not prompt.strip():
-            raise gr.Error("Please enter a prompt.")
-
-        words = prompt.lower().split()
-        unknown = [w for w in words if w not in text_tok.word2idx
-                   and w not in ("<pad>", "<sos>", "<eos>", "<unk>")]
-
-        text_tokens = text_tok.encode(prompt).unsqueeze(0).to(device)
-        images = []
-        try:
-            for _ in range(int(num_samples)):
-                with torch.no_grad():
-                    toks = model.generate(
-                        text_tokens,
-                        sos_token=palette_tok.sos_token,
-                        eos_token=palette_tok.eos_token,
-                        temperature=temperature,
-                        top_k=int(top_k),
-                        top_p=top_p,
-                    )
-                arr = palette_tok.decode_tokens(toks[0].cpu().tolist())
-                img = Image.fromarray(arr, "RGB")
-                s = int(scale)
-                if s > 1:
-                    img = img.resize((32 * s, 32 * s), Image.NEAREST)
-images.append((img, None))
-        except Exception as e:
-            raise gr.Error(f"Generation error: {e}")
-
-        if unknown:
-            gr.Warning("Unknown words (mapped to <unk>): " + ", ".join(unknown))
-        return images
-
-
-    # ─── Vocabulary for the info panel ──────────────────────────────
-    with open(VOCAB_PATH) as _f:
-        _vocab_raw = json.load(_f)
-    VOCAB_DISPLAY = ", ".join(sorted(w for w in _vocab_raw if not w.startswith("<")))
-
-    EXAMPLES = [
-        "a red pixel art sword",
-        "a blue pixel art knight",
-        "a green pixel art dragon",
-        "a purple pixel art wizard",
-        "a gold pixel art crown",
-        "a dark pixel art skeleton",
-        "a teal pixel art ship",
-        "a orange pixel art fox",
-        "a pink pixel art butterfly",
-        "a yellow pixel art beehive",
-    ]
-
-    with gr.Blocks(
-        title="BitPixelLM — Pixel Art Generator",
-        theme=gr.themes.Soft(),
-        css="""
-        .gallery-item img { image-rendering: pixelated !important; }
-        #gallery img      { image-rendering: pixelated !important; }
-        """,
-    ) as demo:
-        gr.Markdown(
-            """
-            # 🎨 BitPixelLM — Pixel Art Generator
-            ### Generate 32×32 pixel art from text prompts
-
-            **BitPixelLM** is a custom 1.58-bit ternary transformer (BitNet b1.58) trained from scratch to
-            produce pixel art sprites. 7.4M parameters, 75% ternary weights, best val loss 0.4015.
-
-            Try prompts like: `a red pixel art sword` · `a blue pixel art knight` · `a green pixel art dragon`
-            """
-        )
-
-        with gr.Row():
-            with gr.Column(scale=1):
-                prompt_box = gr.Textbox(
-                    label="Prompt",
-                    placeholder="a red pixel art sword",
-                    lines=2,
-                )
-                with gr.Row():
-                    btn = gr.Button("Generate ✨", variant="primary", scale=2)
-                    n_samples = gr.Slider(1, 8, value=4, step=1, label="Samples")
-
-                with gr.Accordion("Advanced", open=False):
-                    temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
-                    top_k      = gr.Slider(0, 256, value=40, step=1,  label="Top-K")
-                    top_p      = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
-                    scale      = gr.Slider(1, 16, value=8, step=1,
-                                           label="Upscale factor",
-                                           info="8× = 256×256 output")
-
-                gr.Markdown(f"**Supported words:** {VOCAB_DISPLAY}")
-
-            with gr.Column(scale=2):
-                gallery = gr.Gallery(
-                    label="Generated Pixel Art",
-                    columns=4, rows=2, height=520,
-                    object_fit="contain", elem_id="gallery",
+"""
+BitPixelLM — Hugging Face Gradio Space
+
+Generates 32x32 pixel art from text prompts using BitPixelLM,
+a 1.58-bit ternary transformer built from scratch with BitNet b1.58.
+"""
+
+import sys
+import json
+import torch
+import numpy as np
+import gradio as gr
+from pathlib import Path
+from PIL import Image
+
+HERE = Path(__file__).parent
+sys.path.insert(0, str(HERE))
+
+from model.tokenizer import PaletteTokenizer
+from model.text_encoder import TextTokenizer, TextEncoder
+from model.bit_pixel_decoder import BitPixelLMDecoder, BitPixelLM
+
+# ─── Paths (all relative to this file — bundled in the Space) ──
+PALETTE_PATH    = HERE / "palette_256.npy"
+VOCAB_PATH      = HERE / "vocab.json"
+CHECKPOINT_PATH = HERE / "best.pt"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = None
+palette_tok = None
+text_tok = None
+
+
+def load_all():
+    global model, palette_tok, text_tok
+
+    palette_tok = PaletteTokenizer(palette_path=str(PALETTE_PATH))
+    with open(VOCAB_PATH) as f:
+        vocab = json.load(f)
+    text_tok = TextTokenizer(vocab)
+
+    ckpt = torch.load(str(CHECKPOINT_PATH), map_location=device, weights_only=False)
+    args = ckpt.get("args", {})
+
+    text_encoder = TextEncoder(
+        vocab_size=text_tok.vocab_size,
+        d_model=args.get("d_model", 256),
+        nhead=args.get("nhead", 8),
+        num_layers=args.get("text_layers", 3),
+        dim_feedforward=args.get("dim_ff", 512),
+        max_seq_len=args.get("max_text_len", 32),
+        dropout=args.get("dropout", 0.1),
+    )
+    pixel_decoder = BitPixelLMDecoder(
+        vocab_size=palette_tok.vocab_size,
+        d_model=args.get("d_model", 256),
+        nhead=args.get("nhead", 8),
+        num_layers=args.get("pixel_layers", 6),
+        dim_feedforward=args.get("dim_ff", 512),
+        img_size=32,
+        dropout=args.get("dropout", 0.1),
+    )
+    m = BitPixelLM(text_encoder, pixel_decoder).to(device)
+    m.load_state_dict(ckpt["model_state_dict"])
+    m.eval()
+    model = m
+    print(f"BitPixelLM loaded | device={device} | vocab={text_tok.vocab_size} words")
+
+
+def generate(prompt, temperature, top_k, top_p, num_samples, scale):
+    if not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+
+    words = prompt.lower().split()
+    unknown = [w for w in words if w not in text_tok.word2idx
+               and w not in ("<pad>", "<sos>", "<eos>", "<unk>")]
+
+    text_tokens = text_tok.encode(prompt).unsqueeze(0).to(device)
+    images = []
+    try:
+        for _ in range(int(num_samples)):
+            with torch.no_grad():
+                toks = model.generate(
+                    text_tokens,
+                    sos_token=palette_tok.sos_token,
+                    eos_token=palette_tok.eos_token,
+                    temperature=temperature,
+                    top_k=int(top_k),
+                    top_p=top_p,
                 )
-
-        gr.Examples(examples=EXAMPLES, inputs=[prompt_box], label="Try an example")
-
-        gr.Markdown(
-            """
-            ---
-            **How it works:** pixel art generation is framed as language modelling.
-            Each of the 1,024 pixels in a 32×32 image is predicted as a discrete token
-            from a 256-color palette, left-to-right, top-to-bottom, conditioned on
-            the text prompt via cross-attention.
-            [Model card](https://huggingface.co/BlakePeavy/BitPixelLM) ·
-            [Paper: BitNet b1.58](https://arxiv.org/abs/2402.17764)
-            """
-        )
-
-        inputs = [prompt_box, temperature, top_k, top_p, n_samples, scale]
-        btn.click(fn=generate, inputs=inputs, outputs=gallery)
-        prompt_box.submit(fn=generate, inputs=inputs, outputs=gallery)
-
-    load_all()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+            arr = palette_tok.decode_tokens(toks[0].cpu().tolist())
+            img = Image.fromarray(arr, "RGB")
+            s = int(scale)
+            if s > 1:
+                img = img.resize((32 * s, 32 * s), Image.NEAREST)
+images.append((img, None))
+    except Exception as e:
+        raise gr.Error(f"Generation error: {e}")
+
+    if unknown:
+        gr.Warning("Unknown words (mapped to <unk>): " + ", ".join(unknown))
+    return images
+
+
+# ─── Vocabulary for the info panel ──────────────────────────────
+with open(VOCAB_PATH) as _f:
+    _vocab_raw = json.load(_f)
+VOCAB_DISPLAY = ", ".join(sorted(w for w in _vocab_raw if not w.startswith("<")))
+
+EXAMPLES = [
+    "a red pixel art sword",
+    "a blue pixel art knight",
+    "a green pixel art dragon",
+    "a purple pixel art wizard",
+    "a gold pixel art crown",
+    "a dark pixel art skeleton",
+    "a teal pixel art ship",
+    "a orange pixel art fox",
+    "a pink pixel art butterfly",
+    "a yellow pixel art beehive",
+]
+
+with gr.Blocks(
+    title="BitPixelLM — Pixel Art Generator",
+    theme=gr.themes.Soft(),
+    css="""
+    .gallery-item img { image-rendering: pixelated !important; }
+    #gallery img      { image-rendering: pixelated !important; }
+    """,
+) as demo:
+    gr.Markdown(
+        """
+        # 🎨 BitPixelLM — Pixel Art Generator
+        ### Generate 32×32 pixel art from text prompts
+
+        **BitPixelLM** is a custom 1.58-bit ternary transformer (BitNet b1.58) trained from scratch to
+        produce pixel art sprites. 7.4M parameters, 75% ternary weights, best val loss 0.4015.
+
+        Try prompts like: `a red pixel art sword` · `a blue pixel art knight` · `a green pixel art dragon`
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt_box = gr.Textbox(
+                label="Prompt",
+                placeholder="a red pixel art sword",
+                lines=2,
+            )
+            with gr.Row():
+                btn = gr.Button("Generate ✨", variant="primary", scale=2)
+                n_samples = gr.Slider(1, 8, value=4, step=1, label="Samples")
+
+            with gr.Accordion("Advanced", open=False):
+                temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
+                top_k      = gr.Slider(0, 256, value=40, step=1,  label="Top-K")
+                top_p      = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
+                scale      = gr.Slider(1, 16, value=8, step=1,
+                                       label="Upscale factor",
+                                       info="8× = 256×256 output")
+
+            gr.Markdown(f"**Supported words:** {VOCAB_DISPLAY}")
+
+        with gr.Column(scale=2):
+            gallery = gr.Gallery(
+                label="Generated Pixel Art",
+                columns=4, rows=2, height=520,
+                object_fit="contain", elem_id="gallery",
+            )
+
+    gr.Examples(examples=EXAMPLES, inputs=[prompt_box], label="Try an example")
+
+    gr.Markdown(
+        """
+        ---
+        **How it works:** pixel art generation is framed as language modelling.
+        Each of the 1,024 pixels in a 32×32 image is predicted as a discrete token
+        from a 256-color palette, left-to-right, top-to-bottom, conditioned on
+        the text prompt via cross-attention.
+        [Model card](https://huggingface.co/BlakePeavy/BitPixelLM) ·
+        [Paper: BitNet b1.58](https://arxiv.org/abs/2402.17764)
+        """
+    )
+
+    inputs = [prompt_box, temperature, top_k, top_p, n_samples, scale]
+    btn.click(fn=generate, inputs=inputs, outputs=gallery)
+    prompt_box.submit(fn=generate, inputs=inputs, outputs=gallery)
+
+load_all()
+demo.launch(server_name="0.0.0.0", server_port=7860)