From 54d01a505feebe33b02c2875296b20dd4edca65b Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 24 Jun 2026 19:41:15 -0300 Subject: [PATCH] Add a live per-step generation progress bar with ETA --- studio/backend/core/inference/diffusion.py | 54 ++++++++++++++- studio/backend/models/inference.py | 10 +++ studio/backend/routes/inference.py | 8 +++ .../backend/tests/test_diffusion_backend.py | 16 +++++ studio/frontend/src/features/images/api.ts | 12 ++++ .../src/features/images/images-page.tsx | 67 +++++++++++++++++-- 6 files changed, 160 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 77eb712ee4..03cc0bb865 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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() diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index e8b800cbf8..79539b8ae9 100644 --- a/studio/backend/models/inference.py +++ b/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.""" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 6f7772417a..6eed4a1d08 100644 --- a/studio/backend/routes/inference.py +++ b/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()) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 7c3bef6774..4b3eb1d6bf 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/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)) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 4ab76b0360..6f2e6f623a 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/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 return parseJson(await authFetch("/api/inference/images/load-progress")); } +export async function getGenerateProgress(): Promise { + return parseJson(await authFetch("/api/inference/images/generate-progress")); +} + export async function loadDiffusionModel(body: DiffusionLoadRequest): Promise { return parseJson( await authFetch("/api/inference/images/load", { diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index f22370072e..47ca5a076a 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/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(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(null); + const genPollTimer = useRef | null>(null); const [status, setStatus] = useState(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() { - ) : 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.
-

{busy === "generating" ? "Generating…" : "Loading…"}

+

Loading…

- ) : ( + ) : busy === "generating" ? null : (

@@ -813,6 +844,30 @@ export function ImagesPage() {

)} + + {/* 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" && ( +
+
+ 1 + ? `Generating · run ${genProgress.done + 1}/${genProgress.total}` + : "Generating…" + } + message="Starting…" + progressPercent={genStep ? genStep.fraction * 100 : null} + progressLabel={genStep ? genStepLabel(genStep) : null} + /> +
+
+ )} {(images.length > 0 || busy === "generating") && (