""" Zeta-Chroma Test Bench — Custom Pixel-Space Pipeline ===================================================== Uses our NextDiTPixelSpace model re-implemented from ComfyUI source. No VAE needed — model operates directly on RGB pixels. Text encoder: Qwen3 4B from Tongyi-MAI/Z-Image. Key text encoding details (from ComfyUI z_image.py): - Prompt wrapped in chat template - Uses hidden layer -2 (second to last) - No padding to max length """ import os, gc, random, math, time import spaces import torch import gradio as gr import numpy as np from datetime import datetime, timezone from huggingface_hub import WebhooksServer, WebhookPayload, hf_hub_download from safetensors.torch import load_file from transformers import AutoTokenizer, AutoModel, AutoConfig from zeta_model import NextDiTPixelSpace from PIL import Image # ═══════════════════════════════════════════════════════════════════ # Config # ═══════════════════════════════════════════════════════════════════ FINETUNE_MODEL_ID = "lodestones/Zeta-Chroma" FINETUNE_FILENAME = "zeta-chroma-base-x0-pixel-no-dino.safetensors" # Dynamically list all .safetensors files in the model repo (picks up future checkpoints automatically) from huggingface_hub import list_repo_files AVAILABLE_CHECKPOINTS = sorted([ f for f in list_repo_files(FINETUNE_MODEL_ID) if f.endswith(".safetensors") ]) # Ensure default is first in the list for the dropdown if FINETUNE_FILENAME in AVAILABLE_CHECKPOINTS: AVAILABLE_CHECKPOINTS.remove(FINETUNE_FILENAME) AVAILABLE_CHECKPOINTS.insert(0, FINETUNE_FILENAME) print(f"Available checkpoints: {AVAILABLE_CHECKPOINTS}") BASE_MODEL_ID = "Tongyi-MAI/Z-Image" CHAT_TEMPLATE = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" DEFAULT_SHIFT = 3.0 DEFAULT_CFG = 3.5 DEFAULT_STEPS = 50 DEFAULT_NEG = "low quality, ugly, unfinished, out of focus, underexposed, overly shadowed, overexposed, blurry, noisy, extra limbs, extra fingers, malformed hands, distorted anatomy" DEFAULT_SEED = 42 SCHEDULERS = ["Simple", "Normal", "Beta"] DEFAULT_SCHEDULER = "Normal" # Samplers: Euler (deterministic), Euler Ancestral (stochastic/adds noise each step) SAMPLERS = ["Euler", "Euler Ancestral", "Heun", "DPM++ 2S"] DEFAULT_SAMPLER = "Euler" RESOLUTIONS = ["512\u00d7512","768\u00d7768","1024\u00d71024","1280\u00d71280","1024\u00d7768","768\u00d71024","1024\u00d7576","576\u00d71024","1280\u00d7768","768\u00d71280","1280\u00d71024","1024\u00d71280"] TEST_PROMPTS = [ {"name": "1 \u00b7 Hyper-Realism","prompt": "Photorealistic portrait of an anthropomorphic grey wolf standing under a flickering neon sign spelling HOWL in glowing magenta letters, on a rain-soaked city street at night. Wet fur with individually visible strands, water droplets catching cyan and magenta neon reflections. Worn leather jacket with realistic creases. Shallow depth of field, bokeh city lights, volumetric fog from steam vents, puddle reflections on wet asphalt. Hyperdetailed fur texture, cinematic color grading, photorealistic rendering.","negative_prompt": DEFAULT_NEG + ", cartoon, anime, flat colors, simple"}, {"name": "2 \u00b7 Semi-Realism / Oil Painting","prompt": "A regal anthropomorphic dragon in the style of a Renaissance oil painting, dramatic chiaroscuro lighting from a single candle. Iridescent emerald and gold scales with intricate patterning, amber eyes with slit pupils reflecting candlelight. Velvet-draped throne room, ornate medieval armor with engraved filigree. Visible brushstroke texture, craquelure aging, rich earth tones, deep shadows, classical composition, museum-quality fine art.","negative_prompt": DEFAULT_NEG + ", cartoon, flat, simple, modern, digital art look"}, {"name": "3 \u00b7 Painterly / Ghibli","prompt": "A cheerful anthropomorphic red fox adventurer in a luminous enchanted forest, vibrant Studio Ghibli-inspired painterly style. Cozy hand-knit green scarf, tiny lantern casting warm golden light on bioluminescent mushrooms and floating pollen. Watercolor-style foliage with paint bleeding at edges, dappled sunlight through ancient tree canopy, whimsical warm atmosphere, complementary orange and teal palette, expressive stylized fur, fantasy illustration.","negative_prompt": DEFAULT_NEG + ", photorealistic, dark, gritty, harsh shadows"}, {"name": "4 \u00b7 Toony Cartoon","prompt": "Extremely cartoony anthropomorphic orange tabby cat with huge exaggerated eyes and a giant mischievous grin, classic 1930s rubber-hose cartoon style. Thick bold outlines, flat cel-shaded colors. Over-the-top surprised pose with stretched noodle limbs, motion lines, stars and exclamation marks around the head. Bright primary colors, simple flat background with speed lines, squash-and-stretch animation, clean vector line art, vintage cartoon aesthetic.","negative_prompt": DEFAULT_NEG + ", realistic, detailed fur, photographic, complex shading"}, {"name": "5 \u00b7 Two Characters / Calm","prompt": "Two anthropomorphic characters at a cozy campfire in a forest clearing at dusk. Left: a tall lean anthro deer with brown fur and branching antlers, flannel shirt and round glasses, reading from a leather-bound book. Right: a small chubby anthro raccoon with grey and black markings, wrapped in an oversized knit blanket, roasting a marshmallow on a stick. Firelight from below casting warm orange highlights and soft shadows on both characters. Each character has distinct body language and expression. Pine trees, stars overhead, autumn leaves, backpack and lantern nearby. Semi-realistic digital painting.","negative_prompt": DEFAULT_NEG + ", merged bodies, fused characters, extra limbs, same face, same species"}, {"name": "6 \u00b7 Two Characters / Action","prompt": "Two anthropomorphic characters fighting in a rain-drenched cyberpunk alley. A muscular anthro tiger with striped fur lunges with glowing blue energy gauntlets, teeth bared. A sleek anthro raven with glossy black feathers and cybernetic wings dives from above with neon green plasma daggers. Different body types and species. Motion blur, rain frozen in midair, neon sign spelling Zeta-Chroma with reflections on wet ground, dramatic foreshortening, high-energy anime-influenced illustration.","negative_prompt": DEFAULT_NEG + ", merged bodies, extra limbs, same species, static pose, stiff, calm"}, {"name": "7 \u00b7 Hands / Paws Detail","prompt": "Close-up of an anthropomorphic otter at a workbench, carefully painting a tiny miniature figurine. Left paw holds the figurine between thumb and index finger, five distinct fingers visible with short claws. Right paw grips a fine-tipped brush. Sharp focus on both paws. Sleek brown fur, cream chest, small round ears, whiskers. Workbench with paint pots, magnifying lamp, scattered miniatures. Soft studio lighting, macro detail, individual fur strands visible.","negative_prompt": DEFAULT_NEG + ", extra fingers, fused fingers, missing fingers, blob paws, mitten hands, no detail"}, {"name": "8 \u00b7 Text Rendering","prompt": "An anthropomorphic golden retriever barista behind a coffee shop counter wearing a green apron with the word ZETA in large white embroidered letters on the chest. A chalkboard behind reads COFFEE in big chalk letters. Holding a paper cup. Warm interior with exposed brick, potted plants, neon OPEN sign in the window. Cozy lighting, shallow depth of field on the barista and text.","negative_prompt": DEFAULT_NEG + ", illegible text, garbled letters, misspelled, backwards text, blurry text"}, {"name": "9 \u00b7 Extreme Lighting","prompt": "Silhouette of an anthropomorphic stag with massive branching antlers on a misty hilltop at sunrise. Intense golden-orange rim lighting traces every edge of the antlers and fur against a purple-to-amber gradient sky. Thick volumetric fog catching god rays streaming between antler tines. Flowing tattered cloak billowing in wind. Dark wildflower shapes in foreground. Dramatic contre-jour, extreme dynamic range, lens flare from the sun directly behind the antlers.","negative_prompt": DEFAULT_NEG + ", flat lighting, overexposed, no fog, no atmosphere, midday"}, {"name": "10 \u00b7 Complex Scene","prompt": "Wide shot of a bustling medieval fantasy marketplace with anthropomorphic animals. Foreground: fox merchant with colorful fabrics under a striped awning. Left: bear blacksmith at a forge. Middle: rabbit children running between stalls. Above: hawk perched on a rooftop. Cobblestone street receding toward a castle on a distant hill. Each stall has distinct wares and signs. Golden hour lighting, timber-frame buildings, flower baskets, chimney smoke, birds in the sky. Rich fantasy illustration with layered depth.","negative_prompt": DEFAULT_NEG + ", empty background, flat, single character, close-up, minimal detail, simple"}, {"name": "11 \u00b7 Legoshi (Beastars)","prompt": "Legoshi from Beastars, the tall grey wolf with gentle eyes, wearing a yellow Hawaiian shirt and flip-flops, relaxing on a tropical beach at sunset. He is building an elaborate sandcastle with tiny flags on top, looking peaceful and content. His large ears are relaxed and slightly back. Ocean waves in the background, palm trees, warm golden light on his grey fur. Soft anime-influenced semi-realistic style matching the Beastars aesthetic, gentle atmosphere.","negative_prompt": DEFAULT_NEG + ", aggressive, scary, dark, school uniform, indoor"}, {"name": "12 \u00b7 Robin Hood (Disney)","prompt": "Robin Hood the fox from the 1973 Disney film, recognizable with his orange-red fur and charming smirk, but dressed in a modern tailored navy suit and tie, sitting at a sleek office desk in a glass-walled corner office in a skyscraper. He leans back confidently in an executive chair, one leg crossed, holding a coffee mug. City skyline visible through floor-to-ceiling windows behind him. Clean Disney animation style with smooth lines, warm colors, and expressive features.","negative_prompt": DEFAULT_NEG + ", medieval, tunic, hat with feather, forest, realistic fur"}, {"name": "13 \u00b7 Nick Wilde (Zootopia)","prompt": "Nick Wilde the red fox from Zootopia in a fantasy RPG setting, wearing ornate leather rogue armor with a hooded cloak, perched on a moonlit castle parapet. He holds a glowing enchanted dagger and has his signature sly half-smile. His green eyes catch the moonlight. Stone castle architecture, starry night sky, distant torchlit medieval town below. Painterly fantasy illustration style with rich colors, keeping Nick's recognizable facial features and smug expression.","negative_prompt": DEFAULT_NEG + ", police uniform, modern city, Judy Hopps, Hawaiian shirt, realistic"}, {"name": "14 \u00b7 Sabrina (Sabrina Online)","prompt": "Sabrina the skunk from Sabrina Online, a young adult female anthropomorphic skunk with her distinctive black and white fur pattern and long headfur, attending a busy comic convention. She wears a casual graphic t-shirt and jeans, carrying a tote bag full of comics, looking excited and a bit overwhelmed by the crowd. Convention hall background with colorful banners, cosplayers out of focus behind her, booth displays. Clean webcomic-inspired line art style with cel shading, bright indoor lighting.","negative_prompt": DEFAULT_NEG + ", realistic, photographic, nude, office setting, computer desk"}, {"name": "15 \u00b7 Bojack Horseman","prompt": "Bojack Horseman, the brown horse with the long face and dark mane from the Netflix show, standing outdoors in a sunlit meadow painting on a canvas easel. He wears a paint-splattered smock over his usual teal sweater, holding a palette and brush, looking genuinely focused and at peace for once. Rolling green hills, wildflowers, blue sky with scattered clouds. Flat graphic style matching the show's distinctive aesthetic with bold outlines, muted warm palette, and simple shapes.","negative_prompt": DEFAULT_NEG + ", 3D, realistic, bar, alcohol, sad, depressed, indoor, Hollywoo"}, {"name": "16 \u00b7 Judy & Fox McCloud (Crossover)","prompt": "Judy Hopps the grey rabbit from Zootopia and Fox McCloud the fox from Star Fox cooking together in a cozy home kitchen. Judy stands on a step stool to reach the counter, energetically chopping vegetables with determination. Fox McCloud leans against the counter beside her with a confident grin, casually stirring a pot with one paw, still wearing his red scarf. Size difference clearly visible between the small rabbit and taller fox. Warm kitchen lighting, hanging copper pots, recipe book open on the counter, steam rising from the pot. Colorful crossover illustration style.","negative_prompt": DEFAULT_NEG + ", police uniforms, space suits, dark, realistic, same size, merged"}, ] # Example prompts for the Custom tab EXAMPLES = [ "A cyberpunk anthro cat hacker sitting at a multi-monitor setup in a dark room, neon glow on fur, hoodie, energy drink cans scattered around", "Portrait of an elegant anthro snow leopard in a flowing silk dress at a moonlit ball, baroque architecture, chandeliers, soft focus", "An anthro wolf blacksmith hammering a glowing sword at a forge, sparks flying, muscular arms, leather apron, medieval workshop", ] # ═══════════════════════════════════════════════════════════════════ # Model loading (CUDA emulation at startup — no GPU quota) # ═══════════════════════════════════════════════════════════════════ def make_sigmas(num_steps, shift=3.0, scheduler="Simple"): """Generate sigma schedule. - Simple/Normal: linear spacing with shift applied - Beta: scipy beta distribution spacing (original workflow, may cause darkness) """ if scheduler == "Beta": from scipy.stats import beta as beta_dist ts = 1.0 - np.linspace(0, 1, num_steps, endpoint=False) ts = beta_dist.ppf(ts, 0.6, 0.6) sigmas = shift * ts / (1 + (shift - 1) * ts) elif scheduler == "Normal": # Normal: linspace but with slight cosine weighting for smoother transitions ts = np.linspace(1, 0, num_steps + 1)[:-1] ts = 0.5 * (1 + np.cos(np.pi * (1 - ts))) # cosine remap sigmas = shift * ts / (1 + (shift - 1) * ts) else: # Simple (default) ts = np.linspace(1, 0, num_steps + 1)[:-1] sigmas = shift * ts / (1 + (shift - 1) * ts) return torch.from_numpy(np.append(sigmas, 0.0)).float() print("=" * 60) print("Loading Qwen3 text encoder...") tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, subfolder="tokenizer") text_encoder = AutoModel.from_pretrained(BASE_MODEL_ID, subfolder="text_encoder", torch_dtype=torch.bfloat16, output_hidden_states=True) text_encoder.eval().to("cuda") print("Downloading Zeta-Chroma checkpoint...") ckpt_path = hf_hub_download(repo_id=FINETUNE_MODEL_ID, filename=FINETUNE_FILENAME) # Get checkpoint metadata — commit SHA, last modified date, and age from huggingface_hub import model_info as _mi from datetime import datetime, timezone as _tz _m = _mi(FINETUNE_MODEL_ID) MODEL_COMMIT_SHA = _m.sha[:12] MODEL_FILE_DATE = str(_m.last_modified)[:19] if _m.last_modified else 'unknown' # Compute human-readable age # Record when this Space instance started (= when this checkpoint was loaded) # This is what 'loaded X ago' shows — how long this process has been running from datetime import datetime, timezone as _tz _STARTUP_TIME = datetime.now(_tz.utc) _STARTUP_ISO = _STARTUP_TIME.strftime('%Y-%m-%d %H:%M:%S') MODEL_FILE_DATE = _STARTUP_ISO # used by JS badge to compute live age MODEL_FILE_AGE = 'just now' # initial value, JS will overwrite MODEL_ID_STRING = f'{FINETUNE_MODEL_ID} @ {MODEL_COMMIT_SHA}' print(f"Model ID: {MODEL_ID_STRING} (loaded at {_STARTUP_ISO} UTC)") print("Creating & loading model...") model = NextDiTPixelSpace(patch_size=32, in_channels=3, dim=3840, n_layers=30, n_refiner_layers=2, n_heads=30, n_kv_heads=30, ffn_dim_multiplier=8/3, cap_feat_dim=2560, decoder_hidden_size=3840, decoder_num_res_blocks=4, decoder_max_freqs=8, decoder_in_channels=3072) sd = load_file(ckpt_path, device="cpu") model.load_state_dict(sd, strict=False) del sd; gc.collect() model = model.to(dtype=torch.bfloat16, device="cuda").eval() print(f"Ready! {sum(p.numel() for p in model.parameters()):,} params") print("=" * 60) RESULTS_LOG = [] # ═══════════════════════════════════════════════════════════════════ # Generation core # ═══════════════════════════════════════════════════════════════════ def encode_prompt(prompt): templated = CHAT_TEMPLATE.format(prompt) inputs = tokenizer(templated, return_tensors="pt", padding=False, truncation=True, max_length=4096, return_attention_mask=True).to("cuda") with torch.no_grad(): return text_encoder(**inputs, output_hidden_states=True).hidden_states[-2].to(torch.bfloat16) # Track which checkpoint is currently loaded _current_ckpt = FINETUNE_FILENAME def _load_checkpoint(ckpt_name): """Load a different checkpoint if not already loaded. Returns True if swapped.""" global model, _current_ckpt, MODEL_COMMIT_SHA, MODEL_ID_STRING, MODEL_FILE_DATE, MODEL_FILE_AGE if ckpt_name == _current_ckpt: return False print(f"Swapping checkpoint: {_current_ckpt} -> {ckpt_name}") new_path = hf_hub_download(repo_id=FINETUNE_MODEL_ID, filename=ckpt_name) new_sd = load_file(new_path, device="cpu") model.load_state_dict(new_sd, strict=False) del new_sd; gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() _current_ckpt = ckpt_name print(f"Checkpoint loaded: {ckpt_name}") return True def generate_single(prompt, negative_prompt, width, height, steps, cfg, shift, seed, scheduler=None, sampler=None): generator = torch.Generator("cuda").manual_seed(seed) prompt_embeds = encode_prompt(prompt) neg_embeds = encode_prompt(negative_prompt) if (negative_prompt and negative_prompt.strip()) else encode_prompt("") sigmas = make_sigmas(steps, shift=shift, scheduler=scheduler or DEFAULT_SCHEDULER).to("cuda") x = torch.randn(1, 3, height, width, generator=generator, device="cuda", dtype=torch.bfloat16) _sampler = sampler or DEFAULT_SAMPLER def _model_call(x_in, sigma_in): """Single model forward pass with CFG. Extracted for reuse by higher-order samplers.""" sb = sigma_in.unsqueeze(0) if sigma_in.dim() == 0 else sigma_in if cfg > 1.0: oc = model(x_in, sb, prompt_embeds) ou = model(x_in, sb, neg_embeds) return ou + cfg * (oc - ou) return model(x_in, sb, prompt_embeds) for i in range(steps): sigma, sigma_next = sigmas[i], sigmas[i + 1] dt = sigma_next - sigma output = _model_call(x, sigma) if _sampler == "Euler Ancestral" and sigma_next > 0: # Euler Ancestral: stochastic step with noise injection sigma_up = min(sigma_next, (sigma_next**2 * (sigma**2 - sigma_next**2) / sigma**2) ** 0.5) sigma_down = (sigma_next**2 - sigma_up**2) ** 0.5 x = x + (sigma_down - sigma) * output x = x + sigma_up * torch.randn_like(x) elif _sampler == "Heun" and sigma_next > 0: # Heun (2nd order, like res2s): take Euler step, evaluate at new point, average x_euler = x + dt * output output_2 = _model_call(x_euler, sigma_next) x = x + dt * 0.5 * (output + output_2) elif _sampler == "DPM++ 2S" and sigma_next > 0: # DPM++ 2nd order single-step (similar to res2s/bong_tangent quality) # Midpoint method: evaluate at sigma_mid, use that for the full step sigma_mid = (sigma + sigma_next) / 2 dt_half = sigma_mid - sigma x_mid = x + dt_half * output output_mid = _model_call(x_mid, sigma_mid) x = x + dt * output_mid else: # Euler (deterministic, default) x = x + dt * output return Image.fromarray((x.add(1).div(2).clamp(0,1).squeeze(0).permute(1,2,0).float().cpu().numpy() * 255).round().astype("uint8")) def parse_resolution(s): p = s.replace("x","\u00d7").split("\u00d7") return int(p[0].strip()), int(p[1].strip()) def format_log(): return "\n".join(f"[{e['ts']}] {e['n']} images \u00b7 {e.get('sha','')}" for e in RESULTS_LOG) or "No runs yet." # ═══════════════════════════════════════════════════════════════════ # Duration estimators for @spaces.GPU # ═══════════════════════════════════════════════════════════════════ def _estimate_custom_duration(prompt, neg, res, steps, cfg, shift, scheduler, sampler, ckpt, seed, rand_seed): w, h = parse_resolution(res) return min(max(int((int(steps) * 0.25 * (w*h/(1024*1024)) * (2 if cfg > 1 else 1) + 3) * 1.5), 10), 120) # IMPORTANT: Receives ALL btn.click inputs (22 args). Must use *args. def _estimate_suite_duration(*args, **kwargs): return min(int((len(TEST_PROMPTS) * (30 * 0.5 + 3)) * 1.5), 250) # ═══════════════════════════════════════════════════════════════════ # GPU functions # ═══════════════════════════════════════════════════════════════════ @spaces.GPU(duration=_estimate_suite_duration) def run_test_suite(seed, steps, cfg, shift, res, neg_override, scheduler, sampler, ckpt, *prompt_texts): seed = int(seed) if seed else DEFAULT_SEED steps = int(steps) if steps else DEFAULT_STEPS cfg = float(cfg) if cfg else DEFAULT_CFG shift = float(shift) if shift else DEFAULT_SHIFT scheduler = scheduler if scheduler else DEFAULT_SCHEDULER sampler = sampler if sampler else DEFAULT_SAMPLER if ckpt and ckpt != _current_ckpt: _load_checkpoint(ckpt) w, h = parse_resolution(res) if res else (1024, 1024) neg = neg_override.strip() if (neg_override and neg_override.strip()) else DEFAULT_NEG ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") n = len(TEST_PROMPTS) results = [] for i, pc in enumerate(TEST_PROMPTS): prompt = prompt_texts[i].strip() if (i < len(prompt_texts) and prompt_texts[i] and prompt_texts[i].strip()) else pc["prompt"] prompt_neg = pc["negative_prompt"].replace(DEFAULT_NEG, neg) if DEFAULT_NEG in pc["negative_prompt"] else neg pct = int((i / n) * 100) status_html = f'''
\u23f3 {pc["name"]} {i+1}/{n}
seed={seed} \u00b7 {w}\u00d7{h} \u00b7 {steps} steps \u00b7 CFG {cfg} \u00b7 {scheduler} \u00b7 {sampler}
''' img = generate_single(prompt, prompt_neg, w, h, steps, cfg, shift, seed, scheduler, sampler) results.append((img, pc['name'])) yield results, status_html, format_log() done_html = f'''
\u2705 Complete!
{n} images \u00b7 {MODEL_COMMIT_SHA} \u00b7 seed={seed} \u00b7 {w}\u00d7{h} \u00b7 {steps}s \u00b7 CFG {cfg}
''' RESULTS_LOG.insert(0, {"ts": ts, "n": n, "sha": MODEL_COMMIT_SHA}) yield results, done_html, format_log() @spaces.GPU(duration=_estimate_custom_duration) def run_custom(prompt, neg, res, steps, cfg, shift, scheduler, sampler, ckpt, seed, rand_seed): if not prompt.strip(): raise gr.Error("Please enter a prompt.") w, h = parse_resolution(res) if ckpt and ckpt != _current_ckpt: _load_checkpoint(ckpt) if rand_seed: seed = random.randint(0, 2**32 - 1) return generate_single(prompt, neg, w, h, int(steps), cfg, shift, int(seed), scheduler, sampler), str(seed) # ═══════════════════════════════════════════════════════════════════ # UI # ═══════════════════════════════════════════════════════════════════ CSS = """ .status-box { min-height: 60px; } .model-badge { display: inline-block; padding: 5px 12px; border-radius: 20px; background: #1e293b; color: #94a3b8; font-size: 11px; font-family: monospace; margin: 2px 4px 2px 0; border: 1px solid #334155; } .gallery-item img { border-radius: 8px; } /* Tighter accordion spacing */ .accordion { margin-bottom: 4px !important; } /* Settings group visual */ .settings-group { border-left: 3px solid #6366f1; padding-left: 12px; margin: 8px 0; } """ with gr.Blocks(title="Zeta-Chroma Test Bench", theme=gr.themes.Soft(), css=CSS, js="""function() { setTimeout(function() { var el = document.getElementById("model-age"); if (!el) return; var utc = el.getAttribute("data-utc"); if (!utc || utc === "unknown" || utc === "just now") return; var then = new Date(utc.replace(" ", "T") + "Z").getTime(); if (isNaN(then)) return; var secs = Math.floor((Date.now() - then) / 1000); var txt; if (secs < 60) txt = secs + "s ago"; else if (secs < 3600) txt = Math.floor(secs/60) + " min ago"; else if (secs < 86400) { var h = Math.floor(secs/3600); txt = h + " hour" + (h!=1?"s":"") + " ago"; } else if (secs < 604800) { var d = Math.floor(secs/86400); txt = d + " day" + (d!=1?"s":"") + " ago"; } else { var w = Math.floor(secs/604800); txt = w + " week" + (w!=1?"s":"") + " ago"; } el.textContent = "\u231b loaded " + txt; }, 300); }""") as ui: # ── Header ── gr.Markdown("# \U0001f3a8 Zeta-Chroma Test Bench") gr.HTML(f'
\U0001f516 {FINETUNE_MODEL_ID} @ {MODEL_COMMIT_SHA}\u231b loaded {MODEL_FILE_AGE}\u26a0\ufe0f WIP\U0001f4e1 Pixel-space \u00b7 No VAE
') with gr.Tabs(): # ══════════════════ TEST SUITE TAB ══════════════════ with gr.TabItem("\U0001f9ea Test Suite", id="suite"): # Checkpoint selector at the very top suite_ckpt = gr.Dropdown(label="Checkpoint", choices=AVAILABLE_CHECKPOINTS, value=FINETUNE_FILENAME) # Status + controls st = gr.HTML(value='
\U0001f3a8 Ready to generate \u2014 adjust settings below and click Run All 16 Prompts
', elem_classes="status-box") with gr.Row(equal_height=True): btn = gr.Button("\U0001f680 Run All 16 Prompts", variant="primary", size="lg", scale=3) stop_btn = gr.Button("\u23f9 Stop", variant="stop", size="lg", scale=1) # Gallery (the main content) g = gr.Gallery(label="Results", columns=2, object_fit="contain", format="png", show_download_button=True, show_share_button=True, preview=True) # Settings in a collapsible section with gr.Accordion("\u2699\ufe0f Settings", open=True): with gr.Row(equal_height=True): suite_seed = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0) suite_res = gr.Dropdown(label="Resolution", choices=RESOLUTIONS, value="1024\u00d71024") suite_sampler = gr.Dropdown(label="Sampler", choices=SAMPLERS, value=DEFAULT_SAMPLER) suite_sched = gr.Dropdown(label="Scheduler", choices=SCHEDULERS, value=DEFAULT_SCHEDULER) with gr.Row(equal_height=True): suite_steps = gr.Slider(label="Steps", minimum=10, maximum=100, value=DEFAULT_STEPS, step=1) suite_cfg = gr.Slider(label="CFG", minimum=1.0, maximum=15.0, value=DEFAULT_CFG, step=0.1) suite_shift = gr.Slider(label="Shift", minimum=1.0, maximum=15.0, value=DEFAULT_SHIFT, step=0.5) suite_neg = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEG, lines=1) with gr.Row(): restore_settings_btn = gr.Button("\U0001f504 Reset Settings", size="sm", variant="secondary") restore_prompts_btn = gr.Button("\U0001f504 Reset All Prompts", size="sm", variant="secondary") # Editable prompts in grouped accordions gr.Markdown("#### \U0001f4dd Prompts *(editable \u2014 resets on refresh)*") prompt_boxes = [] # Group 1: Style range (1-4) with gr.Accordion("\U0001f3a8 Style Tests (1\u20134): Realism \u2192 Cartoon", open=False): for i in range(0, 4): tb = gr.Textbox(value=TEST_PROMPTS[i]["prompt"], label=TEST_PROMPTS[i]["name"], lines=2, show_copy_button=True) prompt_boxes.append(tb) # Group 2: Multi-character (5-6) with gr.Accordion("\U0001f46b Multi-Character (5\u20136): Interaction & Action", open=False): for i in range(4, 6): tb = gr.Textbox(value=TEST_PROMPTS[i]["prompt"], label=TEST_PROMPTS[i]["name"], lines=2, show_copy_button=True) prompt_boxes.append(tb) # Group 3: Technical (7-10) with gr.Accordion("\U0001f527 Technical Tests (7\u201310): Hands, Text, Lighting, Scenes", open=False): for i in range(6, 10): tb = gr.Textbox(value=TEST_PROMPTS[i]["prompt"], label=TEST_PROMPTS[i]["name"], lines=2, show_copy_button=True) prompt_boxes.append(tb) # Group 4: Known characters (11-16) with gr.Accordion("\u2b50 Known Characters (11\u201316): Beastars, Disney, Zootopia, etc.", open=False): for i in range(10, 16): tb = gr.Textbox(value=TEST_PROMPTS[i]["prompt"], label=TEST_PROMPTS[i]["name"], lines=2, show_copy_button=True) prompt_boxes.append(tb) with gr.Accordion("\U0001f4dc Run History", open=False): log = gr.Textbox(value=format_log(), interactive=False, lines=3, show_label=False) # Wire events suite_event = btn.click(fn=run_test_suite, inputs=[suite_seed, suite_steps, suite_cfg, suite_shift, suite_res, suite_neg, suite_sched, suite_sampler, suite_ckpt] + prompt_boxes, outputs=[g, st, log]) stop_btn.click(fn=None, cancels=[suite_event]) restore_settings_btn.click(fn=lambda: (DEFAULT_SEED, "1024\u00d71024", DEFAULT_STEPS, DEFAULT_CFG, DEFAULT_SHIFT, DEFAULT_NEG, DEFAULT_SCHEDULER, DEFAULT_SAMPLER, FINETUNE_FILENAME), outputs=[suite_seed, suite_res, suite_steps, suite_cfg, suite_shift, suite_neg, suite_sched, suite_sampler, suite_ckpt]) restore_prompts_btn.click(fn=lambda: tuple(p["prompt"] for p in TEST_PROMPTS), outputs=prompt_boxes) # ══════════════════ CUSTOM PROMPT TAB ══════════════════ with gr.TabItem("\U0001f3a8 Custom", id="custom"): c_ckpt = gr.Dropdown(label="Checkpoint", choices=AVAILABLE_CHECKPOINTS, value=FINETUNE_FILENAME) with gr.Row(): with gr.Column(scale=1): cp = gr.Textbox(label="Prompt", placeholder="Describe your image...", lines=4) cn = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEG, lines=2) with gr.Row(): cr = gr.Dropdown(label="Resolution", choices=RESOLUTIONS, value="1024\u00d71024", scale=2) cse = gr.Number(label="Seed", value=42, precision=0, scale=1) csr = gr.Checkbox(label="Random", value=True, scale=1) with gr.Row(equal_height=True): cs = gr.Slider(label="Steps", minimum=10, maximum=100, value=DEFAULT_STEPS, step=1) cc = gr.Slider(label="CFG", minimum=1.0, maximum=15.0, value=DEFAULT_CFG, step=0.1) csh = gr.Slider(label="Shift", minimum=1.0, maximum=15.0, value=DEFAULT_SHIFT, step=0.5) with gr.Row(equal_height=True): c_sampler = gr.Dropdown(label="Sampler", choices=SAMPLERS, value=DEFAULT_SAMPLER) c_sched = gr.Dropdown(label="Scheduler", choices=SCHEDULERS, value=DEFAULT_SCHEDULER) with gr.Row(): cb = gr.Button("\U0001f3a8 Generate", variant="primary", size="lg", scale=3) cso = gr.Textbox(label="Seed", interactive=False, scale=1) # Example prompts gr.Markdown("#### \U0001f4a1 Examples *(click to use)*") gr.Examples(examples=[[e] for e in EXAMPLES], inputs=[cp], label="") with gr.Column(scale=2): ci = gr.Image(label="Generated Image", type="pil", show_download_button=True, show_share_button=True) cb.click(fn=run_custom, inputs=[cp, cn, cr, cs, cc, csh, c_sched, c_sampler, c_ckpt, cse, csr], outputs=[ci, cso]) # Footer gr.HTML('
') gr.Markdown(f"*[Zeta-Chroma](https://huggingface.co/lodestones/Zeta-Chroma) \u00b7 `{MODEL_COMMIT_SHA}` \u00b7 [Source Code](https://huggingface.co/spaces/chipfoxx/zeta-chroma-test-bench/tree/main)*") # ═══════════════════════════════════════════════════════════════════ # Webhook # ═══════════════════════════════════════════════════════════════════ app = WebhooksServer(ui=ui, webhook_secret=os.getenv("WEBHOOK_SECRET")) @app.add_webhook("/model_updated") async def on_model_updated(payload: WebhookPayload): """When lodestones/Zeta-Chroma gets a new commit, restart the Space to load new weights. Why restart instead of hot-swap: - ZeroGPU's CUDA emulation mode doesn't support reliable in-process weight reload - restart_space() triggers a clean re-run of startup code which re-downloads the latest checkpoint - Downtime is ~60-90s (no rebuild needed, just runtime restart) """ if payload.event.action != "update" or payload.repo.type != "model": return {"processed": False, "reason": "not a model update"} if payload.repo.name != FINETUNE_MODEL_ID: return {"processed": False, "reason": f"wrong repo: {payload.repo.name}"} sha = payload.repo.head_sha[:7] if payload.repo.head_sha else "?" print(f"\U0001f514 Webhook: {payload.repo.name} updated (commit {sha}) — restarting Space to load new checkpoint...") # Trigger a Space restart — the process will be killed and relaunched, # causing the startup code to re-download and load the latest safetensors file. try: from huggingface_hub import HfApi as _HfApi _api = _HfApi(token=os.environ.get("HF_TOKEN")) _api.restart_space(repo_id="chipfoxx/zeta-chroma-test-bench", factory_reboot=False) except Exception as e: print(f"\u274c Failed to restart: {e}") return {"processed": False, "error": str(e)} return {"processed": True, "commit": sha, "action": "restarting"} app.launch()