import glob
import os
import subprocess
import sys

# diffusers from the MiniMax-H3-Image PR branch (private repo, so the wheel ships with the Space).
for _wheel in glob.glob(os.path.join(os.path.dirname(os.path.abspath(__file__)), "diffusers-*.whl")):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-deps", "--force-reinstall", "--no-cache-dir", _wheel])

import spaces  # before torch

import json
import math
import os
import random
import tempfile
import time
from functools import cache

import gradio as gr
import torch
from PIL import Image
from safetensors.torch import load_file

from diffusers import ModularPipeline

REPO = "diffusers-internal-dev/minimax-h3-image"
ENCODER_SPACE = os.environ.get("H3_IMAGE_ENCODER", "multimodalart/minimax-h3-image-encoder")
MAX_IMAGES = 8
MAX_SEED = 2**31 - 1
MULTIPLE = 32

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}
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.
blocks = ModularPipeline.from_pretrained(REPO).blocks
blocks.sub_blocks.pop("text_encoder")
pipe = blocks.init_pipeline(REPO)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")


@cache
def encoder():
    from gradio_client import Client

    return Client(ENCODER_SPACE, token=os.environ.get("HF_TOKEN"), httpx_kwargs={"timeout": 600})


def image_paths(images):
    return [item[0] if isinstance(item, (list, tuple)) else item for item in images or []]


def size_at_area(width, height, pixels):
    scale = math.sqrt(pixels / (width * height))
    return (
        max(MULTIPLE, round(width * scale / MULTIPLE) * MULTIPLE),
        max(MULTIPLE, round(height * scale / MULTIPLE) * MULTIPLE),
    )


def resolve_size(paths, aspect, resolution):
    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):
    """Validate the request, pick the size and seed, and encode the prompt on the encoder Space."""
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt.")
    paths = image_paths(images)
    if len(paths) > MAX_IMAGES:
        raise gr.Error(f"At most {MAX_IMAGES} images.")
    width, height = resolve_size(paths, aspect, resolution)
    if randomize:
        seed = random.randint(0, MAX_SEED)
    from gradio_client import handle_file

    embeds_path = encoder().predict(prompt.strip(), [{"image": handle_file(p), "caption": None} for p in paths], api_name="/encode")
    return embeds_path, int(seed), int(width), int(height)


def _duration(embeds_path, images, width, height, seed):
    # The canvas is at least 2048x2048; references add their own latent rows.
    megapixels = max(int(width) * int(height), 2048 * 2048) / 1e6
    for path in image_paths(images):
        w, h = Image.open(path).size
        megapixels += min(max(w * h / 1e6, 1.0), 4.2)
    return int(30 + 4 * megapixels + 0.15 * megapixels**2)


def _run(embeds_path, images, width, height, seed):
    embeds = load_file(embeds_path)
    kwargs = dict(
        prompt_embeds=embeds["prompt_embeds"].to("cuda", torch.bfloat16),
        text_token_tags=embeds["text_token_tags"],
        width=int(width),
        height=int(height),
        generator=torch.Generator("cpu").manual_seed(int(seed)),
        output="images",
    )
    paths = image_paths(images)
    if paths:
        kwargs["image"] = [Image.open(p).convert("RGB") for p in paths]
    image = pipe(**kwargs)[0]
    path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    image.save(path)
    return path


@spaces.GPU(duration=_duration, size="xlarge")
def _generate(embeds_path, images, width, height, seed):
    return _run(embeds_path, images, width, height, seed)


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.

    `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).
    """
    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, seed, randomize, effort):
    """The composer's Generate, part 1: optional PE (streamed) and encoding. Returns the render plan."""
    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 = decode_pictures(state.get("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 "✨ Enhancing the prompt with MiniMax-H3-Image-PE ...", gr.update(), "", seed, None
        t0 = time.time()
        caption = ratio = None
        for thinking, caption, ratio in enhance_stream(prompt, effort, seed):
            if caption is None:
                yield f"✨ Enhancing the prompt · thinking ({len(thinking.split())} words) ...", gr.update(), thinking, seed, None
        if caption:
            final = caption
            if aspect == AUTO and ratio:
                aspect = ratio
        yield f"✨ Prompt enhanced in {time.time() - t0:.0f}s · encoding ...", final, thinking, seed, None
    width, height = resolve_size(paths, aspect, resolution)
    yield f"Encoding the prompt · {width}×{height} ...", final, thinking, seed, None
    embeds_path, seed, width, height = prepare(final, paths, aspect, resolution, seed, False)
    plan = {"embeds": embeds_path, "images": paths, "width": width, "height": height, "seed": seed}
    yield f"Generating {width}×{height} ...", final, thinking, seed, plan


def _render_duration(plan, progress=None):
    return _duration(plan["embeds"], plan["images"], plan["width"], plan["height"], plan["seed"])


@spaces.GPU(duration=_render_duration, size="xlarge")
def render(plan, progress=gr.Progress(track_tqdm=True)):
    """Part 2: the image, as its own event so the result panel shows the loading UI and the denoising steps."""
    t0 = time.time()
    path = _run(plan["embeds"], plan["images"], plan["width"], plan["height"], plan["seed"])
    return path, f"{plan['width']}×{plan['height']} · seed {plan['seed']} · {time.time() - t0:.0f}s"

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",
]

COMPOSER_DEFAULTS = {"prompt": "", "aspect": AUTO, "resolution": "4 MP", "enhance": True, "images": []}


def decode_pictures(urls):
    """The composer's pictures (data URLs, in <Picture i> order) as temp files."""
    import base64

    paths = []
    for url in (urls or [])[:MAX_IMAGES]:
        header, _, data = str(url).partition(",")
        ext = {"image/png": ".png", "image/webp": ".webp", "image/gif": ".gif"}.get(header[5:].split(";")[0], ".jpg")
        path = tempfile.NamedTemporaryFile(suffix=ext, delete=False).name
        with open(path, "wb") as f:
            f.write(base64.b64decode(data))
        paths.append(path)
    return paths

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>
  <div class="hx-pics" data-role="pics">
    <div class="hx-thumbs" data-role="thumbs"></div>
    <label class="hx-add" title="Optional: up to 8 pictures to edit; drop files here too">
      <input type="file" accept="image/*" multiple data-role="files" hidden>
      <span class="hx-tile" aria-hidden="true">
        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.6" stroke-linecap="round" stroke-linejoin="round">
          <rect x="3" y="4" width="18" height="16" rx="2.5"/><circle cx="8.5" cy="9.5" r="1.7"/><path d="M3.5 17.5l5-5 4 4 2.5-2.5 5.5 5.5"/>
        </svg>
        <span class="hx-tile-plus">+</span>
      </span>
      <span class="hx-addtext"><span class="hx-addtitle">Pictures to edit</span><span class="hx-opt">optional &middot; refer to them as &lt;Picture 1&gt;, &lt;Picture 2&gt;&#8230;</span></span>
    </label>
  </div>
  <div class="hx-pictags" data-role="pictags" hidden><span class="hx-pictags-label">Insert</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, or an edit of the pictures above"></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-pics { display: flex; align-items: center; gap: 8px; flex-wrap: wrap; min-height: 72px; box-sizing: border-box; border: 1px dashed var(--border-color-primary); border-radius: var(--radius-lg, 8px); padding: 9px 10px; transition: border-color .15s, background .15s; }
.hx-pics.hx-drag { border-color: var(--color-accent); background: var(--background-fill-secondary); }
.hx-thumbs { display: flex; gap: 6px; flex-wrap: wrap; }
.hx-thumbs:empty { display: none; }
.hx-thumb { position: relative; width: 52px; height: 52px; border-radius: 6px; overflow: hidden; border: 1px solid var(--border-color-primary); }
.hx-thumb img { width: 100%; height: 100%; object-fit: cover; display: block; }
.hx-thumb .hx-n { position: absolute; left: 2px; bottom: 2px; background: rgba(0,0,0,.65); color: #fff; font-size: 10px; line-height: 1; padding: 2px 4px; border-radius: 4px; }
.hx-thumb .hx-x { position: absolute; right: 1px; top: 1px; width: 16px; height: 16px; border: none; border-radius: 50%; background: rgba(0,0,0,.65); color: #fff; font-size: 11px; line-height: 16px; padding: 0; cursor: pointer; }
.hx-add { display: inline-flex; align-items: center; gap: 10px; color: var(--body-text-color-subdued); cursor: pointer; user-select: none; }
.hx-tile { position: relative; flex: none; width: 52px; height: 52px; box-sizing: border-box; display: flex; align-items: center; justify-content: center; border: 1px dashed var(--border-color-primary); border-radius: 8px; background: var(--background-fill-secondary); transition: color .15s, border-color .15s; }
.hx-tile svg { width: 28px; height: 28px; }
.hx-tile-plus { position: absolute; right: -6px; bottom: -6px; width: 18px; height: 18px; border-radius: 50%; background: var(--button-primary-background-fill); color: var(--button-primary-text-color); font-size: 14px; font-weight: 700; line-height: 18px; text-align: center; }
.hx-add:hover .hx-tile { color: var(--color-accent); border-color: var(--color-accent); }
.hx-addtext { display: flex; flex-direction: column; gap: 2px; }
.hx-addtitle { color: var(--body-text-color); font-size: 13.5px; font-weight: 600; }
.hx-add:hover .hx-addtitle { color: var(--color-accent); }
.hx-opt { font-size: 11.5px; opacity: .85; }
.hx-pictags { display: flex; align-items: center; gap: 5px; flex-wrap: wrap; margin-top: -4px; }
.hx-pictags[hidden] { display: none; }
.hx-pictags-label { color: var(--body-text-color-subdued); font-size: 12px; margin-right: 2px; }
.hx-pictag { background: transparent; color: var(--body-text-color-subdued); border: 1px dashed var(--border-color-primary); border-radius: 6px; padding: 2px 8px; font-family: var(--font-mono, ui-monospace, monospace); font-size: 11.5px; cursor: pointer; transition: color .15s, border-color .15s; }
.hx-pictag:hover { color: var(--color-accent); border-color: var(--color-accent); }
.hx-thumb img { cursor: pointer; }
.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 = '', pics = [];
const MAX_PICS = 8;
function insertTag(tag) {
  const v = prompt.value;
  const at = prompt.selectionStart == null ? v.length : prompt.selectionStart;
  const end = prompt.selectionEnd == null ? at : prompt.selectionEnd;
  const before = v.slice(0, at), after = v.slice(end);
  const ins = (before && !/\\s$/.test(before) ? ' ' : '') + tag + (after && !/^\\s/.test(after) ? ' ' : (after ? '' : ' '));
  prompt.value = before + ins + after;
  const pos = (before + ins).length;
  prompt.focus();
  prompt.setSelectionRange(pos, pos);
}
function renderTags() {
  const row = $('[data-role="pictags"]');
  Array.from(row.querySelectorAll('.hx-pictag')).forEach(function(b) { b.remove(); });
  row.hidden = pics.length === 0;
  pics.forEach(function(_, i) {
    const b = document.createElement('button');
    b.type = 'button'; b.className = 'hx-pictag'; b.textContent = '<Picture ' + (i + 1) + '>';
    b.title = 'Insert into the prompt';
    b.addEventListener('mousedown', function(e) { e.preventDefault(); });  // keep the prompt's cursor
    b.addEventListener('click', function() { insertTag('<Picture ' + (i + 1) + '>'); });
    row.appendChild(b);
  });
}
function renderPics() {
  renderTags();
  const box = $('[data-role="thumbs"]');
  box.innerHTML = '';
  pics.forEach(function(url, i) {
    const t = document.createElement('div'); t.className = 'hx-thumb';
    const img = document.createElement('img'); img.src = url; img.title = 'Insert <Picture ' + (i + 1) + '> into the prompt';
    img.addEventListener('click', function() { insertTag('<Picture ' + (i + 1) + '>'); });
    t.appendChild(img);
    const n = document.createElement('span'); n.className = 'hx-n'; n.textContent = i + 1; t.appendChild(n);
    const x = document.createElement('button'); x.type = 'button'; x.className = 'hx-x'; x.textContent = '\u00d7'; x.title = 'Remove';
    x.addEventListener('click', function() { pics.splice(i, 1); renderPics(); });
    t.appendChild(x);
    box.appendChild(t);
  });
}
function addFiles(files) {
  Array.from(files || []).filter(function(f) { return f.type.indexOf('image/') === 0; }).forEach(function(f) {
    if (pics.length >= MAX_PICS) return;
    const r = new FileReader();
    r.onload = function() { if (pics.length < MAX_PICS) { pics.push(r.result); renderPics(); } };
    r.readAsDataURL(f);
  });
}
const picsBox = $('[data-role="pics"]');
$('[data-role="files"]').addEventListener('change', function(e) { addFiles(e.target.files); e.target.value = ''; });
['dragenter', 'dragover'].forEach(function(ev) { picsBox.addEventListener(ev, function(e) { e.preventDefault(); picsBox.classList.add('hx-drag'); }); });
['dragleave', 'drop'].forEach(function(ev) { picsBox.addEventListener(ev, function(e) { e.preventDefault(); picsBox.classList.remove('hx-drag'); }); });
picsBox.addEventListener('drop', function(e) { addFiles(e.dataTransfer.files); });

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, i) {
  const b = document.createElement('button');
  b.type = 'button'; b.className = 'hx-chip'; b.textContent = t;
  b.addEventListener('click', function() { prompt.value = t; gather({example_key: 'idea_' + i}); trigger('apply'); });
  $('[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, images: pics.slice() }; }
function gather(extra) { const s = Object.assign(readState(), extra || {}); 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');
  if (Array.isArray(v.images)) { pics = v.images.slice(0, MAX_PICS); renderPics(); }
}
$('[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-composer, #hx-composer > div, #hx-composer .html-container { padding: 0 !important; }
.hx-brand { display: flex; align-items: center; justify-content: center; gap: 12px; margin: 12px auto 0; max-width: 320px; }
.hx-brand img { height: 40px; width: auto; max-width: 140px; object-fit: contain; display: block; }
.hx-logo-dark { display: none !important; }
.dark .hx-logo-light { display: none !important; }
.dark .hx-logo-dark { display: block !important; }
.hx-word { font-size: 30px; font-weight: 400; line-height: 1; letter-spacing: -0.01em; color: var(--body-text-color); white-space: nowrap; }
.hx-grad { background: linear-gradient(90deg, #e8197e, #ff6a3d); -webkit-background-clip: text; background-clip: text; color: transparent; font-weight: 500; }
#hx-links { text-align: center; }
"""

import base64

HERE = os.path.dirname(os.path.abspath(__file__))
LOGO_LIGHT = base64.b64encode(open(os.path.join(HERE, "minimax_logo_light.png"), "rb").read()).decode()
LOGO_DARK = base64.b64encode(open(os.path.join(HERE, "minimax_logo_dark.png"), "rb").read()).decode()
EXAMPLES = {}
if os.path.exists(os.path.join(HERE, "examples", "manifest.json")):
    EXAMPLES = json.load(open(os.path.join(HERE, "examples", "manifest.json"), encoding="utf-8"))


def load_example(state):
    """An idea chip: show its precomputed result (PE at low effort, then H3-Image at PE's ratio, 4 MP)."""
    ex = EXAMPLES.get((state or {}).get("example_key") or "")
    if not ex:
        return gr.update(), gr.update(), gr.update(), gr.update()
    stats = f"cached example · enhanced prompt · {ex['aspect']} · {ex['size']} · seed {ex['seed']}"
    return os.path.join(HERE, ex["image"]), stats, ex["prompt_used"], ex["thinking"]


with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="MiniMax H3 Image") as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(
            '<div class="hx-brand">'
            f'<img src="data:image/png;base64,{LOGO_LIGHT}" class="hx-logo-light" alt="MiniMax">'
            f'<img src="data:image/png;base64,{LOGO_DARK}" class="hx-logo-dark" alt="MiniMax">'
            '<span class="hx-word"><span class="hx-grad">H3</span> Image</span></div>',
            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)",
            elem_id="hx-links",
        )
        with gr.Row():
            with gr.Column():
                composer = gr.HTML(
                    value=dict(COMPOSER_DEFAULTS),
                    elem_id="hx-composer",
                    html_template=COMPOSER_HTML,
                    css_template=COMPOSER_CSS,
                    js_on_load=COMPOSER_JS,
                    container=False,
                    padding=False,
                )
                images = gr.Gallery(type="filepath", visible=False)  # API input; the composer carries its own pictures
                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)

    plan = gr.State()
    composer.submit(
        lambda: (None, None), None, [output, plan], show_progress="hidden", api_visibility="private",
    ).then(
        studio, [composer, seed, randomize, effort], [stats, prompt_used, thinking, seed, plan],
        show_progress="minimal", api_visibility="private",
    ).success(
        render, [plan], [output, stats], show_progress="full", show_progress_on=[output], api_visibility="private",
    )

    composer.apply(
        load_example, [composer], [output, stats, prompt_used, thinking],
        show_progress="hidden", api_visibility="private",
    )

    # 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,
        [api_prompt, images, api_aspect, api_resolution, seed, randomize, api_enhance, api_effort],
        [output, seed, prompt_used],
        api_name="generate",
    )

demo.launch()