match gradio UI with js UI

app.py

@@ -1,8 +1,13 @@
 """Gradio entry point for local use and Hugging Face ZeroGPU."""
 import spaces  # Must initialize ZeroGPU before importing torch/diffusers.
 
+import base64
+import html
+import inspect
+from io import BytesIO
 import os
 import secrets
+import time
 
 import gradio as gr
 import numpy as np
@@ -27,6 +32,35 @@ def parse_settings(seed, checkpoint_text):
     return seed, checkpoints
 
 
+BLANK_IMAGE = Image.new("RGB", (256, 256), "#e7e5dc")
+
+
+def checkpoint_slots(checkpoints, frames=None):
+    """Keep every timestep in its own fixed slot, including pending frames."""
+    frames = frames or {}
+    return [(frames.get(step, BLANK_IMAGE), f"t = {step}") for step in checkpoints]
+
+
+def gaussian_noise(seed):
+    noise = np.random.default_rng(int(seed)).standard_normal((256, 256, 3))
+    return Image.fromarray(np.round((noise.clip(-1, 1) + 1) * 127.5).astype(np.uint8))
+
+
+def progress_text(completed=0, started=None):
+    elapsed = time.monotonic() - started if started is not None else 0
+    rate = completed / elapsed if elapsed > 0 and completed else 0
+    def duration(seconds):
+        if seconds is None:
+            return "--:--"
+        seconds = int(seconds)
+        hours, seconds = divmod(seconds, 3600)
+        minutes, seconds = divmod(seconds, 60)
+        return (f"{hours:02d}:" if hours else "") + f"{minutes:02d}:{seconds:02d}"
+    remaining = 0 if completed == 1000 else (1000 - completed) / rate if rate else None
+    return (f"{completed}/1000 [{duration(elapsed)}<{duration(remaining)}, "
+            f"{f'{rate:.2f}' if rate else '?'}it/s]")
+
+
 def diffuse(image, seed, checkpoint_text):
     seed, checkpoints = parse_settings(seed, checkpoint_text)
     if image is None:
@@ -34,15 +68,17 @@ def diffuse(image, seed, checkpoint_text):
     pixels = np.asarray(image.convert("RGB").resize((256, 256)), dtype=np.float32) / 127.5 - 1
     noise = np.random.default_rng(seed).standard_normal(pixels.shape)
     alpha_bars = np.cumprod(1 - np.linspace(0.0001, 0.02, 1000))
-    frames = []
+    started = time.monotonic()
+    frames = {0: image.convert("RGB").resize((256, 256))}
     final_noise = None
+    yield checkpoint_slots(checkpoints, frames), None, progress_html(0, started)
     for step in checkpoints:
         alpha = 1 if step == 0 else alpha_bars[step - 1]
         sample = np.sqrt(alpha) * pixels + np.sqrt(1 - alpha) * noise
         frame = Image.fromarray(np.round((sample.clip(-1, 1) + 1) * 127.5).astype(np.uint8))
-        frames.append((frame, f"t = {step}"))
+        frames[step] = frame
         final_noise = frame
-    return frames, final_noise
+        yield checkpoint_slots(checkpoints, frames), final_noise if step == 1000 else None, progress_html(step, started)
 
 
 @spaces.GPU(duration=int(os.environ.get("GPU_DURATION_SECONDS", "120")))
@@ -51,35 +87,267 @@ def denoise(seed, checkpoint_text, input_mode, forward_noise):
     initial_image = forward_noise if input_mode == "Forward noise" else None
     if input_mode == "Forward noise" and initial_image is None:
         raise gr.Error("Run Diffuse first to create forward noise.")
-    frames = []
+    started = time.monotonic()
+    frames = {1000: initial_image if initial_image is not None else gaussian_noise(seed)}
+    yield checkpoint_slots(checkpoints, frames), progress_html(0, started)
     for step, frame in generate_images(seed, checkpoints, initial_image):
         if frame is not None:
-            frames.append((frame, f"t = {step}"))
-        if frame is not None or step % 25 == 0:
-            yield list(frames), f"Denoising: {1000 - step}/1000 steps · seed {seed}"
-
-
-with gr.Blocks(title="Diffusion Lab") as demo:
-    gr.Markdown("# Diffusion Lab\nExplore forward noising and the 1,000-step DDPM reverse process.")
-    with gr.Row():
-        seed = gr.Textbox(value="42", label="Seed")
-        randomize = gr.Button("Randomize seed")
-        checkpoints = gr.Textbox(value=", ".join(map(str, CHECKPOINTS)), label="Checkpoint timesteps")
-    randomize.click(lambda: str(secrets.randbelow(2**31)), outputs=seed, queue=False)
-    default_image = ROOT / "assets/celeba_hq_validation_91_256.png"
-    source = gr.Image(value=str(default_image) if default_image.exists() else None,
-                      type="pil", label="Source image")
-    diffuse_button = gr.Button("Diffuse")
-    forward_gallery = gr.Gallery(label="Forward process: image to noise", columns=4)
-    forward_noise = gr.State(None)
-    diffuse_button.click(diffuse, [source, seed, checkpoints], [forward_gallery, forward_noise])
-    input_mode = gr.Radio(["Gaussian noise", "Forward noise"], value="Gaussian noise", label="Reverse input")
-    gr.Markdown("The face model is unconditional. Forward noise is an experiment, not a guaranteed reconstruction.")
-    denoise_button = gr.Button("Denoise", variant="primary")
-    reverse_gallery = gr.Gallery(label="Reverse process: noise to face", columns=4)
-    status = gr.Textbox(label="Progress", interactive=False)
-    denoise_button.click(denoise, [seed, checkpoints, input_mode, forward_noise],
-                         [reverse_gallery, status], concurrency_limit=1)
+            frames[step] = frame
+        # Stream every pipeline status; only resend gallery images at checkpoints.
+        yield (checkpoint_slots(checkpoints, frames) if frame is not None else gr.skip(),
+               progress_html(1000 - step, started))
+
+
+# Reuse the original page's palette, typography, header, and panel styling.
+PAGE_CSS = (ROOT / "styles.css").read_text() + """
+.gradio-container { max-width:1440px !important; padding:0 6% !important;
+  background:var(--paper) !important; color:var(--ink) !important;
+  background-image:radial-gradient(#59644d12 .7px,transparent .7px) !important;
+  background-size:8px 8px !important; font-family:var(--sans) !important; }
+/* Keep the brand close to the top and Configure close to its divider. */
+.gradio-container > .main { padding-top:0 !important; }
+#lab-header { padding:0 !important; margin:0 !important; }
+#lab-header .html-container { padding:0 !important; margin:0 !important; }
+#lab-header .topbar { height:38px; align-items:flex-start; padding-top:2px; }
+#lab-header .brand { font-size:24px; line-height:30px; }
+#lab-header .topbar-right { padding-top:12px; }
+#lab-main { gap:22px; padding:0; }
+.gradio-container > .main { padding-bottom:0 !important; flex-grow:0 !important; }
+.gradio-container, .gradio-container > .wrap { min-height:0 !important; }
+.gradio-container > .wrap { flex-grow:0 !important; }
+.gradio-container footer { margin-top:0 !important; padding:4px 0 !important; min-height:0 !important; }
+.gradio-container .gap:has(> #lab-header) { gap:0 !important; }
+.lab-panel { border-radius:0 !important; gap:12px !important; }
+.lab-panel h1,.lab-panel h2 { color:var(--ink); }
+.lab-panel.control-panel { padding-top:8px !important; }
+.control-panel > .html-container, .control-panel > .block { margin-top:0 !important; }
+.control-panel h1 { margin:0 !important; }
+.lab-panel.process-panel { padding:12px 14px !important; gap:8px !important; }
+.config-row { display:flex !important; gap:0 !important; width:100%; max-width:none; flex-wrap:nowrap !important; }
+.config-label { flex:0 0 110px !important; min-width:110px !important;
+  height:40px !important; min-height:40px !important; padding:0 12px !important; display:flex; align-items:center; justify-content:flex-start !important; opacity:1 !important;
+  background:#f0f2eb !important; border:1px solid #cfd3c9 !important;
+  border-radius:3px 0 0 3px !important; color:var(--green) !important;
+  font:10px var(--mono) !important; }
+.config-field { flex:1 1 0 !important; min-width:0 !important; border:0 !important; padding:0 !important; background:transparent !important; }
+/* Gradio nests the textbox inside a label; undo the original page's
+   .config-row label styles so that wrapper cannot become a second box. */
+.config-field label { display:block !important; width:100% !important; min-width:0 !important;
+  border:0 !important; padding:0 !important; margin:0 !important; background:transparent !important; }
+.config-field .input-container { width:100% !important; max-width:none !important; }
+.config-field textarea,.config-field input { width:100% !important; max-width:none !important; box-sizing:border-box !important; height:40px !important; min-height:40px !important;
+  border:1px solid #cfd3c9 !important; border-radius:0 3px 3px 0 !important;
+  padding:10px 12px !important; font:11px var(--mono) !important; background:white !important; white-space:pre !important; overflow-x:auto !important; resize:none !important; }
+#model-field textarea { background:#f8f7f2 !important; color:#687169 !important; }
+.process-button { flex:0 0 auto !important; min-width:76px !important;
+  height:28px !important; min-height:28px !important; padding:0 12px !important;
+  background:var(--green) !important; color:white !important; border-radius:3px !important;
+  font:10px var(--mono) !important; }
+#apply-timesteps { height:40px !important; border-radius:0 3px 3px 0 !important; }
+.process-heading { margin-bottom:0; }
+.process-toolbar { align-items:center !important; gap:12px !important; }
+/* Three columns keep progress centered independently of the controls. */
+.process-toolbar { display:grid !important; grid-template-columns:minmax(0,1fr) max-content minmax(0,1fr);
+  align-items:center !important; }
+.process-toolbar > .process-button { justify-self:start; }
+.process-toolbar > .process-button:last-child { justify-self:end; }
+.progress-counter { justify-self:center; min-width:0 !important; width:max-content !important; max-width:100%;
+  border:0 !important; background:transparent !important; padding:0 !important; }
+.progress-counter .html-container { padding:0 !important; }
+.progress-counter pre { margin:0; padding:3px 5px; color:var(--coral);
+  background:#f5f4ef; border:1px solid #d6dacf; border-radius:3px;
+  font:11px var(--mono); text-align:center; white-space:pre; overflow-x:auto; }
+.checkpoint-track { background:transparent !important; border:0 !important; padding:0 !important; }
+/* Gradio's auto-height galleries still impose a 450px minimum. */
+.checkpoint-track .grid-wrap { overflow-x:auto !important; overflow-y:hidden !important;
+  height:auto !important; min-height:0 !important; max-height:none !important;
+  padding:0 !important; }
+.checkpoint-track .grid-container { height:auto !important; min-height:0 !important; }
+.checkpoint-track .grid-container { gap:10px !important; grid-auto-rows:auto !important; }
+.checkpoint-track .thumbnail-item { min-width:0 !important; aspect-ratio:1 !important; height:auto !important; }
+.checkpoint-track .thumbnail-item { border-radius:0 !important; background:#e7e5dc !important; }
+.checkpoint-track .caption { font:11px var(--mono) !important; color:var(--green) !important; }
+.checkpoint-row { position:relative; gap:0 !important; overflow:visible !important; }
+.checkpoint-preview-host { position:relative; height:0 !important; min-height:0 !important; padding:0 !important;
+  margin:0 !important; border:0 !important; overflow:visible !important; z-index:1000; }
+.checkpoint-preview-host .html-container { padding:0 !important; overflow:visible !important; }
+.checkpoint-preview { position:absolute; bottom:8px; width:min(256px,100%); z-index:1000;
+  border:2px solid #fbfaf6; background:#fbfaf6; box-shadow:0 8px 24px #25312852; }
+.checkpoint-preview summary { cursor:pointer; padding:3px 6px; font:10px var(--mono); }
+.checkpoint-preview img { width:100%; height:auto; display:block; }
+.checkpoint-preview:not([open]) { display:none; }
+#reverse-input { justify-self:end; width:max-content !important; min-width:0 !important;
+  border:0 !important; background:transparent !important; padding:0 !important; }
+#reverse-input .wrap, #reverse-input .form, #reverse-input .container {
+  width:max-content !important; min-width:0 !important; padding:0 !important;
+  gap:6px !important; justify-content:flex-end !important; }
+#reverse-input label { flex:0 0 auto !important; width:max-content !important;
+  margin:0 !important; padding:4px 6px !important; }
+@media(max-width:620px) {
+  .gradio-container { padding:0 4% !important; }
+  .config-label { flex-basis:82px !important; min-width:82px !important; }
+  .process-heading { align-items:flex-start; flex-direction:column; }
+  .process-toolbar { grid-template-columns:1fr !important; }
+  .process-toolbar > :last-child { justify-self:end; }
+}
+"""
+PAGE_HEAD = """<link href="https://fonts.googleapis.com/css2?family=DM+Mono:wght@400;500&family=DM+Sans:wght@400;500;600;700&family=Space+Grotesk:wght@400;500;600;700&display=swap" rel="stylesheet">"""
+# Gradio 6 moved CSS/head from Blocks to launch; support both runtime versions.
+PAGE_OPTIONS = {"css": PAGE_CSS, "head": PAGE_HEAD}
+BLOCK_OPTIONS = PAGE_OPTIONS if "css" not in inspect.signature(gr.Blocks.launch).parameters else {}
+LAUNCH_OPTIONS = PAGE_OPTIONS if not BLOCK_OPTIONS else {}
+
+
+# Install after a native selection event so this also works when Spaces mounts
+# Blocks without using the local launch options. Traverse open shadow roots.
+INSTALL_PREVIEW_DISMISS = """() => {
+  if (window.diffusionDismissInstalled) return;
+  window.diffusionDismissInstalled = true;
+  const closeOutside = (root, path) => {
+    for (const preview of root.querySelectorAll(".checkpoint-preview[open]")) {
+      if (!path.includes(preview)) preview.open = false;
+    }
+    for (const element of root.querySelectorAll("*")) {
+      if (element.shadowRoot) closeOutside(element.shadowRoot, path);
+    }
+  };
+  document.addEventListener("pointerdown", event => {
+    closeOutside(document, event.composedPath());
+  }, true);
+  document.addEventListener("keydown", event => {
+    if (event.key === "Escape") closeOutside(document, []);
+  });
+}"""
+
+
+def progress_html(completed=0, started=None):
+    return "<pre>" + progress_text(completed, started) + "</pre>"
+
+
+def reset_tracks(seed, checkpoint_text, image):
+    seed, steps = parse_settings(seed, checkpoint_text)
+    initial = {0: image.convert("RGB").resize((256, 256))} if image is not None else {}
+    return (", ".join(map(str, steps)),
+            gr.Gallery(value=checkpoint_slots(steps, initial), columns=len(steps), allow_preview=False, interactive=False),
+            gr.Gallery(value=checkpoint_slots(steps, {1000: gaussian_noise(seed)}), columns=len(steps), allow_preview=False, interactive=False),
+            None, progress_html(), progress_html())
+
+
+def change_source(filepath, seed, checkpoint_text):
+    with Image.open(filepath) as uploaded:
+        image = uploaded.convert("RGB")
+    return (image, *reset_tracks(seed, checkpoint_text, image))
+
+
+def show_checkpoint_preview(gallery, event: gr.SelectData):
+    if not gallery:
+        return ""
+    index = int(event.index)
+    image, caption = gallery[index]
+    if isinstance(image, dict):
+        image = image["path"]
+    if not isinstance(image, Image.Image):
+        with Image.open(image) as loaded:
+            image = loaded.convert("RGB")
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+    center = 100 * (index + 0.5) / len(gallery)
+    # Force fresh markup even when the same thumbnail is clicked again, so
+    # a details element closed in the browser is recreated in its open state.
+    preview_id = secrets.token_hex(8)
+    return (f'<details open class="checkpoint-preview" data-preview-id="{preview_id}" ' 
+            f'style="left:clamp(0px, calc({center}% - 128px), max(0px, calc(100% - 256px)))">'
+            '<summary>Close preview</summary>'
+            f'<img src="data:image/png;base64,{encoded}" alt="{html.escape(caption or "Checkpoint", quote=True)}">'
+            '</details>')
+
+
+with gr.Blocks(title="Diffusion — A DDPM in motion", **BLOCK_OPTIONS) as demo:
+    gr.HTML('<header class="topbar"><a class="brand" href="#lab-main">'
+            '<span class="brand-mark" aria-hidden="true"><i></i><i></i><i></i><i></i></span>'
+            '<span>diffusion model</span></a><div class="topbar-right">'
+            '<span class="live-dot"></span></div></header>',
+            elem_id="lab-header", container=False)
+    with gr.Column(elem_id="lab-main"):
+        with gr.Column(elem_classes=["lab-panel", "control-panel"]):
+            gr.HTML("<h1>Configure</h1>")
+            with gr.Row(elem_classes="config-row"):
+                gr.Button("Model", interactive=False, scale=0, min_width=110, elem_classes="config-label")
+                gr.Textbox(value="google/ddpm-celebahq-256", interactive=False,
+                           scale=10, min_width=0, lines=1, max_lines=1, show_label=False, container=False, elem_id="model-field", elem_classes="config-field")
+            with gr.Row(elem_classes="config-row"):
+                gr.Button("Timesteps", interactive=False, scale=0, min_width=110, elem_classes="config-label")
+                checkpoints = gr.Textbox(value=", ".join(map(str, CHECKPOINTS)), show_label=False,
+                                         scale=10, min_width=0, lines=1, max_lines=1, container=False, elem_classes="config-field")
+                apply_button = gr.Button("Apply", scale=0, min_width=76, elem_id="apply-timesteps", elem_classes="process-button")
+            with gr.Row(elem_classes="config-row"):
+                randomize = gr.Button("Seed", scale=0, min_width=110, elem_classes="config-label")
+                seed = gr.Textbox(value="42", show_label=False, scale=10, min_width=0, lines=1, max_lines=1, container=False, elem_classes="config-field")
+            randomize.click(lambda: str(secrets.randbelow(2**31)), outputs=seed, queue=False)
+
+        default_image = ROOT / "assets/celeba_hq_validation_91_256.png"
+        initial_image = Image.open(default_image).convert("RGB") if default_image.exists() else None
+        initial_frames = {0: initial_image.resize((256, 256))} if initial_image is not None else {}
+        source = gr.State(initial_image)
+        forward_noise = gr.State(None)
+        with gr.Column(elem_classes=["lab-panel", "process-panel"]):
+            gr.HTML('<div class="process-heading"><h2>Forward process: image to noise</h2>'
+                    '<span class="equation">x<sub>t</sub> = √α<sub>t</sub> x<sub>t−1</sub>'
+                    ' + √(1 − α<sub>t</sub>) ε<sub>t</sub></span></div>')
+            with gr.Row(elem_classes="process-toolbar"):
+                diffuse_button = gr.Button("Diffuse", elem_classes="process-button")
+                forward_status = gr.HTML(value=progress_html(), container=False,
+                                         elem_classes="progress-counter")
+                upload_source = gr.UploadButton("Change image", file_types=["image"],
+                                                type="filepath", elem_classes="process-button")
+            with gr.Column(elem_classes="checkpoint-row"):
+                forward_preview = gr.HTML("", container=False, elem_classes="checkpoint-preview-host")
+                forward_gallery = gr.Gallery(value=checkpoint_slots(CHECKPOINTS, initial_frames),
+                                             show_label=False, columns=len(CHECKPOINTS), rows=1,
+                                             height="auto", object_fit="cover", allow_preview=False, interactive=False, elem_id="forward-images",
+                                             elem_classes="checkpoint-track")
+            forward_gallery.select(show_checkpoint_preview, [forward_gallery], [forward_preview],
+                                  queue=False, show_progress="hidden").then(
+                                      fn=None, js=INSTALL_PREVIEW_DISMISS, queue=False)
+            diffuse_button.click(diffuse, [source, seed, checkpoints],
+                                 [forward_gallery, forward_noise, forward_status])
+
+        with gr.Column(elem_classes=["lab-panel", "process-panel"]):
+            gr.HTML('<div class="process-heading"><h2>Reverse process: noise to face</h2>'
+                    '<p class="reverse-equation">x<sub>t−1</sub> = 1/√α<sub>t</sub>'
+                    ' (x<sub>t</sub> − β<sub>t</sub>/√(1−α̅<sub>t</sub>)'
+                    ' ε<sub>θ</sub>(x<sub>t</sub>,t)) + σ<sub>t</sub>z</p></div>')
+            with gr.Row(elem_classes="process-toolbar"):
+                denoise_button = gr.Button("Denoise", elem_classes="process-button")
+                status = gr.HTML(value=progress_html(), container=False,
+                                 elem_classes="progress-counter")
+                input_mode = gr.Radio(["Gaussian noise", "Forward noise"], value="Gaussian noise",
+                                      show_label=False, container=False, scale=0, min_width=0, elem_id="reverse-input")
+            with gr.Column(elem_classes="checkpoint-row"):
+                reverse_preview = gr.HTML("", container=False, elem_classes="checkpoint-preview-host")
+                reverse_gallery = gr.Gallery(value=checkpoint_slots(CHECKPOINTS, {1000: gaussian_noise(42)}),
+                                             show_label=False, columns=len(CHECKPOINTS), rows=1,
+                                             height="auto", object_fit="cover", allow_preview=False, interactive=False, elem_classes="checkpoint-track")
+            reverse_gallery.select(show_checkpoint_preview, [reverse_gallery], [reverse_preview],
+                                  queue=False, show_progress="hidden").then(
+                                      fn=None, js=INSTALL_PREVIEW_DISMISS, queue=False)
+            denoise_button.click(denoise, [seed, checkpoints, input_mode, forward_noise],
+                                 [reverse_gallery, status], concurrency_limit=1,
+                                 **({"stream_every": 0} if "stream_every" in
+                                    inspect.signature(denoise_button.click).parameters else {}))
+        def clear_previews():
+            return "", ""
+
+        for control, event_name in [(apply_button, "click"), (upload_source, "upload"),
+                                    (diffuse_button, "click"), (denoise_button, "click")]:
+            getattr(control, event_name)(clear_previews, outputs=[forward_preview, reverse_preview],
+                                         queue=False, show_progress="hidden")
+        reset_outputs = [checkpoints, forward_gallery, reverse_gallery,
+                         forward_noise, forward_status, status]
+        apply_button.click(reset_tracks, [seed, checkpoints, source], reset_outputs, queue=False)
+        upload_source.upload(change_source, [upload_source, seed, checkpoints],
+                             [source, *reset_outputs], queue=False)
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch(**LAUNCH_OPTIONS)