From 2bf9d73b4af719ffd36b2b4b673be7110a8ec2e0 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 24 Jun 2026 19:50:10 -0300
Subject: [PATCH] Reuse formatEta, compute generation ETA once per step, guard
 re-renders

---
 studio/backend/core/inference/diffusion.py    | 30 ++++++++++++-------
 .../backend/tests/test_diffusion_backend.py   | 20 +++++++++----
 .../src/features/images/images-page.tsx       | 21 ++++++-------
 3 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 03cc0bb865..fa3ed47746 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -64,6 +64,18 @@ class _GenState:
     # 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
+    # Computed once per step (in the callback) so it's stable between polls.
+    eta_seconds: Optional[float] = None
+
+
+def _estimate_eta(total_steps: int, step: int, first_step_at: float, now: float) -> Optional[float]:
+    """Seconds remaining, from the average step time measured after the first step.
+    None until at least one step has elapsed since the first."""
+    steps_since_first = step - 1
+    if not first_step_at or steps_since_first <= 0:
+        return None
+    per_step = (now - first_step_at) / steps_since_first
+    return max(0.0, (total_steps - step) * per_step)
 
 
 class DiffusionBackend:
@@ -332,9 +344,11 @@ class DiffusionBackend:
             gen = _GenState(total_steps = steps)
 
             def _on_step(pipe, step_index, timestep, callback_kwargs):
+                now = time.time()
                 gen.step = step_index + 1
                 if gen.first_step_at == 0.0:
-                    gen.first_step_at = time.time()
+                    gen.first_step_at = now
+                gen.eta_seconds = _estimate_eta(gen.total_steps, gen.step, gen.first_step_at, now)
                 return callback_kwargs
 
             # Not every pipeline accepts the callback; only pass it where supported
@@ -356,18 +370,12 @@ class DiffusionBackend:
         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,
+            "step": gen.step,
+            "total_steps": gen.total_steps,
+            "fraction": min(gen.step / gen.total_steps, 1.0),
+            "eta_seconds": gen.eta_seconds,
         }
 
     def unload(self) -> dict[str, Any]:
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index 4b3eb1d6bf..51ef8b246b 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -296,20 +296,28 @@ 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
+def test_estimate_eta():
+    from core.inference.diffusion import _estimate_eta
 
+    # No rate yet until a step has elapsed since the first.
+    assert _estimate_eta(8, 1, first_step_at = 100.0, now = 100.0) is None
+    assert _estimate_eta(8, 0, first_step_at = 0.0, now = 100.0) is None
+    # 3 steps in 3s since the first ⇒ 1s/step ⇒ 4 steps left ⇒ ~4s.
+    assert _estimate_eta(8, 4, first_step_at = 100.0, now = 103.0) == 4.0
+    # Last step ⇒ 0 remaining.
+    assert _estimate_eta(8, 8, first_step_at = 100.0, now = 107.0) == 0.0
+
+
+def test_generate_progress_reads_gen_state():
     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)
+    backend._gen = _GenState(total_steps = 8, step = 4, eta_seconds = 4.0)
     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
+    assert abs(p["fraction"] - 0.5) < 1e-9 and p["eta_seconds"] == 4.0
 
 
 def test_begin_load_rejects_concurrent(monkeypatch):
diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx
index 47ca5a076a..0fef65b2cc 100644
--- a/studio/frontend/src/features/images/images-page.tsx
+++ b/studio/frontend/src/features/images/images-page.tsx
@@ -38,7 +38,7 @@ import type {
   ModelSelectorChangeMeta,
 } from "@/components/assistant-ui/model-selector/types";
 import { ModelLoadDescription } from "@/features/chat/components/model-load-status";
-import { formatBytes } from "@/features/hub/lib/format";
+import { formatBytes, formatEta } from "@/features/hub/lib/format";
 import { cn } from "@/lib/utils";
 import { toast } from "@/lib/toast";
 
@@ -137,16 +137,12 @@ 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.
+// Bar label for an in-flight generation: step count plus an ETA once it's known
+// (formatEta returns "" for non-positive, so the last step shows just the step).
 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;
+  const eta = p.eta_seconds != null ? formatEta(p.eta_seconds) : "";
+  return eta ? `${base} · ~${eta}` : base;
 }
 
 // The chat tab's model-load toast styling, reused verbatim so the diffusion
@@ -626,7 +622,12 @@ export function ImagesPage() {
     genPollTimer.current = setInterval(async () => {
       try {
         const p = await getGenerateProgress();
-        setGenStep(p.active ? p : null);
+        // Skip the state update (and re-render) when nothing the bar shows moved.
+        setGenStep((prev) => {
+          if (!p.active) return null;
+          if (prev && prev.step === p.step && prev.eta_seconds === p.eta_seconds) return prev;
+          return p;
+        });
       } catch {
         // transient; keep polling
       }