Add a live per-step generation progress bar with ETA

studio/backend/core/inference/diffusion.py

@@ -13,7 +13,9 @@ policy lives in the arbiter the routes call, not here.
 
 from __future__ import annotations
 
+import inspect
 import threading
+import time
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Optional
@@ -53,6 +55,17 @@ class _LoadingState:
     error: Optional[str] = None
 
 
+@dataclass
+class _GenState:
+    """An in-flight generation, updated per denoising step for the progress bar."""
+
+    total_steps: int
+    step: int = 0
+    # Set when the first step finishes; the ETA rate is measured from there so the
+    # slower first step (warmup) doesn't skew it.
+    first_step_at: float = 0.0
+
+
 class DiffusionBackend:
     """Holds at most one loaded diffusers pipeline. All mutations are serialised."""
 
@@ -64,6 +77,9 @@ class DiffusionBackend:
         self._lock = threading.Lock()
         self._state: Optional[_LoadState] = None
         self._loading: Optional[_LoadingState] = None
+        # The callback mutates this and generate_progress() reads it, both without
+        # the lock (generate holds it for the whole call), so polling stays live.
+        self._gen: Optional[_GenState] = None
 
     @property
     def is_loaded(self) -> bool:
@@ -313,11 +329,47 @@ class DiffusionBackend:
             if negative_prompt:
                 kwargs["negative_prompt"] = negative_prompt
 
-            images = state.pipe(**kwargs).images
+            gen = _GenState(total_steps = steps)
+
+            def _on_step(pipe, step_index, timestep, callback_kwargs):
+                gen.step = step_index + 1
+                if gen.first_step_at == 0.0:
+                    gen.first_step_at = time.time()
+                return callback_kwargs
+
+            # Not every pipeline accepts the callback; only pass it where supported
+            # so the step counter never breaks generation.
+            if "callback_on_step_end" in inspect.signature(state.pipe.__call__).parameters:
+                kwargs["callback_on_step_end"] = _on_step
+
+            self._gen = gen
+            try:
+                images = state.pipe(**kwargs).images
+            finally:
+                self._gen = None
             # Return the PIL images (not yet encoded): the route embeds each
             # image's recipe and persists it via the gallery.
             return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id}
 
+    def generate_progress(self) -> dict[str, Any]:
+        """Live per-step progress for an in-flight generation (lock-free read)."""
+        gen = self._gen
+        if gen is None or gen.total_steps <= 0:
+            return {"active": False, "step": 0, "total_steps": 0, "fraction": 0.0, "eta_seconds": None}
+        step, total = gen.step, gen.total_steps
+        eta = None
+        steps_since_first = step - 1
+        if gen.first_step_at and steps_since_first > 0:
+            per_step = (time.time() - gen.first_step_at) / steps_since_first
+            eta = max(0.0, (total - step) * per_step)
+        return {
+            "active": True,
+            "step": step,
+            "total_steps": total,
+            "fraction": min(step / total, 1.0),
+            "eta_seconds": eta,
+        }
+
     def unload(self) -> dict[str, Any]:
         with self._lock:
             self._unload_locked()

studio/backend/models/inference.py

@@ -1759,6 +1759,16 @@ class GalleryListResponse(BaseModel):
     images: list[GalleryImage] = Field(default_factory = list)
 
 
+class DiffusionGenerateProgressResponse(BaseModel):
+    """Live per-step progress for an in-flight generation."""
+
+    active: bool = Field(False, description = "Whether a generation is running")
+    step: int = Field(0, description = "Denoising steps completed so far")
+    total_steps: int = Field(0, description = "Total denoising steps for this run")
+    fraction: float = Field(0.0, description = "step / total_steps, clamped to [0,1]")
+    eta_seconds: Optional[float] = Field(None, description = "Estimated seconds remaining")
+
+
 class DiffusionLoadProgressResponse(BaseModel):
     """Download/finalize progress for an in-flight diffusion load."""
 

studio/backend/routes/inference.py

@@ -1052,6 +1052,7 @@ from models.inference import (
     DiffusionLoadRequest,
     DiffusionGenerateRequest,
     DiffusionGenerateResponse,
+    DiffusionGenerateProgressResponse,
     DiffusionStatusResponse,
     DiffusionLoadProgressResponse,
     GalleryImage,
@@ -10213,3 +10214,10 @@ async def diffusion_load_progress(current_subject: str = Depends(get_current_sub
     from core.inference.diffusion import get_diffusion_backend
 
     return DiffusionLoadProgressResponse(**get_diffusion_backend().load_progress())
+
+
+@studio_router.get("/images/generate-progress", response_model = DiffusionGenerateProgressResponse)
+async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)):
+    from core.inference.diffusion import get_diffusion_backend
+
+    return DiffusionGenerateProgressResponse(**get_diffusion_backend().generate_progress())

studio/backend/tests/test_diffusion_backend.py

@@ -296,6 +296,22 @@ def test_load_progress_fraction_clamped(monkeypatch):
     assert p["bytes_downloaded"] == 1000  # clamped to the estimate
 
 
+def test_generate_progress_reports_step_and_eta():
+    import time as _time
+
+    from core.inference.diffusion import _GenState
+
+    backend = DiffusionBackend()
+    assert backend.generate_progress()["active"] is False
+
+    # Halfway through an 8-step run, ~1s/step measured since the first step.
+    backend._gen = _GenState(total_steps = 8, step = 4, first_step_at = _time.time() - 3)
+    p = backend.generate_progress()
+    assert p["active"] is True and p["step"] == 4 and p["total_steps"] == 8
+    assert abs(p["fraction"] - 0.5) < 1e-9
+    assert p["eta_seconds"] is not None and 0 < p["eta_seconds"] < 30  # ~4 steps left
+
+
 def test_begin_load_rejects_concurrent(monkeypatch):
     backend = DiffusionBackend()
     monkeypatch.setattr(DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0))

studio/frontend/src/features/images/api.ts

@@ -14,6 +14,14 @@ export interface DiffusionStatus {
   cpu_offload: boolean;
 }
 
+export interface DiffusionGenerateProgress {
+  active: boolean;
+  step: number;
+  total_steps: number;
+  fraction: number;
+  eta_seconds: number | null;
+}
+
 export interface DiffusionLoadProgress {
   phase: "downloading" | "finalizing" | "ready" | "error" | null;
   bytes_downloaded: number;
@@ -77,6 +85,10 @@ export async function getDiffusionLoadProgress(): Promise<DiffusionLoadProgress>
   return parseJson(await authFetch("/api/inference/images/load-progress"));
 }
 
+export async function getGenerateProgress(): Promise<DiffusionGenerateProgress> {
+  return parseJson(await authFetch("/api/inference/images/generate-progress"));
+}
+
 export async function loadDiffusionModel(body: DiffusionLoadRequest): Promise<DiffusionStatus> {
   return parseJson(
     await authFetch("/api/inference/images/load", {

studio/frontend/src/features/images/images-page.tsx

@@ -43,6 +43,7 @@ import { cn } from "@/lib/utils";
 import { toast } from "@/lib/toast";
 
 import {
+  type DiffusionGenerateProgress,
   type DiffusionLoadProgress,
   type DiffusionStatus,
   type GalleryImage,
@@ -52,6 +53,7 @@ import {
   getDiffusionLoadProgress,
   getDiffusionStatus,
   getGallery,
+  getGenerateProgress,
   loadDiffusionModel,
   unloadDiffusionModel,
 } from "./api";
@@ -135,6 +137,18 @@ function formatTimestamp(epochSeconds: number): string {
   return new Date(epochSeconds * 1000).toLocaleString();
 }
 
+function formatEta(seconds: number): string {
+  if (seconds < 1) return "<1s";
+  if (seconds < 60) return `${Math.round(seconds)}s`;
+  return `${Math.floor(seconds / 60)}m ${Math.round(seconds % 60)}s`;
+}
+
+// Bar label for an in-flight generation: step count plus an ETA once it's known.
+function genStepLabel(p: DiffusionGenerateProgress): string {
+  const base = `Step ${p.step}/${p.total_steps}`;
+  return p.eta_seconds != null ? `${base} · ~${formatEta(p.eta_seconds)} left` : base;
+}
+
 // The chat tab's model-load toast styling, reused verbatim so the diffusion
 // load toast is visually identical (persistent, progress bar, same chrome).
 const LOAD_TOAST_CLASSNAMES = {
@@ -359,6 +373,9 @@ export function ImagesPage() {
   const [busy, setBusy] = useState<Busy>(null);
   // {done, total} while a multi-run generation is in flight (for the button).
   const [genProgress, setGenProgress] = useState<{ done: number; total: number } | null>(null);
+  // Live per-step progress (step / total + ETA) polled during generation.
+  const [genStep, setGenStep] = useState<DiffusionGenerateProgress | null>(null);
+  const genPollTimer = useRef<ReturnType<typeof setInterval> | null>(null);
   const [status, setStatus] = useState<DiffusionStatus | null>(null);
   // Records come from the backend (durable); srcById maps each id to its object
   // URL (loaded images) or data URL (the one just generated).
@@ -491,9 +508,10 @@ export function ImagesPage() {
 
   useEffect(() => {
     void refreshStatus();
-    // Stop polling if the page unmounts mid-load.
+    // Stop polling if the page unmounts mid-load / mid-generate.
     return () => {
       if (pollTimer.current) clearTimeout(pollTimer.current);
+      if (genPollTimer.current) clearInterval(genPollTimer.current);
     };
   }, [refreshStatus]);
 
@@ -602,6 +620,17 @@ export function ImagesPage() {
 
     setBusy("generating");
     setGenProgress({ done: 0, total: count });
+    setGenStep(null);
+    // Poll the backend's per-step progress across the whole run (all sequential
+    // generations), so the bar tracks the live denoising steps.
+    genPollTimer.current = setInterval(async () => {
+      try {
+        const p = await getGenerateProgress();
+        setGenStep(p.active ? p : null);
+      } catch {
+        // transient; keep polling
+      }
+    }, 300);
     try {
       for (let i = 0; i < count; i++) {
         const res = await generateDiffusionImage({
@@ -625,8 +654,11 @@ export function ImagesPage() {
     } catch (err) {
       toast.error(err instanceof Error ? err.message : "Image generation failed");
     } finally {
+      if (genPollTimer.current) clearInterval(genPollTimer.current);
+      genPollTimer.current = null;
       setBusy(null);
       setGenProgress(null);
+      setGenStep(null);
     }
   }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, ensureSrc]);
 
@@ -796,14 +828,13 @@ export function ImagesPage() {
                   </Button>
                 </div>
               </>
-            ) : busy === "generating" || selected ? (
-              // First image generating, or the selected record's blob is still
-              // loading — spin in place rather than flashing the empty state.
+            ) : selected ? (
+              // The selected record's blob is still loading — spin in place.
               <div className="flex flex-col items-center gap-3 text-muted-foreground">
                 <Spinner className="size-8" />
-                <p className="text-sm">{busy === "generating" ? "Generating…" : "Loading…"}</p>
+                <p className="text-sm">Loading…</p>
               </div>
-            ) : (
+            ) : busy === "generating" ? null : (
               <div className="flex flex-col items-center gap-3 text-muted-foreground">
                 <HugeiconsIcon icon={ImageAdd02Icon} className="size-12" strokeWidth={1.5} />
                 <p className="text-sm">
@@ -813,6 +844,30 @@ export function ImagesPage() {
                 </p>
               </div>
             )}
+
+            {/* Live generation progress: a per-step bar with ETA, centered when
+                there's nothing else to show, tucked at the bottom over an image. */}
+            {busy === "generating" && (
+              <div
+                className={cn(
+                  "pointer-events-none absolute flex justify-center px-4",
+                  selectedSrc ? "inset-x-0 bottom-4" : "inset-0 items-center",
+                )}
+              >
+                <div className="w-72 max-w-full rounded-xl bg-background/85 p-3 shadow-lg ring-1 ring-border backdrop-blur">
+                  <ModelLoadDescription
+                    title={
+                      genProgress && genProgress.total > 1
+                        ? `Generating · run ${genProgress.done + 1}/${genProgress.total}`
+                        : "Generating…"
+                    }
+                    message="Starting…"
+                    progressPercent={genStep ? genStep.fraction * 100 : null}
+                    progressLabel={genStep ? genStepLabel(genStep) : null}
+                  />
+                </div>
+              </div>
+            )}
           </div>
 
           {(images.length > 0 || busy === "generating") && (