Music3-style composer UI + optional MiniMax-H3-Image-PE prompt enhancement

app.py

@@ -9,10 +9,12 @@ for _wheel in glob.glob(os.path.join(os.path.dirname(os.path.abspath(__file__)),
 
 import spaces  # before torch
 
+import json
 import math
 import os
 import random
 import tempfile
+import time
 from functools import cache
 
 import gradio as gr
@@ -28,38 +30,10 @@ MAX_IMAGES = 8
 MAX_SEED = 2**31 - 1
 MULTIPLE = 32
 
-SIZES = {
-    "4 MP": {
-        "1:1": (2048, 2048),
-        "4:3": (2368, 1760),
-        "3:4": (1760, 2368),
-        "3:2": (2496, 1664),
-        "2:3": (1664, 2496),
-        "16:9": (2720, 1536),
-        "9:16": (1536, 2720),
-    },
-    "8 MP": {
-        "1:1": (2912, 2912),
-        "4:3": (3360, 2496),
-        "3:4": (2496, 3360),
-        "3:2": (3552, 2368),
-        "2:3": (2368, 3552),
-        "16:9": (3872, 2176),
-        "9:16": (2176, 3872),
-    },
-    "16 MP": {
-        "1:1": (4096, 4096),
-        "4:3": (4736, 3552),
-        "3:4": (3552, 4736),
-        "3:2": (5024, 3360),
-        "2:3": (3360, 5024),
-        "16:9": (5472, 3072),
-        "9:16": (3072, 5472),
-    },
-}
+RATIOS = ["1:1", "5:4", "4:3", "3:2", "16:9", "21:9", "4:5", "3:4", "2:3", "9:16"]  # MiniMax-H3-Image-PE's set
 MEGAPIXELS = {"4 MP": 4, "8 MP": 8, "16 MP": 16}
-MATCH = "Match Picture 1"
-ASPECTS = [MATCH] + list(SIZES["4 MP"])
+AUTO = "auto"  # PE's recommendation for text-to-image, Picture 1's ratio for edits, else 1:1
+PE_SPACE = os.environ.get("H3_IMAGE_PE", "multimodalart/minimax-h3-image-pe")
 
 # Everything but the Qwen3-VL text encoder, which runs on ENCODER_SPACE: `prompt_embeds` + `text_token_tags` become
 # inputs of the pipeline.
@@ -90,12 +64,34 @@ def size_at_area(width, height, pixels):
 
 
 def resolve_size(paths, aspect, resolution):
-    if aspect == MATCH:
-        if not paths:
-            return SIZES[resolution]["1:1"]
-        width, height = Image.open(paths[0]).size
-        return size_at_area(width, height, MEGAPIXELS[resolution] * 1024 * 1024)
-    return SIZES[resolution][aspect]
+    pixels = MEGAPIXELS.get(resolution, 4) * 1024 * 1024
+    if aspect not in RATIOS:
+        if paths:
+            width, height = Image.open(paths[0]).size
+            return size_at_area(width, height, pixels)
+        aspect = "1:1"
+    rw, rh = (int(v) for v in aspect.split(":"))
+    return size_at_area(rw, rh, pixels)
+
+
+@cache
+def pe_client():
+    from gradio_client import Client
+
+    return Client(PE_SPACE, token=os.environ.get("HF_TOKEN"), httpx_kwargs={"timeout": 900})
+
+
+def enhance_stream(prompt, effort, seed):
+    """Streams MiniMax-H3-Image-PE: yields (thinking, None, None) while it reasons, then (thinking, caption, ratio)."""
+    job = pe_client().submit(prompt, effort, int(seed) % (2**31), api_name="/expand")
+    shown = 0
+    for update in job:
+        thinking = update[2] or ""
+        if len(thinking) - shown > 200:
+            shown = len(thinking)
+            yield thinking, None, None
+    caption, ratio, thinking, _ = job.result()
+    yield thinking, caption, ratio
 
 
 def prepare(prompt, images, aspect, resolution, seed, randomize):
@@ -143,51 +139,232 @@ def _generate(embeds_path, images, width, height, seed):
     return path
 
 
-def generate(prompt, images, aspect, resolution, seed, randomize, progress=gr.Progress(track_tqdm=True)):
+def generate(prompt, images=None, aspect=AUTO, resolution="4 MP", seed=0, randomize=True, enhance=False,
+             reasoning_effort="low", progress=gr.Progress(track_tqdm=True)):
     """Text-to-image, or an edit of up to 8 images referred to as `<Picture 1>`, `<Picture 2>`, ... in the prompt.
 
-    Returns the image path, the seed used, and the size.
+    `aspect`: one of 1:1, 5:4, 4:3, 3:2, 16:9, 21:9, 4:5, 3:4, 2:3, 9:16, or "auto". `enhance`: expand the prompt with
+    MiniMax-H3-Image-PE first (text-to-image only). Returns (image, seed, prompt used).
     """
-    progress(0, desc="Encoding the prompt ...")
-    embeds_path, seed, width, height = prepare(prompt, images, aspect, resolution, seed, randomize)
-    progress(0.05, desc=f"Generating {width}x{height} ...")
-    return _generate(embeds_path, images, width, height, seed), seed
-
-
-EXAMPLES = [
-    "A German Expressionist oil painting of a Berlin street at night in the manner of Ernst Ludwig Kirchner, with gaunt figures hurrying past a glowing gas lamp.",
-    "A two-color risograph print of a fox reading a newspaper on a park bench, the headline reads \"FOX NEWS\", no other text.",
-    "A 1970s film frame of a woman waiting at a rainy bus stop in Tokyo, lit by a vending machine.",
-    "A photorealistic overhead food photograph of a rustic farm-to-table feast on a weathered dark wooden table: a scored sourdough loaf on a jute placemat, a copper bowl of seasonal vegetables, a cheese board, heirloom tomatoes, and a glass of red wine beside a lit taper candle, warm sunlight from the left.",
+    if not prompt or not prompt.strip():
+        raise gr.Error("Enter a prompt.")
+    final = prompt.strip()
+    if enhance and not image_paths(images):
+        progress(0, desc="Enhancing the prompt ...")
+        *_, (thinking, caption, ratio) = enhance_stream(final, reasoning_effort, seed)
+        final = caption or final
+        if aspect == AUTO and ratio:
+            aspect = ratio
+    progress(0.05, desc="Encoding the prompt ...")
+    embeds_path, seed, width, height = prepare(final, images, aspect, resolution, seed, randomize)
+    progress(0.1, desc=f"Generating {width}x{height} ...")
+    return _generate(embeds_path, images, width, height, seed), seed, final
+
+
+def studio(state, images, seed, randomize, effort):
+    """The composer's Generate: optional PE (streamed), then the image."""
+    state = state or {}
+    prompt = (state.get("prompt") or "").strip()
+    aspect = state.get("aspect") or AUTO
+    resolution = state.get("resolution") or "4 MP"
+    if not prompt:
+        raise gr.Error("Describe the image first.")
+    paths = image_paths(images)
+    if randomize:
+        seed = random.randint(0, MAX_SEED)
+    final, thinking = prompt, ""
+    if state.get("enhance") and paths:
+        gr.Info("Prompt enhancement is text-to-image only for now; edits use your prompt as written.")
+    elif state.get("enhance"):
+        yield gr.update(), "✨ Enhancing the prompt with MiniMax-H3-Image-PE ...", gr.update(), "", seed
+        t0 = time.time()
+        caption = ratio = None
+        for thinking, caption, ratio in enhance_stream(prompt, effort, seed):
+            if caption is None:
+                yield gr.update(), f"✨ Enhancing the prompt · thinking ({len(thinking.split())} words) ...", gr.update(), thinking, seed
+        if caption:
+            final = caption
+            if aspect == AUTO and ratio:
+                aspect = ratio
+        yield gr.update(), f"✨ Prompt enhanced in {time.time() - t0:.0f}s · encoding ...", final, thinking, seed
+    width, height = resolve_size(paths, aspect, resolution)
+    yield gr.update(), f"Encoding the prompt · {width}×{height} ...", final, thinking, seed
+    embeds_path, seed, width, height = prepare(final, images, aspect, resolution, seed, False)
+    yield gr.update(), f"Generating {width}×{height} ...", final, thinking, seed
+    t0 = time.time()
+    path = _generate(embeds_path, images, width, height, seed)
+    yield path, f"{width}×{height} · seed {seed} · {time.time() - t0:.0f}s", final, thinking, seed
+
+
+IDEAS = [
+    "a two-color risograph print of a fox reading a newspaper on a park bench, headline \"FOX NEWS\"",
+    "a 1970s film frame of a woman waiting at a rainy Tokyo bus stop, lit by a vending machine",
+    "isometric cutaway of a tiny underground bakery",
+    "a poster for a jazz festival in Lisbon",
 ]
 
-with gr.Blocks(title="MiniMax-H3-Image") as demo:
-    gr.Markdown(
-        "# MiniMax-H3-Image\n"
-        "Text-to-image and editing with up to 8 images, in 6 steps, through the diffusers modular pipeline. "
-        "Add images to edit: refer to them in the prompt as `<Picture 1>`, `<Picture 2>`, …"
+COMPOSER_DEFAULTS = {"prompt": "", "aspect": AUTO, "resolution": "4 MP", "enhance": True}
+
+COMPOSER_HTML = """
+<div class="hx-card">
+  <div class="hx-head">
+    <div class="hx-seg" data-role="resolution" role="tablist">
+      <button type="button" class="hx-seg-btn" data-value="4 MP">4 MP</button>
+      <button type="button" class="hx-seg-btn" data-value="8 MP">8 MP</button>
+      <button type="button" class="hx-seg-btn" data-value="16 MP">16 MP</button>
+    </div>
+    <span class="hx-headhint">text-to-image &middot; or edit up to 8 pictures</span>
+  </div>
+  <textarea class="hx-desc" data-role="prompt" rows="4"
+    placeholder="Describe your image&#8230;  e.g. a poster for a jazz festival in Lisbon. To edit, add pictures below and refer to them as <Picture 1>, <Picture 2>&#8230;"></textarea>
+  <div class="hx-chips" data-role="idea-chips"></div>
+  <div class="hx-row">
+    <span class="hx-label">Aspect</span>
+    <div class="hx-ratios" data-role="aspect"></div>
+  </div>
+  <div class="hx-foot">
+    <label class="hx-toggle" title="Expand the prompt into a detailed caption (and pick the ratio when Aspect is Auto) with MiniMax-H3-Image-PE">
+      <input type="checkbox" data-role="enhance">&#10024; Enhance prompt
+    </label>
+    <span class="hx-spacer"></span>
+    <button type="button" class="hx-primary" data-role="generate">Generate</button>
+  </div>
+</div>
+"""
+
+COMPOSER_CSS = """
+.hx-card { background: var(--block-background-fill); border: var(--block-border-width, 1px) solid var(--block-border-color, var(--border-color-primary)); border-radius: var(--block-radius, 12px); box-shadow: var(--block-shadow, none); padding: 16px; display: flex; flex-direction: column; gap: 12px; }
+.hx-head { display: flex; align-items: center; justify-content: space-between; gap: 10px; flex-wrap: wrap; }
+.hx-seg { display: inline-flex; background: var(--background-fill-secondary); border: 1px solid var(--border-color-primary); border-radius: 999px; padding: 3px; gap: 2px; }
+.hx-seg-btn { border: none; background: transparent; color: var(--body-text-color-subdued); padding: 5px 14px; border-radius: 999px; font-family: inherit; font-size: 13.5px; font-weight: 600; cursor: pointer; transition: background .15s, color .15s; }
+.hx-seg-btn[aria-selected="true"] { background: var(--button-primary-background-fill); color: var(--button-primary-text-color); }
+.hx-headhint { color: var(--body-text-color-subdued); font-size: 12.5px; }
+.hx-card textarea { width: 100%; box-sizing: border-box; background: var(--input-background-fill); border: var(--input-border-width, 1px) solid var(--input-border-color, var(--border-color-primary)); border-radius: var(--input-radius, 8px); padding: 10px 12px; color: var(--body-text-color); font-family: inherit; font-size: 16px; line-height: 1.5; min-height: 118px; resize: vertical; transition: border-color .15s; }
+.hx-card textarea::placeholder { color: var(--input-placeholder-color, var(--body-text-color-subdued)); }
+.hx-card textarea:focus { outline: none; border-color: var(--input-border-color-focus, var(--color-accent)); }
+.hx-chips { display: flex; flex-wrap: wrap; gap: 6px; }
+.hx-chip { background: var(--button-secondary-background-fill); color: var(--button-secondary-text-color); border: 1px solid var(--button-secondary-border-color, var(--border-color-primary)); border-radius: 999px; padding: 4px 12px; font-family: inherit; font-size: 12.5px; cursor: pointer; transition: border-color .15s; text-align: left; }
+.hx-chip:hover { border-color: var(--color-accent); }
+.hx-row { display: flex; align-items: center; gap: 10px; flex-wrap: wrap; }
+.hx-label { color: var(--block-title-text-color, var(--body-text-color)); font-size: 13px; font-weight: 600; }
+.hx-ratios { display: flex; flex-wrap: wrap; gap: 4px; }
+.hx-ratio { background: transparent; color: var(--body-text-color-subdued); border: 1px dashed var(--border-color-primary); border-radius: 6px; padding: 3px 9px; font-family: var(--font-mono, ui-monospace, monospace); font-size: 12px; cursor: pointer; transition: color .15s, border-color .15s, background .15s; }
+.hx-ratio:hover { color: var(--color-accent); border-color: var(--color-accent); }
+.hx-ratio[aria-selected="true"] { border-style: solid; border-color: var(--color-accent); color: var(--color-accent); background: var(--background-fill-secondary); }
+.hx-foot { display: flex; align-items: center; gap: 10px; }
+.hx-spacer { flex: 1; }
+.hx-toggle { display: inline-flex; align-items: center; gap: 7px; color: var(--body-text-color); font-size: 14px; cursor: pointer; user-select: none; }
+.hx-toggle input { width: 16px; height: 16px; accent-color: var(--color-accent); cursor: pointer; }
+.hx-primary { background: var(--button-primary-background-fill); color: var(--button-primary-text-color); border: var(--button-border-width, 1px) solid var(--button-primary-border-color, transparent); border-radius: var(--button-large-radius, var(--radius-lg, 8px)); padding: 10px 28px; font-family: inherit; font-size: var(--button-large-text-size, 16px); font-weight: 600; cursor: pointer; box-shadow: var(--button-primary-shadow, none); transition: background .15s, transform .05s; }
+.hx-primary:hover { background: var(--button-primary-background-fill-hover, var(--button-primary-background-fill)); }
+.hx-primary:active { transform: translateY(1px); }
+"""
+
+COMPOSER_JS = """
+const $ = function(s) { return element.querySelector(s); };
+const $$ = function(s) { return Array.from(element.querySelectorAll(s)); };
+const IDEAS = __IDEAS__;
+const RATIOS = ['auto'].concat(__RATIOS__);
+const prompt = $('[data-role="prompt"]');
+const enhance = $('[data-role="enhance"]');
+let aspect = 'auto', resolution = '4 MP', lastPushed = '';
+
+RATIOS.forEach(function(r) {
+  const b = document.createElement('button');
+  b.type = 'button'; b.className = 'hx-ratio'; b.dataset.value = r;
+  b.textContent = r === 'auto' ? 'Auto' : r;
+  if (r === 'auto') b.title = 'Enhanced: the ratio PE recommends. Editing: Picture 1\\'s ratio. Otherwise 1:1.';
+  b.addEventListener('click', function() { setAspect(r); });
+  $('[data-role="aspect"]').appendChild(b);
+});
+IDEAS.forEach(function(t) {
+  const b = document.createElement('button');
+  b.type = 'button'; b.className = 'hx-chip'; b.textContent = t;
+  b.addEventListener('click', function() { prompt.value = t; prompt.focus(); });
+  $('[data-role="idea-chips"]').appendChild(b);
+});
+function setAspect(r) { aspect = r; $$('.hx-ratio').forEach(function(b) { b.setAttribute('aria-selected', String(b.dataset.value === r)); }); }
+function setResolution(r) { resolution = r; $$('.hx-seg-btn').forEach(function(b) { b.setAttribute('aria-selected', String(b.dataset.value === r)); }); }
+$$('.hx-seg-btn').forEach(function(b) { b.addEventListener('click', function() { setResolution(b.dataset.value); }); });
+function readState() { return { prompt: prompt.value, aspect: aspect, resolution: resolution, enhance: enhance.checked }; }
+function gather() { const s = readState(); lastPushed = JSON.stringify(s); props.value = s; }
+function applyState(v) {
+  if (!v) return;
+  if (prompt.value !== (v.prompt || '')) prompt.value = v.prompt || '';
+  enhance.checked = !!v.enhance;
+  setAspect(v.aspect || 'auto');
+  setResolution(v.resolution || '4 MP');
+}
+$('[data-role="generate"]').addEventListener('click', function() { gather(); trigger('submit'); });
+prompt.addEventListener('keydown', function(e) { if (e.key === 'Enter' && (e.metaKey || e.ctrlKey)) { e.preventDefault(); gather(); trigger('submit'); } });
+watch('value', function() { if (JSON.stringify(props.value) !== lastPushed) applyState(props.value); });
+applyState(props.value);
+""".replace("__IDEAS__", json.dumps(IDEAS)).replace("__RATIOS__", json.dumps(RATIOS))
+
+CSS = """
+#col-container { max-width: 1300px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+.html-container { padding: 0; }
+.hx-title { text-align: center; margin: 14px 0 0; font-size: 34px; font-weight: 800; letter-spacing: -0.02em; }
+.hx-title span { background: linear-gradient(90deg, var(--color-accent), var(--secondary-500, var(--color-accent))); -webkit-background-clip: text; background-clip: text; color: transparent; }
+#hx-links { text-align: center; }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="MiniMax H3 Image") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.HTML('<h1 class="hx-title">MiniMax <span>H3 Image</span></h1>', container=False, padding=False)
+        gr.Markdown(
+            "Text-to-image and editing of up to 8 pictures in 6 steps, with optional prompt enhancement by "
+            "MiniMax-H3-Image-PE "
+            "[[model]](https://huggingface.co/MiniMaxAI/MiniMax-H3-Image) | "
+            "[[prompt enhancer]](https://huggingface.co/MiniMaxAI/MiniMax-H3-Image-PE) | "
+            "[[diffusers PR]](https://github.com/huggingface/diffusers-new-model-addition-minimax/pull/1)",
+            elem_id="hx-links",
+        )
+        with gr.Row():
+            with gr.Column():
+                composer = gr.HTML(
+                    value=dict(COMPOSER_DEFAULTS),
+                    html_template=COMPOSER_HTML,
+                    css_template=COMPOSER_CSS,
+                    js_on_load=COMPOSER_JS,
+                    container=False,
+                    padding=False,
+                )
+                images = gr.Gallery(
+                    label="Pictures to edit (optional, up to 8) — refer to them as <Picture 1>, <Picture 2>, …",
+                    type="filepath", columns=4, height=200,
+                )
+                with gr.Accordion("Advanced", open=False):
+                    effort = gr.Radio(["low", "medium", "xhigh"], value="low", label="Prompt enhancer reasoning effort",
+                                      info="xhigh is the released default; low thinks briefly and is several times faster")
+                    with gr.Row():
+                        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                        randomize = gr.Checkbox(label="Randomize seed", value=True)
+            with gr.Column():
+                output = gr.Image(label="Result", type="filepath", format="png", show_label=False, height=620)
+                stats = gr.Markdown()
+                with gr.Accordion("Prompt used", open=False):
+                    prompt_used = gr.Textbox(show_label=False, lines=8)
+                with gr.Accordion("Enhancer thinking", open=False):
+                    thinking = gr.Textbox(show_label=False, lines=10)
+
+    composer.submit(
+        studio, [composer, images, seed, randomize, effort], [output, stats, prompt_used, thinking, seed],
+        show_progress="minimal", api_visibility="private",
     )
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(label="Prompt", lines=4, placeholder="Describe the image, or the edit")
-            images = gr.Gallery(label="Images to edit (optional, up to 8)", type="filepath", columns=4, height=220)
-            with gr.Row():
-                aspect = gr.Dropdown(ASPECTS, value="1:1", label="Aspect ratio")
-                resolution = gr.Radio(list(SIZES), value="4 MP", label="Resolution")
-            with gr.Row():
-                seed = gr.Number(value=0, label="Seed", precision=0)
-                randomize = gr.Checkbox(value=True, label="Randomize seed")
-            button = gr.Button("Generate", variant="primary")
-        with gr.Column():
-            output = gr.Image(label="Result", type="filepath", format="png")
-    gr.Examples(EXAMPLES, inputs=prompt)
-
-    images.change(lambda imgs: gr.update(value=MATCH) if imgs else gr.update(value="1:1"), images, aspect)
-    gr.on(
-        [button.click, prompt.submit],
+
+    # Plain API endpoint (also what the composer runs, minus the streaming UI updates).
+    api_prompt = gr.Textbox(visible=False)
+    api_aspect = gr.Textbox(visible=False, value=AUTO)
+    api_resolution = gr.Textbox(visible=False, value="4 MP")
+    api_enhance = gr.Checkbox(visible=False, value=False)
+    api_effort = gr.Textbox(visible=False, value="low")
+    api_button = gr.Button(visible=False)
+    api_button.click(
         generate,
-        [prompt, images, aspect, resolution, seed, randomize],
-        [output, seed],
+        [api_prompt, images, api_aspect, api_resolution, seed, randomize, api_enhance, api_effort],
+        [output, seed, prompt_used],
         api_name="generate",
     )