UI: min-height 500px auto-sizing output, gr.Progress step reporting

app.py

@@ -11,6 +11,8 @@ os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 import spaces  # MUST be imported before torch / transformers / sensenova_u1
 
 import random
+import threading
+import time
 
 import gradio as gr
 import numpy as np
@@ -211,6 +213,27 @@ def _estimate_duration(pil_inputs, prompt, aspect_ratio, seed, *args, **kwargs):
     return 120
 
 
+def _register_step_hook(counter: dict) -> object | None:
+    """Count denoising steps by hooking the per-step timestep embedder.
+
+    The sampling loops in ``t2i_generate`` / ``it2i_generate`` call
+    ``fm_modules['timestep_embedder']`` exactly once per step, so a forward
+    hook on it is a cheap, non-invasive progress signal.
+    """
+    try:
+        embedder = model.fm_modules["timestep_embedder"]
+    except Exception:  # pragma: no cover - defensive
+        return None
+
+    def _tick(_module, _args, _output):
+        counter["step"] += 1
+
+    try:
+        return embedder.register_forward_hook(_tick)
+    except Exception:  # pragma: no cover - defensive
+        return None
+
+
 @spaces.GPU(duration=_estimate_duration)
 def _generate_on_gpu(
     pil_inputs: list,
@@ -218,52 +241,82 @@ def _generate_on_gpu(
     aspect_ratio: str,
     seed: int,
 ):
-    """GPU worker. Inputs are already validated / preprocessed on CPU."""
+    """GPU worker. Inputs are already validated / preprocessed on CPU.
+
+    Runs sampling on a background thread inside the GPU fork and yields
+    ``("progress", steps_done, total_steps)`` ticks while it works, then a
+    final ``("result", image)``. The caller turns those ticks into a
+    ``gr.Progress`` bar on the output component.
+    """
     torch.cuda.reset_peak_memory_stats()
+    counter = {"step": 0}
+    total_steps = DEFAULT_NUM_STEPS
+    hook = _register_step_hook(counter)
+    result: dict = {}
+
+    def _run():
+        try:
+            with torch.inference_mode():
+                if not pil_inputs:
+                    width, height = T2I_RESOLUTIONS[aspect_ratio]
+                    tensor = model.t2i_generate(
+                        tokenizer,
+                        prompt,
+                        image_size=(width, height),
+                        cfg_scale=DEFAULT_CFG_SCALE,
+                        cfg_norm="none",
+                        timestep_shift=DEFAULT_TIMESTEP_SHIFT,
+                        cfg_interval=(0.0, 1.0),
+                        num_steps=DEFAULT_NUM_STEPS,
+                        batch_size=1,
+                        seed=int(seed),
+                        think_mode=False,
+                    )
+                else:
+                    out_w, out_h = _editing_output_size(pil_inputs[0], EDIT_TARGET_PIXELS)
+                    tensor = model.it2i_generate(
+                        tokenizer,
+                        prompt,
+                        pil_inputs,
+                        image_size=(out_w, out_h),
+                        cfg_scale=DEFAULT_CFG_SCALE,
+                        img_cfg_scale=1.0,
+                        cfg_norm="none",
+                        timestep_shift=DEFAULT_TIMESTEP_SHIFT,
+                        cfg_interval=(0.0, 1.0),
+                        num_steps=DEFAULT_NUM_STEPS,
+                        batch_size=1,
+                        think_mode=False,
+                        seed=int(seed),
+                    )
+            result["image"] = _to_pil(tensor)[0]
+            del tensor
+        except BaseException as exc:  # re-raised on the generator thread below
+            result["error"] = exc
+
+    worker = threading.Thread(target=_run, daemon=True)
+    started = time.time()
     try:
-        with torch.inference_mode():
-            if not pil_inputs:
-                width, height = T2I_RESOLUTIONS[aspect_ratio]
-                tensor = model.t2i_generate(
-                    tokenizer,
-                    prompt,
-                    image_size=(width, height),
-                    cfg_scale=DEFAULT_CFG_SCALE,
-                    cfg_norm="none",
-                    timestep_shift=DEFAULT_TIMESTEP_SHIFT,
-                    cfg_interval=(0.0, 1.0),
-                    num_steps=DEFAULT_NUM_STEPS,
-                    batch_size=1,
-                    seed=int(seed),
-                    think_mode=False,
-                )
-            else:
-                out_w, out_h = _editing_output_size(pil_inputs[0], EDIT_TARGET_PIXELS)
-                tensor = model.it2i_generate(
-                    tokenizer,
-                    prompt,
-                    pil_inputs,
-                    image_size=(out_w, out_h),
-                    cfg_scale=DEFAULT_CFG_SCALE,
-                    img_cfg_scale=1.0,
-                    cfg_norm="none",
-                    timestep_shift=DEFAULT_TIMESTEP_SHIFT,
-                    cfg_interval=(0.0, 1.0),
-                    num_steps=DEFAULT_NUM_STEPS,
-                    batch_size=1,
-                    think_mode=False,
-                    seed=int(seed),
-                )
-        images_out = _to_pil(tensor)
-        del tensor
+        worker.start()
+        while True:
+            worker.join(0.4)
+            done = min(counter["step"], total_steps)
+            if not worker.is_alive():
+                break
+            yield ("progress", done, total_steps, time.time() - started)
         print(f"[generate] peak GPU memory: {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB")
-        return images_out[0]
-    except torch.cuda.OutOfMemoryError as exc:  # pragma: no cover - depends on load
-        raise gr.Error(
-            "Ran out of GPU memory for this request. Try a smaller input image "
-            "or a smaller aspect ratio."
-        ) from exc
+        exc = result.get("error")
+        if isinstance(exc, torch.cuda.OutOfMemoryError):  # pragma: no cover
+            raise gr.Error(
+                "Ran out of GPU memory for this request. Try a smaller input image "
+                "or a smaller aspect ratio."
+            ) from exc
+        if exc is not None:
+            raise exc
+        yield ("result", result["image"], total_steps, time.time() - started)
     finally:
+        if hook is not None:
+            hook.remove()
         # Never leave a failed call's tensors behind for the next request.
         torch.cuda.empty_cache()
 
@@ -274,7 +327,7 @@ def generate(
     aspect_ratio: str = "1:1",
     seed: int = 42,
     randomize_seed: bool = True,
-    progress=gr.Progress(track_tqdm=True),
+    progress=gr.Progress(),
 ):
     """Generate an image from a text prompt, or edit an uploaded image.
 
@@ -307,7 +360,31 @@ def generate(
         sizes = ", ".join(f"{im.width}x{im.height}" for im in pil_inputs)
         print(f"[generate] editing with {len(pil_inputs)} input image(s) at {sizes}")
 
-    image_out = _generate_on_gpu(pil_inputs, prompt, aspect_ratio, seed)
+    # Drive a progress bar on the output component from the GPU worker's ticks.
+    estimated = _estimate_duration(pil_inputs, prompt, aspect_ratio, seed)
+    progress(0.0, desc="Starting generation…")
+    image_out = None
+    for kind, payload, total_steps, elapsed in _generate_on_gpu(
+        pil_inputs, prompt, aspect_ratio, seed
+    ):
+        if kind == "progress":
+            steps_done = int(payload)
+            if steps_done > 0:
+                progress(
+                    min(steps_done / max(total_steps, 1), 1.0),
+                    desc=f"Denoising step {steps_done}/{total_steps}",
+                )
+            else:
+                # Prefix / conditioning pass, before the first denoising step.
+                progress(
+                    min(elapsed / max(estimated, 1), 0.05),
+                    desc="Encoding prompt…",
+                )
+        else:
+            image_out = payload
+            progress(1.0, desc="Done")
+    if image_out is None:  # pragma: no cover - worker always yields a result
+        raise gr.Error("Generation produced no image. Please try again.")
     return image_out, seed
 
 
@@ -316,6 +393,7 @@ def generate_t2i(
     aspect_ratio: str = "1:1",
     seed: int = 42,
     randomize_seed: bool = True,
+    progress=gr.Progress(),
 ):
     """Text-to-image only: generate an image from a prompt (no input image).
 
@@ -328,7 +406,7 @@ def generate_t2i(
     # Wired to the text-to-image gr.Examples: those rows carry (prompt,
     # aspect_ratio) only, so they must NOT be bound straight to `generate`,
     # whose first argument is the input-image list.
-    return generate(None, prompt, aspect_ratio, seed, randomize_seed)
+    return generate(None, prompt, aspect_ratio, seed, randomize_seed, progress=progress)
 
 
 def generate_edit(
@@ -336,6 +414,7 @@ def generate_edit(
     prompt: str,
     seed: int = 42,
     randomize_seed: bool = True,
+    progress=gr.Progress(),
 ):
     """Image editing: apply an edit instruction to the uploaded image(s).
 
@@ -345,7 +424,7 @@ def generate_edit(
         seed: RNG seed for reproducible sampling.
         randomize_seed: if True, pick a fresh random seed each run.
     """
-    return generate(images, prompt, "1:1", seed, randomize_seed)
+    return generate(images, prompt, "1:1", seed, randomize_seed, progress=progress)
 
 
 # T2I examples: prompt + aspect ratio
@@ -384,6 +463,18 @@ EDIT_EXAMPLES = [
 CSS = """
 #col-container { max-width: 1100px; margin: 0 auto; }
 .dark .gradio-container { color: var(--body-text-color); }
+
+/* Output: let the frame grow to the natural height of the displayed image,
+   with a 500px floor so an empty / small result never collapses. */
+#output-image { min-height: 500px; }
+#output-image .image-container,
+#output-image .image-frame,
+#output-image .empty { min-height: 500px; height: auto; }
+#output-image img {
+    height: auto;
+    max-height: none;
+    object-fit: contain;
+}
 """
 
 with gr.Blocks(title="SenseNova-U1.5-8B-MoT") as demo:
@@ -423,7 +514,15 @@ with gr.Blocks(title="SenseNova-U1.5-8B-MoT") as demo:
                 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             generate_button = gr.Button("Generate", variant="primary")
         with gr.Column(scale=1):
-            output_image = gr.Image(label="Output", type="pil", format="png", interactive=False)
+            # The output component sizes itself to the generated image (no fixed
+            # `height`), with a 500px minimum enforced in CSS via this elem_id.
+            output_image = gr.Image(
+                label="Output",
+                type="pil",
+                format="png",
+                interactive=False,
+                elem_id="output-image",
+            )
             used_seed = gr.Number(label="Seed used", interactive=False)
 
     with gr.Accordion("Text-to-Image Examples", open=True):