app.py

"""
🔮 Cryptid Cannon — summon a pack of collectible cryptids.

One click fires off a 4-card pack: each card is a one-of-a-kind creature with
its own name, type, rarity and stats, rendered by FLUX.2 [klein] 4B. A
deliberately delightful Build Small toy (Thousand Token Wood track). Built on
the klein starter's verified ZeroGPU + pipeline pattern.
"""
from __future__ import annotations

import random
import time

# --- ZeroGPU shim: import `spaces` BEFORE torch -----------------------------
try:
    import spaces  # type: ignore

    GPU = spaces.GPU
except Exception:  # local / non-ZeroGPU fallback

    def GPU(*dargs, **dkwargs):  # noqa: N802
        if len(dargs) == 1 and callable(dargs[0]) and not dkwargs:
            return dargs[0]

        def wrap(fn):
            return fn

        return wrap


import gradio as gr
import torch
from diffusers import Flux2KleinPipeline

MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"  # 4B, Apache 2.0, ungated
STEPS = 4
GUIDANCE = 1.0
MAX_SEED = 2**31 - 1
PACK_SIZE = 4

# Build on CPU at module scope (ZeroGPU pattern), guarded so the UI still boots
# on tiny hardware instead of a blank runtime error.
pipe = None
LOAD_ERR = ""
try:
    print(f"Loading {MODEL_ID} on CPU…")
    pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
    print("  loaded.")
except Exception as e:  # noqa: BLE001
    LOAD_ERR = str(e)
    print("Model load failed:", e)

# --- the chaos tables -------------------------------------------------------
CREATURES = [
    "mushroom knight", "cloud cat", "disgruntled potato wizard", "sock goblin",
    "philosophical jellyfish", "tiny bureaucrat dragon", "moss yeti",
    "origami werewolf", "caffeinated snail", "Victorian frog detective",
    "marshmallow golem", "interdimensional library owl", "grumpy lighthouse crab",
    "pocket-sized thundercloud", "retired circus moth",
]
TRAITS = [
    "wearing a tiny hat", "made of stained glass", "on a unicycle",
    "knitting a scarf", "looking very smug", "mid-sneeze", "holding a lantern",
    "deeply confused", "in formal attire", "glowing faintly", "drinking tea",
    "surrounded by fireflies", "guarding a single sock",
]
STYLES = [
    "children's storybook illustration, soft watercolor",
    "claymation still, studio lighting",
    "bold flat vector sticker art",
    "cozy painted scene, warm light",
    "1970s fantasy paperback cover",
    "cute isometric 3D render",
]

# --- card-flavor tables -----------------------------------------------------
NAME_PREFIX = ["Sir", "Lord", "Tiny", "Old", "Captain", "Mega", "Lil", "Grand",
               "Baby", "Professor", "Madam", "Saint"]
NAME_ROOT = ["Mossbottom", "Snagglet", "Borp", "Wuffle", "Crumb", "Noodle",
             "Pib", "Glorp", "Thistle", "Bumbershoot", "Quark", "Sprout",
             "Fenwick", "Pumpkin", "Wobble"]
NAME_SUFFIX = ["the Damp", "of the Bog", "McFluff", "the Untidy", "VonCrumb",
               "the Eternal", "Snacksworth", "the Mild", "of Aisle 7",
               "the Confused", "the Third", "the Slightly Cursed"]
TYPES = ["🍄 Fungal", "☁️ Vapor", "🔧 Clockwork", "🫧 Goo", "🌙 Lunar",
         "📚 Arcane", "🥔 Root", "⚡ Static", "🕯️ Spectral", "🌾 Feral"]
RARITY = [("Common", "⚪", 60), ("Uncommon", "🟢", 25), ("Rare", "🔵", 9),
          ("Epic", "🟣", 4), ("Legendary", "🟡", 1.5), ("Cursed", "🔴", 0.5)]
FLAVORS = [
    "Smells faintly of rain and regret.",
    "Has filed three complaints with management.",
    "Believes it is much larger than it is.",
    "Collects buttons. Will not explain why.",
    "Last seen unionizing the local pigeons.",
    "Powered entirely by spite and snacks.",
    "Legally distinct from a houseplant.",
    "Hums a tune no one else can hear.",
    "Once won a staring contest with the moon.",
]


def _prompt(twist: str) -> str:
    subject = twist.strip() if twist and twist.strip() else (
        f"{random.choice(CREATURES)} {random.choice(TRAITS)}"
    )
    return f"A {subject}, {random.choice(STYLES)}, centered, whimsical, no text"


def _name() -> str:
    return f"{random.choice(NAME_PREFIX)} {random.choice(NAME_ROOT)} {random.choice(NAME_SUFFIX)}"


def _card_stats():
    rname, remoji, _ = random.choices(RARITY, weights=[r[2] for r in RARITY])[0]
    # rarer creatures hit a little harder
    boost = {"Common": 0, "Uncommon": 8, "Rare": 16, "Epic": 26,
             "Legendary": 38, "Cursed": 45}[rname]
    hp = random.randint(20, 55) + boost
    atk = random.randint(10, 45) + boost
    dfn = random.randint(10, 45) + boost
    return rname, remoji, hp, atk, dfn


@GPU(duration=120)
def summon_pack(twist: str):
    if pipe is None:
        raise gr.Error(f"Model isn't loaded (this Space needs a GPU). {LOAD_ERR[:200]}")
    pipe.to("cuda")
    gallery, cards = [], []
    t = time.time()
    for _ in range(PACK_SIZE):
        prompt = _prompt(twist)
        seed = random.randint(0, MAX_SEED)
        img = pipe(
            prompt=prompt,
            width=1024,
            height=1024,
            num_inference_steps=STEPS,
            guidance_scale=GUIDANCE,
            generator=torch.Generator(device="cuda").manual_seed(seed),
        ).images[0]
        name = _name()
        typ = random.choice(TYPES)
        rname, remoji, hp, atk, dfn = _card_stats()
        flavor = random.choice(FLAVORS)
        gallery.append((img, f"{remoji} {name}"))
        cards.append(
            f"### {remoji} {name}\n"
            f"**{typ}** · *{rname}*\n\n"
            f"❤️ HP {hp}   ⚔️ ATK {atk}   🛡️ DEF {dfn}\n\n"
            f"> {flavor}"
        )
    details = "## 🃏 Your pack\n\n" + "\n\n---\n\n".join(cards)
    details += f"\n\n<sub>Summoned in {time.time() - t:.1f}s · klein 4B · {STEPS} steps</sub>"
    return gallery, details


# --- styling: system fonts + no SSR, so it always renders on Spaces ----------
THEME = gr.themes.Soft(
    font=["system-ui", "-apple-system", "Segoe UI", "Roboto", "Helvetica", "Arial", "sans-serif"],
    font_mono=["ui-monospace", "SFMono-Regular", "Consolas", "monospace"],
)
CSS = """
footer {visibility: hidden;}
.gradio-container, .gradio-container .prose, .gradio-container p,
.gradio-container h1, .gradio-container h2, .gradio-container h3 {
  font-family: system-ui, -apple-system, "Segoe UI", Roboto, Helvetica, Arial, sans-serif !important;
}
"""

with gr.Blocks(title="Cryptid Cannon", theme=THEME, css=CSS) as demo:
    gr.Markdown(
        "# 🔮 Cryptid Cannon\n"
        "Fire the cannon and summon a **pack of four collectible cryptids** — "
        "each a one-of-a-kind creature with its own name, type, rarity and stats. "
        "Powered by **FLUX.2 [klein] 4B** (4B params, Apache 2.0, runs locally on "
        "the Space GPU). Leave the box empty for pure chaos, or give the whole "
        "pack a theme."
    )
    with gr.Row():
        twist = gr.Textbox(
            label="Pack theme (optional)",
            placeholder="haunted bakery",
            lines=1,
            scale=4,
        )
        btn = gr.Button("🔮 Fire the cannon", variant="primary", scale=1)
    gr.Examples(
        ["", "haunted bakery", "deep-sea disco", "cottagecore goblins", "office supplies that gained sentience"],
        twist,
        label="Try a theme",
    )
    gallery = gr.Gallery(label="Your pack", columns=4, height=300, object_fit="contain")
    cards = gr.Markdown()
    btn.click(summon_pack, twist, [gallery, cards])

if __name__ == "__main__":
    # ssr_mode=False: Gradio-5 server-side rendering is a common cause of an app
    # rendering as unstyled raw HTML on HF Spaces; the client bundle is reliable.
    demo.queue(max_size=8).launch(
        server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False
    )