Fix example clicks (cache_examples=False); Custom output size (2048² T2I / 1536² edit cap); UI: short header, Advanced settings, examples reproduce with their seed, Viggle theme

README.md

@@ -56,7 +56,8 @@ app produced for it (fixed seed; 6 steps and prompt enhancement on unless the ro
 otherwise — the launch-graphic row is 8 steps at the 3:2 2048² size with the 3k-character brief sent as
 written, and the five `qwen_*` rows are sent as written at 2048² / 1536²), so you can see what the model
 does without spending GPU time — click a row to load the prompt, the references, the result, and the
-steps / size / enhancement settings that produced it. `release/render_examples.py` regenerates the rows
+steps / size / enhancement settings and seed that produced it (with *Randomize seed* switched off), so
+**Generate** reproduces the row. `release/render_examples.py` regenerates the rows
 and `examples/manifest.json` whenever the weights change. The reference photos
 `woman1/woman2/cat/cat_window/bird.webp` and the three-reference prompt come from the
 [black-forest-labs/flux-klein-9b-kv](https://huggingface.co/spaces/black-forest-labs/flux-klein-9b-kv) Space;
@@ -146,6 +147,12 @@ reference at 1024² area, which is what the pipeline does when `height`/`width`
 dimensions onto the 1536² bucket of the nearest aspect ratio, and says so in the status line, so an
 API client or a stale dropdown cannot push a three-reference call out of budget.
 
+**Custom** (last entry of both menus) shows Width / Height sliders, 512 to 2720 in steps of 32, pre-filled
+with the last preset picked. A custom size above **2048² area for text-to-image** or **1536² area for editing**
+(the largest trained bucket of each mode) is scaled down to that area keeping its ratio, floored to multiples
+of 32, and the status line says so; smaller sizes are used as given. Sizes off the trained buckets work but
+were not evaluated.
+
 Reference images are always resized to **1024² area** before encoding regardless of the output size,
 so extra references cost VRAM and time but not output resolution.
 

app.py

@@ -78,6 +78,7 @@ EXAMPLES = json.loads((EXAMPLES_DIR / "manifest.json").read_text(encoding="utf-8
 # different menus. Reference images are always encoded at 1024^2 area regardless.
 RATIOS = [("1:1", 1.0), ("16:9", 16 / 9), ("9:16", 9 / 16), ("4:3", 4 / 3), ("3:4", 3 / 4), ("3:2", 3 / 2), ("2:3", 2 / 3)]
 AUTO = "Auto · match the last reference (1024² area)"
+CUSTOM = "Custom · set width and height below"
 
 
 def _bucket(area):
@@ -88,11 +89,16 @@ def _bucket(area):
     return sizes
 
 
-SIZES = {AUTO: None, **_bucket(1024), **_bucket(1536), **_bucket(2048)}
-T2I_CHOICES = [*_bucket(1024), *_bucket(2048)]
-EDIT_CHOICES = [AUTO, *_bucket(1024), *_bucket(1536)]
+SIZES = {AUTO: None, CUSTOM: None, **_bucket(1024), **_bucket(1536), **_bucket(2048)}
+T2I_CHOICES = [*_bucket(1024), *_bucket(2048), CUSTOM]
+EDIT_CHOICES = [AUTO, *_bucket(1024), *_bucket(1536), CUSTOM]
 # the (width, height) pairs the editing menu actually offers, in menu order
 EDIT_DIMS = [SIZES[label] for label in EDIT_CHOICES if SIZES[label]]
+# Custom sizes keep their aspect ratio but are scaled into the mode's largest distilled bucket, which is also what the
+# memory table in the README measured: 2048² area for text-to-image, 1536² for editing. The sliders reach the longest
+# side any preset has.
+CUSTOM_CAP = {False: 2048, True: 1536}  # keyed by "has references"
+MAX_SIDE = max(max(size) for size in SIZES.values() if size)
 
 if STUDENT == "full":
     transformer = QwenImage21Transformer2DModel.from_pretrained(
@@ -168,27 +174,36 @@ def enhance_prompt(prompt, images, ratio_name):
 
 
 @gpu
-def generate(prompt, image_1, image_2, image_3, size_label, seed, randomize_seed, enhance=True, steps=STEPS):
+def generate(prompt, image_1, image_2, image_3, size_label, seed, randomize_seed, enhance=True, steps=STEPS,
+             custom_width=1024, custom_height=1024):
     steps = min(max(int(steps), 3), 8)  # 6 (RAW_NODES) is the validated schedule; 3-8 are accepted, fewer or more degrade quickly
     images = [image for image in (image_1, image_2, image_3) if image is not None]
     # unchecked "Randomize seed" always honours the box, so the UI never says one thing and does another
     seed = random.randint(0, MAX_SEED) if (randomize_seed or seed is None) else max(int(seed), 0)
     size = SIZES.get(size_label)
     width, height = size if size else (None, None)
+    ratio_name = size_label.split(" · ")[0]
+    clamped = ""
+    if size_label == CUSTOM:
+        cap = CUSTOM_CAP[bool(images)]
+        scale = min(1.0, cap / (custom_width * custom_height) ** 0.5)
+        width, height = int(custom_width * scale) // 32 * 32, int(custom_height * scale) // 32 * 32
+        ratio_name = f"{width}:{height}"
+        if scale < 1:
+            clamped = f" · scaled to {width}×{height} ({'editing' if images else 'text-to-image'} caps at {cap}² area)"
     # Belt and braces: an editing call must land on one of the editing buckets even if the caller
     # reaches generate() with a text-to-image size (stale dropdown, gr.Examples, API client). Test
     # membership in the offered dims rather than a raw area threshold: calculate_dimensions(1536**2,
     # 4/3) is 1760x1344 = 2_365_440 px, slightly *above* 1536**2, so an area test would flag two
     # sizes the editing menu itself offers. Snap to the 1536² bucket of the nearest named ratio.
-    clamped = ""
-    if images and width and (width, height) not in EDIT_DIMS:
+    elif images and width and (width, height) not in EDIT_DIMS:
         ratio = min(RATIOS, key=lambda name_ratio: abs(name_ratio[1] - width / height))[1]
         width, height, _ = calculate_dimensions(1536 * 1536, ratio)
         clamped = f" · clamped to {width}×{height} (editing caps at the 1536² bucket)"
     used_prompt, enhance_note = prompt, ""
     if enhance and prompt.strip():
         started = time.perf_counter()
-        used_prompt = enhance_prompt(prompt, images, size_label.split(" · ")[0])
+        used_prompt = enhance_prompt(prompt, images, ratio_name)
         enhance_note = f" · enhance {time.perf_counter() - started:.1f}s"
         if used_prompt == prompt:
             enhance_note += " (rewrite failed to parse, original prompt used)"
@@ -217,20 +232,44 @@ def refresh_sizes(image_1, image_2, image_3, current):
     return gr.update(choices=choices, value=current if current in choices else choices[0])
 
 
-with gr.Blocks(title="Viggle Turbo v0.2.1 · Qwen-Image-2.1 6-step") as demo:
+def custom_size_controls(size_label, width, height):
+    """The width/height sliders show only for Custom; a preset pre-fills them, so Custom starts from the last size picked."""
+    size = SIZES.get(size_label)
+    return gr.update(visible=size_label == CUSTOM), size[0] if size else width, size[1] if size else height
+
+
+# Viggle brand tokens (viggle-brandbook): jolt-green calls to action with ink text, warm neutrals, Satoshi (from Fontshare)
+JOLT = gr.themes.Color(c50="#e6fdec", c100="#c2f9d0", c200="#8ff2a9", c300="#52ea7c", c400="#1fe456", c500="#00e13f",
+                       c600="#00b833", c700="#008f28", c800="#006b1e", c900="#004d16", c950="#002e0d", name="jolt")
+WARM = gr.themes.Color(c50="#fcfcfc", c100="#f1eeea", c200="#e1e1e1", c300="#d9d9d9", c400="#9a9a9a", c500="#74685a",
+                       c600="#4d453d", c700="#2e2a25", c800="#221d18", c900="#1b1614", c950="#151110", name="warm")
+THEME = gr.themes.Default(primary_hue=JOLT, secondary_hue=JOLT, neutral_hue=WARM,
+                          font=[gr.themes.Font("Satoshi"), "ui-sans-serif", "system-ui", "sans-serif"]).set(
+    body_text_color="#29231e", body_text_color_dark="#f4f1ec",
+    button_primary_background_fill="#00e13f", button_primary_background_fill_dark="#00e13f",
+    button_primary_background_fill_hover="#00c937", button_primary_background_fill_hover_dark="#00c937",
+    button_primary_border_color="#00e13f", button_primary_border_color_dark="#00e13f",
+    button_primary_text_color="#29231e", button_primary_text_color_dark="#29231e",
+)
+FONT_HEAD = '<link rel="stylesheet" href="https://api.fontshare.com/v2/css?f[]=satoshi@400,500,700&display=swap">'
+
+with gr.Blocks(title="Viggle Turbo v0.2.1 · Qwen-Image-2.1 6-step", theme=THEME, head=FONT_HEAD) as demo:
     gr.Markdown(
         "# Viggle Turbo v0.2.1 — 6-step Qwen-Image-2.1\n"
-        "A DMD-distilled student of **Qwen-Image-2.1** that generates and edits in **6 sampling steps**, with no "
-        "classifier-free guidance: about **5× faster** than the 40-step base model and very competitive with it in quality, "
-        "and on some prompts we prefer its output. The **Comparison** tab has 33 examples of the official Qwen Space side by "
-        "side, turbo vs base, same seed. Complicated edits (multi-reference composition, face swaps, identity-preserving "
-        "edits) can still fall short of the base model.\n\n"
-        "Leave the reference images empty for text-to-image; add one to three of them to edit, compose or transfer style. "
-        "The size menu switches to the editing buckets as soon as a reference is attached. "
-        "**v0.2.1 (2026-09-24):** the step-700 checkpoint of the v0.2 run on the 6-step schedule — intra-prompt diversity "
-        "0.98× the base model, 0% composition drift; earlier versions and the numbers are on the model card. "
-        f"Weights: **{STUDENT_TAG}** from [{STUDENT_REPO}](https://huggingface.co/{STUDENT_REPO})."
+        "Text-to-image and image editing in **6 steps**: about **5× faster** than the 40-step Qwen-Image-2.1 and very "
+        "competitive with it in quality — see the **Comparison** tab. Leave the references empty for text-to-image, "
+        "or add one to three to edit, compose or transfer style."
     )
+    with gr.Accordion("About this model", open=False):
+        gr.Markdown(
+            "A DMD-distilled student of **Qwen-Image-2.1**, sampled with no classifier-free guidance. On some prompts we "
+            "prefer its output to the base model's; complicated edits (multi-reference composition, face swaps, "
+            "identity-preserving edits) can still fall short of it. The **Comparison** tab has 33 examples of the official "
+            "Qwen Space side by side, turbo vs base, same seed.\n\n"
+            "**v0.2.1 (2026-09-24):** the step-700 checkpoint of the v0.2 run on the 6-step schedule — intra-prompt diversity "
+            "0.98× the base model, 0% composition drift; earlier versions and the numbers are on the model card. "
+            f"Weights: **{STUDENT_TAG}** from [{STUDENT_REPO}](https://huggingface.co/{STUDENT_REPO})."
+        )
     with gr.Tab("Generate"):
         with gr.Row():
             with gr.Column(scale=3):
@@ -239,23 +278,28 @@ with gr.Blocks(title="Viggle Turbo v0.2.1 · Qwen-Image-2.1 6-step") as demo:
                     image_1 = gr.Image(label="Reference 1 (optional)", type="pil", image_mode="RGB", height=200)
                     image_2 = gr.Image(label="Reference 2 (optional)", type="pil", image_mode="RGB", height=200)
                     image_3 = gr.Image(label="Reference 3 (optional)", type="pil", image_mode="RGB", height=200)
-                with gr.Row():
-                    # allow_custom_value: gradio validates API calls against the *initial* choices (the text-to-image
-                    # menu), which would reject every editing size sent through /generate; generate() snaps
-                    # anything off-menu itself.
-                    size_label = gr.Dropdown(label="Output size", choices=T2I_CHOICES, value=T2I_CHOICES[0], allow_custom_value=True)
-                    steps = gr.Slider(label="Steps", minimum=3, maximum=8, step=1, value=STEPS,
-                                      info="6 is the validated schedule. Extra steps are always added at the high-noise end (7 also validated, 8 untested); fewer steps drop them again (4 = training nodes, 3 = uniform spacing, off-contract).")
-                with gr.Row():
-                    seed = gr.Number(label="Seed", value=0, precision=0, interactive=False)
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                # allow_custom_value: gradio validates API calls against the *initial* choices (the text-to-image
+                # menu), which would reject every editing size sent through /generate; generate() snaps
+                # anything off-menu itself.
+                size_label = gr.Dropdown(label="Output size", choices=T2I_CHOICES, value=T2I_CHOICES[0], allow_custom_value=True,
+                                         info="The menu switches to the editing sizes as soon as a reference is attached. Custom "
+                                         "sizes above 2048² area (1536² when editing) are scaled down, keeping the ratio.")
+                with gr.Row(visible=False) as custom_row:
+                    custom_width = gr.Slider(label="Width", minimum=512, maximum=MAX_SIDE, step=32, value=1024)
+                    custom_height = gr.Slider(label="Height", minimum=512, maximum=MAX_SIDE, step=32, value=1024)
                 enhance = gr.Checkbox(
                     label="Enhance prompt",
                     value=True,
                     info="Rewrite the prompt with the official Qwen-Image prompt-enhancement instructions "
                     "(runs on the built-in Qwen3-VL text encoder, +4–15 s). The text actually sent is shown under the result.",
                 )
-                run = gr.Button("Generate", variant="primary")
+                with gr.Accordion("Advanced settings", open=False) as advanced:
+                    steps = gr.Slider(label="Steps", minimum=3, maximum=8, step=1, value=STEPS,
+                                      info="6 is the validated schedule; extra steps are added at the high-noise end.")
+                    with gr.Row():
+                        seed = gr.Number(label="Seed", value=0, precision=0, interactive=False)
+                        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                run = gr.Button("Generate", variant="primary", size="lg")
             with gr.Column(scale=4):
                 output_image = gr.Image(label="Result", type="pil", height=560)
                 info = gr.Markdown()
@@ -263,11 +307,13 @@ with gr.Blocks(title="Viggle Turbo v0.2.1 · Qwen-Image-2.1 6-step") as demo:
 
         def example_details(prompt, image_1, image_2, image_3, result):
             """Runs on an example click (no GPU): fills the status line and the prompt stored with the row, and sets the
-            steps / size / enhancement controls to what produced it (rows differ; Run then reproduces the result). It also sets
-            the size menu's mode itself: loading an example does not fire the reference images' .input listeners."""
+            steps / size / enhancement / seed controls to what produced it, so Generate reproduces the result. It opens
+            Advanced settings to show that the seed is now fixed. It also sets the size menu's mode itself: loading an
+            example does not fire the reference images' .input listeners."""
             row = next(row for row in EXAMPLES if row["prompt"] == prompt)
             choices = EDIT_CHOICES if any(row["refs"]) else T2I_CHOICES
-            return row["info"], row["used_prompt"], row["steps"], gr.update(choices=choices, value=row["size_label"]), row["enhance"]
+            return (row["info"], row["used_prompt"], row["steps"], gr.update(choices=choices, value=row["size_label"]),
+                    row["enhance"], row["seed"], False, gr.update(open=True))
 
         if EXAMPLES:
             gr.Examples(
@@ -277,10 +323,13 @@ with gr.Blocks(title="Viggle Turbo v0.2.1 · Qwen-Image-2.1 6-step") as demo:
                 ],
                 inputs=[prompt, image_1, image_2, image_3, output_image],
                 fn=example_details,
-                outputs=[info, used_prompt, steps, size_label, enhance],
+                outputs=[info, used_prompt, steps, size_label, enhance, seed, randomize_seed, advanced],
                 run_on_click=True,
+                # On Spaces gr.Examples caches by default (lazily on ZeroGPU), and gradio 5.50 crashes caching an output that
+                # updates a Dropdown's choices ('Dropdown' object has no attribute 'proxy_url'). example_details needs no GPU.
+                cache_examples=False,
                 examples_per_page=12,
-                label="Examples · results pre-rendered by this model — click a row to load it (6 steps, prompt enhancement on, unless the status line says otherwise)",
+                label="Examples · results pre-rendered by this model — click a row to load it with its settings; Generate reproduces it",
             )
 
     with gr.Tab("Comparison"):
@@ -303,10 +352,12 @@ with gr.Blocks(title="Viggle Turbo v0.2.1 · Qwen-Image-2.1 6-step") as demo:
             inputs=[image_1, image_2, image_3, size_label],
             outputs=size_label,
         )
+    size_label.change(custom_size_controls, inputs=[size_label, custom_width, custom_height],
+                      outputs=[custom_row, custom_width, custom_height])
     randomize_seed.change(lambda on: gr.update(interactive=not on), inputs=randomize_seed, outputs=seed)
     run.click(
         generate,
-        inputs=[prompt, image_1, image_2, image_3, size_label, seed, randomize_seed, enhance, steps],
+        inputs=[prompt, image_1, image_2, image_3, size_label, seed, randomize_seed, enhance, steps, custom_width, custom_height],
         outputs=[output_image, info, used_prompt],
     )